多核对齐切分:让矩阵乘法在多核上"各干各的"
【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言,原生支持C和C++标准规范,主要由类库和语言扩展层构成,提供多层级API,满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit
矩阵越算越慢,单核资源明明在空转,跑一次大矩阵的 Matmul 却要等上好半天——这是很多算子开发新手第一次接触多核并行时的真实感受。答案其实就藏在"多核对齐切分"这四个字里:把一个大矩阵切成若干小块,分给不同的 AI Core 并行计算。本文以 CANN asc-devkit 中的 Matmul 为例,带你搞懂多核切分的两种策略、Tiling 参数和核数设置的避坑点。
切蛋糕的智慧:把大活拆成小活
想象你在筹备一场 100 人的晚宴,只有一位厨师。他做完 100 份菜需要一整天。但如果把菜谱切成 10 份,请 10 位厨师同时开工,理论上 1 小时就能全部搞定。唯一的代价是:需要先设计好"怎么切",让每位厨师的工作量均衡、且互不干扰。
矩阵乘法(C = A × B)的多核并行就是同一件事。C 矩阵的每个元素都可以独立计算,天然适合并行。问题只在于:怎么把 A、B 矩阵切开,才能让每个核各算一块、最后拼起来正好等于完整结果?这就是"多核对齐切分"要解决的核心问题,而切分方案会被记录在多核 Tiling 参数(SingleCoreM / SingleCoreN / SingleCoreK)里。
两种切分策略,一张图看懂
Matmul 的切分策略有两种:不切 K 轴和切 K 轴。它们的本质区别在于:要不要因为"中间累加"而引入跨核通信。
策略一:只切 M、N 轴(不切 K)
这是最简单直观的做法:
- A 矩阵沿 M 轴切,每个核拿到
SingleCoreM × K的 A 分块; - B 矩阵沿 N 轴切,每个核拿到
K × SingleCoreN的 B 分块; - 每个核独立计算
SingleCoreM × SingleCoreN的 C 分块,互不依赖、无需通信。
每个核都持有完整的 K 维度数据,相当于每位厨师都拿到完整的"配方",自己从头做到尾。缺点也很明显:K 很大时,每个核都要反复加载整条 K 维数据,数据搬运开销大。
策略二:M、N、K 三个轴都切
当 K 维度很大时,把 K 也切开:
- A 矩阵切成
SingleCoreM × SingleCoreK的块; - B 矩阵切成
SingleCoreK × SingleCoreN的块; - 每个核只算其中一块,部分结果= A1×B1(或 A2×B2),最后把多个核的部分结果累加成完整的 C 分块。
形象地说,这就像把"从头做到尾"改成了"流水线分工":每个核只负责其中一段工序,但最后需要把几段工序的产出合起来。因此切 K 轴引入了跨核累加,需要额外的同步与归约开销。
💡怎么选?K 不大、核数够用 → 只切 M、N,简单可靠;K 特别大导致单核搬运 K 维数据成为瓶颈 → 切 K,用通信换内存带宽。CANN 的 MultiCoreMatmulTiling 会自动帮你算出每种策略下的 SingleCoreM / SingleCoreN / SingleCoreK,你只需关心"设多少核"。
关键参数与核数设置:SetDim vs SetBlockDim
多核 Tiling 有四个关键参数,前三个描述"怎么切",最后一个描述"用几个核":
| 参数 | 含义 | 谁来算 |
|---|---|---|
| SingleCoreM | 单核处理的 M 方向大小 | Tiling 自动计算 |
| SingleCoreN | 单核处理的 N 方向大小 | Tiling 自动计算 |
| SingleCoreK | 单核处理的 K 方向大小(仅切 K 场景出现) | Tiling 自动计算 |
| SetDim / SetBlockDim | 设置参与计算的核数 | 开发者手动设置 |
参数自动算,但核数必须你亲自设,而且这里藏着新手最容易踩的坑:
⚠️SetDim 和 SetBlockDim 完全不是一回事。
- SetDim:设置"可用的核数",即告诉 Tiling 计算器"我有这么多核可以用",Tiling 会据此推导 SingleCoreM/N/K。它不决定实际加载哪些核。
- SetBlockDim:设置"整个算子实际加载的核数",这是真正会生效、会被启动的核数,必须设置。
简单记:
SetDim是"允许用多少核",SetBlockDim是"真的用多少核"。纯 Cube 模式下,你应根据 Tiling 实际计算出的核数来配 SetBlockDim;而 MIX 模式(Cube + 矢量计算)的核数规则更复杂,建议参考算子实践章节中"矩阵编程"的核数设置说明。
最小可运行示例:把 Tiling 串起来
下面这段骨架代码展示了多核 Tiling 的完整流程,核心步骤只有五步:
// ① 创建多核 Tiling 对象(单核场景用 MatmulTiling,多核用 MultiCoreMatmulTiling) auto platform = platform_ascendc::PlatformAscendCManager::GetInstance(socVersion); matmul_tiling::MultiCoreMatmulTiling tiling(*platform); // ② 设置可用核数:把当前 AI 处理器的 Cube 核数全部告诉 Tiling tiling.SetDim(platform.GetCoreNumAic()); // ③ 声明 A、B、C 及可选的 Bias 的类型、存储位置和格式 tiling.SetAType(TPosition::GM, CubeFormat::ND, DataType::DT_FLOAT16); tiling.SetBType(TPosition::GM, CubeFormat::ND, DataType::DT_FLOAT16); tiling.SetCType(TPosition::GM, CubeFormat::ND, DataType::DT_FLOAT); tiling.SetBiasType(TPosition::GM, CubeFormat::ND, DataType::DT_FLOAT); // ④ 传入矩阵形状(OrgShape 是原始形状,Shape 可用于带 padding 的场景) tiling.SetOrgShape(M, N, K); tiling.SetShape(M, N, K); tiling.EnableBias(isBias); // ⑤ 生成 Tiling 数据,返回值 -1 表示生成失败 optiling::TCubeTiling tilingData; int ret = tiling.GetTiling(tilingData); // if ret = -1, gen tiling failed逐段解释:
- ①②:先建对象、报出"家底"(可用核数),让 Tiling 计算器知道资源上限;
- ③④:声明数据长什么样(类型、格式、形状),Tiling 据此推演每种切分策略下的分块大小;
- ⑤:
GetTiling一次性把 SingleCoreM / SingleCoreN / SingleCoreK 等参数填进tilingData,供 host 侧使用;之后记得按实际使用核数设置SetBlockDim。
完整可运行样例可在项目 examples 目录中找到:只切 M、N 的场景看01_simd_cpp_api/00_introduction/02_matrix/matmul;切 K 的场景看01_simd_cpp_api/03_libraries/00_matrix/matmul_splitk。
常见疑问 Q&A
Q1:什么时候应该切 K 轴?A:当 K 特别大、单核反复搬运整条 K 维数据成为性能瓶颈时,切 K 能降低单核内存压力;代价是需要跨核累加。K 不大时优先只切 M、N。
Q2:核数到底怎么定?A:纯 Cube 场景,用SetDim报出可用核数,Tiling 会给出实际使用的核数,再据此配置SetBlockDim。记住二者别混淆。
Q3:核数设得越多就越快吗?A:不一定。矩阵尺寸固定时,切分越细、单核负载越小,可能造成大量核空转甚至通信开销超过收益。设置原则是"核数尽量整除矩阵维度、负载均衡"。
Q4:SingleCoreK 只在切 K 场景才有吗?A:是的。不切 K 时每个核持有完整的 K 维,只有切 K 后才有"单核处理的 K 分块大小"这一概念。
总结与延伸
多核对齐切分的本质,是把"一位厨师做一百份菜"变成"多位厨师分工协作":只切 M、N 轴简单无通信,切 K 轴以通信换带宽;分块大小由 Tiling 自动算出,核数则靠你通过 SetDim 与 SetBlockDim 正确设置。搞懂这两对关系,多核 Matmul 的性能优化就有了方向。
想继续深入?可以接着读项目中"矩阵编程(高阶 API)"的算子实现章节,了解 MIX 模式的核数设置规则;动手跑一遍上面提到的 matmul 与 matmul_splitk 样例,把单核改造为多核,你会对 Tiling 参数的流向有更直观的感受。
【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言,原生支持C和C++标准规范,主要由类库和语言扩展层构成,提供多层级API,满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考