make_tiled_mma

Create a tiled_mma policy for subsequent MMA operations.

This involves the atomic MMA method input (as mma_atom) and the tiling shape input (as atom_shape), where the product of the matrix size calculated by mma_atom and the corresponding m/n/k factors in atom_shape represents the actual matrix size calculated by the tiled_mma method.

Currently, m and n in atom_shape do not support arbitrary tiling and must be set to 1, while k supports arbitrary sizes.

Interface Definition

template<typename MmaAtom, typename Shape>

TiledMma<MmaAtom, Shape> make_tiled_mma(MmaAtom mma_atom, Shape atom_shape);

Template Parameters

**Table 1** Template parameter definition
Parameter	Type	Description
MmaAtom	typename	MMA atomic policy type.
Shape	typename	Shape type.

Parameters

**Table 2** Parameter definition
Parameter	Type	Description	Input/Output
mma_atom	MmaAtom, Ops<enumeration mma_atom_t>	MMA atomic policy. Currently, mma_atom_t can be set to: KP36_32x16x1_F64F64F64: indicates the MMA atomic behavior with a computation dimension of 32x16x1, where matrix A is column-major with a computation stride of Stride<1, 32>, matrix B is row-major with a computation stride of Stride<16, 1>, matrix C is row-major with a computation stride of Stride<16, 1>; and the computation precision is *float64 float64 = float64. KP36_32x16x512_F64F64F64: indicates the MMA atomic behavior with a computation dimension of 32x16x512, where matrix A is column-major with a computation stride of Stride<1, 32>, matrix B is row-major with a computation stride of Stride<16, 1>, matrix C is row-major with a computation stride of Stride<16, 1>; and the computation precision is float64 * float64 = float64. KP36_16x64x2_BF16BF16F32: indicates the MMA atomic behavior with a computation dimension of 16x64x2, where matrix A is a ZZ-peak matrix with a computation stride of Stride<2, Stride<1, 32>>, matrix B is an NN-peak matrix with a computation stride of Stride<Stride<1, 128>, 2>, matrix C is row-major with a computation stride of Stride<64, 1>, and the computation precision is bfloat16 * bfloat16 = float32. KP36_16x64x1_BF16BF16F32: indicates the MMA atomic behavior with a computation dimension of 16x64x1, where matrix A is column-major with a computation stride of Stride<1, 16>, matrix B is row-major with a computation stride of Stride<64, 1>, matrix C is row-major with a computation stride of Stride<64, 1>; and the computation precision is bfloat16 * bfloat16 = float32. KP36_16x64x4_INT8INT8INT32: indicates the MMA atomic behavior with a computation dimension of 16x64x4, where matrix A is a ZZ-peak matrix with a computation stride of Stride<4, Stride<1, 64>>, matrix B is an NN-peak matrix with a computation stride of Stride<Stride<1, 256>, 4>, matrix C is row-major with a computation stride of Stride<64, 1>; and the computation precision is int8 * int8 = int32. KP36_32x32x4_INT8INT8INT32: indicates the MMA atomic behavior with a computation dimension of 32x32x4, where matrix A is a ZZ-peak matrix with a computation stride of Stride<4, Stride<1, 128>>, matrix B is an NN-peak matrix with a computation stride of Stride<Stride<1, 128>, 4>, matrix C is row-major with a computation stride of Stride<32, 1>; and the computation precision is int8 * int8 = int32**.	Input
atom_shape	Shape	Number of times that the atom policy is executed in each dimension, including m, n, and k.	Input

Return Value

Returns the TiledMma<MmaAtom, Shape> object.

Examples

#include "stdlib.h"
#include "kupl_mma.h"
using namespace kupl::tensor;

int main()
{
    constexpr int MATRIX_M  = 32;
    constexpr int MATRIX_N  = 16;
    constexpr int MATRIX_K = 512;
    double *data_a = (double *)malloc(sizeof(double) * MATRIX_M * MATRIX_K);
    double *data_b = (double *)malloc(sizeof(double) * MATRIX_K * MATRIX_N);
    double *data_c = (double *)malloc(sizeof(double) * MATRIX_M * MATRIX_N);

    auto shape_a = make_shape(Int<32>{}, Int<512>{});
    auto shape_b = make_shape(Int<512>{}, Int<16>{});
    auto shape_c = make_shape(Int<32>{}, Int<16>{});

    auto stride_a = make_stride(Int<1>{}, Int<32>{});
    auto stride_b = make_stride(Int<16>{}, Int<1>{});
    auto stride_c = make_stride(Int<16>{}, Int<1>{});

    auto layout_a = make_layout(shape_a, stride_a);
    auto layout_b = make_layout(shape_b, stride_b);
    auto layout_c = make_layout(shape_c, stride_c);

    auto mma_atom_shape = make_shape(Int<1>{}, Int<1>{}, Int<1>{});
    auto tiled_mma = make_tiled_mma(Ops<KP36_32x16x512_F64F64F64>{}, mma_atom_shape);
    auto store_atom_shape = make_shape(Int<1>{}, Int<1>{});
    auto tile_store = make_tiled_store(Ops<KP36_32x16_F64_STORE>{}, store_atom_shape);

    auto tensor_a = make_tensor(data_a, layout_a);
    auto tensor_b = make_tensor(data_b, layout_b);
    auto tensor_c = make_tensor(data_c, layout_c);

    tensor_tiled_mma(tiled_mma, tensor_c, tensor_a, tensor_b, tensor_c);
    tensor_tiled_store(tile_store, tensor_c);

    free(data_a);
    free(data_b);
    free(data_c);
    return 0;
}

The preceding example demonstrates the MMA process based on the 32*16*512_F64F64F64 matrix shape, where the tiled_mma policy is created using make_tiled_mma.

Parent topic: Matrix Programming API Functions