cubecl_cpp/hip/
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 HipMmaProcessor;
9
10impl Processor for HipMmaProcessor {
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 scope =
19                        Scope::root(false).with_global_state(processing.global_state.clone());
20                    let row_idx: Value =
21                        row_index::expand(&scope, lane_id.into(), i.into(), matrix.ident).into();
22                    let tmp_processing = scope.process([]);
23                    for inst in tmp_processing.instructions {
24                        processing.instructions.push(inst);
25                    }
26
27                    processing.instructions.push(Instruction::new(
28                        Operation::Copy(row_idx),
29                        instruction.out(),
30                    ));
31                }
32                Operation::CoopMma(CoopMma::ColIndex { lane_id, i, matrix }) => {
33                    let scope =
34                        Scope::root(false).with_global_state(processing.global_state.clone());
35                    let row_idx: Value =
36                        col_index::expand(&scope, lane_id.into(), i.into(), matrix.ident).into();
37                    let tmp_processing = scope.process([]);
38                    for inst in tmp_processing.instructions {
39                        processing.instructions.push(inst);
40                    }
41
42                    processing.instructions.push(Instruction::new(
43                        Operation::Copy(row_idx),
44                        instruction.out(),
45                    ));
46                }
47                _ => {
48                    processing.instructions.push(instruction);
49                }
50            }
51        }
52
53        processing
54    }
55}
56
57#[cube]
58fn row_index(lane_id: u32, i: u32, #[comptime] ident: MatrixIdent) -> u32 {
59    match ident {
60        MatrixIdent::A => lane_id % 16,
61        MatrixIdent::B => i,
62        // 2 * i, offset by 1 if lane_id >= 16
63        MatrixIdent::Accumulator => i * 2 + (lane_id / 16),
64    }
65}
66
67#[cube]
68fn col_index(lane_id: u32, i: u32, #[comptime] ident: MatrixIdent) -> u32 {
69    match ident {
70        MatrixIdent::A => i,
71        MatrixIdent::B => lane_id % 16,
72        MatrixIdent::Accumulator => lane_id % 16,
73    }
74}
cubecl_cpp/hip/processors.rs

cubecl_cpp/hip/
processors.rs
