cubecl_cpp/shared/
signature.rs

1use core::marker::PhantomData;
2
3use cubecl_core::ir::{
4    ContextExt,
5    attributes::{ATTR_BUFFER_IO, BufferIOAttr, FuncInterface, buffer_io_by_position},
6    cube_op,
7    dialect::OperationPtrExt,
8    interfaces::{AlignedType, HasElementType},
9    prelude::*,
10    rewrite::visit_all_ops_of_type,
11    types::{ArrayType, PointerType, RuntimeArrayType, VectorType, scalar::IndexType},
12};
13use cubecl_opt::passes::alloc_shared_memory::AllocSharedOp;
14use cubecl_runtime::kernel::Visibility;
15use itertools::Itertools;
16use pliron::{
17    builtin::{
18        attributes::{StringAttr, TypeAttr},
19        op_interfaces::SingleBlockRegionInterface,
20        ops::{FuncOp, ModuleOp},
21    },
22    graph::walkers::uninterruptible::immutable::walk_op,
23    pass::Pass,
24    utils::table::{IMap, ISet},
25};
26
27use crate::{
28    shared::{
29        CompilationState, CppValue, shared_op, shared_op_with_out,
30        ty::{InfoStructType, TypeExtCPP, UniformPointerType},
31        type_info_definition_sized,
32    },
33    target::*,
34};
35
36#[cube_op(name = "cpp.declare_types")]
37#[result_ty(none)]
38pub struct DeclareInfoTypeOp {}
39
40shared_op!(DeclareInfoTypeOp, |_, ctx| {
41    let state = ctx.aux_ty::<CompilationState>();
42    let mut out = String::new();
43    type_info_definition_sized(&mut out, ctx, &state.info).unwrap();
44    out
45});
46
47#[cube_op(name = "cpp.declare_complex_helpers")]
48#[result_ty(none)]
49pub struct DeclareComplexHelpersOp {}
50
51shared_op!(DeclareComplexHelpersOp, |_, _| {
52    crate::cuda::dialect::COMPLEX_HELPERS.into()
53});
54
55#[cube_op(name = "cpp.load_info")]
56#[result_ty(fixed = InfoStructType::get(ctx).into())]
57pub struct LoadInfoOp {
58    pub ptr: Value,
59}
60
61#[cube_op(name = "cpp.load_dynamic_meta")]
62#[result_ty(fixed = UniformPointerType::get(
63    ctx,
64    IndexType::get(ctx).into(),
65).into())]
66pub struct LoadDynMetaOp {
67    ptr: Value,
68}
69
70shared_op!(LoadInfoOp, |op, ctx| {
71    let ptr = op.ptr(ctx).name(ctx);
72    let out = op.get_result(ctx);
73    let out_ty = out.get_type(ctx).to_cpp(ctx);
74    format!("const {out_ty}& {} = *{ptr};", out.name(ctx))
75});
76
77shared_op_with_out!(LoadDynMetaOp, |op, ctx| {
78    let ptr = op.ptr(ctx).name(ctx);
79    let out_ty = op.get_result(ctx).get_type(ctx).to_cpp(ctx);
80    format!("reinterpret_cast<{out_ty}>({ptr} + 1)")
81});
82
83#[cube_op(name = "cpp.declare_vector", format = "attr($vector_ty, $TypeAttr)")]
84#[result_ty(none)]
85pub struct DeclareVectorOp {
86    vector_ty: TypeAttr,
87}
88
89shared_op!(DeclareVectorOp, |op, ctx| {
90    let vector = op.vector_ty(ctx).get_type(ctx).deref(ctx);
91    let vector = vector.downcast_ref::<VectorType>().unwrap();
92    let align = vector.align(ctx);
93    let inner_ty = vector.inner.to_cpp(ctx);
94    let vec = vector.vectorization;
95    let fields = (0..vec).map(|i| format!("{inner_ty} i_{i};")).join(" ");
96    format!("struct alignas({align}) {inner_ty}_{vec} {{ {fields} }};\n")
97});
98
99#[cube_op(name = "cpp.include", format = "attr($header, $StringAttr)")]
100#[result_ty(none)]
101pub struct IncludeOp {
102    pub header: StringAttr,
103}
104
105shared_op!(IncludeOp, |op, ctx| {
106    format!("#include <{}>\n", op.header(ctx).as_str())
107});
108
109shared_op!(AllocSharedOp, |op, ctx| {
110    let name = op.get_result(ctx).name(ctx);
111    let align = op.alignment(ctx).0;
112    format!("extern __shared__ __align__({align}) char {name}[];")
113});
114
115/// Run on module
116#[derive(Default)]
117pub struct DeclareVectorTypesPass;
118
119#[pass_name]
120impl Pass for DeclareVectorTypesPass {
121    fn run(
122        &mut self,
123        op: Ptr<Operation>,
124        ctx: &mut Context,
125        _analyses: &mut AnalysisManager,
126    ) -> Result<PassResult> {
127        let module = op.as_op::<ModuleOp>(ctx).expect("Should be run on module");
128        // Deduplicate by type name because some types are semantic only (i.e. tf32 is the same as f32)
129        let mut vectors = IMap::default();
130
131        walk_op(
132            ctx,
133            &mut vectors,
134            &WALKCONFIG_PREORDER_FORWARD,
135            op,
136            |ctx, vectors, node| {
137                let mut ins = |val: Value| {
138                    if let Some(vector) = try_get_vector_type(ctx, val) {
139                        vectors.insert((vector.inner.to_cpp(ctx), vector.vectorization), vector);
140                    }
141                };
142                match node {
143                    IRNode::Operation(op) if let Some(res) = op.opt_result(ctx) => ins(res),
144                    IRNode::BasicBlock(ptr) => {
145                        for arg in ptr.deref(ctx).arguments() {
146                            ins(arg);
147                        }
148                    }
149                    _ => {}
150                }
151            },
152        );
153
154        let mut res = PassResult::default();
155        for &vector in vectors.values() {
156            let decl = DeclareVectorOp::new(ctx, vector.get_self_handle(ctx));
157            decl.get_operation()
158                .insert_at_front(module.get_body(ctx, 0), ctx);
159            res.ir_changed |= IRStatus::Changed;
160        }
161        Ok(res)
162    }
163}
164
165fn try_get_vector_type(ctx: &Context, value: Value) -> Option<VectorType> {
166    let ty = value.get_type(ctx).deref(ctx);
167    let elem = type_cast::<dyn HasElementType>(&*ty)?.element_type(ctx)?;
168    elem.deref(ctx).downcast_ref().copied()
169}
170
171#[op_interface]
172pub trait RequiresIncludesOp<T> {
173    verify_op_succ!();
174    fn includes(&self, ctx: &Context) -> Vec<String>;
175}
176
177#[type_interface]
178pub trait RequiresIncludesType<T> {
179    verify_ty_succ!();
180    fn includes(&self, ctx: &Context) -> Vec<String>;
181}
182
183macro_rules! op_includes {
184    ($target: ty, [$($ty: ty),*] => $inc: expr) => {
185        $(#[pliron::derive::op_interface_impl]
186        impl crate::shared::signature::RequiresIncludesOp<$target> for $ty {
187            fn includes(&self, _ctx: &pliron::context::Context) -> Vec<String> {
188                vec![$inc.into()]
189            }
190        })*
191    };
192}
193pub(crate) use op_includes;
194
195macro_rules! ty_includes {
196    ($target: ty, [$($ty: ty),*] => $inc: expr) => {
197        $(#[pliron::derive::type_interface_impl]
198        impl crate::shared::signature::RequiresIncludesType<$target> for $ty {
199            fn includes(&self, _ctx: &pliron::context::Context) -> Vec<String> {
200                vec![$inc.into()]
201            }
202        })*
203    };
204}
205pub(crate) use ty_includes;
206
207macro_rules! nested_include_types {
208    ($target: ty) => {
209        #[type_interface_impl]
210        impl RequiresIncludesType<$target> for PointerType {
211            fn includes(&self, ctx: &Context) -> Vec<String> {
212                let inner = self.inner.deref(ctx);
213                if let Some(includes) = type_cast::<dyn RequiresIncludesType<$target>>(&*inner) {
214                    includes.includes(ctx)
215                } else {
216                    vec![]
217                }
218            }
219        }
220
221        #[type_interface_impl]
222        impl RequiresIncludesType<$target> for ArrayType {
223            fn includes(&self, ctx: &Context) -> Vec<String> {
224                let inner = self.inner.deref(ctx);
225                if let Some(includes) = type_cast::<dyn RequiresIncludesType<$target>>(&*inner) {
226                    includes.includes(ctx)
227                } else {
228                    vec![]
229                }
230            }
231        }
232
233        #[type_interface_impl]
234        impl RequiresIncludesType<$target> for RuntimeArrayType {
235            fn includes(&self, ctx: &Context) -> Vec<String> {
236                let inner = self.inner.deref(ctx);
237                if let Some(includes) = type_cast::<dyn RequiresIncludesType<$target>>(&*inner) {
238                    includes.includes(ctx)
239                } else {
240                    vec![]
241                }
242            }
243        }
244
245        #[type_interface_impl]
246        impl RequiresIncludesType<$target> for VectorType {
247            fn includes(&self, ctx: &Context) -> Vec<String> {
248                let inner = self.inner.deref(ctx);
249                if let Some(includes) = type_cast::<dyn RequiresIncludesType<$target>>(&*inner) {
250                    includes.includes(ctx)
251                } else {
252                    vec![]
253                }
254            }
255        }
256    };
257}
258
259nested_include_types!(Cuda);
260nested_include_types!(Hip);
261nested_include_types!(Metal);
262
263/// Run on module
264#[derive(Default)]
265pub struct CollectIncludesPass<T: CppTarget> {
266    _ty: PhantomData<T>,
267}
268
269#[pass_name]
270impl<T: CppTarget> Pass for CollectIncludesPass<T> {
271    fn run(
272        &mut self,
273        op: Ptr<Operation>,
274        ctx: &mut Context,
275        _analyses: &mut AnalysisManager,
276    ) -> Result<PassResult> {
277        let module = op.as_op::<ModuleOp>(ctx).expect("Should be run on module");
278        let mut includes = ISet::default();
279
280        walk_op(
281            ctx,
282            &mut includes,
283            &WALKCONFIG_PREORDER_FORWARD,
284            op,
285            |ctx, includes, node| {
286                let mut ins = |val: Value| {
287                    let ty = val.get_type(ctx).deref(ctx);
288                    if let Some(includes_ty) = type_cast::<dyn RequiresIncludesType<T>>(&*ty) {
289                        includes.extend(includes_ty.includes(ctx));
290                    }
291                };
292                match node {
293                    IRNode::Operation(op) => {
294                        if let Some(res) = op.opt_result(ctx) {
295                            ins(res);
296                        }
297                        let dyn_op = op.dyn_op(ctx);
298                        if let Some(includes_op) = op_cast::<dyn RequiresIncludesOp<T>>(&*dyn_op) {
299                            includes.extend(includes_op.includes(ctx));
300                        }
301                    }
302                    IRNode::BasicBlock(ptr) => {
303                        for arg in ptr.deref(ctx).arguments() {
304                            ins(arg);
305                        }
306                    }
307                    _ => {}
308                }
309            },
310        );
311
312        let mut res = PassResult::default();
313        for include in includes {
314            let decl = IncludeOp::new(ctx, include);
315            decl.get_operation()
316                .insert_at_front(module.get_body(ctx, 0), ctx);
317            res.ir_changed |= IRStatus::Changed;
318        }
319        Ok(res)
320    }
321}
322
323pub fn shared_memory_size(ctx: &Context, module: Ptr<Operation>) -> usize {
324    let mut size = 0;
325    visit_all_ops_of_type::<AllocSharedOp, _>(ctx, &mut size, module, |ctx, size, op| {
326        *size += op.size(ctx).0;
327    });
328    size
329}
330
331/// The four-state per-buffer IO, by buffer position — see
332/// [`buffer_io_by_position`].
333pub fn buffer_io(ctx: &Context, entry_func: FuncOp) -> Vec<BufferIOAttr> {
334    buffer_io_by_position(ctx, entry_func).into_iter().collect()
335}
336
337pub fn buffers(ctx: &Context, entry_func: FuncOp) -> Vec<Visibility> {
338    let entry = entry_func.get_entry_block(ctx);
339    let num_args = entry.deref(ctx).get_num_arguments();
340    let mut out = vec![];
341    for i in 0..num_args {
342        if let Some(io) = entry_func.get_arg_attr::<BufferIOAttr>(ctx, i, &ATTR_BUFFER_IO) {
343            out.push(match io.is_writable() {
344                true => Visibility::ReadWrite,
345                false => Visibility::Read,
346            });
347        }
348    }
349    out
350}
cubecl_cpp/shared/signature.rs

cubecl_cpp/shared/
signature.rs
