CANN/catlass GEMM恒等块调度

Gemm Identity Block Swizzle

【免费下载链接】catlass本项目是CANN的算子模板库，提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass

代码位置

功能说明

Swizzle策略决定了AI Core对基本任务块的分配关系和计算顺序，详见Swizzle解释。GemmIdentityBlockSwizzle策略对结果矩阵C在$M$、$N$方向切基本块，根据$SwizzleOffset$和$SwizzleDirection$决定基本块排布顺序，并在排布上按照物理核序号依次分配基本块。

如下图示例，AI Core共20个，基本块沿箭头方向依次分配给AI Core。箭头方向则由$SwizzleOffset = 1$和$SwizzleDirection = 0$来决定：

常用Methods说明

调用示例

Block组装

参考basic_matmul

// Swizzle offset is 3 and direction is 0. using BlockScheduler = typename Gemm::Block::GemmIdentityBlockSwizzle<3, 0>; // kernel level using MatmulKernel = Gemm::Kernel::BasicMatmul<BlockMmad, BlockEpilogue, BlockScheduler>;

Matmul的kernel使用

参考basic_matmul，在kernel代码的void operator()<AscendC::AIC>函数中。

实例化BlockScheduler

BlockScheduler matmulBlockScheduler(params.problemShape, MakeCoord(L1TileShape::M, L1TileShape::N));

获取基本块总数

uint32_t coreLoops = matmulBlockScheduler.GetCoreLoops();

基本块遍历中，获取当前block的block坐标和实际shape

for (uint32_t loopIdx = AscendC::GetBlockIdx(); loopIdx < coreLoops; loopIdx += AscendC::GetBlockNum()) { // Compute block location GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx); GemmCoord actualBlockShape = matmulBlockScheduler.GetActualBlockShape(blockCoord); ... }

GroupMatmul的kernel使用

参考grouped_matmul_slice_m_per_token_dequant_multistage_workspace，在kernel代码的void operator()<AscendC::AIC>函数中。

实例化BlockScheduler

BlockScheduler blockScheduler;

group的遍历中，根据每个group的shape更新blockScheduler，并获取单个group的基本块总数

for (uint32_t groupIdx = 0; groupIdx < params.problemCount; ++groupIdx) { ... blockScheduler.Update(inGroupProblemShape, L1TileShape::ToCoordMN()); uint32_t coreLoops = blockScheduler.GetCoreLoops(); ... }

单个group的基本块遍历中，获取当前block的block坐标和实际shape

for (uint32_t loopIdx = startLoopIdx; loopIdx < coreLoops; loopIdx += coreNum) { GemmCoord blockCoordMNK = blockScheduler.GetBlockCoord(loopIdx); GemmCoord actualBlockShapeMNK = blockScheduler.GetActualBlockShape(blockCoordMNK); ... }

【免费下载链接】catlass本项目是CANN的算子模板库，提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass