已关闭
[Requirement|需求建议]: SoftplusV2Grad算子Ascend C实现 #229
cxj01创建于  2025年12月12日关闭于  5月6日
cxj01
2025年12月12日 创建

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 (价值/作用)

  1. 参考内置算子实现,需实现相同功能,包含相同数据类型和相同数据格式,其中int64,double数据类型可暂不支持,广播操作可暂不支持。
  2. 算子性能要求持平原有TBE实现算子,当前可暂只支持在使用所有核进行计算时性能需不低于原有TBE算子95%。
  3. 适用于Atlas 800I A2推理产品

Design(设计方案)

3.2 需求整体设计

3.2.1 host侧设计

3.2.1.1 tiling 策略

tiling 策略的核心是将输入数据划分为适合 AI Core 硬件特性(如 UB 容量、计算单元吞吐量)的小块(tile),确保计算过程中数据高效复用,减少全局内存访问。本代码的 tiling 策略设计如下:
基础参数定义

  • BLOCK_SIZE = 32:Datacopy等算子需要32字节对齐;
  • BUFFER_NUM = 2:UB 中用于双缓冲(Double Buffer)的缓冲区数量,通过双缓冲隐藏数据搬运 latency;
  • UB_DATA_NUMBER = 16:根据需要存储的数据块的数量确定。【暂定是16,基于以下考量:1、冗余设计:避免实际计算过程中占满UB,导致不必要溢出; 2、保证有足够的暂存结果可以放到UB中,3、目前并没有在kernel计算时进行缓冲区的复用,比如数据类型转换时、中间exp计算等数据,后续调优可以再进行优化】
    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算子流程见下图:
AscendC执行逻辑.png
以下按 “初始化→主流程→数据搬运→核心计算→结果输出” 的逻辑,详细拆解计算步骤:

3.2.2.1 初始化阶段(Init方法):配置硬件与任务参数

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 字节对齐的计算量)决定,确保适配硬件访问粒度。

3.2.2.2 主流程控制(Process方法):循环处理 Tile 数据

Process是 kernel 的主调度逻辑,按 host 侧划分的tileNum(总 tile 数)循环处理每个 tile,核心是 “分批次处理数据,避免 UB 溢出”,步骤如下:

  1. 循环次数确定:循环tileNum次,前tileNum-1次处理 “完整 tile”(大小为tileDataNum),最后 1 次处理 “尾 tile”(大小为tailDataNum,因总数据量可能不被tileDataNum整除);

  2. 每轮调度逻辑:每次循环依次调用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):

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),需先生成掩码区分 “线性分支” 和 “普通分支”:

  1. 初始化thresholdTensor(全为thresholdVal,host 侧传递的阈值,默认 20.0)和constantOneTensor(全为 1.0);

  2. 调用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)),为避免精度损失,代码做了两处关键优化

  1. 溢出保护:当-betaSelf过大时,exp(-betaSelf)可能溢出,设置安全阈值expSafeThreshold(BF16 为 80.0,其他为 50.0),通过Min(AscendC 取小 API)限制-betaSelf最大值:
Min(safeNegBetaSelf, negBetaSelf, maxThresholdTensor, ...); // negBetaSelf = -betaSelf
  1. 高精度除法替代倒数:传统用Reciprocal(1+exp(...))计算 sigmoid 会损失精度,此处改为Div(AscendC 向量除法 API):
  • 先计算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转回目标数

likedislike
cxj01
2025年12月12日 评论:

/assign

likedislike
胡一航成员
2025年12月15日 评论:

感谢您的反馈,可以关联上您的PR,在验证后由commiter审视合入。

likedislike
chenqi317成员
2025年12月25日 评论:

经sig组织评审,您的方案已通过,请关联相关贡献的代码

likedislike
Cchenqi317成员
2025年12月25日 issue状态由 进行中 改变为 待办的
Cchenqi317成员
3月13日 将 zhajianqing123 设为负责人
Cchenqi317成员
3月13日 将 liujie12345678 设为负责人
Cchenqi317成员
4月13日 将 fullt 设为负责人
Cchenqi317成员
4月13日 移除了负责人 zhajianqing123
Cchenqi317成员
4月13日 移除了负责人 liujie12345678
Ffulltower成员
4月20日 关联了pull request:SoftplusV2Grad算子Ascend C实现
Ffulltower成员
4月20日 关联了pull request:提交Ascend C实现的GeluGrad算子
Cchenqi317成员
4月24日 将 chenqi317 设为负责人
CANN-robotCANN-robot成员
5月6日 关闭了 issue