Closed
[Requirement|需求建议]: InplaceRsqrt 新增整型数据类型支持 #2156
Nice tryAdd on  Jul 7Closed at  Jul 30
Nice try
Nice tryCollaborators
Jul 7 Create

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

一、背景信息 (必填)

当前 Rsqrt / InplaceRsqrt 算子仅支持浮点数据类型(float16、float、bfloat16),不支持整型和布尔类型。在 AscendC 算子的通用生态中,越来越多的模型和场景需要 Rsqrt 算子能够处理 int8、int16、int32、uint8、bool 等数据类型,以实现端到端的整型推理链路,避免频繁的类型转换开销。

二、价值/作用 (必填)

本需求为 Rsqrt / InplaceRsqrt 算子新增 int8、int16、int32、uint8、bool 五种数据类型的支持,带来以下价值:

  1. 补齐算子数据类型覆盖:Rsqrt 对齐同类数学算子(如 Sqrt、Acos)的数据类型支持范围,覆盖 AscendC 算子的全部通用数据类型。
  2. 减少推理链路类型转换:整型模型可直接调用 Rsqrt 无需插入 Cast 节点,简化计算图,降低调度开销。
  3. 增强算子通用性:满足社区开发者在混合精度训练、量化推理、布尔掩码处理等场景中对 Rsqrt 算子更广泛的数据类型需求。

三、设计方案 (必填)

总体思路

Rsqrt 算子的核心计算为 y = 1 / sqrt(x),对于浮点类型直接通过 AscendC 内置的 Sqrt + Div 原语完成;对于整型/布尔类型,通过类型提升(promotion)到 half/float → 执行 Rsqrt 计算 → CAST_RINT 截断回原类型的方式实现,同时对 int16/int8 引入分段线性校正保证结果精度。

改动范围

涉及 5 个层次的修改:

1. 算子定义层(op_def / op_proto)

在 Input 和 Output 的 DataType 列表中新增 DT_INT8、DT_INT16、DT_INT32、DT_UINT8、DT_BOOL,注册格式为 ND。

2. Tiling 层(op_host/rsqrt_tiling.cpp)

新增各数据类型在 UB 空间中的 Buffer 分配常量:

数据类型 DATA_NUM 含义
INT32 2 每个元素需 2 倍 UB 数据量
INT16 2 同上
INT8 6 有符号 8 位整数需较大 UB
UINT8 4 无符号 8 位整数
BOOL 2 布尔类型

Tiling 函数 RsqrtTilingFunc 的 switch 分支扩展至上述类型,其余逻辑(coreNum 计算、buffer 数量选择、tile 切分)保持通用。

3. ACLNN API 层(op_api/aclnn_rsqrt.cpp)

  • 统一 DTYPE_SUPPORT_LIST 覆盖全部 8 种类型(移除旧版独立仅浮点的 DTYPE_OUT_LIST)
  • 新增 CheckDtypeValid 校验:self 和 out 的数据类型必须一致,避免 ViewCopy 出现未定义行为

4. Kernel 层(op_kernel/rsqrt.h / rsqrt.cpp)

KernelRsqrt<DTYPE_X, BUFFER_NUM> 模板类按 TYPE_X 分发至以下计算路径:

  • float / half:Sqrt → Div(1, result)(保持原逻辑)
  • bfloat16:Cast 到 float → Sqrt → Div → Cast 回 bfloat16
  • int32:Cast 到 float → Rsqrt → Cast(CAST_RINT) 写回 int32
  • int16:Cast 到 half → Rsqrt → ComputeIntCorrection(分段校正)→ Cast(CAST_RINT)
  • int8:同 int16 流程(使用 tmp1/tmp2 half 缓冲区)
  • uint8:Cast 到 half → Rsqrt → Cast(CAST_RINT)
  • bool:Duplicate(0x0101) 填充 int16_t 模式,每个输入元素输出 true

其中 ComputeIntCorrection 用于修正 int16/int8 的 Rsqrt 结果:

p1 = min(p1, 2.0f)
p2 = p1 - 1.0f
p2 = max(p2, 0.0f)
p2 = p2 * 3.0f
p1 = p1 - p2

此校正将 [0, 2] 范围的 Rsqrt 结果压缩到正确范围,对齐 pytorch/tbe 的语义。

rsqrt_tiling_data.h 中将 tile/tail 相关字段从 uint32_t 升级为 uint64_t,避免大 shape 下溢出。

5. 测试层(tests/ut)

op_api 测试(16 个用例):覆盖全部支持类型的正常入参、空 tensor、非连续、高维、dtype 不匹配、shape 不一致、nullptr 等场景。

op_host 测试(8 个用例):infershape 推导 + tiling 参数计算验证(float、float16、bf16)。

op_kernel 测试(7 个用例):每个新增类型独立的端到端 ICPU 仿真测试(int8/int16/int32/uint8/bool 各 1 个,float32 覆盖单流水 + 双流水)。使用 rsqrt_test_entry.h 共享入口模板通过宏拼接生成唯一函数名(如 rsqrt_int8),避免多 DTYPE_X 文件的 ODR 冲突。数据生成脚本 gen_data.py 和对比脚本 compare_data.py 同时支持 8 种数据类型。

伪代码示意

def rsqrt_kernel(x: Tensor[AnyDType], y: Tensor[AnyDType]):
    tiling = compute_tiling(x.shape, x.dtype)
    if x.dtype in {float16, float32}:
        y = 1.0 / sqrt(x)
    elif x.dtype == bfloat16:
        x_f32 = cast(x, float32)
        y_f32 = 1.0 / sqrt(x_f32)
        y = cast(y_f32, bfloat16, round_mode=CAST_RINT)
    elif x.dtype == int32:
        x_f32 = cast(x, float32)
        y_f32 = rsqrt(x_f32)
        y = cast(y_f32, int32, round_mode=CAST_RINT)
    elif x.dtype in {int16, int8}:
        x_f16 = cast(x, float16)
        y_f16 = rsqrt(x_f16)
        y_f16 = int_correction(y_f16)  # min(2)→add(-1)→max(0)→mul(3)→sub
        y = cast(y_f16, original_dtype, round_mode=CAST_RINT)
    elif x.dtype == uint8:
        x_f16 = cast(x, float16)
        y_f16 = rsqrt(x_f16)
        y = cast(y_f16, uint8, round_mode=CAST_RINT)
    elif x.dtype == bool:
        y = fill(0x0101)  # map to int16_t pattern, output True

💡 备注(选填)

  • 本次改动仅影响 experimental/math/rsqrt 路径
  • Tiling 层适配 ascend910b / ascend910_93 两个 SoC 版本
  • Kernel 层支持 AIV 模式调度
  • 所有 UT 均通过 ICPU 仿真器验证
likedislike
Nice tryNice tryCollaborators
Jul 7 Associated pull request
陈思
陈思Collaborators
Jul 7 Comments:

/assign @Nice_try

likedislike
CANN-robotCANN-robotCollaborators
Jul 7 Assign Nice_try as responsible person
Nice tryNice tryCollaborators
Jul 7 Modify description of issue
Nice tryNice tryCollaborators
Jul 7 Associated pull request
llimwangCollaborators
Jul 30 Comments:

该 Issue 已超过 20 天未更新,先行关闭以便归档;如仍需跟进请评论或重新打开。

likedislike
LllimwangCollaborators
Jul 30 Close issue