Module cmma

This module exposes cooperative matrix-multiply and accumulate operations.

Most of the functions are actually unsafe, since they mutate their input, even if they are passed as reference.

Example

This is a basic 16x16x16 matrix multiplication example.

#[cube(launch)]
pub fn example(lhs: &Array<F16>, rhs: &Array<F16>, out: &mut Array<F32>) {
    let a = cmma::Matrix::<F16>::new(
        cmma::MatrixIdent::A,
        16,
        16,
        16,
        cmma::MatrixLayout::RowMajor,
    );
    let b = cmma::Matrix::<F16>::new(
        cmma::MatrixIdent::B,
        16,
        16,
        16,
        cmma::MatrixLayout::ColMajor,
    );
    let c = cmma::Matrix::<F32>::new(
        cmma::MatrixIdent::Accumulator,
        16,
        16,
        16,
        cmma::MatrixLayout::Undefined,
    );
    cmma::fill::<F32>(&c, F32::new(0.0));
    cmma::load::<F16>(&a, lhs.as_slice(), u32::new(16));
    cmma::load::<F16>(&b, rhs.as_slice(), u32::new(16));

    cmma::execute::<F16, F16, F32, F32>(&a, &b, &c, &c);

    cmma::store::<F32>(
        out.as_slice_mut(),
        &c,
        u32::new(16),
        cmma::MatrixLayout::RowMajor,
    );
}

Modules

cast

Module containing the expand function for store().

execute

Module containing the expand function for execute().

fill

Module containing the expand function for fill().

load

Module containing the expand function for load().

load_with_layout

Module containing the expand function for load_with_layout().

store

Module containing the expand function for store().

Structs

Matrix

A matrix represent a 2D grid of numbers.

MatrixExpand

Expand type of Matrix.

Enums

MatrixIdent

MatrixLayout

Functions

cast

Store the matrix in the given array following the given stride and layout.

execute

Execute the matrix-multiply and accumulate operation on the given matrices.

fill

Fill the matrix with the provided value.

load

Load the matrix with the provided array using the stride.

load_with_layout

Load the matrix with the provided array using the stride with an explicit layout. Explicit layouts are required when loading accumulators.

store

Store the matrix in the given array following the given stride and layout.