cubecl_ir/dialect/
matrix.rs

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

cubecl_ir/dialect/
matrix.rs
