Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.
cann训练营二期算子开发任务。
参考版本内置SoftplusV2Grad算子的TBE实现,在昇腾NPU上使用Ascend C编程语言实现相同功能的算子。
亚信科技算子团队
tiling 策略的核心是将输入数据划分为适合 AI Core 硬件特性(如 UB 容量、计算单元吞吐量)的小块(tile),确保计算过程中数据高效复用,减少全局内存访问。本代码的 tiling 策略设计如下: 基础参数定义
算子属性集成
将 SoftplusBackward 的关键属性(beta和threshold)存入 tiling 数据中,供 device 侧 kernel 使用。这些参数直接影响计算逻辑(如softplus(x) = (1/beta)log(1 + exp(betax)),反向计算需依赖这些参数),通过 tiling 传递避免 device 侧额外读取全局内存。
分核策略的目标是将总任务均匀分配到多个 AI Core 上,最大化并行计算效率 核心数量限制 首先根据输入数据量和单个 tile 的大小判断是否需要多核心:
若单个 tile 即可覆盖所有输入数据(tileDataNum >= inputNum),则仅使用 1 个核心(避免多核心调度开销);
否则,核心数量不超过输入数据按BLOCK_SIZE对齐后的总块数(避免核心闲置): 任务均衡分配 输入数据按BLOCK_SIZE对齐后(inputLengthAlgin32),总块数为inputLengthAlgin32 / BLOCK_SIZE。将这些块分配到coreNum个核心上,可能存在 “尾块”(总块数不能被核心数整除),因此分为两种核心:
小核:处理everyCoreInputBlockNum个块(everyCoreInputBlockNum = 总块数 / coreNum);
大核:处理everyCoreInputBlockNum + 1个块(用于处理尾块,数量为tailBlockNum = 总块数 % coreNum)。
分核参数记录
为每个核心记录其处理的数据量(smallCoreDataNum/bigCoreDataNum)、包含的 tile 数量(finalSmallTileNum/finalBigTileNum)、最后一个 tile 的大小(smallTailDataNum/bigTailDataNum)等,确保 device 侧每个核心能准确执行分配的任务。
数据分块和内存优化的核心是减少全局内存访问次数,提高 UB 利用率,同时满足硬件的内存访问对齐要求。
数据分块逻辑
输入数据对齐:输入数据量按BLOCK_SIZE=32对齐(inputLengthAlgin32),确保分块无碎片,避免边界处理的额外开销;
核心内部分块:每个核心处理的数据进一步划分为多个 tile(如小核的finalSmallTileNum个 tile),每个 tile 大小为tileDataNum(由 UB 容量决定),最后一个 tile 可能为 “尾 tile”(smallTailDataNum/bigTailDataNum),确保数据无遗漏。
内存访问优化
UB 复用:通过tileDataNum限制单个 tile 大小不超过 UB 容量,确保计算时数据可从 UB 直接读取,减少全局内存访问;
双缓冲机制:通过BUFFER_NUM=2在 UB 中分配两个缓冲区,一个用于当前计算,另一个用于预加载下一个 tile 的数据,隐藏数据搬运耗时;
内存对齐:计算每个 tile 的实际处理数据量时,按 256 字节对齐(硬件要求的内存访问粒度):
smallprocessDataNum_computes = (((tileDataNum * sizeof(float) + 256 - 1) / 256) * 256) / sizeof(float);
确保内存访问地址对齐,避免硬件访问效率损失。
Ascend C的SoftplusBackward算子流程见下图: 以下按 “初始化→主流程→数据搬运→核心计算→结果输出” 的逻辑,详细拆解计算步骤:
Init是 kernel 启动后的首个步骤,主要作用是读取 host 侧传递的 tiling 参数、划分核心任务角色、初始化全局内存(GM)与统一缓冲区(UB),为后续计算做准备,具体步骤如下:
1. 核心角色与任务参数初始化
核心类型判断:根据当前核心索引(coreIdx = GetBlockIdx())与 host 侧计算的 “尾块数量”(tailBlockNum),判断当前核心是 “大核”(处理尾块,任务量稍多)还是 “小核”(处理普通块):
若coreIdx < tailBlockNum:为大核,加载大核参数(bigCoreDataNum、finalBigTileNum等);
否则为小核,加载小核参数(smallCoreDataNum、finalSmallTileNum等)。
全局偏移量计算:根据核心类型计算当前核心处理数据在 GM 中的起始地址(globalOffset),确保多核心不重复处理数据:
// 大核偏移:核心索引 × 大核单核心数据量 // 小核偏移:尾块总数据量 + (核心索引-尾块数) × 小核单核心数据量 globalOffset = (coreIdx < tailBlockNum) ? (coreIdx * bigCoreDataNum) : (tailBlockNum * bigCoreDataNum + (coreIdx - tailBlockNum) * smallCoreDataNum);
2. GM 缓冲区绑定(全局内存映射)
根据数据类型(dataType:0=FP16、1=FP32、2=BF16),将 GM 中的输入(gradOutput、self)和输出(gradInput)地址映射为对应类型的GlobalTensor(AscendC 中用于访问 GM 的封装类),确保数据读写类型匹配:
例如 FP16 类型:reinterpret_cast<gm half*>(gradOutput)后绑定到gradOutputGmHalf;
BF16 类型:通过bfloat16(本质uint16_t)映射,后续需通过自定义函数转换为 float 计算。
3. U 缓冲区初始化(双缓冲与计算缓冲区)
UB 是 AI Core 的高速缓冲区,Init需为 “数据搬运” 和 “核心计算” 分配 UB 空间,核心设计包括:
双缓冲(Double Buffer):为gradOutput、self、gradInput分别初始化TQue(AscendC 队列类),缓冲区数量BUFFER_NUM=2,实现 “数据搬运与计算并行”(一个缓冲区用于计算,另一个预加载下一批数据);
计算缓冲区:初始化 16 个TBuf(UB 计算缓冲区),用于存储中间计算结果(如betaSelfBuf存储beta*self、sigmoidBuf存储 sigmoid 结果),缓冲区大小由processDataNum_computes(按 256 字节对齐的计算量)决定,确保适配硬件访问粒度。
Process是 kernel 的主调度逻辑,按 host 侧划分的tileNum(总 tile 数)循环处理每个 tile,核心是 “分批次处理数据,避免 UB 溢出”,步骤如下:
循环次数确定:循环tileNum次,前tileNum-1次处理 “完整 tile”(大小为tileDataNum),最后 1 次处理 “尾 tile”(大小为tailDataNum,因总数据量可能不被tileDataNum整除);
每轮调度逻辑:每次循环依次调用CopyIn(GM→UB 搬运数据)、Compute(UB 内核心计算)、CopyOut(UB→GM 输出结果),形成 “搬运 - 计算 - 输出” 的流水线。
CopyIn的作用是将当前 tile 的 GM 数据(gradOutput、self)搬运到 UB 的双缓冲队列(gradOutputQueue、selfQueue),按数据类型差异化处理:
1. 双缓冲分配与数据拷贝
首先调用AllocTensor从队列中分配一个LocalTensor(AscendC 中访问 UB 的封装类);
通过DataCopy(AscendC 数据拷贝 API)将GlobalTensor中当前 tile 的数据(起始地址为progress * tileDataNum,progress为当前循环轮次)拷贝到LocalTensor;
调用EnQue将加载好数据的LocalTensor放入队列,等待Compute读取。
2. 数据类型适配
FP16/FP32:直接通过DataCopy拷贝,类型天然匹配;
BF16:因 BF16 无原生计算支持,需先以bfloat16类型搬运到 UB,后续在Compute中转换为 float 计算。
Compute是 kernel 的核心,实现 SoftplusBackward 的数学计算逻辑(Softplus 正向:softplus(x) = (1/beta)log(1+exp(betax)),反向:gradInput = gradOutput * sigmoid(beta*x),需处理阈值分支),按数学步骤拆解为 7 个关键环节,同时适配多数据类型并优化精度:
环节 1:输入数据类型统一(转为 float 计算)
因 BF16/FP16 无原生复杂计算支持,需先将 UB 中的输入数据转为 float(AI Core 对 float 计算支持更完善):
FP16→float:调用Cast(AscendC 类型转换 API),RoundMode::CAST_NONE(无舍入);
FP32:直接DataCopy,无需转换;
BF16→float:通过自定义Bf16ToFloat函数转换(将 BF16 的 16 位数据左移 16 位补零,再强转为 float):
uint32_t tmp = (static_cast<uint32_t>(bf_val) << 16); // BF16→32位整数补零 return *reinterpret_cast<float*>(&tmp); // 强转为float
环节 2:计算beta * self(Softplus 反向的基础项)
通过Muls(AscendC 标量乘向量 API)计算beta * self,结果存入betaSelfBuf:
公式:betaSelf[i] = selfFloat[i] * betaVal(betaVal为 host 侧传递的算子属性);
作用:betaSelf是后续计算 sigmoid 和阈值分支的核心输入。
环节 3:生成线性分支掩码(阈值保护逻辑)
Softplus 反向在betax > threshold时,为避免exp(betax)溢出,采用线性近似(gradInput ≈ gradOutput),需先生成掩码区分 “线性分支” 和 “普通分支”:
初始化thresholdTensor(全为thresholdVal,host 侧传递的阈值,默认 20.0)和constantOneTensor(全为 1.0);
调用Compare(AscendC 比较 API),按betaSelf[i] > thresholdVal生成linearMask(uint8_t类型,1 表示线性分支,0 表示普通分支):
Compare(linearMask, betaSelf, thresholdTensor, CMPMODE::GT, ...);
环节 4:高精度 Sigmoid 计算(普通分支核心)
普通分支需计算gradFactor = sigmoid(betaSelf) = 1/(1+exp(-betaSelf)),为避免精度损失,代码做了两处关键优化:
Min(safeNegBetaSelf, negBetaSelf, maxThresholdTensor, ...); // negBetaSelf = -betaSelf
先计算expNegBetaSelf = exp(safeNegBetaSelf)(Exp API);
再计算onePlusExp = expNegBetaSelf + 1.0(Adds API);
最后用Div(numerator, denominator)计算:numerator全为 1.0,denominator = onePlusExp,结果存入sigmoidBuf,精度比Reciprocal提升约 1e-5。
环节 5:梯度因子选择(线性 / 普通分支融合)
通过Select(AscendC 分支选择 API)根据linearMask融合两个分支的梯度因子gradFactor:
当linearMask[i] = 1(线性分支):gradFactor[i] = 1.0(对应constantOneTensor);
当linearMask[i] = 0(普通分支):gradFactor[i] = sigmoidBuf[i];
公式:gradFactor = linearMask ? 1.0 : sigmoid(betaSelf),结果存入selectResultBuf。
环节 6:计算最终梯度(gradInput = gradOutput * gradFactor)
通过Mul(AscendC 向量乘向量 API)计算输出梯度:
公式:gradInputFloat[i] = gradOutputFloat[i] * gradFactor[i](gradOutputFloat为输入梯度,gradFactor为分支选择后的因子);
结果存入gradInputFloatBuf,为后续输出做准备。
环节 7:计算结果类型转回(float→目标类型)
将 float 类型的gradInputFloat转回目标数
/assign
感谢您的反馈,可以关联上您的PR,在验证后由commiter审视合入。
经sig组织评审,您的方案已通过,请关联相关贡献的代码
Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.
Backgroud(背景信息)
cann训练营二期算子开发任务。
参考版本内置SoftplusV2Grad算子的TBE实现,在昇腾NPU上使用Ascend C编程语言实现相同功能的算子。
Origin(信息来源)
亚信科技算子团队
Benefit / Necessity (价值/作用)
Design(设计方案)
3.2 需求整体设计
3.2.1 host侧设计
3.2.1.1 tiling 策略
tiling 策略的核心是将输入数据划分为适合 AI Core 硬件特性(如 UB 容量、计算单元吞吐量)的小块(tile),确保计算过程中数据高效复用,减少全局内存访问。本代码的 tiling 策略设计如下:
基础参数定义
UB 容量适配
代码通过ascendcPlatform.GetCoreMemSize获取 UB 总大小,计算 UB 可容纳的元素数:uint64_t ubElementCapacity = ubSize / sizeof(float); // 以float为单位计算UB可容纳的元素数
再结合基础参数计算单个 tile 的大小(tileDataNum):
uint32_t tileBlockNum = (ubElementCapacity / BLOCK_SIZE / BUFFER_NUM) / UB_DATA_NUMBER; // 每个tile包含的BLOCK数量
uint32_t tileDataNum = tileBlockNum * BLOCK_SIZE; // 单个tile的元素数(最终tile大小)
该设计确保单个 tile 能完全放入 UB,且通过双缓冲(BUFFER_NUM=2)实现计算与数据搬运并行。
数据类型适配
针对不同数据类型(FP16、FP32、BF16)的内存占用差异,通过dataTypeFlag标识数据类型,并限制 BF16 仅在支持的芯片(如 Ascend910B、310B)上使用。同时,通过typeLength记录数据类型的字节长度(如 FP16 为 2 字节,FP32 为 4 字节),确保分块时内存计算准确。
算子属性集成
将 SoftplusBackward 的关键属性(beta和threshold)存入 tiling 数据中,供 device 侧 kernel 使用。这些参数直接影响计算逻辑(如softplus(x) = (1/beta)log(1 + exp(betax)),反向计算需依赖这些参数),通过 tiling 传递避免 device 侧额外读取全局内存。
3.2.1.2 分核策略
分核策略的目标是将总任务均匀分配到多个 AI Core 上,最大化并行计算效率
核心数量限制
首先根据输入数据量和单个 tile 的大小判断是否需要多核心:
若单个 tile 即可覆盖所有输入数据(tileDataNum >= inputNum),则仅使用 1 个核心(避免多核心调度开销);
否则,核心数量不超过输入数据按BLOCK_SIZE对齐后的总块数(避免核心闲置):
任务均衡分配
输入数据按BLOCK_SIZE对齐后(inputLengthAlgin32),总块数为inputLengthAlgin32 / BLOCK_SIZE。将这些块分配到coreNum个核心上,可能存在 “尾块”(总块数不能被核心数整除),因此分为两种核心:
小核:处理everyCoreInputBlockNum个块(everyCoreInputBlockNum = 总块数 / coreNum);
大核:处理everyCoreInputBlockNum + 1个块(用于处理尾块,数量为tailBlockNum = 总块数 % coreNum)。
分核参数记录
为每个核心记录其处理的数据量(smallCoreDataNum/bigCoreDataNum)、包含的 tile 数量(finalSmallTileNum/finalBigTileNum)、最后一个 tile 的大小(smallTailDataNum/bigTailDataNum)等,确保 device 侧每个核心能准确执行分配的任务。
3.2.1.3 数据分块和内存优化策略
数据分块和内存优化的核心是减少全局内存访问次数,提高 UB 利用率,同时满足硬件的内存访问对齐要求。
数据分块逻辑
输入数据对齐:输入数据量按BLOCK_SIZE=32对齐(inputLengthAlgin32),确保分块无碎片,避免边界处理的额外开销;
核心内部分块:每个核心处理的数据进一步划分为多个 tile(如小核的finalSmallTileNum个 tile),每个 tile 大小为tileDataNum(由 UB 容量决定),最后一个 tile 可能为 “尾 tile”(smallTailDataNum/bigTailDataNum),确保数据无遗漏。
内存访问优化
UB 复用:通过tileDataNum限制单个 tile 大小不超过 UB 容量,确保计算时数据可从 UB 直接读取,减少全局内存访问;
双缓冲机制:通过BUFFER_NUM=2在 UB 中分配两个缓冲区,一个用于当前计算,另一个用于预加载下一个 tile 的数据,隐藏数据搬运耗时;
内存对齐:计算每个 tile 的实际处理数据量时,按 256 字节对齐(硬件要求的内存访问粒度):
smallprocessDataNum_computes = (((tileDataNum * sizeof(float) + 256 - 1) / 256) * 256) / sizeof(float);
确保内存访问地址对齐,避免硬件访问效率损失。
3.2.2 kernel侧设计
Ascend C的SoftplusBackward算子流程见下图:

以下按 “初始化→主流程→数据搬运→核心计算→结果输出” 的逻辑,详细拆解计算步骤:
3.2.2.1 初始化阶段(Init方法):配置硬件与任务参数
Init是 kernel 启动后的首个步骤,主要作用是读取 host 侧传递的 tiling 参数、划分核心任务角色、初始化全局内存(GM)与统一缓冲区(UB),为后续计算做准备,具体步骤如下:
1. 核心角色与任务参数初始化
核心类型判断:根据当前核心索引(coreIdx = GetBlockIdx())与 host 侧计算的 “尾块数量”(tailBlockNum),判断当前核心是 “大核”(处理尾块,任务量稍多)还是 “小核”(处理普通块):
若coreIdx < tailBlockNum:为大核,加载大核参数(bigCoreDataNum、finalBigTileNum等);
否则为小核,加载小核参数(smallCoreDataNum、finalSmallTileNum等)。
全局偏移量计算:根据核心类型计算当前核心处理数据在 GM 中的起始地址(globalOffset),确保多核心不重复处理数据:
2. GM 缓冲区绑定(全局内存映射)
根据数据类型(dataType:0=FP16、1=FP32、2=BF16),将 GM 中的输入(gradOutput、self)和输出(gradInput)地址映射为对应类型的GlobalTensor(AscendC 中用于访问 GM 的封装类),确保数据读写类型匹配:
例如 FP16 类型:reinterpret_cast<gm half*>(gradOutput)后绑定到gradOutputGmHalf;
BF16 类型:通过bfloat16(本质uint16_t)映射,后续需通过自定义函数转换为 float 计算。
3. U 缓冲区初始化(双缓冲与计算缓冲区)
UB 是 AI Core 的高速缓冲区,Init需为 “数据搬运” 和 “核心计算” 分配 UB 空间,核心设计包括:
双缓冲(Double Buffer):为gradOutput、self、gradInput分别初始化TQue(AscendC 队列类),缓冲区数量BUFFER_NUM=2,实现 “数据搬运与计算并行”(一个缓冲区用于计算,另一个预加载下一批数据);
计算缓冲区:初始化 16 个TBuf(UB 计算缓冲区),用于存储中间计算结果(如betaSelfBuf存储beta*self、sigmoidBuf存储 sigmoid 结果),缓冲区大小由processDataNum_computes(按 256 字节对齐的计算量)决定,确保适配硬件访问粒度。
3.2.2.2 主流程控制(Process方法):循环处理 Tile 数据
Process是 kernel 的主调度逻辑,按 host 侧划分的tileNum(总 tile 数)循环处理每个 tile,核心是 “分批次处理数据,避免 UB 溢出”,步骤如下:
循环次数确定:循环tileNum次,前tileNum-1次处理 “完整 tile”(大小为tileDataNum),最后 1 次处理 “尾 tile”(大小为tailDataNum,因总数据量可能不被tileDataNum整除);
每轮调度逻辑:每次循环依次调用CopyIn(GM→UB 搬运数据)、Compute(UB 内核心计算)、CopyOut(UB→GM 输出结果),形成 “搬运 - 计算 - 输出” 的流水线。
3.2.2.3 数据输入搬运(CopyIn方法):GM→UB 数据加载
CopyIn的作用是将当前 tile 的 GM 数据(gradOutput、self)搬运到 UB 的双缓冲队列(gradOutputQueue、selfQueue),按数据类型差异化处理:
1. 双缓冲分配与数据拷贝
首先调用AllocTensor从队列中分配一个LocalTensor(AscendC 中访问 UB 的封装类);
通过DataCopy(AscendC 数据拷贝 API)将GlobalTensor中当前 tile 的数据(起始地址为progress * tileDataNum,progress为当前循环轮次)拷贝到LocalTensor;
调用EnQue将加载好数据的LocalTensor放入队列,等待Compute读取。
2. 数据类型适配
FP16/FP32:直接通过DataCopy拷贝,类型天然匹配;
BF16:因 BF16 无原生计算支持,需先以bfloat16类型搬运到 UB,后续在Compute中转换为 float 计算。
3.2.2.4 核心计算流程(Compute方法):Softplus 反向传播逻辑
Compute是 kernel 的核心,实现 SoftplusBackward 的数学计算逻辑(Softplus 正向:softplus(x) = (1/beta)log(1+exp(betax)),反向:gradInput = gradOutput * sigmoid(beta*x),需处理阈值分支),按数学步骤拆解为 7 个关键环节,同时适配多数据类型并优化精度:
环节 1:输入数据类型统一(转为 float 计算)
因 BF16/FP16 无原生复杂计算支持,需先将 UB 中的输入数据转为 float(AI Core 对 float 计算支持更完善):
FP16→float:调用Cast(AscendC 类型转换 API),RoundMode::CAST_NONE(无舍入);
FP32:直接DataCopy,无需转换;
BF16→float:通过自定义Bf16ToFloat函数转换(将 BF16 的 16 位数据左移 16 位补零,再强转为 float):
环节 2:计算beta * self(Softplus 反向的基础项)
通过Muls(AscendC 标量乘向量 API)计算beta * self,结果存入betaSelfBuf:
公式:betaSelf[i] = selfFloat[i] * betaVal(betaVal为 host 侧传递的算子属性);
作用:betaSelf是后续计算 sigmoid 和阈值分支的核心输入。
环节 3:生成线性分支掩码(阈值保护逻辑)
Softplus 反向在betax > threshold时,为避免exp(betax)溢出,采用线性近似(gradInput ≈ gradOutput),需先生成掩码区分 “线性分支” 和 “普通分支”:
初始化thresholdTensor(全为thresholdVal,host 侧传递的阈值,默认 20.0)和constantOneTensor(全为 1.0);
调用Compare(AscendC 比较 API),按betaSelf[i] > thresholdVal生成linearMask(uint8_t类型,1 表示线性分支,0 表示普通分支):
环节 4:高精度 Sigmoid 计算(普通分支核心)
普通分支需计算gradFactor = sigmoid(betaSelf) = 1/(1+exp(-betaSelf)),为避免精度损失,代码做了两处关键优化:
先计算expNegBetaSelf = exp(safeNegBetaSelf)(Exp API);
再计算onePlusExp = expNegBetaSelf + 1.0(Adds API);
最后用Div(numerator, denominator)计算:numerator全为 1.0,denominator = onePlusExp,结果存入sigmoidBuf,精度比Reciprocal提升约 1e-5。
环节 5:梯度因子选择(线性 / 普通分支融合)
通过Select(AscendC 分支选择 API)根据linearMask融合两个分支的梯度因子gradFactor:
当linearMask[i] = 1(线性分支):gradFactor[i] = 1.0(对应constantOneTensor);
当linearMask[i] = 0(普通分支):gradFactor[i] = sigmoidBuf[i];
公式:gradFactor = linearMask ? 1.0 : sigmoid(betaSelf),结果存入selectResultBuf。
环节 6:计算最终梯度(gradInput = gradOutput * gradFactor)
通过Mul(AscendC 向量乘向量 API)计算输出梯度:
公式:gradInputFloat[i] = gradOutputFloat[i] * gradFactor[i](gradOutputFloat为输入梯度,gradFactor为分支选择后的因子);
结果存入gradInputFloatBuf,为后续输出做准备。
环节 7:计算结果类型转回(float→目标类型)
将 float 类型的gradInputFloat转回目标数