make_stride
创建Stride,用于规定矩阵行/列方向元素跨度。
接口定义
template<typename... Args>
Stride<Args...> make_stride(Args ... args);
参数
参数名 |
类型 |
描述 |
输入/输出 |
|---|---|---|---|
args... |
Args...,支持任意Int<>类型对象及任意Stride<>类型对象作为输入。 |
用于定义矩阵行/列方向元素跨度。对于非峰形矩阵只需输入Int<>类型对象,对于峰形矩阵需要输入Stride<>类型对象进行嵌套描述。 |
输入 |
返回值
返回Stride<Args...>对象
示例
#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<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_stride创建stride,用于规定矩阵行/列方向的元素跨度。
父主题: 矩阵编程接口函数