make_tiled_mma

创建tiled_mma策略。MmaAtom Ops计算矩阵大小与atom_shape相应m、n和k相乘所得乘积表示实际计算矩阵大小。

当前atom_shape中m/n不支持任意扩展，只能设置为1，k支持任意大小。

接口定义

template<typename MmaAtom, typename Shape>

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

模板参数

表1 模板参数定义
参数名	类型	描述
MmaAtom	typename	mma原子策略类型
Shape	typename	形状类型

参数

表2 参数定义
参数名	类型	描述	输入/输出
mma_atom	MmaAtom，Ops<枚举mma_atom_t>	mma原子策略，当前mma_atom_t可设置为 MMA_32x16x1_F64F64F64：表示计算尺寸为32161，计算步长为A矩阵Stride<1, 32>、B矩阵Stride<16, 1>、C矩阵Stride<16, 1>，计算精度为float64float64=float64的mma原子行为 MMA_32x16x512_F64F64F64：表示计算尺寸为3216512，计算步长为A矩阵Stride<1, 32>、B矩阵Stride<16, 1>、C矩阵Stride<16, 1>，计算精度为float64float64=float64的mma原子行为 MMA_16x64x2_BF16BF16F32：表示计算尺寸为16642，计算步长为A矩阵Stride<2, Stride<1, 32>>、B矩阵Stride<Stride<1, 128>, 2>、C矩阵Stride<64, 1>，计算精度为bfloat16bfloat16=float32的mma原子行为 MMA_16x64x1_BF16BF16F32：表示计算尺寸为16641，计算步长为A矩阵Stride<1, 16>、B矩阵Stride<64, 1>、C矩阵Stride<64, 1>，计算精度为bfloat16bfloat16=float32的mma原子行为	输入
atom_shape	Shape	原子策略在各维度执行次数，包含m、n和k三个数值	输入

返回值

返回TiledMma<MmaAtom, Shape>对象

示例

#include "stdlib.h"
#include "kupl.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<MMA_32x16x512_F64F64F64>{}, mma_atom_shape);
    auto store_atom_shape = make_shape(Int<1>{}, Int<1>{});
    auto tile_store = make_tiled_store(Ops<STORE_32x16_F64>{}, 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;
}
 

上述示例演示了基于32*16*512_F64F64F64矩阵形状的mma流程，其中通过make_tiled_mma创建mma tile策略。

父主题： 矩阵编程接口函数