CublasLt 极简入门

news2024/11/14 11:32:24

CublasLt 入门

官方文档： cublasLt文档

网上搜了挺多，Cublas的文档有不少，但是对于cuda10以上轻量化的cublasLt 的文档或博客却很少。

cuBLASLt使用启发式方法根据问题大小、GPU配置和其他参数选择最合适的内核来执行。这需要在主机CPU上执行一些计算，这可能需要几十微秒。为了克服这种开销，建议使用cublasLtMatmulAlgoGetHeuristic（）查询启发式算法一次，然后使用cublasLtMatmul（）在后续计算中重用结果。

1. 初始化 cuBLASLt 句柄

和cublas的使用一样，在使用cublasLt之前需要先创建一个cublasLtHandle_t 句柄

cublasLtHandle_t ltHandle;
cublasLtCreate(&ltHandle)

2. 设置operation描述符

cublasLtMatmulDescOpaque_t operationDesc = {};
cublasOperation_t transa；
cublasOperation_t transb；

cublasLtMatmulDescInit(&operationDesc, 		CUBLAS_COMPUTE_64F, CUDA_R_64F);
// 设置 A和B转置信息
cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa));

cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transb));

3.设置矩阵描述符

为乜咯矩阵初始化和设置矩阵尺寸，类型，信息等。

在执行矩阵乘法前，需要定义输入矩阵和输出矩阵的描述符。每个描述符都包含矩阵的尺寸、数据类型、布局等信息。cublasLtMatrixLayout_t 用于描述每个矩阵的布局，包括行数、列数、数据类型等。

cublasLtMatrixLayoutOpaque_t Adesc = {}, Bdesc = {}, Cdesc = {};

cublasLtMatrixLayoutInit(&Adesc, CUDA_R_64F, m , k, lda);
cublasLtMatrixLayoutInit(&Bdesc, CUDA_R_64F, k, n, ldb);
cublasLtMatrixLayoutInit(&Cdesc, CUDA_R_64F, m, n, ldc);

通过cublasLtMatrixLayoutInit 函数初始化之前创建的矩阵布局描述符。

cublasStatus_t cublasLtMatrixLayoutInit( 		 cublasLtMatrixLayout_t matLayout,
                                        cudaDataType type,
                                        uint64_t rows,
                                        uint64_t cols,
                                        int64_t ld);

/*
matLayout  : 输出参数， 指向将要被本函数初始化的 matrix layout descriptor结构体的指针
type       : 具体的数据精度类型
rows, cols : 矩阵行和列的维度
Ld         : 矩阵的主维。 如果是列主，就是行数
*/

4 初始化矩阵乘法算法结构

对具体的矩阵乘法算法以及输入矩阵A，B，C以及输出矩阵D通过函数cublasLtMatmulAlgoInit() 为矩阵乘法算法(cublasLtMatmul)结构进行初始化。

cublasLtMatmulAlgoInit(ltHandle,  //
                       CUBLAS_COMPUTE_64F,   // compute
                       CUDA_R_64F,   // scale
                       CUDA_R_64F,   // A
                       CUDA_R_64F,   // B
                       CUDA_R_64F,   // C
                       CUDA_R_64F,   // D
                       algoId,
                       &algo));
// 设置tile大小 tileId = CUBLASLT_MATMUL_TILE_16x16
cublasLtMatmulAlgoConfigSetAttribute(&algo, 	     CUBLASLT_ALGO_CONFIG_TILE_ID, &tileId, sizeof(tileId));

// 设置reduction结构，当split大于1时候设置
// reductionMode = CUBLASLT_REDUCTION_SCHEME_INPLACE
cublasLtMatmulAlgoConfigSetAttribute(&algo, CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME, &reductionMode, sizeof(reductionMode));

// 设置splits数量  splitKFactor = 256
cublasLtMatmulAlgoConfigSetAttribute(&algo, CUBLASLT_ALGO_CONFIG_SPLITK_NUM, &splitKFactor, sizeof(splitKFactor));

其中cublasLtMatmulAlgoInit() 函数作用：

cublasStatus_t cublasLtMatmulAlgoInit(
      cublasLtHandle_t lightHandle,
      cublasComputeType_t computeType,
      cudaDataType_t scaleType,
      cudaDataType_t Atype,
      cudaDataType_t Btype,
      cudaDataType_t Ctype,
      cudaDataType_t Dtype,
      int algoId,
      cublasLtMatmulAlgo_t *algo);

/*
初始化矩阵乘法算法结构。
lightHandle   : 输入参数，cublasLtHandle_t的上下文句柄
computeType   : 计算类型
scaleType     : scale类型，一般和computeType一样（就是缩放因子alpha,beta的数据类型）
Atype, Btype, Ctype, and Dtype : 输入和输出矩阵的数据类型
algoId        : 指定初始化的算法类型，应该为cublasLtMatmulAlgoGetIds的有效返回值
algo          : 指向待初始化的不透明结构体指针
*/

通过cublasLtMatmulAlgoInit 初始化了algo之后，通过cublasLtMatmulAlgoConfigSetAttribute() 为 algo 设置指定配置属性的值。

cublasStatus_t cublasLtMatmulAlgoConfigSetAttribute(
      cublasLtMatmulAlgo_t *algo,
      cublasLtMatmulAlgoConfigAttributes_t attr,
      const void *buf,
      size_t sizeInBytes);

/*
algo   :  经过了初始化的保存矩阵乘法算法描述符的不透明数据结构类型
attr   :  需要通过本函数设置的配置属性
buf    :  要被设置的实际值
sizeInBytes ： buf的大小（字节），用于验证。
*/

属性设置值

进行矩阵乘法

// C = alpha * op(A) + beta * op(B)
cublasLtMatmul(ltHandle,
               &operationDesc,
               alpha,
               A,
               &Adesc,
               B,
               &Bdesc,
               beta,
               C,
               &Cdesc,
               C,
               &Cdesc,
               &algo,
               workspace,
               workspaceSize,
               stream)