已关闭
[Requirement|需求建议]: 【社区任务】Logspace算子开发交付(任务编号 529-13) #2029
bububu创建于  6月23日关闭于  7月30日
bububu
6月23日 创建

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

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

一、背景信息 (必填)

参考昇腾版本内置 LogSpace 算子的 TBE 实现,在昇腾 NPU 上基于 Ascend C 编程语言实现功能一致的算子,完成算子设计、开发、测试全流程工作,验收通过后将算子提交至昇腾算子开源仓。LogSpace 为纯生成类算子,按对数空间生成等比序列:result[i] = base^(start + i * step),其中 step = (end - start) / (steps - 1)steps == 1 时仅生成 base^startsteps == 0 时输出空张量)。

二、价值/作用 (必填)

与原 TBE 算子核心功能完全对齐,支持原算子对应的所有数据类型、数据格式;在原 TBE 仅支持 FLOAT/FLOAT16/BFLOAT16 的基础上,特别地需要新增 INT8、UINT8、INT16、INT32 整型输出支持。说明:整型输出与 torch.logspace(...).to(dtype) 一致——先按 float32 计算 base^x,再向零截断(CAST_TRUNC)落整型;当结果超出整型范围时按该整型位宽取模回绕(例如 uint8 等价于 value mod 256int8 等价于有符号 8 bit 回绕),而不是饱和。
必须实现算子泛化功能,满足各类合法输入场景(任意 start/end/base/steps、大指数溢出、整数幂网格点等)的计算需求,验收阶段将采用泛化数据进行验收。

三、设计方案 (必填)

host 侧设计

  1. 参数解析与校验

    • 输入:无输入张量(纯生成算子)。
    • 属性:start(aclScalar→float)、end(aclScalar→float)、steps(int64_t,>= 0)、base(double→float,> 0)。
    • 输出:result,必选张量,其 dtype 决定输出类型。
    • 校验 result 非空。
    • 校验 result 元素数与 steps 一致。
    • AICore 路径支持 FLOAT/FLOAT16/BFLOAT16/INT8/UINT8/INT16/INT32,其中 INT8/UINT8/INT16/INT32 为本次新增支持类型。
  2. tiling 策略

    • LogSpace 是无输入、无广播的纯生成逐元素算子,可复用通用逐元素 tiling。
    • tiling 输入只需要输出元素总量 steps、输出 dtype size、UB 大小、AIV core 数等基础信息。
    • host 侧预先算出序列步长 step 与对数底 logBase = ln(base),下发到 tiling。
    • 按元素总量切分 core 任务(满核优先,单核最少处理 MIN_PER_CORE = 64 个元素,不足则缩减核数),并在单 core 内按 UB 容量(UB_CHUNK_ELEMS = 2048 fp32 元素/次)继续切分 tile。
    • steps == 0 走空执行路径,不需要额外 workspace。
    • 整型大值溢出场景:host 侧预计算整数幂网格点精确值表(baseN[],double pow 精确算出)与 double-float 递推常数,随 tiling 下发,供 kernel 修正端点 / 网格点精度(见差异点)。
  3. tilingKey 规划策略

    • tilingKey(D_T_Y, MODE) 二元组下发:D_T_Y 为输出 dtype,MODE0 = NORMAL(steps >= 2) / 1 = SINGLE(steps == 0 / 1)
    • 共 14 条路径(7 dtype × NORMAL/SINGLE)全部启用。
    • 浮点类型(float32/float16/bfloat16)走向量直接计算路径,bfloat16 使用 float32 中间计算后转回。
    • 整型类型(int8/int16/int32/uint8)在 kernel 内以 float32 计算 base^x,末步 Cast(CAST_TRUNC) 向零取整落整型,匹配 torch。
    • int8/uint8 大值溢出(max(base^start, base^end) > 2^24)由 host 下发 useDfV 门控,走向量 double-float 高精度指数路径(见 kernel 侧)。

kernel 侧设计

  1. kernel 侧实现描述

    • 进行 InitProcess 两个阶段,其中 Process 包括计算(Compute)、搬出(CopyOut)。LogSpace 无输入张量,故无 CopyIn。
    • 每个核按其认领的输出区间 [coreOffset, coreOffset + coreNum) 逐 chunk 生成 base^(start + i * step)
    • 根据 dtype 选择计算路径:
      • float16/float32/int16/int32:向量计算 exp((start + i*step) * ln(base)),浮点直接落,整型末步 Cast(CAST_TRUNC)
      • bfloat16:以 float32 中间计算,再 Castbfloat16
      • int8/uint8(大值溢出路径):向量 double-float 指数(Cody-Waite 范围规约 k = round(arg/ln2)r = arg - k*ln2,整数运算 + Sterbenz 保证标量也精确;向量单元上做 Horner 泰勒展开,1/k! double-float 常数由 host 下发),并以 expBase[c] = expBase[c-1] * const 递推(整核只做 1 次 df-exp、每 chunk 仅 1 次向量 df 乘,const = base^(chunkElems * step) 由 host 下发);随后对单个 float32 值取 mod 256value >= 2^31 时用 keep 掩码 1 - min(1, floor(value/2^31)) 置零(此时 fp32 为 256 的倍数,取模为 0)。
    • 端点 / 整数幂网格点精度:用 host pow 精确算出的值在 UB float32 buffer 上 SetValue 打补丁,消除设备 exp 对整数幂的舍入误差。
    • 将 UB 中的结果搬出到 GM。
  2. AscendC 实现流程图

graph TD
    A["开始"] --> B["Host 校验 result / 计算 step、logBase 并生成 tiling"]
    B --> C["Kernel Init: 读取 GM 地址与 tiling,按 MODE 选路"]
    C --> D{"steps == 0?"}
    D -->|是| Z["空执行,直接结束"]
    D -->|否| E["Compute: 按区间生成 base^(start + i*step)"]
    E --> F{"输出 dtype?"}
    F -->|float32/float16| G["向量 exp 直接落"]
    F -->|bfloat16| H["float32 中间计算后 Cast 回 bf16"]
    F -->|int16/int32| I["float32 计算后 Cast(CAST_TRUNC)"]
    F -->|int8/uint8| J{"大值溢出 useDfV?"}
    J -->|否| I
    J -->|是| K["向量 df-exp + 递推 + keep 掩码 + fp32 mod256"]
    G --> L["端点/整数幂网格点 host pow 精确 SetValue 打补丁"]
    H --> L
    I --> L
    K --> L
    L --> M["CopyOut: 写回 result"]
    M --> N["结束"]

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

  1. TBE 使用 classify + variable_shape + auto_schedule 自动处理动态 shape;AscendC 需要在 host tiling 中显式完成元素数统计、分核和 tile 大小规划。
  2. TBE 仅支持 float32/float16/bfloat16;AscendC 新增 int8/int16/int32/uint8 整型路径,统一以 float32 计算 base^xCast(CAST_TRUNC) 向零取整,并对 int8/uint8 实现与 torch.logspace(...).to(dtype) 一致的位宽回绕(取模)语义。
  3. NPU 标量单元精度不足:double-float 的误差自由变换(TwoSum/TwoProd)在标量域会被编译器代数重排塌缩成单精度(#pragma clang fp contract(off) / volatile 屏障均无效)。AscendC 将 double-float 高精度指数放在向量单元计算(向量为 IEEE 单精,忠实匹配 numpy float32 golden),解决 int8/uint8 大指数取模的精度问题。
  4. 等比序列 base^x 的逐 chunk 参考值采用递推(每 chunk 乘常数 base^(chunkElems*step)),把一次性昂贵的 df-exp 摊薄到整核一次,兼顾精度与性能。
  5. 端点与整数幂网格点用 host double pow 精确写出,规避设备 exp 对整数幂的舍入误差导致首尾元素偏差。
  6. bfloat16 在 AscendC 侧使用 float32 中间计算再转回,避免不同后端对直接 bf16 计算的支持差异。

支持硬件

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

算子约束限制

参数名 类别 描述 数据类型 数据格式
start 属性 对数序列的起始指数。 FLOAT(host 侧由 aclScalar 转入) -
end 属性 对数序列的结束指数。 FLOAT(host 侧由 aclScalar 转入) -
steps 属性 序列中的元素数量(>= 0)。 INT64 -
base 属性 对数空间的底数(> 0)。 FLOAT(host 侧由 double 转入) -
result 输出张量 输出的对数间隔序列张量。 FLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16、INT32 ND
  1. base 必须大于 0;steps 必须大于等于 0。
  2. 整型输出与 torch.logspace(...).to(dtype) 一致:float32 计算后向零截断,超范围按位宽取模回绕。
  3. LogSpace 无输入张量、无广播语义,不需要处理多输入 shape 对齐。
likedislike
Bbububu
6月23日 关联了pull request:【CANN开源开放社区任务】【社区任务】AscendC实现Logspace算子贡献
sunchun成员
6月24日 评论:

/assign @bububu

likedislike
CANN-robot
CANN-robot成员
6月24日 评论:

Notice

This issue can not be assigned to bububu. Please try to assign to the repository members.

likedislike
陈思
陈思成员
6月24日 评论:

/assign @LiJianhao2

likedislike
CANN-robotCANN-robot成员
6月24日 将 LiJianhao2 设为负责人
llimwang成员
7月30日 评论:

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

likedislike
Lllimwang成员
7月30日 关闭了 issue
CANN-robotCANN-robot成员
8月6日 添加了label:resolved