cubecl_core/frontend/wgmma/
base.rs

1use super::layout::{Major, WgmmaTileLayout};
2
3use crate::frontend::container::inner_array_ty;
4use crate::frontend::pending::{GroupToken, PendingExpand};
5use crate::{self as cubecl, prelude::*, unexpanded};
6use core::marker::PhantomData;
7use cubecl_ir::{
8    dialect::matrix::{
9        WgmmaDescriptorOp, WgmmaFenceOp, WgmmaFenceOperandOp, WgmmaOp, warpgroup_elems_per_unit,
10    },
11    features::WgmmaElems,
12    types::{MatrixShape, pending::AsyncGroup},
13};
14use cubecl_macros::{comptime_type, cube, intrinsic};
15use pliron::builtin::attributes::TypeAttr;
16
17use alloc::format;
18
19pub use cubecl_ir::dialect::matrix::{WARPGROUP_M, WARPGROUP_UNITS};
20
21/// The accumulator's elements come in runs of two adjacent columns of a row.
22const ACCUMULATOR_RUN: u32 = 2;
23/// The fragment's elements come in runs of one 32-bit register.
24const FRAGMENT_RUN_BYTES: usize = 4;
25
26/// The 64-bit descriptor of a tile in shared memory, which a warpgroup MMA reads `A` or `B`
27/// through.
28#[derive(CubeType, Clone, Copy)]
29pub struct MatrixDescriptor<E: Scalar> {
30    #[allow(unused)]
31    descriptor: u64,
32    #[allow(unused)]
33    #[cube(comptime)]
34    layout: WgmmaTileLayout,
35    #[cube(comptime)]
36    _elem: PhantomData<E>,
37}
38
39#[cube]
40impl<E: Scalar> MatrixDescriptor<E> {
41    /// Describes `tile`, laid out as `layout`, whose lines hold its elements one or several at a
42    /// time. `tile` must be in shared memory and start at a multiple of
43    /// [`WgmmaTileLayout::alignment`]: a swizzled tile anywhere else is read with its lines
44    /// permuted wrong, silently.
45    #[allow(unused_variables)]
46    pub fn new<L: CubePrimitive<Scalar = E>>(
47        tile: &[L],
48        #[comptime] layout: WgmmaTileLayout,
49    ) -> Self {
50        intrinsic!(|scope| {
51            if let Err(err) = layout.validate(E::elem_type(scope).size()) {
52                scope.push_error(format!("a warpgroup MMA can't read the tile: {err}"));
53            }
54            let ptr = unsafe { *tile.__expand_as_ptr_method(scope) }.value(scope);
55            let op = WgmmaDescriptorOp::new(
56                scope.ctx_mut(),
57                ptr,
58                layout.leading_byte_offset(),
59                layout.stride_byte_offset(),
60                layout.swizzle,
61            );
62            MatrixDescriptorExpand {
63                descriptor: scope.register_with_result(&op).into(),
64                layout,
65                _elem: PhantomData,
66            }
67        })
68    }
69
70    /// The descriptor of the sub-tile that starts at `(row, k)`, which an MMA reads its 64 rows
71    /// of `A`, or its `n` rows of `B`, and [`WgmmaTileLayout::k_step`] of K from. `k` is a
72    /// multiple of the step. `row` is a multiple of 8 in a K-major tile, and of
73    /// [`WgmmaTileLayout::panel_elems`] in an MN-major one, so the sub-tile starts a panel: the
74    /// MMA reads whole panels, and one started inside a panel runs into the next line.
75    pub fn at(&self, row: usize, k: usize) -> Self {
76        let elem_size = E::size();
77        let offset = comptime![self.layout].byte_offset(row, k, elem_size);
78        // The address field holds the address in 16-byte units, and no shared memory address
79        // carries out of it.
80        MatrixDescriptor::<E> {
81            descriptor: self.descriptor + (offset as u64 >> 4),
82            layout: comptime![self.layout],
83            _elem: PhantomData,
84        }
85    }
86
87    /// The layout of the tile.
88    pub fn layout(&self) -> comptime_type!(WgmmaTileLayout) {
89        intrinsic!(|_| self.layout)
90    }
91}
92
93/// The registers each unit of a warpgroup holds of a `64 x n` accumulator.
94#[derive(CubeType, Clone)]
95pub struct Accumulator<CD: Numeric> {
96    registers: Array<CD>,
97    #[allow(unused)]
98    #[cube(comptime)]
99    n: usize,
100}
101
102#[cube]
103impl<CD: Numeric> Accumulator<CD> {
104    /// A zeroed `64 x n` accumulator. `n` is a multiple of 8 up to 256, or of 16 for the integer
105    /// types.
106    pub fn new(#[comptime] n: usize) -> Self {
107        let len = comptime![warpgroup_elems_per_unit(n)];
108        let mut registers = Array::new(len);
109        #[unroll]
110        for i in 0..len {
111            registers[i] = CD::from_int(0);
112        }
113        Accumulator::<CD> { registers, n }
114    }
115
116    /// The `n` of the accumulator.
117    pub fn n(&self) -> comptime_type!(usize) {
118        intrinsic!(|_| self.n)
119    }
120
121    /// The elements each unit holds.
122    #[allow(clippy::len_without_is_empty)]
123    pub fn len(&self) -> comptime_type!(usize) {
124        intrinsic!(|_| warpgroup_elems_per_unit(self.n))
125    }
126
127    /// The `nth` element this unit holds.
128    pub fn get(&self, nth: usize) -> CD {
129        self.registers[nth]
130    }
131
132    /// The `(row, col)` of the `nth` element `unit` holds. `unit` is the unit's position in its
133    /// warpgroup, `0..128`.
134    ///
135    /// Each plane holds 16 rows. A unit holds two adjacent elements of its row `lane / 4` of
136    /// them, then the same in row `lane / 4 + 8`, then moves on 8 columns.
137    pub fn position_of_nth(&self, unit: u32, nth: u32) -> (u32, u32) {
138        position_of_nth(unit, nth, ACCUMULATOR_RUN)
139    }
140
141    /// Hands the accumulator to the MMAs that [`Pending::execute`] issues.
142    pub fn start(self) -> Pending<Accumulator<CD>> {
143        intrinsic!(|scope| {
144            fence_operand(scope, &self.registers);
145            scope.register(&WgmmaFenceOp::new(scope.ctx_mut()));
146            PendingExpand::ready(scope, AsyncGroup::Warpgroup, self)
147        })
148    }
149}
150
151// A handle to the same registers, as an array's is.
152impl<CD: Numeric> Clone for AccumulatorExpand<CD> {
153    fn clone(&self) -> Self {
154        Self {
155            registers: self.registers,
156            n: self.n,
157        }
158    }
159}
160
161impl<CD: Numeric> Pending<Accumulator<CD>> {
162    /// This accumulator named at `O`, panicking unless `O` is the element it holds.
163    ///
164    /// Like a slice's `downcast`, this only satisfies the Rust type system: the registers and the
165    /// wait are the same, so it costs nothing.
166    pub fn downcast<O: Numeric>(&self) -> Pending<Accumulator<O>> {
167        unexpanded!()
168    }
169
170    /// This accumulator named at `O`, whatever element it holds.
171    ///
172    /// # Safety
173    ///
174    /// Only for keeping the accumulator at a dummy type, to be named again with
175    /// [`downcast`](Self::downcast) where it is used, as a slice's `downcast_unchecked` is.
176    pub unsafe fn downcast_unchecked<O: Numeric>(&self) -> Pending<Accumulator<O>> {
177        unexpanded!()
178    }
179}
180
181impl<CD: Numeric> PendingExpand<Accumulator<CD>> {
182    pub fn __expand_downcast_method<O: Numeric>(
183        &self,
184        scope: &Scope,
185    ) -> PendingExpand<Accumulator<O>> {
186        let named = O::__expand_as_type(scope);
187        let held = inner_array_ty(scope, self.value.registers.value(scope))
188            .deref(scope.ctx())
189            .inner;
190        assert_eq!(
191            held, named,
192            "Pending<Accumulator>::downcast: the accumulator holds {held:?}, not {named:?}"
193        );
194        self.__expand_downcast_unchecked_method(scope)
195    }
196
197    #[doc(hidden)]
198    pub fn __expand_downcast_unchecked_method<O: Numeric>(
199        &self,
200        _scope: &Scope,
201    ) -> PendingExpand<Accumulator<O>> {
202        PendingExpand {
203            value: AccumulatorExpand {
204                registers: self.value.registers.expand.into(),
205                n: self.value.n,
206            },
207            token: self.token,
208        }
209    }
210}
211
212impl<CD: Numeric> PendingValue for Accumulator<CD> {
213    fn wait_for(_pending: Pending<Self>) -> Self {
214        unexpanded!()
215    }
216
217    /// Waits for every MMA issued into the accumulator, committed or not. It commits them and
218    /// waits on that last group, so the token [`Accumulator::start`] made is never read, and the
219    /// wait covers every warpgroup MMA group of the unit.
220    fn __expand_wait_for(scope: &Scope, pending: PendingExpand<Self>) -> AccumulatorExpand<CD> {
221        let token = GroupToken::commit(scope, AsyncGroup::Warpgroup);
222        token.__expand_wait_method(scope);
223        // The MMAs wrote the registers behind the compiler's back: no read moves above the wait.
224        fence_operand(scope, &pending.value.registers);
225        pending.value
226    }
227}
228
229#[cube]
230impl<CD: Numeric> Pending<Accumulator<CD>> {
231    /// Issues `D = A * B + D`, with `A` and `B` in shared memory. Each MMA reads 64 rows of `A`
232    /// and `n` of `B`, and [`WgmmaTileLayout::k_step`] elements of K, from where the
233    /// descriptors point.
234    #[allow(unused_variables)]
235    pub fn execute<A: Scalar, B: Scalar>(
236        &mut self,
237        a: &MatrixDescriptor<A>,
238        b: &MatrixDescriptor<B>,
239    ) {
240        intrinsic!(|scope| {
241            let a_layout = a.layout;
242            let a = a.descriptor.read_value(scope);
243            issue::<A, B, CD>(scope, &self.value, a, Some(a_layout), b)
244        })
245    }
246
247    /// Issues `D = A * B + D`, with `A` in registers and `B` in shared memory. The fragment
248    /// must stay unchanged until the MMA's group completes.
249    #[allow(unused_variables)]
250    pub fn execute_registers<A: Scalar, B: Scalar>(
251        &mut self,
252        a: &Fragment<A>,
253        b: &MatrixDescriptor<B>,
254    ) {
255        intrinsic!(|scope| {
256            // The fragment was written by the units, and the MMA reads it out of the same
257            // registers: settle the writes, and order them before the MMA.
258            fence_operand(scope, &a.registers);
259            scope.register(&WgmmaFenceOp::new(scope.ctx_mut()));
260            let a = a.registers.read_value(scope);
261            issue::<A, B, CD>(scope, &self.value, a, None, b)
262        })
263    }
264
265    /// Waits for every MMA issued into the accumulator, zeroes it, and hands it back to the MMAs
266    /// to come, in the same registers: what [`Accumulator::new`] then [`Accumulator::start`] give,
267    /// with no copy written after the fence the MMAs to come read past.
268    pub fn reset(&mut self) {
269        intrinsic!(|scope| {
270            let token = GroupToken::commit(scope, AsyncGroup::Warpgroup);
271            token.__expand_wait_method(scope);
272            fence_operand(scope, &self.value.registers);
273            let len = warpgroup_elems_per_unit(self.value.n);
274            zero_registers::expand::<CD>(scope, &mut self.value.registers, len);
275            fence_operand(scope, &self.value.registers);
276            scope.register(&WgmmaFenceOp::new(scope.ctx_mut()));
277        })
278    }
279
280    /// Commits the MMAs issued since the last commit into a group, and returns its completion.
281    /// Once it resolves, the MMAs are done reading their tiles and fragments.
282    pub fn commit(&mut self) -> Pending<()> {
283        intrinsic!(|scope| PendingExpand::commit(scope, AsyncGroup::Warpgroup, ()))
284    }
285}
286
287/// Sets the `len` registers to zero.
288#[cube]
289fn zero_registers<CD: Numeric>(registers: &mut Array<CD>, #[comptime] len: usize) {
290    #[unroll]
291    for i in 0..len {
292        registers[i] = CD::from_int(0);
293    }
294}
295
296/// Pins `registers` in place: the compiler moves no access to them across this point.
297fn fence_operand<E: CubePrimitive>(scope: &Scope, registers: &NativeExpand<Array<E>>) {
298    let registers = registers.__extract_list(scope);
299    scope.register(&WgmmaFenceOperandOp::new(scope.ctx_mut(), registers))
300}
301
302/// Issues one MMA into `acc`, whose registers it writes asynchronously. `a_layout` is the layout
303/// of `A`'s tile, or `None` for `A` in registers.
304fn issue<A: Scalar, B: Scalar, CD: Numeric>(
305    scope: &Scope,
306    acc: &AccumulatorExpand<CD>,
307    a: cubecl_ir::pliron::value::Value,
308    a_layout: Option<WgmmaTileLayout>,
309    b: &MatrixDescriptorExpand<B>,
310) {
311    let elems = WgmmaElems {
312        a: A::elem_type(scope),
313        b: B::elem_type(scope),
314        cd: CD::elem_type(scope),
315    };
316    let shape = MatrixShape {
317        m: WARPGROUP_M,
318        n: acc.n,
319        k: WgmmaTileLayout::k_step(elems.a.size()),
320    };
321    check_supported(scope, elems, shape);
322    if let Some(a_layout) = a_layout {
323        check_covers(scope, "A", a_layout, shape.m, shape.k);
324    }
325    check_covers(scope, "B", b.layout, shape.n, shape.k);
326    let a_major = a_layout.map_or(Major::K, |layout| layout.major);
327    let b_major = b.layout.major;
328    let b = b.descriptor.read_value(scope);
329    let registers = acc.registers.__extract_list(scope);
330    let a_ty = TypeAttr::new(A::__expand_as_type(scope));
331    let b_ty = TypeAttr::new(B::__expand_as_type(scope));
332    scope.register(&WgmmaOp::new(
333        scope.ctx_mut(),
334        a,
335        b,
336        registers,
337        a_ty,
338        b_ty,
339        shape,
340        a_major,
341        b_major,
342    ));
343}
344
345/// Checks that `layout`'s tile holds the `rows x k` the MMA reads of `operand`. A sub-tile from
346/// [`MatrixDescriptor::at`] keeps the whole tile's layout, so this only catches a tile too small
347/// for one MMA, not an offset that runs past its end.
348fn check_covers(scope: &Scope, operand: &str, layout: WgmmaTileLayout, rows: usize, k: usize) {
349    if layout.rows < rows || layout.k < k {
350        scope.push_error(format!(
351            "a warpgroup MMA reads {rows} rows and {k} of K of {operand}, but its tile is {} x {}",
352            layout.rows, layout.k
353        ));
354    }
355}
356
357fn check_supported(scope: &Scope, elems: WgmmaElems, shape: MatrixShape) {
358    let Some(props) = scope.state().device_properties.clone() else {
359        return;
360    };
361    let wgmma = &props.features.matmul.wgmma;
362    if !wgmma.iter().any(|config| config.matches(elems, shape)) {
363        let WgmmaElems { a, b, cd } = elems;
364        let MatrixShape { m, n, k } = shape;
365        scope.push_error(format!(
366            "the device doesn't support a {m}x{n}x{k} warpgroup MMA of {a:?} x {b:?} into \
367             {cd:?}; supported configurations: {wgmma:?}"
368        ));
369    }
370}
371
372/// The registers each unit of a warpgroup holds of a `64 x k` tile of `A`, for
373/// [`Pending::execute_registers`]. `k` is [`WgmmaTileLayout::k_step`].
374#[derive(CubeType)]
375pub struct Fragment<A: Scalar> {
376    registers: Array<A>,
377}
378
379#[cube]
380impl<A: Scalar> Fragment<A> {
381    /// An uninitialized fragment.
382    #[allow(clippy::new_without_default)]
383    pub fn new() -> Self {
384        intrinsic!(|scope| {
385            let len = fragment_elems(A::elem_type(scope).size());
386            FragmentExpand {
387                registers: Array::__expand_new(scope, len),
388            }
389        })
390    }
391
392    /// The elements each unit holds.
393    #[allow(clippy::len_without_is_empty)]
394    pub fn len(&self) -> comptime_type!(usize) {
395        intrinsic!(|scope| fragment_elems(A::elem_type(scope).size()))
396    }
397
398    /// The `nth` element this unit holds.
399    pub fn get(&self, nth: usize) -> A {
400        self.registers[nth]
401    }
402
403    /// Sets the `nth` element this unit holds.
404    pub fn set(&mut self, nth: usize, value: A) {
405        self.registers[nth] = value;
406    }
407
408    /// The `(row, col)` of the `nth` element `unit` holds, laid out as the accumulator, but in
409    /// runs of 32 bits rather than of two elements.
410    pub fn position_of_nth(&self, unit: u32, nth: u32) -> (u32, u32) {
411        intrinsic!(|scope| {
412            let run = (FRAGMENT_RUN_BYTES / A::elem_type(scope).size()) as u32;
413            position_of_nth::expand(scope, unit, nth, run)
414        })
415    }
416}
417
418/// The elements of `A` of `elem_size` bytes each unit holds for one MMA.
419fn fragment_elems(elem_size: usize) -> usize {
420    warpgroup_elems_per_unit(WgmmaTileLayout::k_step(elem_size))
421}
422
423/// The `(row, col)` of the `nth` element `unit` holds, in runs of `run` adjacent elements.
424#[cube]
425fn position_of_nth(unit: u32, nth: u32, #[comptime] run: u32) -> (u32, u32) {
426    let plane = unit / 32;
427    let lane = unit % 32;
428    let chunk = nth / run;
429    let row = plane * 16 + lane / 4 + (chunk % 2) * 8;
430    let col = (chunk / 2) * (run * 4) + (lane % 4) * run + nth % run;
431    (row, col)
432}
cubecl_core/frontend/wgmma/base.rs

cubecl_core/frontend/wgmma/
base.rs
