1use cubecl_core::ir::{
2 CanMaterialize, Pure,
3 attributes::{FuncInterface, IndexAttr},
4 dialect::general::{BufferLenOp, ReadScalarOp, ShapeOp, StrideOp},
5 ident,
6 interfaces::ScalarType,
7 prelude::*,
8 rewrite::RewriteOp,
9 try_cast_ty,
10 types::scalar::IndexType,
11};
12use pliron::{
13 builtin::{
14 attributes::TypeAttr,
15 given_names::{insert_block_arg_name, insert_operation_result_name},
16 ops::FuncOp,
17 },
18 irbuild::match_rewrite::apply_match_rewrite,
19};
20
21use crate::{
22 cuda::signature::ATTR_GRID_CONSTANT,
23 shared::{
24 CompilationOptions, CompilationState, shared_op_with_out,
25 signature::{LoadDynMetaOp, LoadInfoOp},
26 ty::{InfoStructType, UniformPointerType},
27 },
28};
29
30#[cube_op(name = "cpp.read_scalar", format = "$0 `[` attr($id, $IndexAttr) `]`")]
31#[result_ty(from_inputs = |ctx, _, _, ty: &TypeAttr| ty.get_type(ctx))]
32#[op_traits(Pure, CanMaterialize)]
33pub struct CppReadScalarOp {
34 pub base: Value,
35 pub id: IndexAttr,
36 pub ty: TypeAttr,
37}
38
39#[cube_op(
40 name = "cpp.read_static_meta",
41 format = "$0 `[` attr($offset, $IndexAttr) `]`"
42)]
43#[result_ty(fixed = IndexType::get(ctx).to_handle())]
44#[op_traits(Pure, CanMaterialize)]
45pub struct CppReadStaticMetaOp {
46 pub base: Value,
47 pub offset: IndexAttr,
48}
49
50#[cube_op(name = "cpp.read_dynamic_meta", format = "$0 `[` $1 ` + ` $2 `]`")]
51#[result_ty(fixed = IndexType::get(ctx).to_handle())]
52#[op_traits(Pure, CanMaterialize)]
53pub struct CppReadDynamicMetaOp {
54 pub base: Value,
55 pub offset: Value,
56 pub dim: Value,
57}
58
59shared_op_with_out!(CppReadScalarOp, |op, ctx| {
60 let ty = op.ty(ctx).get_type(ctx).deref(ctx);
61 let elem = try_cast_ty!(ty, ctx, dyn ScalarType).elem_type(ctx);
62 let base = op.base(ctx).name(ctx);
63 let offset = op.id(ctx).0;
64 format!("{base}.scalars_{elem}[{offset}]")
65});
66
67shared_op_with_out!(CppReadStaticMetaOp, |op, ctx| {
68 let base = op.base(ctx).name(ctx);
69 let offset = op.offset(ctx).0;
70 format!("{base}.static_meta[{offset}]")
71});
72
73shared_op_with_out!(CppReadDynamicMetaOp, |op, ctx| {
74 let base = op.base(ctx).name(ctx);
75 let offset = op.offset(ctx).name(ctx);
76 let dim = op.dim(ctx).name(ctx);
77 format!("{base}[{offset} + {dim}]")
78});
79
80#[derive(Default)]
81pub struct LowerInfoPass;
82
83#[pass_name]
84impl Pass for LowerInfoPass {
85 fn run(
86 &mut self,
87 op: Ptr<Operation>,
88 ctx: &mut Context,
89 _analyses: &mut AnalysisManager,
90 ) -> Result<PassResult> {
91 let (has_info, has_dynamic_meta) = {
92 let info = &ctx.aux_ty::<CompilationState>().info;
93 (info.has_info(), info.has_dynamic_meta)
94 };
95 let func = op.as_op::<FuncOp>(ctx).unwrap();
96 let entry_block = func.get_entry_block(ctx);
97 let supports_features = ctx.aux_ty::<CompilationOptions>().supports_features;
98
99 let info_name = ident("info");
100 let dyn_meta_name = ident("dynamic_meta");
101
102 let mut info_st = None;
103 let mut dyn_meta = None;
104
105 if supports_features.grid_constants {
106 if has_dynamic_meta {
107 let usize = IndexType::get(ctx).to_handle();
108 let index_ptr = UniformPointerType::get(ctx, usize);
109 let id = func.push_argument(ctx, index_ptr.to_handle());
110 insert_block_arg_name(ctx, entry_block, id, Some(dyn_meta_name));
111 let value = entry_block.deref(ctx).get_argument(id);
112 dyn_meta = Some(value);
113 }
114
115 if has_info {
116 let id = func.push_argument(ctx, InfoStructType::get(ctx).into());
117 func.set_arg_attr_unit(ctx, id, &ATTR_GRID_CONSTANT);
118 insert_block_arg_name(ctx, entry_block, id, Some(info_name));
119 let value = entry_block.deref(ctx).get_argument(id);
120 info_st = Some(value);
121 }
122 } else if has_info {
123 let info_st_ty = InfoStructType::get(ctx).to_handle();
124 let info_ptr = UniformPointerType::get(ctx, info_st_ty);
125 let id = func.push_argument(ctx, info_ptr.to_handle());
126
127 let ptr = entry_block.deref(ctx).get_argument(id);
128
129 let load_info = LoadInfoOp::new(ctx, ptr);
130 info_st = Some(load_info.get_result(ctx));
131 insert_operation_result_name(ctx, load_info.get_operation(), 0, Some(info_name));
132 load_info.get_operation().insert_at_front(entry_block, ctx);
133
134 let load_dyn = LoadDynMetaOp::new(ctx, ptr);
135 dyn_meta = Some(load_dyn.get_result(ctx));
136 insert_operation_result_name(ctx, load_dyn.get_operation(), 0, Some(dyn_meta_name));
137 load_dyn.get_operation().insert_at_front(entry_block, ctx);
138 }
139
140 if let Some(info_st) = info_st {
141 let mut rewrite = MatchRewriteOp::new(ReplaceScalars(info_st));
142 apply_match_rewrite(ctx, &mut rewrite, Default::default(), op)?;
143 let mut rewrite = MatchRewriteOp::new(ReplaceBufferLen(info_st));
144 apply_match_rewrite(ctx, &mut rewrite, Default::default(), op)?;
145 }
146 if let Some((info_st, dyn_meta)) = info_st.zip(dyn_meta) {
147 let mut rewrite = MatchRewriteOp::new(ReplaceShape(info_st, dyn_meta));
148 apply_match_rewrite(ctx, &mut rewrite, Default::default(), op)?;
149 let mut rewrite = MatchRewriteOp::new(ReplaceStride(info_st, dyn_meta));
150 apply_match_rewrite(ctx, &mut rewrite, Default::default(), op)?;
151 }
152
153 let mut res = PassResult::default();
154 res.ir_changed |= IRStatus::Changed;
155 Ok(res)
156 }
157}
158
159struct ReplaceScalars(Value);
160impl RewriteOp<ReadScalarOp> for ReplaceScalars {
161 fn rewrite(&mut self, ctx: &mut Context, rewriter: &mut MatchRewriter, op: ReadScalarOp) {
162 let id = *op.id(ctx);
163 let ty = op.ty(ctx).clone();
164 let new_op = CppReadScalarOp::new(ctx, self.0, id, ty);
165 rewriter.replace_op_with(ctx, op.get_operation(), new_op.get_operation());
166 }
167}
168
169struct ReplaceBufferLen(Value);
170impl RewriteOp<BufferLenOp> for ReplaceBufferLen {
171 fn rewrite(&mut self, ctx: &mut Context, rewriter: &mut MatchRewriter, op: BufferLenOp) {
172 let buffer_idx = op.buffer_idx(ctx).0;
173 let meta = ctx.aux_ty::<CompilationState>().info.metadata;
174 let offset = meta.buffer_len_index(buffer_idx);
175 let new_op = CppReadStaticMetaOp::new(ctx, self.0, offset);
176 rewriter.replace_op_with(ctx, op.get_operation(), new_op.get_operation());
177 }
178}
179
180struct ReplaceShape(Value, Value);
181impl RewriteOp<ShapeOp> for ReplaceShape {
182 fn rewrite(&mut self, ctx: &mut Context, rewriter: &mut MatchRewriter, op: ShapeOp) {
183 let (static_, dynamic) = (self.0, self.1);
184 let buffer_idx = op.buffer_idx(ctx).0;
185 let meta = ctx.aux_ty::<CompilationState>().info.metadata;
186 let offset = meta.shape_offset_index(buffer_idx);
187 let dim = op.dim(ctx);
188
189 let offs = CppReadStaticMetaOp::new(ctx, static_, offset);
190 offs.get_operation().insert_before(ctx, op.get_operation());
191 let new_op = CppReadDynamicMetaOp::new(ctx, dynamic, offs.get_result(ctx), dim);
192 rewriter.replace_op_with(ctx, op.get_operation(), new_op.get_operation());
193 }
194}
195
196struct ReplaceStride(Value, Value);
197impl RewriteOp<StrideOp> for ReplaceStride {
198 fn rewrite(&mut self, ctx: &mut Context, rewriter: &mut MatchRewriter, op: StrideOp) {
199 let (static_, dynamic) = (self.0, self.1);
200 let buffer_idx = op.buffer_idx(ctx).0;
201 let meta = ctx.aux_ty::<CompilationState>().info.metadata;
202 let offset = meta.stride_offset_index(buffer_idx);
203 let dim = op.dim(ctx);
204
205 let offs = CppReadStaticMetaOp::new(ctx, static_, offset);
206 offs.get_operation().insert_before(ctx, op.get_operation());
207 let new_op = CppReadDynamicMetaOp::new(ctx, dynamic, offs.get_result(ctx), dim);
208 rewriter.replace_op_with(ctx, op.get_operation(), new_op.get_operation());
209 }
210}