cubecl_ir/dialect/
ssa_matrix.rs

1use core::fmt;
2
3use alloc::{format, string::ToString, vec::Vec};
4
5use cubecl_macros_internal::{NamedRewrite, cube_op};
6use pliron::{
7    builtin::{
8        attributes::IdentifierAttr,
9        ops::FuncOp,
10        types::{FunctionType, IntegerType, Signedness},
11    },
12    combine::{Parser, parser::char::char},
13    identifier::Identifier,
14    input_err,
15    irfmt::parsers::{process_parsed_ssa_defs, spaced, ssa_opd_parse},
16    location::Location,
17    op::OpObj,
18    parsable::{self, IntoParseResult, Parsable, ParseResult},
19    printable::{self, Printable},
20    symbol_table::SymbolTableCollection,
21    verify_err,
22};
23
24use crate::{
25    CanMaterialize,
26    dialect::{
27        general::SymbolUserOpVerifyErr,
28        matrix::{self, MatrixLayoutAttr, parse_closure, print_closure},
29        memory,
30        synchronization::SyncScope,
31    },
32    interfaces::{TypedExt, synchronizes},
33    prelude::*,
34};
35
36/// Fill a matrix with a scalar value.
37/// Note: Unlike most matrix ops, this does not have implicit synchronization because there's no
38/// coordination between threads.
39#[cube_op(name = "ssa_matrix.fill")]
40#[result_ty(argument)]
41#[op_traits(CanMaterialize)]
42pub struct FillOp {
43    pub value: Value,
44}
45
46#[op_interface_impl]
47impl MatrixToSSAOp for matrix::FillOp {
48    fn to_owned_matrix(
49        &self,
50        ctx: &mut Context,
51        rewriter: &mut DialectConversionRewriter,
52        _: &OperandsInfo,
53    ) -> Result<()> {
54        let value = self.value(ctx);
55        let out_ty = self.matrix(ctx).unwrap_ptr(ctx);
56        let op = FillOp::new(ctx, out_ty, value);
57        let matrix = rewriter.append_op_with_result(ctx, &op);
58        store_value(self.matrix(ctx), matrix, ctx, rewriter);
59        rewriter.erase_operation(ctx, self.get_operation());
60
61        Ok(())
62    }
63}
64
65#[cube_op(name = "ssa_matrix.load")]
66#[result_ty(argument)]
67#[op_traits(CanMaterialize)]
68pub struct LoadOp {
69    #[operand(ptr_read)]
70    pub source: Value,
71    pub stride: Value,
72    pub layout: MatrixLayoutAttr,
73}
74synchronizes!(LoadOp, SyncScope::Plane);
75
76#[op_interface_impl]
77impl MatrixToSSAOp for matrix::LoadOp {
78    fn to_owned_matrix(
79        &self,
80        ctx: &mut Context,
81        rewriter: &mut DialectConversionRewriter,
82        _: &OperandsInfo,
83    ) -> Result<()> {
84        let source = self.source(ctx);
85        let stride = self.stride(ctx);
86        let layout = self.layout(ctx).0;
87        let out_ty = self.matrix(ctx).unwrap_ptr(ctx);
88        let op = LoadOp::new(ctx, out_ty, source, stride, layout);
89        let matrix = rewriter.append_op_with_result(ctx, &op);
90        store_value(self.matrix(ctx), matrix, ctx, rewriter);
91        rewriter.erase_operation(ctx, self.get_operation());
92
93        Ok(())
94    }
95}
96
97#[cube_op(name = "ssa_matrix.store")]
98#[result_ty(none)]
99#[op_traits(CanMaterialize)]
100pub struct StoreOp {
101    pub matrix: Value,
102    #[operand(ptr_write)]
103    pub destination: Value,
104    pub stride: Value,
105    pub layout: MatrixLayoutAttr,
106}
107synchronizes!(StoreOp, SyncScope::Plane);
108
109#[op_interface_impl]
110impl MatrixToSSAOp for matrix::StoreOp {
111    fn to_owned_matrix(
112        &self,
113        ctx: &mut Context,
114        rewriter: &mut DialectConversionRewriter,
115        _: &OperandsInfo,
116    ) -> Result<()> {
117        let matrix = load_value(self.matrix(ctx), ctx, rewriter);
118        let dest = self.destination(ctx);
119        let stride = self.stride(ctx);
120        let layout = self.layout(ctx).0;
121        let op = StoreOp::new(ctx, matrix, dest, stride, layout);
122        rewriter.append_op(ctx, &op);
123        rewriter.erase_operation(ctx, self.get_operation());
124
125        Ok(())
126    }
127}
128
129#[cube_op(name = "ssa_matrix.multiply_accumulate")]
130#[result_ty(argument)]
131#[op_traits(CanMaterialize)]
132pub struct MultiplyAccumulateOp {
133    pub mat_a: Value,
134    pub mat_b: Value,
135    pub mat_c: Value,
136}
137synchronizes!(MultiplyAccumulateOp, SyncScope::Plane);
138
139#[op_interface_impl]
140impl MatrixToSSAOp for matrix::MultiplyAccumulateOp {
141    fn to_owned_matrix(
142        &self,
143        ctx: &mut Context,
144        rewriter: &mut DialectConversionRewriter,
145        _: &OperandsInfo,
146    ) -> Result<()> {
147        let mat_a = load_value(self.mat_a(ctx), ctx, rewriter);
148        let mat_b = load_value(self.mat_b(ctx), ctx, rewriter);
149        let mat_c = load_value(self.mat_c(ctx), ctx, rewriter);
150        let out_ty = self.mat_d(ctx).unwrap_ptr(ctx);
151        let op = MultiplyAccumulateOp::new(ctx, out_ty, mat_a, mat_b, mat_c);
152        let matrix_out = rewriter.append_op_with_result(ctx, &op);
153        store_value(self.mat_d(ctx), matrix_out, ctx, rewriter);
154        rewriter.erase_operation(ctx, self.get_operation());
155
156        Ok(())
157    }
158}
159
160/// Cast a matrix from one type to another.
161/// Note: Unlike most matrix ops, this does not have implicit synchronization because there's no
162/// coordination between threads.
163#[cube_op(name = "ssa_matrix.cast")]
164#[result_ty(argument)]
165#[op_traits(CanMaterialize)]
166pub struct CastOp {
167    pub input: Value,
168}
169
170#[op_interface_impl]
171impl MatrixToSSAOp for matrix::CastOp {
172    fn to_owned_matrix(
173        &self,
174        ctx: &mut Context,
175        rewriter: &mut DialectConversionRewriter,
176        _: &OperandsInfo,
177    ) -> Result<()> {
178        let matrix_in = load_value(self.input(ctx), ctx, rewriter);
179        let out_ty = self.output(ctx).unwrap_ptr(ctx);
180        let op = CastOp::new(ctx, out_ty, matrix_in);
181        let matrix_out = rewriter.append_op_with_result(ctx, &op);
182        store_value(self.output(ctx), matrix_out, ctx, rewriter);
183        rewriter.erase_operation(ctx, self.get_operation());
184
185        Ok(())
186    }
187}
188
189/// Executes a closure for each element in the matrix.
190/// Note: Unlike most matrix ops, this does not have implicit synchronization because there's no
191/// coordination between threads.
192#[pliron_op(
193    name = "ssa_matrix.elementwise",
194    attributes = (ssa_matrix_elementwise_closure: IdentifierAttr),
195    verifier = "succ"
196)]
197#[op_interfaces(OneResultInterface)]
198#[op_traits(CanMaterialize)]
199pub struct ElementwiseOp;
200
201impl ElementwiseOp {
202    pub fn new(
203        ctx: &mut Context,
204        matrix_in: Value,
205        closure: Identifier,
206        captures: Vec<Value>,
207    ) -> Self {
208        let out_ty = vec![matrix_in.get_type(ctx)];
209        let mut opds = vec![matrix_in];
210        opds.extend(captures);
211        let op = Self {
212            op: Operation::new(ctx, Self::get_concrete_op_info(), out_ty, opds, vec![], 0),
213        };
214        op.set_attr_ssa_matrix_elementwise_closure(ctx, IdentifierAttr::new(closure));
215        op
216    }
217
218    pub fn matrix_in(&self, ctx: &Context) -> Value {
219        self.get_operation().operand(ctx, 0)
220    }
221
222    pub fn closure(&self, ctx: &Context) -> Identifier {
223        let attr = self.get_attr_ssa_matrix_elementwise_closure(ctx).unwrap();
224        attr.clone().into()
225    }
226
227    pub fn closure_captures(&self, ctx: &Context) -> Vec<Value> {
228        self.get_operation().deref(ctx).operands().skip(1).collect()
229    }
230}
231
232impl Printable for ElementwiseOp {
233    fn fmt(
234        &self,
235        ctx: &Context,
236        _state: &printable::State,
237        f: &mut fmt::Formatter<'_>,
238    ) -> fmt::Result {
239        write!(
240            f,
241            "{} = {} ({}) ",
242            self.get_result(ctx).disp(ctx),
243            self.get_opid().disp(ctx),
244            self.matrix_in(ctx).disp(ctx)
245        )?;
246        print_closure(ctx, &self.closure(ctx), &self.closure_captures(ctx), f)
247    }
248}
249impl Parsable for ElementwiseOp {
250    type Arg = Vec<(Identifier, Location)>;
251    type Parsed = OpObj;
252    fn parse<'a>(
253        input: &mut parsable::StateStream<'a>,
254        arg: Self::Arg,
255    ) -> ParseResult<'a, Self::Parsed> {
256        let cur_loc = input.loc();
257
258        spaced(char('(')).parse_stream(input).into_result()?;
259        let mat_in = ssa_opd_parse(input, ())?.0;
260        spaced(char(')')).parse_stream(input).into_result()?;
261
262        let (closure, captures) = parse_closure(input)?.0;
263        let ctx = &mut input.state.ctx;
264
265        if arg.len() != 1 {
266            input_err!(
267                cur_loc,
268                "Expected 1 result, got {} during parsing",
269                arg.len()
270            )?;
271        }
272
273        let op = ElementwiseOp::new(ctx, mat_in, closure, captures);
274        process_parsed_ssa_defs(input, &arg, op.get_operation())?;
275        Ok(OpObj::new(op)).into_parse_result()
276    }
277}
278
279impl ElementwiseOp {
280    /// Callee type, including implicit args
281    fn callee_type(&self, ctx: &Context) -> TypeHandle {
282        let u32 = IntegerType::get(ctx, 32, Signedness::Unsigned).to_handle();
283        let elem_ty = self.matrix_in(ctx).element_ty(ctx).scalar_ty(ctx);
284        let mut args = vec![u32, u32, elem_ty];
285        args.extend(self.closure_captures(ctx).iter().map(|it| it.get_type(ctx)));
286        FunctionType::get(ctx, args, vec![elem_ty]).to_handle()
287    }
288}
289
290#[op_interface_impl]
291impl SymbolUserOpInterface for ElementwiseOp {
292    fn verify_symbol_uses(
293        &self,
294        ctx: &Context,
295        symbol_tables: &mut SymbolTableCollection,
296    ) -> Result<()> {
297        let callee_sym = self.closure(ctx);
298        let Some(callee) =
299            symbol_tables.lookup_symbol_in_nearest_table(ctx, self.get_operation(), &callee_sym)
300        else {
301            return verify_err!(
302                self.loc(ctx),
303                SymbolUserOpVerifyErr::SymbolNotFound(callee_sym.to_string())
304            );
305        };
306        let Some(func_op) = (&*callee as &dyn Op).downcast_ref::<FuncOp>() else {
307            return verify_err!(
308                self.loc(ctx),
309                SymbolUserOpVerifyErr::NotFunc(callee_sym.to_string())
310            );
311        };
312        let func_op_ty = func_op.get_type(ctx);
313
314        if func_op_ty != self.callee_type(ctx) {
315            return verify_err!(
316                self.loc(ctx),
317                SymbolUserOpVerifyErr::FuncTypeErr(format!(
318                    "expected {}, got {}",
319                    func_op_ty.disp(ctx),
320                    self.callee_type(ctx).disp(ctx)
321                ))
322            );
323        }
324        Ok(())
325    }
326
327    fn used_symbols(&self, ctx: &Context) -> Vec<Identifier> {
328        vec![self.closure(ctx)]
329    }
330}
331
332#[op_interface_impl]
333impl MatrixToSSAOp for matrix::ElementwiseOp {
334    fn to_owned_matrix(
335        &self,
336        ctx: &mut Context,
337        rewriter: &mut DialectConversionRewriter,
338        _: &OperandsInfo,
339    ) -> Result<()> {
340        let matrix_in = load_value(self.matrix_in(ctx), ctx, rewriter);
341        let closure = self.closure(ctx);
342        let captures = self.closure_captures(ctx);
343
344        let op = ElementwiseOp::new(ctx, matrix_in, closure, captures);
345        let matrix_out = rewriter.append_op_with_result(ctx, &op);
346        store_value(self.matrix_out(ctx), matrix_out, ctx, rewriter);
347        rewriter.erase_operation(ctx, self.get_operation());
348
349        Ok(())
350    }
351}
352
353fn load_value(ptr: Value, ctx: &mut Context, rewriter: &mut DialectConversionRewriter) -> Value {
354    let load = memory::LoadOp::new(ctx, ptr);
355    rewriter.append_op_with_result(ctx, &load)
356}
357
358fn store_value(
359    ptr: Value,
360    value: Value,
361    ctx: &mut Context,
362    rewriter: &mut DialectConversionRewriter,
363) {
364    let store = memory::StoreOp::new(ctx, ptr, value);
365    rewriter.append_op(ctx, &store);
366}
367
368#[op_interface]
369trait MatrixToSSAOp {
370    verify_op_succ!();
371    fn to_owned_matrix(
372        &self,
373        ctx: &mut Context,
374        rewriter: &mut DialectConversionRewriter,
375        operands_info: &OperandsInfo,
376    ) -> Result<()>;
377}
378
379pub type MatrixToSSAPass = DialectConversionPass<MatrixToSSAConversion>;
380
381#[derive(Default, NamedRewrite)]
382pub struct MatrixToSSAConversion;
383
384impl DialectConversion for MatrixToSSAConversion {
385    fn can_convert_op(&self, ctx: &Context, op: Ptr<Operation>) -> bool {
386        op.impls::<dyn MatrixToSSAOp>(ctx)
387    }
388
389    fn rewrite(
390        &mut self,
391        ctx: &mut Context,
392        rewriter: &mut DialectConversionRewriter,
393        op: Ptr<Operation>,
394        operands_info: &OperandsInfo,
395    ) -> Result<()> {
396        let dyn_op = op.dyn_op(ctx);
397        let to_owned = op_cast::<dyn MatrixToSSAOp>(&*dyn_op).unwrap();
398        to_owned.to_owned_matrix(ctx, rewriter, operands_info)
399    }
400}
cubecl_ir/dialect/ssa_matrix.rs

cubecl_ir/dialect/
ssa_matrix.rs
