make_tiled_store

Create a tiled_store policy for subsequent matrix writeback operations.

The operation involves the matrix writeback of store_atom and shape extension of atom_shape. The size of the tiled_store writeback matrix is determined by multiplying the store_atom matrix size by the m or n dimension specified in atom_shape.

Currently, m or n in atom_shape cannot be extended randomly and can only be set to 1.

Interface Definition

template<typename StoreAtom, typename Shape>

TiledStore<StoreAtom, Shape> make_tiled_store( StoreAtom store_atom, Shape atom_shape);

Template Parameters

**Table 1** Template parameter definition
Parameter	Type	Description
StoreAtom	typename	Store atom policy type.
Shape	typename	Shape type.

Parameters

**Table 2** Parameter definition
Parameter	Type	Description	Input/Output
store_atom	StoreAtom, Ops<store_atom_t>	Store atom policy. Currently, store_atom_t can be set to: KP36_32x16_F64_STORE: store atom operation with the writeback size of 32 × 16 and writeback precision of float64, utilizing a row-major matrix with a computation stride of Stride<16, 1>. KP36_16x64_F32_STORE: store atom operation with the writeback size of 16 × 64 and writeback precision of float32, utilizing a row-major matrix with a computation stride of Stride<64, 1>. KP36_16x64_INT32_STORE: store atom operation with the writeback size of 16 × 64 and writeback precision of int32, utilizing a row-major matrix with a computation stride of Stride<64, 1>. KP36_32x32_INT32_STORE: store atom operation with the writeback size of 32 × 32 and writeback precision of int32, utilizing a row-major matrix with a computation stride of Stride<32, 1>.	Input
atom_shape	Shape	Number of times that the atom policy is executed in each dimension, including m and n.	Input

Return Value

TiledStore<StoreAtom, Shape> object returned.

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 store tiling policy is created using make_tiled_store.

Parent topic: Matrix Programming API Functions