cubecl_cpp/shared/
operation.rs

1use cubecl_core::{
2    cmma::MatrixType,
3    ir::{
4        AddressSpace, Scope, cube_op,
5        dialect::{
6            base::OperationPtrExt,
7            general::{CommentOp, PoisonOp, PrintfOp},
8            math::FmaOp,
9            memory::{DeclareVariableOp, UnrelatedAllocInfo},
10        },
11        prelude::*,
12        types::PointerType,
13    },
14};
15use cubecl_opt::passes::alloc_shared_memory::SliceSharedOp;
16use itertools::Itertools;
17use pliron::{
18    arg_err,
19    attribute::{AttrObj, boxed_attr_cast},
20    builtin::{attributes::TypeAttr, ops::ConstantOp},
21    opts::mem2reg::{AllocInfo, PromotableAllocationInterface},
22};
23
24use crate::{
25    error::{CompileError, Result},
26    shared::{
27        CppValue, format_const,
28        lowering::LowerOp,
29        ty::{TypeExtCPP, TypedExtCPP},
30        unroll::unrolling,
31    },
32    target::{CtxTarget, Shared, Target, dispatch_target},
33};
34
35#[op_interface]
36pub trait OpToCPP<T> {
37    verify_op_succ!();
38    fn to_cpp(&self, ctx: &Context) -> String;
39}
40
41macro_rules! shared_op {
42    ($ty: ty, $impl: expr) => {
43        #[pliron::derive::op_interface_impl]
44        impl $crate::shared::operation::OpToCPP<$crate::target::Shared> for $ty {
45            fn to_cpp(&self, ctx: &pliron::context::Context) -> String {
46                $crate::shared::closure_inference_hack::<$ty, String>(self, ctx, $impl)
47            }
48        }
49    };
50}
51pub(crate) use shared_op;
52
53macro_rules! shared_op_with_out {
54    ($ty: ty, $impl: expr) => {
55        #[pliron::derive::op_interface_impl]
56        impl $crate::shared::operation::OpToCPP<$crate::target::Shared> for $ty {
57            fn to_cpp(&self, ctx: &pliron::context::Context) -> String {
58                use cubecl_core::ir::prelude::*;
59                use $crate::shared::CppValue;
60                let op = $crate::shared::closure_inference_hack::<$ty, String>(self, ctx, $impl);
61                let out = self.get_result(ctx).fmt_left(ctx);
62                format!("{out} = {op};\n")
63            }
64        }
65    };
66}
67pub(crate) use shared_op_with_out;
68
69pub trait OpExtCPP {
70    fn to_cpp(&self, ctx: &Context) -> Result<String>;
71}
72
73impl OpExtCPP for Ptr<Operation> {
74    fn to_cpp(&self, ctx: &Context) -> Result<String> {
75        let op_dyn = self.dyn_op(ctx);
76        let target_cpp = dispatch_target!(ctx, {
77            let inner_cpp = op_cast::<dyn OpToCPP<Target>>(op_dyn.as_ref());
78            inner_cpp.map(|it| it.to_cpp(ctx))
79        });
80        if let Some(cpp) = target_cpp {
81            return Ok(cpp);
82        }
83        let shared = op_cast::<dyn OpToCPP<Shared>>(op_dyn.as_ref())
84            .ok_or_else(|| CompileError::UnsupportedOp(op_dyn.disp(ctx).to_string()))?;
85        Ok(shared.to_cpp(ctx))
86    }
87}
88
89#[cube_op(name = "cpp.declare_local")]
90#[result_ty(from_inputs = variable_ptr_ty)]
91pub struct DeclareLocalOp {
92    pub value_ty: TypeAttr,
93    #[attribute(optional, untyped)]
94    pub initializer: AttrObj,
95}
96
97#[op_interface_impl]
98impl PromotableAllocationInterface for DeclareLocalOp {
99    fn alloc_info(&self, ctx: &Context) -> Vec<AllocInfo> {
100        vec![AllocInfo {
101            ptr: self.get_result(ctx),
102            ty: self.value_ty(ctx).get_type(ctx),
103        }]
104    }
105
106    fn default_value(
107        &self,
108        ctx: &mut Context,
109        inserter: &mut dyn Inserter,
110        alloc_info: &AllocInfo,
111    ) -> cubecl_core::ir::prelude::Result<Value> {
112        if alloc_info.ptr != self.get_result(ctx) {
113            return arg_err!(self.loc(ctx), UnrelatedAllocInfo);
114        }
115        if let Some(initializer) = self.initializer(ctx).map(|it| it.clone()) {
116            let initializer = boxed_attr_cast(initializer).unwrap();
117            let constant = ConstantOp::new(ctx, initializer);
118            inserter.insert_op(ctx, &constant);
119            Ok(constant.get_result(ctx))
120        } else {
121            let poison = PoisonOp::new(ctx, alloc_info.ty);
122            inserter.insert_op(ctx, &poison);
123            Ok(poison.get_result(ctx))
124        }
125    }
126
127    fn promote(
128        &self,
129        ctx: &mut Context,
130        rewriter: &mut dyn Rewriter,
131        alloc_infos: &[AllocInfo],
132    ) -> cubecl_core::ir::prelude::Result<()> {
133        if alloc_infos.len() != 1 || alloc_infos[0].ptr != self.get_result(ctx) {
134            return arg_err!(self.loc(ctx), UnrelatedAllocInfo);
135        }
136        rewriter.erase_operation(ctx, self.get_operation());
137        Ok(())
138    }
139}
140
141shared_op!(DeclareLocalOp, |op, ctx| {
142    let ty = op.value_ty(ctx).get_type(ctx);
143    let name = op.get_result(ctx).name(ctx);
144    let value_ty = ty.to_cpp(ctx);
145    let out_ty = op.get_result(ctx).get_type(ctx).to_cpp(ctx);
146    let init = op.initializer(ctx).map(|init| format_const(ctx, &init, ty));
147    if let Some(init) = init {
148        format!("{value_ty} {name}_store = {init};\n{out_ty} {name} = &{name}_store;")
149    } else {
150        format!("{value_ty} {name}_store;\n{out_ty} {name} = &{name}_store;")
151    }
152});
153
154#[cube_op(name = "cpp.declare_matrix")]
155#[result_ty(from_inputs = variable_ptr_ty)]
156pub struct DeclareMatrixOp {
157    pub value_ty: TypeAttr,
158}
159
160fn variable_ptr_ty(ctx: &Context, value_ty: &TypeAttr) -> TypeHandle {
161    let value_ty = value_ty.get_type(ctx);
162    PointerType::get(ctx, value_ty, AddressSpace::Local).into()
163}
164
165shared_op_with_out!(SliceSharedOp, |op, ctx| {
166    let ptr_ty = op.get_result(ctx).get_type(ctx).to_cpp(ctx);
167    let block = op.block(ctx).name(ctx);
168    let offset = op.offset(ctx).0;
169    format!("reinterpret_cast<{ptr_ty}>(&{block}[{offset}])")
170});
171
172// These are the same in IR semantics, but have very different codegen in C++. So split them up.
173#[op_interface_impl]
174impl LowerOp for DeclareVariableOp {
175    fn should_lower(&self, _ctx: &Context) -> bool {
176        true
177    }
178
179    fn lower(&self, scope: &Scope) -> Vec<Value> {
180        let ctx = scope.ctx_mut();
181        let value_ty = self.value_ty(ctx).clone();
182        let addr_space = self.addr_space(ctx).0;
183        vec![match addr_space {
184            AddressSpace::Global(_) => panic!("Unsupported address space for declaration"),
185            AddressSpace::Shared => panic!("Should be lowered to block allocation"),
186            AddressSpace::Local => {
187                if value_ty.get_type(ctx).deref(ctx).is::<MatrixType>() {
188                    let op = DeclareMatrixOp::new(ctx, value_ty);
189                    scope.register_with_result(&op)
190                } else {
191                    let init = self.initializer(ctx).map(|it| it.clone());
192                    let op = DeclareLocalOp::new(ctx, value_ty, init);
193                    scope.register_with_result(&op)
194                }
195            }
196        }]
197    }
198}
199
200shared_op_with_out!(FmaOp, |op, ctx| {
201    let a = op.a(ctx).name(ctx);
202    let b = op.b(ctx).name(ctx);
203    let c = op.c(ctx).name(ctx);
204    let res = op.get_result(ctx);
205    // CUDA/HIP headers have no `fma` overload for `__half`/bf16 — the call is
206    // ambiguous between the double/float/_Float16 candidates — so the half
207    // types go through the dedicated intrinsics. Metal's `fma` is generic over
208    // its scalar and vector float types, including half.
209    let f = if matches!(ctx.target(), Target::Metal) {
210        "fma"
211    } else if res.is_half(ctx) {
212        "__hfma"
213    } else if res.is_half2(ctx) {
214        "__hfma2"
215    } else {
216        "fma"
217    };
218    format!("{f}({a}, {b}, {c})")
219});
220// `fma` has no vector overloads in CUDA/HIP headers, so a vector fma must be
221// scalarized lane by lane like the other math functions.
222unrolling!(FmaOp);
223
224shared_op!(CommentOp, |op, ctx| {
225    let content = String::from(op.comment(ctx).clone());
226    if content.contains("\n") {
227        format!("/* {content} */\n")
228    } else {
229        format!("// {content}\n")
230    }
231});
232
233shared_op!(PrintfOp, |op, ctx| {
234    let format_string = op.format_string(ctx);
235    let args = op.args(ctx);
236    let args = args.iter().map(|it| format!(", {}", it.name(ctx))).join("");
237    format!("printf({:?}{args});", format_string.as_str())
238});
cubecl_cpp/shared/operation.rs

cubecl_cpp/shared/
operation.rs
