已关闭
[Requirement|需求建议]: 需求 ReduceStdV2Update 算子支持 Ascend950 ascendc 实现 #2327
zhongheng创建于  7月27日关闭于  8月3日
zhongheng
7月27日 创建

Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.

一、背景信息 (必填)

需求 ReduceStdV2Update 算子支持 Ascend950 ascendc 实现

二、价值/作用 (必填)

  1. 使 aclnnVar/aclnnVarMean 等上层接口在 Ascend 950 上具备 AscendC 原生 kernel 实现能力,替代旧版 TBE binary,获得更好的性能可维护性和功能扩展性。
  2. 使 aclnnVarMean 可复用外部传入的 mean 结果,避免重复求均值,提升方差/均值联合计算场景的性能。
  3. 作为 torch.var/torch.std 的底层实现组件,支撑 PyTorch 训练和推理框架在 Ascend 950 上的归约统计计算。

三、设计方案 (必填)

3.1 使能方式(涉及哪些框架:如Aclnn直调、Pytorch训练等)
  • Aclnn 间接调用:本算子是 l0op 内部子算子(ACLNNTYPE=aclnn_exclude),不对外暴露独立 aclnn L2 接口。aclnn 调用经上层 aclnnVar/aclnnVarCorrection/aclnnVarMean 间接触发。
  • GE IR 图模式:通过 if_std 属性可控制输出方差(if_std=false)或标准差(if_std=true)。
  • PyTorch 训练/推理torch.var/torch.std → torch_npu → aclnnVar/aclnnStdl0op::ReduceStdV2Update(kernel)

调用链:

aclnnVar(x, dim, correction, keepdim)
  → ReduceMean(x) → Expand(mean, x.shape)
  → l0op::ReduceStdV2Update(x, mean, dim, unbiased, keepdim)   ← 本算子
  → output_var
3.2 总体设计
3.2.1 算子支持的数据类型
输入/输出 数据类型
x(self) FLOAT、FLOAT16、BFLOAT16
mean FLOAT、FLOAT16、BFLOAT16(dtype 与 x 一致)
output_var FLOAT、FLOAT16、BFLOAT16(dtype 与 x 一致)

属性:

  • dim:ListInt,归约维度列表,支持负索引,dim 为空时归约所有维度
  • if_std:Bool(默认 false),false=输出方差,true=输出标准差(开方)
  • unbiased:Bool(默认 true),true=无偏(N-1),false=有偏(N)
  • keepdim:Bool(默认 false),是否保留归约维度
  • correction:Int(默认 1),0=有偏(N),1=无偏(N-1)
3.2.2 host侧设计
  • OpDef(def.cpp):定义输入输出 dtype/format/属性配置,三路 dtype 索引对齐组合(x[i]→mean[i] 同 dtype),支持 DynamicRank 和 DynamicShape。
  • InferShape(infershape.cpp):output shape = x 沿 dim 归约后的 shape(keepdim=true 维度设 1,false 移除)。复用 InferShape4ReduceStdFunc 逻辑,支持负索引归一化和全归约。format 限制由 def.cpp 的 .Format({ge::FORMAT_ND}) 保证,infershape 不做 format 检查。
  • Tiling(tiling_arch35.cpp):标准 4 步 pattern 预处理(去 size=1 轴 → 合轴 → 补 leading A → 补 R 增广)+ 双切分策略(A 切分 → R 切分 → R 全载扩 A)+ 多核切分(fuse aLoop)+ group 切分(可选)。cof = 1/(N-correction) 在 host 高精度计算后下发,含除零保护。TilingKey 编码:templateType(bit0) | isEmptyTensor(bit1) | isTailR(bit2),动态分发 normal/group/empty 三套模板。
3.2.3 kernel侧设计

实现 normal/group/empty 三套模板,共 15 份 binary(5 模板组合 × 3 dtype):

  • normal 模板:处理常规 A×R 2D 归约。三级流水线(CopyIn → Compute → CopyOut),热循环内零 PipeBarrier。FP16/BF16 输入固定提升到 FP32 累加(Cast→compute→Cast),结果转回原 dtype。if_std=true 时使用 high_precision vsqrt 开方,含 sqrt 负值保护(max(var,0))。
  • group 模板:处理多核 2D 分核场景(A 轴切分多核 + R 轴组内归约),通过 SetScheduleMode(1) 启用 group 调度。
  • empty 模板:处理空 Tensor 边界。EMPTY_A(∃ A 轴 size=0)→ output 空 tensor,kernel 零操作;EMPTY_R(∃ R 轴 size=0)→ output=0,kernel 写零。

双输入并行载入:x + mean 两个 GM tensor 策略 B 并行载入(preIn 并行槽不开 DoubleBuffer)。

3.3 支持硬件
硬件平台 支持情况 说明
Ascend 950PR/950DT 本次新增 AscendC Kernel 实现(arch35)
Atlas A3 训练/推理 已有 TBE binary 覆盖
Atlas A2 训练/推理 已有 TBE binary 覆盖
Atlas 200I/500 A2 推理 × 不支持
Atlas 推理系列 已有 TBE binary 覆盖
Atlas 训练系列 已有 TBE binary 覆盖

3.4 算子约束限制

  1. 数据类型限制xmean 仅支持 FLOAT、FLOAT16、BFLOAT16,且二者数据类型须一致;不支持 DOUBLE、复数、整型。
  2. shape 约束mean 必须已通过 Expand 广播到 x 的 shape;output shape 为 x 沿 dim 归约后的 shape(keepdim=true 维度设 1,false 移除)。
  3. correction/unbiased 约束:仅支持 correction=0(有偏)和 correction=1(无偏)两种语义;correction>1 由上层 aclnnVarCorrection 处理。
  4. 接口约束:aclnn 接口固定输出方差(if_std=false);GE IR 图模式可通过 if_std 属性控制输出方差或标准差。本算子为 l0op 内部子算子,不对外暴露独立 aclnn L2 接口。
  5. 精度约束:FP16/BF16 输入固定提升到 FP32 累加,结果转回原 dtype;BF16 输出回转使用 round 模式;输出 dtype 与输入 x 一致,不在接口层做类型提升。
  6. 边界情况:空 Tensor、单元素且 correction≥1 的场景由上层 aclnnVar/aclnnVarCorrection 提前拦截返回 NAN/INF,kernel 实现含防御性处理。
    💡 备注(选填)
likedislike
zhongheng
7月27日 评论:

/assign

likedislike
CANN-robotCANN-robot成员
7月27日 将 zhongheng 设为负责人
Zzhongheng
7月27日 关联了pull request:[CANNBot]ReduceStdV2Update算子新增AscendC实现
Zzhongheng
7月30日 修改了issue 的描述
Zzhongheng
8月3日 issue状态由 进行中 改变为 已完成
Zzhongheng
8月3日 关闭了 issue
CANN-robotCANN-robot成员
8月3日 添加了label:Accepted