cubecl_cpp/metal/
builtin.rs

1use std::collections::{HashMap, HashSet};
2
3use cubecl_core::{
4    frontend::HasValue,
5    ir::{
6        Builtin, Scope,
7        attributes::{FuncInterface, IndexAttr},
8        dialect::general::ReadBuiltinOp,
9        ident,
10        prelude::*,
11    },
12};
13use pliron::{
14    builtin::{
15        given_names::set_block_arg_name,
16        ops::FuncOp,
17        types::{IntegerType, Signedness},
18    },
19    irbuild::{listener::DummyListener, match_rewrite::MatchRewrite},
20    value::Value,
21};
22
23use crate::{
24    metal::{BuiltInAttr, metal_op_with_out, signature::ATTR_BUILTIN_ATTRIBUTE},
25    shared::{
26        CompilationState,
27        builtin::{LowerBuiltins, absolute_pos, constant, cube_count, cube_pos},
28    },
29    target::Metal,
30};
31
32#[cube_op(name = "msl.read_dim3_builtin")]
33#[result_ty(argument)]
34pub struct ReadDim3BuiltinOp {
35    pub builtin: Value,
36    pub dim: IndexAttr,
37}
38
39metal_op_with_out!(ReadDim3BuiltinOp, |op, ctx| {
40    format!("{}[{}]", op.builtin(ctx).name(ctx), op.dim(ctx).0)
41});
42
43impl MatchRewrite for LowerBuiltins<Metal> {
44    fn r#match(&mut self, ctx: &Context, op: Ptr<Operation>) -> bool {
45        op.is_op::<ReadBuiltinOp>(ctx)
46    }
47
48    fn rewrite(
49        &mut self,
50        ctx: &mut Context,
51        rewriter: &mut MatchRewriter,
52        op: Ptr<Operation>,
53    ) -> Result<()> {
54        let builtin = op.as_op::<ReadBuiltinOp>(ctx).unwrap().builtin(ctx).0;
55        let scope = Scope::from_context_and_inserter(ctx, rewriter);
56        if let Some(new_value) = builtin.maybe_lower_metal(&scope) {
57            rewriter.replace_operation_with_values(ctx, op, vec![new_value]);
58        }
59        Ok(())
60    }
61}
62
63trait MetalBuiltin {
64    fn maybe_lower_metal(&self, scope: &Scope) -> Option<Value>;
65}
66
67impl MetalBuiltin for Builtin {
68    fn maybe_lower_metal(&self, scope: &Scope) -> Option<Value> {
69        let cube_dim = scope.ctx().aux_ty::<CompilationState>().cube_dim;
70        match self {
71            Builtin::UnitPos => None,
72            // This is common enough to be worth replacing. Z is almost always 1, and Y is often 1.
73            // Replacing it with a constant allows simplifying the positional math
74            Builtin::UnitPosX if cube_dim.x == 1 => Some(constant::expand(scope, 0).value(scope)),
75            Builtin::UnitPosY if cube_dim.y == 1 => Some(constant::expand(scope, 0).value(scope)),
76            Builtin::UnitPosZ if cube_dim.z == 1 => Some(constant::expand(scope, 0).value(scope)),
77            Builtin::UnitPosX | Builtin::UnitPosY | Builtin::UnitPosZ => None,
78            Builtin::CubePosCluster => Some(constant::expand(scope, 0).value(scope)),
79            Builtin::CubePosClusterX => Some(constant::expand(scope, 0).value(scope)),
80            Builtin::CubePosClusterY => Some(constant::expand(scope, 0).value(scope)),
81            Builtin::CubePosClusterZ => Some(constant::expand(scope, 0).value(scope)),
82            Builtin::CubePos => Some(cube_pos::expand(scope).value(scope)),
83            Builtin::CubePosX | Builtin::CubePosY | Builtin::CubePosZ => None,
84            Builtin::CubeDim => Some(constant::expand(scope, cube_dim.num_elems()).value(scope)),
85            Builtin::CubeDimX => Some(constant::expand(scope, cube_dim.x).value(scope)),
86            Builtin::CubeDimY => Some(constant::expand(scope, cube_dim.y).value(scope)),
87            Builtin::CubeDimZ => Some(constant::expand(scope, cube_dim.z).value(scope)),
88            Builtin::CubeClusterDim => Some(constant::expand(scope, 1).value(scope)),
89            Builtin::CubeClusterDimX => Some(constant::expand(scope, 1).value(scope)),
90            Builtin::CubeClusterDimY => Some(constant::expand(scope, 1).value(scope)),
91            Builtin::CubeClusterDimZ => Some(constant::expand(scope, 1).value(scope)),
92            Builtin::CubeCount => Some(cube_count::expand(scope).value(scope)),
93            Builtin::CubeCountX | Builtin::CubeCountY | Builtin::CubeCountZ => None,
94            Builtin::PlaneDim => None,
95            Builtin::PlanePos => None,
96            Builtin::UnitPosPlane => None,
97            Builtin::AbsolutePos => Some(absolute_pos::expand(scope).value(scope)),
98            Builtin::AbsolutePosX | Builtin::AbsolutePosY | Builtin::AbsolutePosZ => None,
99        }
100    }
101}
102
103pub fn append_msl_builtins(ctx: &mut Context, entry_func: FuncOp) {
104    let op = entry_func.get_operation();
105    let mut used = HashSet::new();
106    let mut read_ops = HashSet::new();
107    let state = &mut (&mut used, &mut read_ops);
108    visit_all_ops_of_type::<ReadBuiltinOp, _>(ctx, state, op, |ctx, (used, ops), op| {
109        ops.insert(op);
110        used.insert(built_in_attr(op.builtin(ctx).0));
111    });
112
113    let values = used
114        .into_iter()
115        .map(|attr| {
116            let entry = entry_func.get_entry_block(ctx);
117            let i = entry_func.push_argument(ctx, attr.ty(ctx));
118            entry_func.set_arg_attr(ctx, i, &ATTR_BUILTIN_ATTRIBUTE, Box::new(attr));
119            set_block_arg_name(ctx, entry, i, Some(ident(attr.to_string())));
120            let value = entry.deref(ctx).get_argument(i);
121            (attr, value)
122        })
123        .collect::<HashMap<_, _>>();
124
125    let mut rewriter = IRRewriter::<DummyListener>::default();
126
127    for op in read_ops {
128        rewriter.set_insertion_point_before_operation(op.get_operation());
129
130        let builtin = op.builtin(ctx).0;
131        let attr = built_in_attr(builtin);
132        let value = values[&attr];
133        match builtin {
134            Builtin::UnitPos | Builtin::PlaneDim | Builtin::PlanePos | Builtin::UnitPosPlane => {
135                rewriter.replace_operation_with_values(ctx, op.get_operation(), vec![value]);
136            }
137            Builtin::UnitPosX | Builtin::CubePosX | Builtin::CubeCountX | Builtin::AbsolutePosX => {
138                read_dim3(ctx, &mut rewriter, op, value, 0);
139            }
140            Builtin::UnitPosY | Builtin::CubePosY | Builtin::CubeCountY | Builtin::AbsolutePosY => {
141                read_dim3(ctx, &mut rewriter, op, value, 1);
142            }
143            Builtin::UnitPosZ | Builtin::CubePosZ | Builtin::CubeCountZ | Builtin::AbsolutePosZ => {
144                read_dim3(ctx, &mut rewriter, op, value, 2);
145            }
146            other => unreachable!("{other:?} should be lowered"),
147        }
148    }
149}
150
151fn read_dim3(
152    ctx: &mut Context,
153    rewriter: &mut impl Rewriter,
154    op: ReadBuiltinOp,
155    value: Value,
156    dim: usize,
157) {
158    let u32 = IntegerType::get(ctx, 32, Signedness::Unsigned).to_handle();
159    let new_op = ReadDim3BuiltinOp::new(ctx, u32, value, dim);
160    rewriter.append_op(ctx, &new_op);
161    rewriter.replace_operation(ctx, op.get_operation(), new_op.get_operation());
162}
163
164fn built_in_attr(builtin: Builtin) -> BuiltInAttr {
165    match builtin {
166        Builtin::UnitPos => BuiltInAttr::ThreadIndexInThreadgroup,
167        Builtin::UnitPosX | Builtin::UnitPosY | Builtin::UnitPosZ => {
168            BuiltInAttr::ThreadPositionInThreadgroup
169        }
170        Builtin::CubePosX | Builtin::CubePosY | Builtin::CubePosZ => {
171            BuiltInAttr::ThreadgroupPositionInGrid
172        }
173        Builtin::CubeCountX | Builtin::CubeCountY | Builtin::CubeCountZ => {
174            BuiltInAttr::ThreadgroupsPerGrid
175        }
176        Builtin::PlaneDim => BuiltInAttr::ThreadsPerSIMDgroup,
177        Builtin::PlanePos => BuiltInAttr::SIMDgroupIndexInThreadgroup,
178        Builtin::UnitPosPlane => BuiltInAttr::ThreadIndexInSIMDgroup,
179        Builtin::AbsolutePosX | Builtin::AbsolutePosY | Builtin::AbsolutePosZ => {
180            BuiltInAttr::ThreadPositionInGrid
181        }
182        other => unreachable!("{other:?} should be lowered"),
183    }
184}
cubecl_cpp/metal/builtin.rs

cubecl_cpp/metal/
builtin.rs
