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

cubecl_ir/dialect/
matrix.rs
