ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

CANNBot-DSL 源码精读①:非量化 Matmul 的 host 侧 tiling 推导与自适应多核滑动窗口调度深度拆解

CANNBot-DSL 源码精读①:非量化 Matmul 的 host 侧 tiling 推导与自适应多核滑动窗口调度深度拆解 CANNBot-DSL 源码精读①非量化 Matmul 的 host 侧 tiling 推导与自适应多核滑动窗口调度深度拆解【免费下载链接】cannbot-dsl基于 CANNBot-DSL 的 Ascend NPU 复杂算子示例集合。项目地址: https://gitcode.com/cann/cannbot-dslCANNBot-DSL 是面向昇腾AscendNPU 的算子开发 DSL本仓库开源了基于它生成的复杂算子示例集合。本篇源码精读以非量化 Matmul矩阵乘算子为样本拆解两大核心机制host 侧 tiling 推导流水线与device 端自适应滑动窗口多核调度帮你看懂 Ascend NPU 算子开发的关键路径。项目速览CANNBot-DSL 是什么Matmul 算子做了什么CANNBot-DSL 示例仓库面向 NPU ARCH 3510Ascend 950PR / Ascend 950DT开源了 VoxelConv、Flash Attention、Flash KDA、Matmul、PointNet SA、RMS Norm 六个复杂算子样例代码由 CANNBot 基于 CANNBot-DSL 生成源码统一放在 samples/ 目录下。本期主角 Matmul 的计算目标非常经典C[M,N] A[M,K] B[N,K]^Tfp16/bf16 输入fp32 累加。关键约束一览来自 samples/matmul/README.md维度说明数据类型float16 / bfloat16A、B、C 一致转置模式仅支持 transpose_aFalse、transpose_bTrue并行策略自适应滑动窗口多核调度L2 缓存按矩阵复用情况自适应开关整个算子由一份 600 多行的 Python 文件完成结构清晰是理解NPU 算子开发流程的理想入口。代码地图Matmul 算子源码结构与硬件参数速查全部源码集中在 samples/matmul/matmul.py分为三个层次模块位置职责MatmulTilingmatmul.py#L68host 侧 tiling决定 tile 尺寸、核数、缓存策略MatmulKernelmatmul.py#L545device 核滑动窗口调度 L1/L0 流水线matmul()matmul.py#L641torch 接口参数校验、内核发射硬件数据通路为GM → L1MTE2 搬运→ L0A/L0BMTE1 搬运→ MMAD 矩阵乘 → L0C → GMFIXPIPE 写回。tiling 决策依赖的硬件常量如下资源容量用途L1512 KBA/B 分块的中间缓存L0A / L0B各 64 KBMMAD 输入缓存L0C256 KBfp32 结果累加器L2128 MB片上共享缓存AIC 核32 个矩阵乘计算核Host 侧 tiling 推导四步流水线baseM/baseN/baseK 是怎么定出来的MatmulTiling在构造时通过_compute()一次性跑完整个推导matmul.py#L177顺序执行四个阶段。第一步 ResetBase给 baseM/baseN/baseK 一个保守起点_reset_base()matmul.py#L189baseM / baseN min(256M/N 向上对齐到 16 的倍数)baseK 128B ÷ 元素字节数 64fp16/bf16 场景第二步 RebalanceBlockcube-bound 判型 穷举搜索最优 baseM/baseN这是 tiling 中信息量最大的部分matmul.py#L195分三步走。1cube-bound 判型先用核频、DDR/L2 速率估算整芯片 HBM 带宽与 L2 带宽再结合 (MN)/MN 的数据复用率算出边界值cubeBoundEdgematmul.py#L216。若1/baseM 1/baseN超过该边界判定为访存受限对齐单位放宽到 128 以降低搬运次数否则视为计算受限使用 64。2穷举搜索在 L0C 容量约束baseM×baseN×4B 双缓冲放得下内遍历 baseM×baseN 的全部候选组合matmul.py#L281用两个指标打分cube-bound 参数1/baseM 1/baseN刻画数据复用强度越小越好负载均衡率matmul.py#L502把 32 个核分到尾部 tile上的不均衡也纳入考量越高越好。只有双指标综合占优的组合同步更新最优解——这就是自适应的含义不同 M×N×K 形状会自动得到不同的 baseM/baseN。3baseK 受 L0A 约束最大 baseK L0A 容量 ÷ 2(双缓冲) ÷ dtype 字节 ÷ max(baseM, baseN)matmul.py#L332K 能整段放下就直接用放不下再按 128/64/32/16 顺序回退。第三步 CalL1TilingL1 预算约束下的 K 方向步长_cal_l1_tiling()matmul.py#L349从 1 到 8 逐个尝试 K 步数step_k只有当 A/B 分块baseM×baseK×step_k在双缓冲 4 级缓冲下放得进 512KB L1、且单侧 MTE 搬运量不超 48KB 时才算合法取最大合法步数作为stepKa/stepKb。K 步越大GM 往返越少L1 复用率越高。第四步 Finalize核数、缓冲深度与 L2 缓存模式一锤定音_finalize()matmul.py#L435输出 device 内核真正消费的全部参数参数推导规则含义mL1 / nL1 / kL1形状对齐值与 baseM/N 取小每 block 的 L1 tile 尺寸usedCoreNummin(⌈M/baseM⌉×⌈N/baseN⌉, 32)实际发射的 AIC 核数l1BufferNum4 级缓冲放得下 L1 取 4否则 2L1 ping-pong 深度l0cDB双份 baseM×baseN×fp32 放得下 L0C 取 2否则 1L0C 双缓冲l2CacheDisable见下方自适应开关L2 缓存模式L2 缓存自适应开关matmul.py#L389是个值得学习的设计A/B/C 总量 128MB L2 时保持默认否则对跨 tile 不复用的矩阵例如 baseN 已覆盖整个 N 维时A 被每个核只读一次关闭 L2 缓存把容量让给会被反复复用的那个矩阵。四种状态对应常量L2_CACHE_DEFAULT / A_L2_CACHE_DISABLE / B_L2_CACHE_DISABLE / ALL_L2_CACHE_DISABLE。Device 端自适应滑动窗口多核调度如何分配 tile调度逻辑集中在内核代码的 40 行里信息密度很高。轮转条纹分配32 核负载均衡不靠额外算法核数usedCoreNum由 host 按形状动态决定小矩阵可能远小于 32。内核用一行代码完成任务分发matmul.py#L598for tile_idx in range(block_idx, total_tiles, block_num):第 i 个核从total_tiles个输出 tile 中取第 i 个之后每隔block_num取一个。相比连续区间划分条纹式轮转让每个核的天然负载差仅为 0/1 个 tile负载均衡免费获得。滑动窗口M 维按 4 行一组推进若按简单行列序遍历 tile换到下一行 M 时刚加载的 B 数据就会被逐出缓存。内核把 M 维按WINDOW_LEN 4行组成滑动窗口matmul.py#L543每个窗口内先算完 4 行 M × 全部 N 列——4 个 M tile 共享同一批 B tileB 数据在整个窗口期间保持 L2 热窗口滑向下一组 4 行时才换入新 A。尾部不足 4 行按tail_window单独处理matmul.py#L563窗口长度还会自适应收缩为min(4, m_tiles)M 很小时不浪费。蛇形遍历奇数窗口反向扫描 N 列更巧的细节在 matmul.py#L611奇数窗口把 N 列遍历方向反转。这样上一个窗口的终点与下一个窗口的起点在 N 维上相邻前窗尾部的 B tile 还在 L2 里时后窗就接上了跨窗口的缓存复用首尾相连。L1/L0 ping-pong 流水线拷贝与计算重叠拿到 tile 后每个核跑两层循环matmul.py#L613外层kL1 粒度MTE2 把 A/B 的 kL1 段从 GM 搬进 L1nd2nz 拷贝引擎按 host 决策写入逐矩阵 L2 控制位内层baseK 粒度MTE1 从 L1 切 baseK 段到 L0A/L0B执行 MMADfp32 累加进 L0C首个 k 用initTrue清累加器K 方向累加完成后L0C 经 FIXPIPE 写回 GM。L12 或 4 级与 L02 级双缓冲让搬运和计算完全重叠这正是ping-pong 流水线。性能实测CANNBot-DSL matmul 与 CANN 内置模板对比FP16/BF16项目在 10 组 shape 上对比了 CANNBot-DSL 实现与 CANN 包内置 matmul 基础模板精度用例见 test/matmul/test_matmul.pyFP16 下几个代表性数据用例M×K×NCANNBot-DSLμsCANN 内置μs4096×3840×38432.833.25120×2049×2048109.0110.67168×2048×1535137.0116.36144×8193×640226.6188.6可以看到中小 shape 上两者基本打平、互有胜负长条形 shape如 6144×8193×640内置模板的缓存策略更激进。作为示例集这份实现的价值更在于公开了一条完整可读的host tiling device 调度参考实现每个决策都可复现、可学习是研究NPU 算子性能调优的好教材。新手上手一键运行 Matmul 算子精度测试1️⃣ 克隆仓库git clone https://gitcode.com/cann/cannbot-dsl2️⃣ 运行精度测试pytest 驱动覆盖 10 组 shapepytest test/matmul/test_matmul.py -v3️⃣ 在 torch 中调用节选自 samples/matmul/README.mdfrom matmul import matmul a torch.randn(1024, 1024, dtypetorch.float16).npu() b torch.randn(1024, 1024, dtypetorch.float16).npu() c matmul(a, b, transpose_aFalse, transpose_bTrue)环境依赖可通过 install_deps.sh 一键安装。小结Matmul 算子教会我们的三件事tiling 是折中艺术baseM/baseN 由cube-bound 参数 × 负载均衡率共同寻优L1/L0A/L0C/L2 各级容量都参与约束调度是局部性之战滑动窗口 蛇形遍历目的都只有一个——让数据在 L2 里保持热自适应来自 host 侧核数、缓冲深度、L2 开关全部按形状在发射前推导完成device 内核代码本身零分支。仓库中还沉淀了 Flash Attention、Flash KDA 等更复杂的融合算子samples/后续源码精读将继续拆解它们的多维 tiling 与调度策略值得关注。【免费下载链接】cannbot-dsl基于 CANNBot-DSL 的 Ascend NPU 复杂算子示例集合。项目地址: https://gitcode.com/cann/cannbot-dsl创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表