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^start,steps == 0 时输出空张量)。
result[i] = base^(start + i * step)
step = (end - start) / (steps - 1)
steps == 1
base^start
steps == 0
与原 TBE 算子核心功能完全对齐,支持原算子对应的所有数据类型、数据格式;在原 TBE 仅支持 FLOAT/FLOAT16/BFLOAT16 的基础上,特别地需要新增 INT8、UINT8、INT16、INT32 整型输出支持。说明:整型输出与 torch.logspace(...).to(dtype) 一致——先按 float32 计算 base^x,再向零截断(CAST_TRUNC)落整型;当结果超出整型范围时按该整型位宽取模回绕(例如 uint8 等价于 value mod 256、int8 等价于有符号 8 bit 回绕),而不是饱和。 必须实现算子泛化功能,满足各类合法输入场景(任意 start/end/base/steps、大指数溢出、整数幂网格点等)的计算需求,验收阶段将采用泛化数据进行验收。
FLOAT/FLOAT16/BFLOAT16
INT8、UINT8、INT16、INT32
torch.logspace(...).to(dtype)
float32
base^x
uint8
value mod 256
int8
start/end/base/steps
参数解析与校验
start
end
steps
base
result
FLOAT/FLOAT16/BFLOAT16/INT8/UINT8/INT16/INT32
INT8/UINT8/INT16/INT32
tiling 策略
LogSpace
step
logBase = ln(base)
MIN_PER_CORE = 64
UB_CHUNK_ELEMS = 2048
baseN[]
pow
tilingKey 规划策略
tilingKey
(D_T_Y, MODE)
D_T_Y
MODE
0 = NORMAL(steps >= 2)
1 = SINGLE(steps == 0 / 1)
float32/float16/bfloat16
bfloat16
int8/int16/int32/uint8
Cast(CAST_TRUNC)
int8/uint8
max(base^start, base^end) > 2^24
useDfV
kernel 侧实现描述
Init
Process
[coreOffset, coreOffset + coreNum)
base^(start + i * step)
float16/float32/int16/int32
exp((start + i*step) * ln(base))
Cast
k = round(arg/ln2)
r = arg - k*ln2
1/k!
expBase[c] = expBase[c-1] * const
const = base^(chunkElems * step)
mod 256
value >= 2^31
1 - min(1, floor(value/2^31))
SetValue
exp
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["结束"]
classify + variable_shape + auto_schedule
#pragma clang fp contract(off)
volatile
base^(chunkElems*step)
double pow
/assign @bububu
This issue can not be assigned to bububu. Please try to assign to the repository members.
/assign @LiJianhao2
该 Issue 已超过 20 天未更新,先行关闭以便归档;如仍需跟进请评论或重新打开。
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^start,steps == 0时输出空张量)。二、价值/作用 (必填)
与原 TBE 算子核心功能完全对齐,支持原算子对应的所有数据类型、数据格式;在原 TBE 仅支持
FLOAT/FLOAT16/BFLOAT16的基础上,特别地需要新增INT8、UINT8、INT16、INT32整型输出支持。说明:整型输出与torch.logspace(...).to(dtype)一致——先按float32计算base^x,再向零截断(CAST_TRUNC)落整型;当结果超出整型范围时按该整型位宽取模回绕(例如uint8等价于value mod 256、int8等价于有符号 8 bit 回绕),而不是饱和。必须实现算子泛化功能,满足各类合法输入场景(任意
start/end/base/steps、大指数溢出、整数幂网格点等)的计算需求,验收阶段将采用泛化数据进行验收。三、设计方案 (必填)
host 侧设计
参数解析与校验
start(aclScalar→float)、end(aclScalar→float)、steps(int64_t,>= 0)、base(double→float,> 0)。result,必选张量,其 dtype 决定输出类型。result非空。result元素数与steps一致。FLOAT/FLOAT16/BFLOAT16/INT8/UINT8/INT16/INT32,其中INT8/UINT8/INT16/INT32为本次新增支持类型。tiling 策略
LogSpace是无输入、无广播的纯生成逐元素算子,可复用通用逐元素 tiling。steps、输出 dtype size、UB 大小、AIV core 数等基础信息。step与对数底logBase = ln(base),下发到 tiling。MIN_PER_CORE = 64个元素,不足则缩减核数),并在单 core 内按 UB 容量(UB_CHUNK_ELEMS = 2048fp32 元素/次)继续切分 tile。steps == 0走空执行路径,不需要额外 workspace。baseN[],doublepow精确算出)与 double-float 递推常数,随 tiling 下发,供 kernel 修正端点 / 网格点精度(见差异点)。tilingKey 规划策略
tilingKey由(D_T_Y, MODE)二元组下发:D_T_Y为输出 dtype,MODE为0 = NORMAL(steps >= 2)/1 = SINGLE(steps == 0 / 1)。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 侧设计
kernel 侧实现描述
Init和Process两个阶段,其中Process包括计算(Compute)、搬出(CopyOut)。LogSpace 无输入张量,故无 CopyIn。[coreOffset, coreOffset + coreNum)逐 chunk 生成base^(start + i * step)。float16/float32/int16/int32:向量计算exp((start + i*step) * ln(base)),浮点直接落,整型末步Cast(CAST_TRUNC)。bfloat16:以float32中间计算,再Cast回bfloat16。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 256,value >= 2^31时用 keep 掩码1 - min(1, floor(value/2^31))置零(此时 fp32 为 256 的倍数,取模为 0)。pow精确算出的值在 UBfloat32buffer 上SetValue打补丁,消除设备exp对整数幂的舍入误差。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 实现差异点和原因
classify + variable_shape + auto_schedule自动处理动态 shape;AscendC 需要在 host tiling 中显式完成元素数统计、分核和 tile 大小规划。float32/float16/bfloat16;AscendC 新增int8/int16/int32/uint8整型路径,统一以float32计算base^x后Cast(CAST_TRUNC)向零取整,并对int8/uint8实现与torch.logspace(...).to(dtype)一致的位宽回绕(取模)语义。#pragma clang fp contract(off)/volatile屏障均无效)。AscendC 将 double-float 高精度指数放在向量单元计算(向量为 IEEE 单精,忠实匹配 numpyfloat32golden),解决int8/uint8大指数取模的精度问题。base^x的逐 chunk 参考值采用递推(每 chunk 乘常数base^(chunkElems*step)),把一次性昂贵的 df-exp 摊薄到整核一次,兼顾精度与性能。double pow精确写出,规避设备exp对整数幂的舍入误差导致首尾元素偏差。bfloat16在 AscendC 侧使用float32中间计算再转回,避免不同后端对直接 bf16 计算的支持差异。支持硬件
算子约束限制
base必须大于 0;steps必须大于等于 0。torch.logspace(...).to(dtype)一致:float32计算后向零截断,超范围按位宽取模回绕。LogSpace无输入张量、无广播语义,不需要处理多输入 shape 对齐。