ruda-kernel 0.1.50

Ruda Rust kernel DSL and device tensor operations.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

use crate::dsl::{cmma, prelude::*};

use crate::tiling::{
    MatrixLayout, StageIdent, as_cmma_layout,
    tile::{
        compute::matmul::cmma::{CmmaFragmentReader as _, CmmaStageReader, CmmaStageWriter},
        data::{Strided, StridedTile},
    },
};

#[ruda]
pub fn cmma_execute<L: Numeric, R: Numeric, A: Numeric, S: cmma::MatrixScope>(
    lhs: &cmma::Matrix<L, S>,
    rhs: &cmma::Matrix<R, S>,
    acc: &mut cmma::Matrix<A, S>,
) {
    cmma::execute::<L, R, A, A, _>(lhs, rhs, acc, acc);
}

#[ruda]
pub fn cmma_load_from_shared<E: Numeric, ES: Size, N: Numeric, IO: SliceVisibility, S: cmma::MatrixScope>(
    shared: &StridedTile<E, ES, IO>,
    matrix: &mut cmma::Matrix<N, S>,
    #[comptime] ident: StageIdent,
    #[comptime] matrix_layout: MatrixLayout,
) {
    let shared = shared.to_read_only();
    match ident {
        StageIdent::Lhs | StageIdent::Rhs => {
            CmmaStageReader::<Strided>::load_fragment(&shared, matrix, ComptimeOption::new_None());
        }
        StageIdent::Acc => {
            CmmaStageReader::<Strided>::load_fragment(
                &shared,
                matrix,
                ComptimeOption::new_Some(as_cmma_layout(matrix_layout)),
            );
        }
        _ => panic!("Invalid ident for CMMA load"),
    }
}

#[ruda]
pub fn cmma_load_zeros<N: Numeric, S: cmma::MatrixScope>(matrix: &mut cmma::Matrix<N, S>) {
    cmma::fill(matrix, N::from_int(0));
}

#[ruda]
pub fn cmma_write_to_shared<E: Numeric, ES: Size, A: Numeric, S: cmma::MatrixScope>(
    shared: &mut StridedTile<E, ES, ReadWrite>,
    matrix: &cmma::Matrix<A, S>,
) {
    let casted = cmma::cast::<A, E, _>(matrix);
    CmmaStageWriter::store_fragment(shared, &casted);
}