ops-nn Tiling策略详解:张量切分与并行调度的核心机制和3大交付件
ops-nn Tiling策略详解:张量切分与并行调度的核心机制和3大交付件
【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库,实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn
ops-nn 是 CANN 提供的神经网络类计算算子库,用于实现网络在 NPU 上的加速计算。而在每个算子跑上 NPU 之前,都绕不开一个关键环节——Tiling(张量切分):把大张量切成小块、分配给成百上千个 AI Core 并行计算。本文将带你快速搞懂 Tiling 的核心机制,并盘点它的3 大交付件。
为什么NPU算子必须要做Tiling?
NPU 由多个AI Core组成,每个核内部都有片上存储空间(UB,Unified Buffer),但它比显存小得多,无法一次性把整个张量加载进来计算。因此必须把输入张量切分为多个小块(Tile),逐块搬运、逐块计算。
用一句话概括 Tiling 的职责:
Tiling 策略 = 决定「怎么切」+「谁来切」+「每块多大」,并把这些决策告诉 Kernel。
Tiling 运行在 Host 侧(CPU),Kernel 运行在 Device 侧(NPU)。两者之间不直接通信,而是通过TilingData结构体传递信息——Tiling 算好参数写进去,Kernel 读出来执行。
一条流水线:从Tiling到Kernel并行执行
整个协作流程可以拆成 4 步:
- 查硬件家底:Tiling 入口函数(
TilingFunc)先获取平台信息——可用 AI Core 核数(coreNum)和单核 UB 大小(ubSize)。 - 核间切分(Block Tiling):把总数据量平均分配给各核,「优先用更多的核并行计算」。
- 核内切分(UB Tiling):根据 UB 容量和双缓冲等流水需求,算出每核每次可搬运的数据块大小。
- 下发决策:把切分参数写入
TilingData,通过SetBlockDim设置实际启用的核数,并通过TilingKey标识本次选择的具体算法路径。
Kernel 侧收到后,用GetBlockIdx()找到自己的数据段,循环执行「搬入 → 计算 → 搬出」流水线,即可并行跑满全卡。
3大交付件清单:一次看懂Tiling要写什么
一个标准算子的 Tiling 实现,一共需要3 个交付件:
| # | 交付件 | 所在目录 | 职责 |
|---|---|---|---|
| 1 | ${op_name}_tiling.cpp | op_host/ | Tiling 主切分逻辑,核心计算都在这里 |
| 2 | ${op_name}_tiling_key.h | op_kernel/ | 定义 TilingKey,标识不同的切分/算法分支 |
| 3 | ${op_name}_tiling_data.h | op_kernel/ | 定义 TilingData 结构体,存储块大小、并行度等配置 |
三个文件各有分工,缺一不可:
tiling.cpp是「大脑」:完成平台信息获取、shape 校验、核切分与 UB 切分计算,最终注册 Tiling 函数。tiling_key.h是「路标」:同一个算子可能有 float 路径、int 路径、不同架构路径,Kernel 侧依据 TilingKey 走不同的if constexpr分支,选择对应的实现。tiling_data.h是「契约」:Tiling 侧写入、Kernel 侧读取,是两侧唯一的参数传递通道,例如:
// ${op_name}_tiling_data.h struct ${op_name}TilingData { int64_t totalLength; // 总数据量大小 int64_t tileNum; // 每个核的数据切块数量 };完整目录说明可参考 项目目录结构文档,其中详细标注了${op_name}_tiling.cpp、${op_name}_tiling_key.h、${op_name}_tiling_data.h各自的位置与用途。
实战拆解:以AddExample为例的Tiling流程
官方在 examples/add_example 目录给出了完整的加法算子样例,是学习 Tiling 的最佳起点。它的核心 Tiling 函数 add_example_tiling.cpp 流程如下:
GetPlatformInfo:获取ubSize(UB 大小)与coreNum(AIV 核数);GetShapeAttrsInfo:校验输入输出 shape 均为 4 维,提取总元素数totalIdx与数据类型;- 核切分:
blockFactor = CeilDiv(totalIdx, coreNum),算出每个核处理的元素数,并用SetBlockDim(usedCoreNum)启用对应核数; - UB 切分:考虑 2 输入 + 1 输出、双缓冲共 6 块 UB 缓冲区,计算
ubFactor; - 设置 TilingKey:
FLOAT类型走SCH_MODE_0,INT32类型走SCH_MODE_1。
最终生成的 AddExampleTilingData 只有三个字段,麻雀虽小五脏俱全:
struct AddExampleTilingData { int64_t totalNum = 0; // 总元素数 int64_t blockFactor = 0; // 每核元素数(核间切分) int64_t ubFactor = 0; // 单核单次搬运块大小(核内切分) };对应的分支路由定义在 add_example_tiling_key.h 中,Kernel 入口 add_example.cpp 通过if constexpr (schMode == ...)选择 float/int32 不同实现,完成调度。
新手上手建议:如何验证你的Tiling写对了
- 先读样例再动手:AI Core算子开发指南 中的 Tiling 章节给出了完整的代码骨架,包括
TilingParse(图模式交付件,无逻辑时可置空)与TilingFunc(主入口)的标准写法。 - Tiling 与 Kernel 严格对齐:
tiling_data.h中的每个字段都要在 Kernel 的Init中被消费,字段不一致是新手最常见的踩坑点。 - 补齐 Tiling UT:在
tests/ut/op_host/test_${op_name}_tiling.cpp中构造输入 shape、属性与 compileInfo,断言输出的 TilingKey、TilingData 与 Workspace 是否符合预期,无需 NPU 环境即可快速验证。
掌握 Tiling 策略,就掌握了 NPU 算子性能的第一把钥匙:切得巧,核核都不闲着;切得糙,再强的硬件也白搭。
【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库,实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
