cutlass入门: 调用cutlass做通用矩阵乘法Gemm（附代码）

cutlass是CUDA C++模板抽象的集合，用于实现CUDA中所有级别和规模的高性能矩阵乘法（GEMM）和相关计算。相较于cuBLAS和cuDNN，cutlass中包含了更多可重用的模块化软件组件，这使得cutlass相较于前两者更为灵活。

cutlass项目官方网站：GitHub - NVIDIA/cutlass: CUDA Templates for Linear Algebra Subroutines

本文将展示如何用cutlass实现最基本的矩阵计算。

cutlass的使用流程与普通kernel大致相同：先在host端分配空间生成数据，再将host端的数据传入device端的buffer中，输入参数调用cutlass模块进行运算，最后将device端的数据传回Host端。

本代码实现的矩阵计算为D = α * A * B + β * C。其中，标量α = β = 1，矩阵A的大小为3840*4096，矩阵B的大小为4096*4096，矩阵C的大小为3840*4096。这三个矩阵均被初始化为全部元素为1的矩阵。那么经过简单计算我们可以知道矩阵D所包含的全部元素应该为4097。

代码及注释如下：

#include <iostream>                                           // 标准输入输出流
#include "cutlass/gemm/device/gemm.h"                         // 引入cutlass头文件

using ColumnMajor = cutlass::layout::ColumnMajor;             // 列主序存储方式
using RowMajor    = cutlass::layout::RowMajor;                // 行主序存储方式

using CutlassGemm = cutlass::gemm::device::Gemm<float,        // A矩阵数据类型
                                                RowMajor,     // A矩阵存储方式
                                                float,        // B矩阵数据类型
                                                RowMajor,     // B矩阵存储方式
                                                float,        // C矩阵数据类型
                                                RowMajor>;    // C矩阵存储方式

void generate_tensor_2D(float *ptr, int i_M, int i_N){        // 二维矩阵填充函数（此处全部填充1）
    for(int i = 0; i < i_M; i++){
        for(int j = 0; j < i_N; j++){
            *(ptr + i*i_N + j ) = 1.0;
        }
    }
}

int main(int argc, const char *arg[]) {

    int M = 3840;           //M
    int N = 4096;           //N
    int K = 4096;           //K

    int lda = K;
    int ldb = K;
    int ldc = N;
    int ldd = N;

    float alpha = 1.0;      //alpha
    float beta = 1.0;       //beta

    float *A;               //申明A矩阵host端指针
    float *B;               //申明B矩阵host端指针
    float *C;               //申明C矩阵host端指针
    float *D;               //申明D矩阵host端指针

    size_t A_mem_size = sizeof(float) * M * K; //memory size of matrix A = M * K * sizeof(float)
    size_t B_mem_size = sizeof(float) * K * N; //memory size of matrix B = K * N * sizeof(float)
    size_t C_mem_size = sizeof(float) * M * N; //memory size of matrix C = M * N * sizeof(float)
    size_t D_mem_size = sizeof(float) * M * N; //memory size of matrix C = M * N * sizeof(float)

    A = (float*)malloc(A_mem_size);  // host端A矩阵分配内存
    B = (float*)malloc(B_mem_size);  // host端B矩阵分配内存
    C = (float*)malloc(C_mem_size);  // host端C矩阵分配内存
    D = (float*)malloc(D_mem_size);  // host端D矩阵分配内存

    generate_tensor_2D(A, M, K);     // 填充A矩阵
    generate_tensor_2D(B, K, N);     // 填充B矩阵  
    generate_tensor_2D(C, M, N);     // 填充C矩阵  

    float *d_A;            // 申明device端A矩阵的指针
    float *d_B;            // 申明device端B矩阵的指针
    float *d_C;            // 申明device端C矩阵的指针
    float *d_D;            // 申明device端D矩阵的指针

    cudaMalloc((void**)&d_A, A_mem_size);  // device端为A矩阵分配内存
    cudaMalloc((void**)&d_B, B_mem_size);  // device端为B矩阵分配内存
    cudaMalloc((void**)&d_C, C_mem_size);  // device端为C矩阵分配内存
    cudaMalloc((void**)&d_D, D_mem_size);  // device端为D矩阵分配内存

    cudaMemcpy(d_A, A, A_mem_size, cudaMemcpyHostToDevice); // 将矩阵A的数据传递到device端
    cudaMemcpy(d_B, B, B_mem_size, cudaMemcpyHostToDevice); // 将矩阵B的数据传递到device端
    cudaMemcpy(d_C, C, C_mem_size, cudaMemcpyHostToDevice); // 将矩阵C的数据传递到device端

    CutlassGemm gemm_operator;                  // 申明cutlassgemm类
    CutlassGemm::Arguments args({M, N, K},      // Gemm Problem dimensions
                                {d_A, lda},     // source matrix A
                                {d_B, ldb},     // source matrix B
                                {d_C, ldc},     // source matrix C
                                {d_D, ldd},     // destination matrix D
                                {alpha, beta}); // alpha & beta
    gemm_operator(args); //运行Gemm

    cudaMemcpy(D, d_D, D_mem_size, cudaMemcpyDeviceToHost);  //将运行结果D矩阵传回host端
    std::cout << D[0] << std::endl;                          //打印D中第一行第一个数据
    std::cout << D[M * N - 1] << std::endl;                  //打印D中最后一行最后一个数据

    return 0;
}   

代码编译：

$nvcc -I <PATH TO CUTLASS>/include <YOUR SOURCE FILE>

运行结果如下：

 可见，本程序运行成功利用cutlass实现了D = α * A * B + β * C的矩阵计算。

*本文章及程序仅供交流学习，请勿用作商业用途。转载请告知作者。