多核切分下的Matmul矩阵计算从分块到对齐【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit为什么要做多核对齐切分在昇腾AI处理器的算子开发中矩阵乘Matmul是出镜率最高的计算类型之一。它要处理的往往是动辄成千上万行列的大矩阵如果把所有计算任务都压在一个核上不仅算力吃紧、耗时变长也等于把处理器上多核并行的能力白白闲置。解决办法说起来很朴素把大矩阵拆成若干小块分发给多个核同时计算。这就好比一个大型工程单靠一个人干会旷日持久把它拆成一批小任务交给一个团队大家各管一摊、并行推进最后再汇总效率立刻不一样。在Matmul的世界里这种把大块数据切成小瓷砖片再铺给多个核的做法就是本文要讲的多核对齐切分业内术语常写作矩阵Tiling——Tiling可以理解成按规则分块。切分要动刀的位置有讲究。矩阵乘涉及三张表A矩阵提供行方向的数据B矩阵提供列方向的数据C矩阵存放相乘的结果。沿哪条轴切、切几刀直接决定了每个核分到多大的活、最后怎么拼回完整结果。下面先看一套门槛更低的入门方案。入门方案只沿M、N两轴下刀矩阵的维度可以拆成三条轴M、N、K。初次上手时我们可以先不动K轴只在另外两条轴上做划分逻辑最直观。A矩阵沿M轴切成若干份每份叫SingleCoreM单核只需处理SingleCoreM × K大小的数据B矩阵沿N轴切成若干份每份叫SingleCoreN单核对应处理K × SingleCoreN大小的数据两个小块相乘自然得到SingleCoreM × SingleCoreN大小的C矩阵分块这正是单个核要产出的那部分结果。用数字感受一下假设8个核一起上阵A沿M轴分成4块、B沿N轴分成2块整个计算就变成了4×28个互不干扰的小任务每个核各领一份。比如5号核core5拿到的是A的第几块乘B的第几块独立算出属于自己的那块C谁也不挨着谁。这套方案的优点是清爽各核之间没有数据依赖做完即走无需额外汇总。它的代价是并行上限受限于M、N方向的天然维度当矩阵的M、N不够胖时再多核也派不上用场。进阶方案M、N、K三轴全切如果矩阵规模特别大或者希望调动更多核可以升级策略——把K轴也切开让并行度再上一个台阶。此时每个核手里的数据变成了A矩阵沿M轴、K轴分别切分单核拿到SingleCoreM × SingleCoreK的小块B矩阵沿K轴、N轴分别切分单核拿到SingleCoreK × SingleCoreN的小块C矩阵的分块大小依然是SingleCoreM × SingleCoreN但它的来源不再是一次相乘而是多次相乘的累加。累加的过程很关键设K轴被切成了3段那么某个C分块是通过A1×B1 A2×B2 A3×B3拼出来的其中每一项Ai×Bi都可以分给不同核并行算算完再把三份部分和加起来。这套方案相当于在M、N之外又开辟了K这条并行通道能容纳的核数上限更高代价则是多了跨核的累加环节各核先各自算部分和再沿K方向汇总逻辑上比入门方案多一环。两种策略各有用武之地选择依据主要是矩阵的形状和可用的核数。不过无论选哪套真正落到代码上都绕不开一件事——把切分方案告诉Tiling参数。Tiling参数与核数设置SetDim和SetBlockDim的分工切分规则最终要固化到一组Tiling参数里比如SingleCoreM、SingleCoreN、SingleCoreK分别记录每条轴上每核分得的尺寸。开发者不需要手算这些值host侧主机侧调用API就能自动获取。与单核场景不同多核场景需要先用MultiCoreMatmulTiling构造一个多核Tiling对象随后通过SetDim接口告诉Tiling计算Matmul能用多少个核。这里有两个容易混的接口必须分清接口管什么是否必填SetDimMatmul计算可用的核数只在多核场景使用用来算Tiling参数多核时设置SetBlockDim整个算子实际加载的核数决定真正被调用的核必须设置SetBlockDim的取值有规则可循而SetDim在纯Cube模式下遵循这样的逻辑SetDim设置当前AI处理器可用的核数 → Tiling计算会算出Matmul实际使用的核数 → 实际使用的核数一定小于等于可用核数 → SetBlockDim按实际使用核数来配置。一句话记忆SetDim是给Tiling报个家底SetBlockDim是拍板到底用几个核。如果算子除了矩阵计算还混有矢量计算MIX模式核数设置规则会略有不同需要额外参考MIX场景的核数设置约定这里不展开。一段看得懂的关键代码下面这段示意代码串起了整个流程构造多核Tiling对象、声明各矩阵的格式与类型、给出原始形状和切分形状、最后取出Tiling参数。注释里写了每一步在干什么。// 1. 拿到当前AI处理器的平台信息 auto ascendcPlatform platform_ascendc::PlatformAscendCManager::GetInstance(socVersion); // 2. 构造多核Tiling对象Matmul的切分方案都由它生成 matmul_tiling::MultiCoreMatmulTiling cubeTiling(*ascendcPlatform); // 3. 纯Cube场景把当前处理器上可用的Cube核数报给Tiling cubeTiling.SetDim(ascendcPlatform.GetCoreNumAic()); // 4. 声明A、B、C以及Bias在内存中的位置、格式和数据类型 cubeTiling.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16); cubeTiling.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16); cubeTiling.SetCType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT); cubeTiling.SetBiasType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT); // 5. 告诉Tiling原始形状以及经过切分后的计算形状 cubeTiling.SetOrgShape(M, N, K); cubeTiling.SetShape(M, N, K); // 6. 可选按需打开Bias cubeTiling.EnableBias(isBias); // 7. 取出最终Tiling参数ret为-1表示生成失败 optiling::TCubeTiling tilingData; int ret cubeTiling.GetTiling(tilingData);配合完整的算子样例如多核切M、N和多核切K两类样例一起阅读理解会更深。常见疑问Q1SetDim和SetBlockDim什么时候数值相同当Tiling计算出的实际使用核数恰好等于可用核数时两者数值一致若实际用不了那么多核SetBlockDim会按实际核数设得比SetDim更小。Q2只切M、N轴时各核之间需要通信吗不需要。每个核只做自己的那一片乘加输出也各自独立这是入门方案最大的优点。Q3切了K轴之后谁负责把部分和加起来每个C分块的累加由对应的那组核协同完成多个Ai×Bi并行算出部分和后再按K方向求和汇总。小结与思考回到开头的问题多核对齐切分的本质是把一个大矩阵乘重构为若干个小矩阵乘的集合。是只切M、N两轴还是三轴全切本质是在并行的上限和累加的开销之间做权衡。前者代码简单、核间解耦适合M、N较胖的场景后者并行度更高但要额外处理K方向的累加适合大矩阵、多核场景。顺着这个思路再往前想一步如果K轴切出的块数不能均分怎么办矩阵边缘不足一个SingleCoreK的尾块如何对齐、如何分配这些正是Tiling计算在背后默默替你解决的细节。理解了切分—分配—累加—对齐这条主线再去看具体的Tiling参数和核数配置就不会觉得它们是一堆冷冰冰的数值了。【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
