tensor_tiled_mma

An MMA interface for calculating A x B + C.

This interface involves the TiledMma policy and the tensor D/A/B/C matrix inputs, where the shape and stride of tensor D/A/B/C must be consistent with those defined in the TiledMma policy.

Interface Definition

template<typename TiledMma,

typename dtypeD, typename LayoutD,

typename dtypeA, typename LayoutA,

typename dtypeB, typename LayoutB,

typename dtypeC, typename LayoutC>

void tensor_tiled_mma(TiledMma mma, Tensor<dtypeD, LayoutD> D, Tensor<dtypeA, LayoutA> A, Tensor<dtypeB, LayoutB> B, Tensor<dtypeC, LayoutC> C);

Template Parameters

**Table 1** Template parameter definition
Parameter	Type	Description
TiledMma	typename	MMA tiling policy type.
dtypeD, dtypeA, dtypeB, dtypeC	typename	Precision types of D, A, B, and C.
LayoutD, LayoutA, LayoutB, LayoutC	typename	Layout types of D, A, B, and C.

Parameters

**Table 2** Parameter definition
Parameter	Type	Description	Input/Output
mma	TiledMma	MMA tiling policy.	Input
D	Tensor<dtypeD, LayoutD>	Matrix object D.	Input
A	Tensor<dtypeA, LayoutA>	Matrix object A.	Input
B	Tensor<dtypeB, LayoutB>	Matrix object B.	Input
C	Tensor<dtypeC, LayoutC>	Matrix object C.	Input

Return Value

void

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 tensor_tiled_mma receives tensor objects to perform matrix multiplication.

Parent topic: Matrix Programming API Functions