3.2 矢量编程实战:从零实现一个 Add 算子
通过完整的 Add 算子实现走通 Ascend C 开发闭环——Host 侧 Tiling、DataCopy 数据搬运、矢量计算接口、写回 Global Memory,掌握昇腾算子开发的最小完整范例
理论(3.1)讲完,直接动手。这一章用昇腾算子开发的「Hello World」——Add 算子(向量加法)——带你走通 Ascend C 的完整开发闭环。写完这个算子,你就掌握了昇腾算子开发 80% 的骨架。
📑 目录
- 1. Add 算子的完整构成
- 2. Host 侧:Tiling 参数计算
- 3. Device 侧:核函数主体
- 4. DataCopy 与矢量计算接口
- 5. 完整可参考代码
- 6. 从 Add 到任意 Elementwise 算子的扩展
- 7. 常见坑与调试思路
- 总结
- 延伸阅读
1. Add 算子的完整构成
一个 Ascend C 算子通常由两部分组成:
┌─ Host 侧(在 CPU 上运行)─────────────────────┐
│ · 计算 Tiling 参数(每块搬多少、循环几轮) │
│ · 分配 GM 内存 │
│ · 下发核函数(ACLRT_LAUNCH_KERNEL) │
│ · 同步等待 │
└──────────────────────────────────────────────┘
┌─ Device 侧(核函数,在 AI Core 上运行)────────┐
│ · 根据 Tiling 参数切分数据 │
│ · DataCopy 把块搬进 UB │
│ · 矢量接口 Add 计算 │
│ · DataCopy 把结果搬回 GM │
└──────────────────────────────────────────────┘
目标:dst[i] = src1[i] + src2[i],长度 N(例如 8192 个元素)。
2. Host 侧:Tiling 参数计算
Tiling 的核心是回答三个问题:
- UB 能装多少:UB 容量已知(例如 192KB),算出每个 Block 最多能处理多少元素(
blockLen)。 - 要切多少块:
numBlocks = ceil(N / blockLen)。 - 多核怎么分:把
numBlocks均分给多个 AI Core(每个核处理一部分块)。
简化版伪代码:
// Host 侧 Tiling
const int32_t totalLen = 8192;
// 假定单核单轮最多处理 1024 个元素(实际取决于 UB 容量与数据类型)
const int32_t blockLen = 1024;
const int32_t numBlocks = (totalLen + blockLen - 1) / blockLen;
// 假设 4 个核并行,每个核负责 numBlocks / 4 轮
// 把 numBlocks、blockLen 等参数打包成 TilingData 传给核函数
关键思想:Tiling 参数在 CPU 上算好,通过参数结构体传给核函数——核函数只负责「按参数执行」,不做复杂决策。这种「Host 决策、Device 执行」的分工是昇腾算子的通用范式。
3. Device 侧:核函数主体
核函数的主体结构如下:
// gm 指针 = Global Memory(HBM)
__global__ void addKernel(const __gm__ float *src1,
const __gm__ float *src2,
__gm__ float *dst,
const int32_t totalLen,
const int32_t blockLen) {
// 1. 确定本核要处理的起始块(多核按核号切分)
int32_t startIdx = /* 本核起始位置 */;
// 2. 循环处理本核负责的每个 Block
for (int32_t i = 0; i < /* 本核负责的块数 */; i++) {
int32_t offset = startIdx + i * blockLen; // 当前块的起始元素下标
int32_t curLen = min(blockLen, totalLen - offset); // 边界处理
// 2.1 在 UB 上分配缓冲区
// 2.2 DataCopy:src1、src2 的当前块 → UB
// 2.3 矢量计算:UB 上执行 Add
// 2.4 DataCopy:结果 UB → dst
}
}
注意 min(blockLen, totalLen - offset) 这句——边界处理是 Tiling 最容易写错的地方:最后一块可能不满,取小了会漏算,取大了会越界。
4. DataCopy 与矢量计算接口
DataCopy(数据搬运)
Ascend C 通过 DataCopy 接口在 GM 与 UB 之间搬数据:
DataCopy(dstUb, src1Gm + offset, curLen); // GM → UB
DataCopy(dstGm + offset, resUb, curLen); // UB → GM
特点:
- 一次
DataCopy搬一块连续数据,比逐元素访问高效得多。 - 追求对齐与连续访问,零散访问会损失带宽。
矢量计算接口
Ascend C 提供一套矢量计算 API,操作对象是 UB 上的缓冲区:
Add(aUb, bUb, cUb, curLen); // cUb = aUb + bUb(逐元素)
还有其他常用接口:Sub、Mul、Exp、Relu、SoftMax、ReduceSum 等,命名直观,都遵循「输入缓冲区 + 输出缓冲区 + 长度」的签名风格。
5. 完整可参考代码
下面是核函数主体(示意,省略工程细节,完整工程见官方示例):
#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BLOCK_LEN = 1024;
class AddKernel {
public:
__aicore__ inline AddKernel(const __gm__ float *src1,
const __gm__ float *src2,
__gm__ float *dst,
int32_t totalLen)
: src1Global(src1), src2Global(src2), dstGlobal(dst),
totalLen(totalLen) {}
__aicore__ inline void Process() {
int32_t perBlock = BLOCK_LEN;
// 多核:本核负责从 blockId 开始、间隔核数的块
int32_t numBlocks = (totalLen + perBlock - 1) / perBlock;
int32_t coreNum = GetBlockNum();
int32_t coreIdx = GetBlockIdx();
for (int32_t block = coreIdx; block < numBlocks; block += coreNum) {
int32_t offset = block * perBlock;
int32_t curLen = (offset + perBlock <= totalLen)
? perBlock
: (totalLen - offset);
// 分配 UB 缓冲区
LocalTensor<float> aUb, bUb, cUb;
// 申请 UB 空间(省略 alloc 细节)
// ...
// 搬运入 UB
DataCopy(aUb, src1Global[offset], curLen);
DataCopy(bUb, src2Global[offset], curLen);
// 矢量计算
Add(cUb, aUb, bUb, curLen);
// 搬回 GM
DataCopy(dstGlobal[offset], cUb, curLen);
}
}
private:
const __gm__ float *src1Global;
const __gm__ float *src2Global;
__gm__ float *dstGlobal;
int32_t totalLen;
};
extern "C" __global__ __aicore__ void add_kernel(
const __gm__ float *src1, const __gm__ float *src2,
__gm__ float *dst, int32_t totalLen) {
AddKernel op(src1, src2, dst, totalLen);
op.Process();
}
说明:以上为教学示意代码,省略了 UB 内存申请(
AllocTensor)、同步(FreeTensor)与 Host 侧下发细节。可直接对照昇腾官方Add算子示例工程运行验证。
6. 从 Add 到任意 Elementwise 算子的扩展
Add 是 Elementwise 算子的代表。掌握了它,就能套用到一整个算子族:
| 算子 | 变化点 |
|---|---|
Sub / Mul / Div | 换一个矢量接口 |
Exp / Log | 换接口 + 注意数值范围 |
ReLU | 单输入算子,简化搬运 |
Scale(x * s + b) | 矢量接口 + 标量参数 |
LayerNorm 的逐元素部分 | 矢量计算 + 块内规约 |
套路不变:Tiling 切块 → 搬运入 UB → 矢量计算 → 搬回 GM。变的是接口与边界条件。
7. 常见坑与调试思路
- 越界:最后一块长度算错。用
min(blockLen, totalLen - offset)兜底。 - UB 溢出:申请缓冲区超过 UB 容量,编译或运行时报错。检查 Tiling 块大小。
- 搬运未完成就计算:DataCopy 是异步的,需要确保数据就位再算。使用同步机制(
EnQue/DeQue或SyncAll)——这就是 3.3 的内容。 - 结果错位:多核切分时 offset 算错,导致重复或遗漏。先在 Host 侧打印 Tiling 参数核对。
- 调优起点:先用
msprof看「搬运占比 vs 计算占比」,再用双缓冲优化(3.3/3.5)。
总结
- Ascend C 算子开发的最小闭环:Host Tiling → DataCopy 入 UB → 矢量计算 → DataCopy 回 GM。
- Tiling 三要素:块大小、循环组织、边界处理;Host 决策、Device 执行。
- 矢量接口命名直观,Elementwise 算子族可以「一招鲜」迁移。
- 注意异步搬运的同步问题——这是通往高性能算子的第一道坎。