已关闭
[Requirement|需求建议]: Neg算子新增数据类型支持 #1510
CyndiZ创建于  5月11日关闭于  6月27日
CyndiZ
5月11日 创建

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

一、背景信息 (必填)

参考昇腾版本内置neg算子的 TBE 实现,在昇腾 NPU 上基于 Ascend C 编程语言实现功能一致的算子,完成算子设计、开发、测试全流程工作,验收通过后将算子提交至昇腾算子开源仓。

二、价值/作用 (必填)

与原 TBE 算子核心功能完全对齐,支持原算子对应的所有数据类型、数据格式;特别地需要支持int16、uint8、int64类型输入。说明:当输入类型为 uint8时,其行为和torch.neg一致,torch.neg(uint8) 不会返回负数,而是返回一个 uint8 类型的张量,其值等于 256 - x(对于非零值)或 0(对于零值),即发生“回绕”(wrap-around)或“截断”效果,而不是数学上的取负。
必须实现算子泛化功能,满足各类合法输入场景的计算需求,验收阶段将采用泛化数据进行验收。

三、设计方案 (必填)

host 侧设计

  1. 参数解析与校验

    • 输入:x,必选张量。
    • 输出:y,必选张量。
    • 校验 x/y 非空。
    • 校验 x.dtype == y.dtype
    • 校验 x.shape == y.shape
    • AICore 路径支持 BFLOAT16/FLOAT16/FLOAT32/INT8/UINT8/INT16/INT32/INT64,其中 INT16/UINT8 为本次新增支持类型。
  2. tiling 策略

    • Neg 是无属性、无广播的一元逐元素算子,可复用通用逐元素 tiling。
    • tiling 输入只需要元素总量、dtype size、UB 大小、AIV core 数等基础信息。
    • 按元素总量切分 core 任务,并在单 core 内按 UB 容量继续切分 tile。
    • 空 tensor 直接走空执行路径,不需要额外 workspace。
  3. tilingKey 规划策略

    • tilingKey 用于区分调度模式和 dtype 模板。
    • 浮点类型走直接取负路径,bfloat16 可使用 float32 中间计算后转回。
    • int8/uint8 需要保留与 TBE 一致的低 bit 整数回绕语义。

kernel 侧设计

  1. kernel 侧实现描述

    • 进行 InitProcess 两个阶段,其中 Process 包括搬入、计算、搬出。
    • 从 GM 搬入 x 到 UB。
    • 根据 dtype 选择计算路径:
      • float16/float32/int16/int32/int64:直接执行向量取负。
      • bfloat16:转换为 float32 取负,再转回 bfloat16
      • int8:按有符号 8 bit 回绕语义处理,重点验证 -128
      • uint8:按无符号 8 bit 回绕语义处理,等价于 (-x) mod 256
    • 将 UB 中的结果搬出到 GM。
  2. AscendC 实现流程图

graph TD
    A["开始"] --> B["Host 校验 shape / dtype 并生成 tiling"]
    B --> C["Kernel Init: 读取 GM 地址与 tiling"]
    C --> D["CopyIn: 搬入 x tile 到 UB"]
    D --> E{"dtype 是否 bfloat16?"}
    E -->|是| F["Cast 到 float32"]
    E -->|否| G["保持原 dtype"]
    F --> H["执行向量取负"]
    G --> H
    H --> I{"dtype 是否 int8/uint8?"}
    I -->|是| J["按 8 bit 回绕语义修正"]
    I -->|否| K["保持计算结果"]
    J --> L["CopyOut: 写回 y"]
    K --> L
    L --> M["结束"]

AscendC 实现与 TBE 实现差异点和原因

  1. TBE 使用 classify + variable_shape + auto_schedule 自动处理动态 shape;AscendC 需要在 host tiling 中显式完成元素数统计、分核和 tile 大小规划。
  2. TBE 中 int32/int64 通过乘以 -1 实现;扩展 int16 时也应保持整数回绕语义。
  3. TBE 中 int8 经过 float16 中间计算和溢出修正;扩展 uint8 时同样需要按 uint8 取模语义修正,避免饱和转换导致结果不一致。
  4. bfloat16 建议在 AscendC 侧使用 float32 中间计算再转回,避免不同后端对直接 bf16 取负的支持差异。

支持硬件

产品 是否支持
Atlas A2 训练系列产品/Atlas A3 系列产品

算子约束限制

参数名 类别 描述 数据类型 数据格式
x 输入张量 输入张量。 BFLOAT16、FLOAT16、FLOAT32、INT8、UINT8、INT16、INT32、INT64 ND
y 输出张量 输出张量。 BFLOAT16、FLOAT16、FLOAT32、INT8、UINT8、INT16、INT32、INT64 ND
  1. 输入和输出 dtype 必须一致。
  2. 输入和输出 shape 必须一致。
  3. Neg 没有广播语义,不需要处理多输入 shape 对齐。
likedislike
CCyndiZ
5月11日 关联了pull request:【CANN开源开放社区任务】【社区任务】AscendC实现Neg算子贡献
sunchun成员
5月11日 评论:

/assign @CyndiZ

likedislike
CANN-robotCANN-robot成员
5月11日 将 CyndiZ 设为负责人
CANN-robotCANN-robot成员
6月27日 关闭了 issue
CANN-robotCANN-robot成员
6月27日 添加了label:resolved