尧图建网站 尧图建网站 YAOTU WEB BUILD 免费咨询
ARTICLE DETAIL

资讯详情

深耕网站建设与建站编程的一线实战洞察。

三步吃透 Matmul 多核切分:从“一个核干到底“到“八个核分头干活“

三步吃透 Matmul 多核切分:从“一个核干到底“到“八个核分头干活“ 三步吃透 Matmul 多核切分从一个核干到底到八个核分头干活【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit刚接触昇腾算子开发的朋友十有八九都撞过这么一堵墙在 CANN asc-devkit 里用 Matmul 高级 API 写矩阵乘法单核跑通、结果正确、心满意足可一上大矩阵就傻眼了——动辄 M、N 上万的乘法一个核吭哧吭哧算完耗时感人跟AI 算力四个字完全不沾边。问题出在哪很简单你只用了处理器上一个核。这颗 AI 处理器上明明躺着几十个 Cube 核凭什么让它们集体围观今天这篇实战笔记就带你把 Matmul 多核切分多核矩阵切分与 Tiling 参数生成这件事彻底搞明白并给出可直接照抄的 host 侧 Tiling 代码。一个核是单人厨房多个核才是流水线中央厨房先别急着看代码我们用生活打个比方。假设你要给 2000 人做工作餐。一种做法是一个大厨从洗菜、切配到炒菜全部包圆2000 份做完天都黑了——这就是单核 Matmul。另一种做法是把 2000 人分成 8 个就餐区每个区一个大厨负责八个灶台同时开火——这就是多核 Matmul。但分餐不能乱分。菜谱是米饭 × 菜的搭配对应到矩阵乘法就是A 矩阵 × B 矩阵 C 矩阵。切分时你得想清楚每个大厨到底负责哪一块菜这就是矩阵沿M、N、K 三个轴切分的本质。核心思路一句话C 矩阵的每一小块只依赖 A、B 矩阵中与之相关的对应小块所以可以大胆把大矩阵拆开让不同核并行算各自的分块。先看最朴素的切法只动 M 和 NK 轴不动最省事的策略是不切 K 轴只切 M、N 轴。它有三种切分动作记住三个变量名即可把 A 矩阵沿 M 轴切成一摞singleCoreM每个核拿一份singleCoreM × K的数据把 B 矩阵沿 N 轴切成一列singleCoreN每个核拿一份K × singleCoreN的数据两个分块相乘恰好得到singleCoreM × singleCoreN的 C 矩阵分块这就是该核的输出。举个具体数字8 个核A 沿 M 轴切 4 块B 沿 N 轴切 2 块正好 4×28 个核各司其职谁也不碰谁的数据。这种切法逻辑最简单但有个隐患如果 K 很大比如 4096每个核仍要串行把 K 这一整条算完算力依然有浪费。于是就有了第二种玩法。再上难度M、N、K 三轴全切让累加也并行切 K 轴的思路等于把一道大菜再拆成三道半成品最后拼盘。A 矩阵沿 M 轴切singleCoreM再沿 K 轴切singleCoreK每核处理singleCoreM × singleCoreKB 矩阵沿 K 轴切singleCoreK再沿 N 轴切singleCoreN每核处理singleCoreK × singleCoreNC 矩阵的某个 R 块 A1×B1 A2×B2 A3×B3这三次乘加可以拆给多个核并行算最后把部分和累加在一起。注意这里的变化前一种策略每个核独立产出完整 C 分块互不通信而切 K 之后同一个 C 分块需要多个核分别算部分积再汇总累加这一步是切 K 独有的开销。它适合 K 特别大、单纯切 M、N 已经无法把核用满的场景。两种策略都跑不了数据分不完的尴尬当 M、N、K 无法被块大小整除时会冒出尾块多核非对齐切分场景。好在 asc-devkit 提供了SetTail接口在 kernel 侧迭代前把本次的 singleCoreM/singleCoreN/singleCoreK 临时改成尾块大小即可无需推翻重设 tiling。动手实践host 侧如何申请多核 Tiling切分策略最终要落到 Tiling 参数上。单核场景用普通的 MatmulTiling 就行多核场景必须换成MultiCoreMatmulTiling这是最容易踩的第一个坑。完整代码骨架如下// 1. 拿到当前平台的核信息 auto ascendcPlatform platform_ascendc::PlatformAscendCManager::GetInstance(); // 2. 构造多核 Tiling 对象 matmul_tiling::MultiCoreMatmulTiling cubeTiling(*ascendcPlatform); // 3. 告诉 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. 输入真实形状 cubeTiling.SetOrgShape(M, N, K); cubeTiling.SetShape(M, N, K); cubeTiling.EnableBias(isBias); // 6. 产出 tiling 数据 optiling::TCubeTiling tilingData; int ret cubeTiling.GetTiling(tilingData); // ret -1 说明生成失败逐段看其实每步都在回答一个问题步骤代码它在做什么获取平台PlatformAscendCManager::GetInstance()拿到当前处理器的核数等硬件信息是所有 Tiling 的前提构造对象MultiCoreMatmulTiling只有多核场景用它单核场景用普通 MatmulTiling 即可设可用核数SetDim(GetCoreNumAic())声明最多有多少核能干活它是 Tiling 计算的上限声明类型SetAType/SetBType/SetCTypeA、B 是 FP16、C 是 FP32这是矩阵乘常见的精度组合填形状SetOrgShape/SetShape原始形状和当前形状涉及 reshape 场景时两者不同取结果GetTiling(tilingData)生成 singleCoreM/singleCoreN/singleCoreK 等切分参数注意第 6 步返回-1表示 tiling 生成失败记得做错误判断别闷头往下跑。避坑指南两个核数千万别搞混新手最容易在核数设置上翻车因为这里藏着两个长得像、作用完全不同的参数SetDim设的是可用的核数上限它告诉 Tiling 计算器你最多可以切多少个核实际用几个由GetTiling决定。对纯 Cube 场景一般填GetCoreNumAic()也就是当前处理器上全部的 Cube 核数即可。另一个设的是实际加载的核数它决定整个算子到底启动几个核跑这是必须设置的不同版本该接口名略有差异请以你所用的 asc-devkit 版本文档为准常见名为 SetSimdNumBlocks 或 SetBlockDim并应按照GetTiling算出来的实际使用核数来填。一句话记忆SetDim是预算上限另一个是实际花销。前者设大了没关系Tiling 会按需裁剪后者设错了轻则性能打折重则核间数据对不上、结果出错。纯 Cube 场景里SetDim通常比实际使用核数设得更大或相等GetTiling会自动算出一个合理值我们照着把它填进核数设置接口就行。而如果算子里除了矩阵计算还混了矢量计算MIX 模式核数分配规则会更复杂建议直接看官方文档中 MIX 场景的专门说明这里不展开。一句话总结Matmul 多核切分本质就是把一个大矩阵乘法拆成每个核一份小作业要么只切 M、N 轴各算各的要么连 K 轴一起切、多核接力累加。记住MultiCoreMatmulTiling SetDim 正确设置实际核数这个组合你的算子就能从单核独舞升级为全核合唱吞吐量往往就是成倍的差别。想上手跑通完整样例可以 clone 仓库 https://gitcode.com/cann/asc-devkit 后去examples/01_simd_cpp_api/00_introduction/02_matrix下找多核切 M、N 的 matmul 样例在examples/01_simd_cpp_api/04_advanced_api/00_matmul/matmul_splitk里看多核切 K 的实现尾块处理则可以参考matmul_multi_core_unaligned样例。下一步试着把 M、N、K 改成不能被核数整除的值亲手体会一下非对齐切分的乐趣吧。【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表