cubecl_ir/dialect/
matrix.rs

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

cubecl_ir/dialect/
matrix.rs
