cubecl_ir/dialect/
ssa_matrix.rs

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

cubecl_ir/dialect/
ssa_matrix.rs
