cubecl_cpp/cuda/
processors.rs

1use cubecl_core::{
2    self as cubecl,
3    ir::{CoopMma, Instruction, MatrixIdent, Operation, Processor, Scope, ScopeProcessing, Value},
4    prelude::*,
5};
6
7#[derive(new, Debug)]
8pub struct CudaMmaProcessor;
9
10impl Processor for CudaMmaProcessor {
11    fn transform(&self, mut processing: ScopeProcessing) -> ScopeProcessing {
12        let mut instructions = Vec::new();
13        core::mem::swap(&mut processing.instructions, &mut instructions);
14
15        for instruction in instructions {
16            match instruction.operation {
17                Operation::CoopMma(CoopMma::RowIndex { lane_id, i, matrix }) => {
18                    let elems_per_reg = 32 / matrix.storage.elem_type().size_bits();
19                    let scope =
20                        Scope::root(false).with_global_state(processing.global_state.clone());
21                    let row_idx: Value = row_index::expand(
22                        &scope,
23                        lane_id.into(),
24                        i.into(),
25                        elems_per_reg as u32,
26                        matrix.ident,
27                    )
28                    .into();
29                    let tmp_processing = scope.process([]);
30                    for inst in tmp_processing.instructions {
31                        processing.instructions.push(inst);
32                    }
33
34                    processing.instructions.push(Instruction::new(
35                        Operation::Copy(row_idx),
36                        instruction.out(),
37                    ));
38                }
39                Operation::CoopMma(CoopMma::ColIndex { lane_id, i, matrix }) => {
40                    let elems_per_reg = 32 / matrix.storage.elem_type().size_bits();
41                    let scope =
42                        Scope::root(false).with_global_state(processing.global_state.clone());
43                    let col_idx: Value = col_index::expand(
44                        &scope,
45                        lane_id.into(),
46                        i.into(),
47                        elems_per_reg as u32,
48                        matrix.ident,
49                    )
50                    .into();
51                    let tmp_processing = scope.process([]);
52                    for inst in tmp_processing.instructions {
53                        processing.instructions.push(inst);
54                    }
55
56                    processing.instructions.push(Instruction::new(
57                        Operation::Copy(col_idx),
58                        instruction.out(),
59                    ));
60                }
61                _ => {
62                    processing.instructions.push(instruction);
63                }
64            }
65        }
66
67        processing
68    }
69}
70
71/// Derived from PTX shape documentation
72/// <https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-for-mma>
73#[cube]
74fn row_index(
75    lane_id: u32,
76    i: u32,
77    #[comptime] elems_per_reg: u32,
78    #[comptime] ident: MatrixIdent,
79) -> u32 {
80    match ident {
81        MatrixIdent::A => {
82            let group_id = lane_id / 4;
83            let odd_register = (i / elems_per_reg) & 1;
84            group_id + odd_register * 8
85        }
86        MatrixIdent::B => {
87            let thread_id_in_group = lane_id % 4;
88            let offset = thread_id_in_group * elems_per_reg + (i % elems_per_reg);
89            let reg = i / elems_per_reg;
90            offset + elems_per_reg * 4 * reg
91        }
92        MatrixIdent::Accumulator => {
93            let group_id = lane_id / 4;
94            let offset = (i << 2) & 8;
95            group_id + offset
96        }
97    }
98}
99
100/// Derived from PTX shape documentation
101/// <https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-for-mma>
102#[cube]
103fn col_index(
104    lane_id: u32,
105    i: u32,
106    #[comptime] elems_per_reg: u32,
107    #[comptime] ident: MatrixIdent,
108) -> u32 {
109    match ident {
110        MatrixIdent::A => {
111            let thread_id_in_group = lane_id % 4;
112            let offset = thread_id_in_group * elems_per_reg + (i % elems_per_reg);
113            let group_2 = (i / (2 * elems_per_reg)) & 1;
114            offset + 4 * elems_per_reg * group_2
115        }
116        MatrixIdent::B => lane_id >> 2,
117        MatrixIdent::Accumulator => {
118            let thread_id_in_group = lane_id % 4;
119            (thread_id_in_group * 2) + (i % 2)
120        }
121    }
122}
cubecl_cpp/cuda/processors.rs

cubecl_cpp/cuda/
processors.rs
