KuDNN的矩阵乘之对接KML_BLAS_GEMM接口实现详解

KuDNN的矩阵乘之对接KML_BLAS GEMM接口实现详解 本文将围绕KuDNN的GEMM基本原理展开介绍GEMM接口分别对接了KML BLAS的cblas_gemm接口和JIT code本次主要介绍KML BLAS这一算法分支的原理和实现。 1. 概述 本文将围绕KuDNN的GEMM基本原理展开介绍GEMM接口分别对接了KML BLAS的cblas_gemm接口和JIT code本次主要介绍KML BLAS这一算法分支的原理和实现。 KuDNN源码可参考 https://gitcode.com/kunpengcompute/kudnn 2. 接口介绍include/operations/kudnn_gemm.hpp中声明了Gemm接口类方法主要是Gemm构造函数和Run执行计算两个成员函数以及指向真正Gemm实现GemmImpl类的pImpl指针。class KUDNN_API_PUBLIC Gemm final { public: Gemm(const TensorInfo aInfo, const TensorInfo bInfo, const TensorInfo cInfo, const TensorInfo biasInfo, int numThreads 0) noexcept(false); void Run(const void *a, const void *b, void *c, const void *bias, float alpha 1.0f, float beta 0.0f, int numThreads 0) const noexcept(false); private: std::unique_ptrDetail::GemmImpl pImpl; };kudnn的所有接口实现均采用Pointer to Implement机制将类的实现细节从其对象表示中移除通过不透明指针将它们放置在单独的类中。此技术用于构建具有稳定 ABI 的 C 库接口并减少编译时依赖。因为类的私有数据成员参与其对象表示影响大小和布局并且因为类的私有成员函数参与重载决议在成员访问检查之前进行所以对这些实现细节的任何更改都需要重新编译类的所有用户。PImpl 消除了这种编译依赖对实现的更改不会导致重新编译。因此如果库在其 ABI 中使用 PImpl则新版本的库可以更改实现同时与旧版本保持 ABI 兼容。Gemm的必要参数为A,B,C张量计算过程为CA*BC。tensorInfo描述了张量信息包括张量的维度大小(dims)数据类型布局(layout)和步长(stride); 观察接口会发现相比BLAS的标准GEMM接口少量LDA、LDB和LDC参数其实这些参数可以用stride表示实现。# 3. 算法流程## 3.1 实例化Gemm从接口层的Gemm类只提供接口定义其实现完全由Detail::GemmImpl来完成// all public methods just call corresponding pImpl implemnentation Gemm::Gemm(const TensorInfo aInfo, const TensorInfo bInfo, const TensorInfo cInfo, const TensorInfo biasInfo, int numThreads) noexcept(false): pImpl(new Detail::GemmImpl(aInfo, bInfo, cInfo, biasInfo, numThreads)){}Detail::GemmImpl类中有成员GemmInfo记录gemm参数信息其类成员就是srcweidstbias等参数; 同时将impl设置为nullptrGemmImpl(const TensorInfo srcInfo, const TensorInfo weiInfo, const TensorInfo dstInfo, const TensorInfo biaInfo, int numThreads) noexcept(false) : gemmImplInfo(srcInfo, weiInfo, dstInfo, biaInfo), impl(nullptr)在GemmImpl构造函数中①先校验输入参数的合法性②如果是鲲鹏920判断输入输出张量类型是否满足使用JIT的条件如果满足会进一步调用FindSolution查找合适的解决方案 并赋值给类成员impl指针否则保持impl为nullptr表示后续回退到BLAS计算gemm。{ Service::ThrowOnStatus(Gemm::ValidateInput(srcInfo, weiInfo, dstInfo, biaInfo, numThreads), “GEMM”); bool srcJIT weiJIT dstJIT false; #ifdef KUDNNL_920Pro srcJIT srcJIT || (gemmImplInfo.srcInfo.GetType() Element::TypeT::S8); weiJIT weiJIT || (gemmImplInfo.weiInfo.GetType() Element::TypeT::S8); … dstJIT dstJIT || (gemmImplInfo.dstInfo.GetType() Element::TypeT::S32); #endif // KUDNNL_920Pro … if (srcJIT weiJIT dstJIT) { impl FindSolution(gemmImplInfo, Threading::GetMaxNumThreads()); } }FinSolution函数会生成JITcode具体分析将在另一篇《kudnn的JIT》文章中详细展开这里不再赘述。至此不用JIT的情景下的Gemm实例构造的过程结束下面看Gemm的计算过程。## 3.2. Gemm计算Gemm计算的整体架构是GemmImpl::Run (入口) ↓ BatchedExtendedGemm (回退BLAS GEMM算法) ↓ ChooseImpl (类型匹配与转换) ↓ GEMMCaller::Call (参数准备) ↓ GEMMWrapper (底层BLAS调用)### 3.2.1. GemmImpl::Run调用Gemm.Run接口输入需要计算的张量数据指针a,b,c实现计算。Gemm类的Run函数同样只是简单转调GemmImpl的Run函数。GemmImpl的Run实现如下判断impl是否为nullptr如果非空指针则跳转到JIT模块执行。否则调用BatchedExtendedGemm函数采用传统BLAS接口计算。void Run(const void *a, const void *b, void *c, const void *bias, float alpha, float beta, int numThreads){ if (impl ! nullptr) { if (impl-GetNThreads() ! Threading::GetMaxNumThreads()) { impl FindSolution(impl-GetTask(), Threading::GetMaxNumThreads()); } impl-Run({a, b, c, bias, alpha, beta}); } else { GemmHelpers::BatchedExtendedGemm(gemmImplInfo, a, b, c, bias, alpha, beta, numThreads); } }gemmImplInfo是中保存有GEMM的维度大小和布局等元数据信息结合Run接口输入的计算数据可组合参数调用KBLAS完成GEMM计算。### 3.2.2. BatchedExtendedGemmBatchedExtendedGemm函数完成3件事依次介绍①布局标准化 判断输入张量的布局是否符合满足标准ABX布局先通过ReorderLayer转换。下面介绍一下什么是ABX以及为什么要转换。 ABX布局是矩阵乘法中的一种标准内存布局约定它规定了输入矩阵A、B和输出矩阵X在内存中的排列方式。从ABX的代码可知其采用行优先布局,维度从高到低依次排列。Layout GetStandardABXLayout() const { switch (dims.GetNumDims()) { case DIM_1: { return Layout::A; } … case DIM_5: { return Layout::ABCDE; } } }ABX布局转换的核心目的是消除布局差异所有矩阵统一内存排列最大化BLAS性能使用最优的NoTrans路径简化地址计算在batch处理中避免stride语义混乱提升cache效率保证连续内存访问② 5D广播auto broadcastedDimsA Service::BroadcastTo5D( Service::GetShapeAccordingToLayout(gemmInfoExec.srcInfo.GetDims(), gemmInfoExec.srcInfo.GetLayout())); auto broadcastedDimsB Service::BroadcastTo5D( Service::GetShapeAccordingToLayout(gemmInfoExec.weiInfo.GetDims(), gemmInfoExec.weiInfo.GetLayout())); Shape broadcastedDims {std::max(broadcastedDimsA[IDX_0], broadcastedDimsB[IDX_0]), std::max(broadcastedDimsA[IDX_1], broadcastedDimsB[IDX_1]), std::max(broadcastedDimsA[IDX_2], broadcastedDimsB[IDX_2]), 1, 1}; auto reorderedStridesA Service::BroadcastTo5D(Service::GetShapeAccordingToLayout(gemmInfoExec.srcInfo.GetStrides(), gemmInfoExec.srcInfo.GetLayout()), false); auto reorderedStridesB Service::BroadcastTo5D(Service::GetShapeAccordingToLayout(gemmInfoExec.weiInfo.GetStrides(), gemmInfoExec.weiInfo.GetLayout()), false); …③ 三重循环批处理 因为BLAS的gemm一次计算一个2D的矩阵运算对于kudnn最高支持5维张量的情景需要个三重循环依次遍历计算GEMM。for (SizeType i0 0; i0 broadcastedDims[IDX_0]; i0) { for (SizeType i1 0;i1 broadcastedDims[IDX_1]; i1) { for (SizeType i2 0; i2 broadcastedDims[IDX_2]; i2) { const std::byte *aPtr static_castconst std::byte *(aExec) (i0 * reorderedStridesA[IDX_0] i1 * reorderedStridesA[IDX_1] i2 * reorderedStridesA[IDX_2]) * gemmInfoExec.srcInfo.GetType().GetSize(); const std::byte *bPtr … std::byte *cPtr … const std::byte *biasPtr static_castconst std::byte *(biasExec); if (biasExec) { biasPtr … } ChooseImpl(gemmInfoExec, aPtr, bPtr, cPtr, biasPtr, alpha, beta, numThreads); } } }### 3.2.3. ChooseImpl实现了三级回退策略第1级精确类型匹配 ↓ (失败) 第2级扩展精度匹配 ↓ (失败) 第3级全FP32转换所有的Gemm精度类型组合列表如下,由floatfp16bf16int8uint8之间组合而成每一项都是一个特化的GEMMCaller模板类类中实现了对应数据类型的cblas_gemm函数调用共有11种类型组合。static std::vectorstd::shared_ptr g_gemmImpls { std::make_sharedGEMMCallerfloat, float, float, float(), std::make_sharedGEMMCaller__fp16, __fp16, __fp16, __fp16(), std::make_sharedGEMMCaller__fp16, __fp16, float, float(), … std::make_sharedGEMMCallerstd::int8_t, std::int8_t, std::int32_t, std::int32_t(), … };第1级 - 精确匹配如果输入的src、wei、dst和bias的类型与既定的g_gemmImpls里的数据类型完全匹配就直接调用对应的GEMMCaller方法。// check for exact match for (auto gemmImpl : g_gemmImpls) { bool isTypeMatch (gemmImpl-GetSrcDt() srcType) (gemmImpl-GetWeiDt() weiType) (gemmImpl-GetDstDt() dstType) (gemmImpl-GetBiaDt() biaType); if (isTypeMatch) { gemmImpl-Call(gemmInfo, src, wei, dst, bia, alpha, beta, numThreads); return; } }第2级 - 扩展精度匹配 通过GetWiderType部分获取扩展的目标数据类型然后选择dst为widerType类型的GEMMCaller作为目标函数将其余参数扩展转换为目标函数的参数类型然后调用Call函数继续执行。如果其余参数有比目标函数的参数更宽的数据类型则继续回退到全FP32类型。// try to convert to wider type and check if there’s implementation with wider type auto widerType GetWiderType(srcType, weiType, dstType, biaType);for (auto gemmImpl : g_gemmImpls) { bool isWiderType ((gemmImpl-GetDstDt() widerType) (dstType widerType)) ((srcType.GetSize() gemmImpl-GetSrcDt().GetSize()) (srcType.IsSigned() gemmImpl-GetSrcDt().IsSigned())) … if (isWiderType) { MatrixConverter s(srcType, gemmImpl-GetSrcDt(), src, m * k); MatrixConverter w(weiType, gemmImpl-GetWeiDt(), wei, k * n); MatrixConverter b(biaType, gemmImpl-GetBiaDt(), bia, biasM * biasN); gemmImpl-Call(gemmInfo, s.Get(), w.Get(), dst, b.Get(), alpha, beta, numThreads); return; } }GetWiderType函数实现了类型扩展如果存在BF16则返回BF16如果输入全是FP或integer类型就返回位数最宽的那个参数的类型其他情况返回FP32类型至于为什么是FP32是因为FP32是kudnn支持的最宽数据类型对其他数据类型兼容性最好。第3级 - 全FP32回退 将所有输入类型全部转换为FP32类型调用FP32的GEMMCaller计算最后将计算输出dst转换为原来转换前的数据类型作为最终输出。// some integer combinations can’t be supported when all the types are converted to wider type MatrixConverter s(srcType, Element::TypeT::F32, src, m * k); … GEMMCallerfloat, float, float, float {}.Call(gemmInfo, s.Get(), w.Get(), d.Get(), b.Get(), alpha, beta, numThreads); if (dstType ! Element::TypeT::F32) { switch (dstType) { case Element::TypeT::F16: { Service::ConvertFp32ToFp16(static_castconst float *(d.Get()), static_cast__fp16 *(dst), m * n); break; } … } }### 3.2.4. GEMMCaller::CallGEMMCaller负责从gemmInfo中获取矩阵乘所需的M、N、K、lda、ldb、ldc、transa、transb等参数结合src、wei、dst、bias和offsetC得到cblas_gemm所需的全部参数再通过调用BlasSetNumThreadsLocal设置BLAS线程数最后调用GEMMWrapper完成计算### 3.2.5. GEMMWrapperGEMMWrapper特化模板函数调用具体的cblas_?gemm接口完成计算。template typename SrcDt, typename WeiDt, typename DstDt, typename BiaDt void GEMMWrapper(…)以void GEMMWrapperfloat, float, float, float为例将矩阵乘法分为三种情况优化 ① n1: GEMV矩阵-向量乘选取合适的布局后调用 cblas_sgemv ② m1: GEMV向量-矩阵乘 ③ 其他: GEMM直接调用 cblas_sgemm最后通过 AddBias 添加偏置项