拓冰网站定制 
CANN ops-math 算子开发指南Arange 等差序列算子的 aclnn 接口实现与两段式调用实战【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读本文以 CANN ops-math 开源仓库中 experimental/math/arange/README.md 为核心系统讲解实验性数学算子 Arange等差序列生成的接口定义、参数约束、两段式 aclnn 调用流程与底层实现原理。该算子在昇腾 NPU 上实现与 PyTorchtorch.arange语义一致的一维等差序列生成覆盖 FLOAT 至 INT64 共 8 种数据类型。读完本文你将掌握aclnnArangeGetWorkspaceSizeaclnnArange的完整调用套路、调用方前置约束特别是 N 的计算契约以及算子 host 侧 tiling 与 kernel 侧的动态多核切分实现。一、Arange 算子功能与产品支持1.1 功能说明Arange 从start起始、以step为步长、到end结束左闭右开不含end生成一个一维等差序列张量并写入out。其中start、end、step均为 Host 侧标量aclScalarout为一维输出张量aclTensor功能与昇腾内置aclnnArange、PyTorchtorch.arange一致。计算公式如下序列元素满足$$ \text{out}_i \text{start} i \times \text{step}, \quad i 0, 1, \dots, N-1 $$输出元素个数 N左闭右开向上取整$$ N \left\lceil \frac{\text{end} - \text{start}}{\text{step}} \right\rceil $$取整口径为ceil左闭右开与昇腾内置aclnnArange/ PyTorchtorch.arange一致。关键契约out的元素个数 N 由调用方按上式计算并据此分配、构造out张量shape 为[N]算子本身不重新计算或校验 N详见下文调用方前置约束。1.2 产品支持情况产品是否支持Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas A3 训练系列产品/Atlas A3 推理系列产品√产品支持范围在源码中同样有明确体现算子原型通过 arange_def.cpp 中this-AICore().AddConfig(ascend910b).AddConfig(ascend910_93)仅向 A2ascend910b与 A3ascend910_93两个目标平台注册 kernel 变体。从源码结构看其余昇腾产品线如 Ascend 950 系列、Atlas 200I/500 A2 等不在该实验性实现的支持范围内。二、参数说明与数据类型start、end、step、out四个入参的详细说明如下参数名输入/输出/属性描述数据类型数据格式start输入Host 侧的 aclScalar取值范围的起始位置对应公式中的 startFLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16、INT32、INT64NDend输入Host 侧的 aclScalar取值范围的结束位置左闭右开不含 end对应公式中的 endFLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16、INT32、INT64NDstep输入Host 侧的 aclScalar取值的步长对应公式中的 stepFLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16、INT32、INT64NDout输出一维输出张量存放等差序列shape 为 [N]对应公式中的 outFLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16、INT32、INT64ND需要特别强调的两点dtype 一致性start、end、step、out四者的数据类型必须保持一致不做跨数据类型推导。INT32 / INT64 为兼容保留项算子原型实际注册 8 种数据类型FLOAT / FLOAT16 / BFLOAT16 / INT8 / UINT8 / INT16 / INT32 / INT64其中 INT8 / UINT8 / INT16 为必测数据类型INT32 / INT64 为兼容保留项INT32 同时是 examples 与性能对标的主用例。INT32 / INT64 同样走 FP32 中间域计算由于 FP32 尾数仅 24 位当序列值的绝对值超过 2^2416777216时存在精度损失无法精确表示该量级的整数调用方应在此约束内使用或避免对超大值域使用 INT32 / INT64。2.1 原型注册的源码佐证在 arange_def.cpp 中ARANGE_SCALAR_DTYPE_LIST宏定义了 8 种数据类型的注册列表start/end/step三个输入通过.Scalar()标记为标量输入out输出镜像同一 dtype 列表四个操作数的Format与UnknownShapeFormat均为纯 ND。原型层面即保证了四者 dtype 必须一致、格式只支持 ND的约束。三、约束说明3.1 算子约束start、end、step、out四者的数据类型必须保持一致且数据格式只支持 ND。out不支持空 Tensor要求 N ≥ 1。整数类型INT8、UINT8、INT16输出当序列值超出对应类型值域时按硬件 Cast饱和clamp语义处理例如 INT8 越界值截断到 [-128, 127]。该行为已在 NPU 上实测确认调用方应保证序列值落在目标类型值域内以获得与 CPU 标杆一致的结果。确定性算子为纯逐元素等差序列生成out[i] start i*step无 Reduce、无核间累加相同输入恒产生相同输出默认确定性实现。3.2 调用方前置约束值级由调用方保证接口不做值级校验aclnnArange接口的入参校验仅覆盖数据类型白名单 四者一致性与空指针以下值级约束属于调用方前置条件接口不做值级校验。调用方须在调用前自行保证否则行为未定义前置约束调用方须保证step ≠ 0step 非零step 符号匹配step 0 时 start endstep 0 时 start end即 (end - start) 与 step 同号N ≥ 1UINT8 非负out 为 UINT8 时start / end / step 均需为非负且需 step 0、start endUINT8 不可表示负值N 由调用方计算out 的元素个数 N ceil((end - start) / step)由调用方按该公式计算并据此分配、构造 out 张量算子不重新计算或校验 NN ≥ 1不支持空 TensorN ≤ 0 为非法输入四、两段式 aclnn 接口调用说明4.1 函数原型每个算子分为两段式接口必须先调用aclnnArangeGetWorkspaceSize接口获取计算所需 workspace 大小以及包含了算子计算流程的执行器再调用aclnnArange接口执行计算。aclnnStatus aclnnArangeGetWorkspaceSize( const aclScalar *start, const aclScalar *end, const aclScalar *step, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnArange( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)4.2 aclnnArangeGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)startaclScalar*输入Host 侧标量序列起始值对应公式中 startstep 大于 0 时需满足 start 小于 endstep 小于 0 时需满足 start 大于 end数据类型需与 end、step、out 一致FLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16ND-endaclScalar*输入Host 侧标量序列结束值左闭右开不含 end对应公式中 end取值约束同 start数据类型需与 start、step、out 一致FLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16ND-stepaclScalar*输入Host 侧标量步长对应公式中 stepstep 不等于 0数据类型需与 start、end、out 一致FLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16ND-outaclTensor*输出一维输出张量存放等差序列对应公式中 out不支持空 Tensorshape 为一维 [N]Nceil((end-start)/step)由调用方按该公式计算并构造数据类型需与 start、end、step 一致FLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16ND1workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程----说明Atlas A2 / A3 系列产品下start、end、step、out 支持 FLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16四者数据类型须保持一致UINT8 不可表示负值UINT8 场景下 start、end、step 均需为非负且需满足 step 大于 0、start 小于 end。4.3 返回值与错误码第一段接口aclnnArangeGetWorkspaceSize完成入参校验返回aclnnStatus状态码出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001start、end、step、out 存在空指针。ACLNN_ERR_PARAM_INVALID161002start、end、step 或 out 的数据类型不在支持的范围之内。ACLNN_ERR_PARAM_INVALID161002start、end、step、out 的数据类型不一致。ACLNN_ERR_PARAM_INVALID161002step 等于 0或 step 与 (end-start) 的符号关系不满足约束。4.4 aclnnArange 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnArangeGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream五、调用示例test_aclnn_arange.cpp 全流程解析完整可运行示例位于 examples/test_aclnn_arange.cpp可通过以下命令编译并运行bash build.sh --run_example arange eager cust --vendor_namecustom --experimental该命令的详细说明可参考 build.sh 调用说明。示例覆盖 FLOAT 升序、FLOAT 负 step 降序、INT8 窄整型升序以及 FLOAT 非有限值inf / nan传播共五组用例核心流程如下。5.1 关键步骤拆解第 1 步device / stream 初始化固定写法int Init(int32_t deviceId, aclrtStream* stream) { auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return 0; }第 2 步调用方计算 N算子侧不计算、不校验 Nint64_t ComputeN(double start, double end, double step) { return static_castint64_t(std::ceil((end - start) / step)); }第 3 步创建一维连续输出 aclTensor仅分配 device 内存无需拷入初值int CreateOutTensor(int64_t n, size_t elemSize, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size static_castsize_t(n) * elemSize; auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); std::vectorint64_t shape {n}; // 一维 [N] std::vectorint64_t strides {1}; *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); CHECK_RET(*tensor ! nullptr, return -1); return 0; }第 4 步两段式调用核心流程FLOAT 路径// 4.1 构造 start/end/step 三个 Host 侧标量aclScalardtype 须四者一致 aclScalar* sStart aclCreateScalar(start, aclDataType::ACL_FLOAT); aclScalar* sEnd aclCreateScalar(end, aclDataType::ACL_FLOAT); aclScalar* sStep aclCreateScalar(step, aclDataType::ACL_FLOAT); // 4.2 构造一维输出张量 outshape[N]dtype 与标量一致 void* outDeviceAddr nullptr; aclTensor* out nullptr; auto ret CreateOutTensor(n, sizeof(float), outDeviceAddr, aclDataType::ACL_FLOAT, out); // 4.3 第一段获取 workspace 大小与执行器 uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; ret aclnnArangeGetWorkspaceSize(sStart, sEnd, sStep, out, workspaceSize, executor); // 4.4 按需申请 workspace本算子 workspaceSize 通常为 0 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 4.5 第二段执行计算 同步等待 ret aclnnArange(workspaceAddr, workspaceSize, executor, stream); ret aclrtSynchronizeStream(stream); // 4.6 拷回 host 并打印结果 std::vectorfloat result(n); ret aclrtMemcpy(result.data(), result.size() * sizeof(float), outDeviceAddr, static_castsize_t(n) * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); // 4.7 释放标量 / 张量 / device 内存 aclDestroyScalar(sStart); aclDestroyScalar(sEnd); aclDestroyScalar(sStep); aclDestroyTensor(out); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); }5.2 用例清单与预期输出用例路径/参数预期输出case1-FLOAT-ascstart0, end10, step1[0, 1, 2, ..., 9]N10case2-FLOAT-neg-stepstart5, end-5, step-2[5, 3, 1, -1, -3]N5case3-INT8-ascstart-3, end12, step3[-3, 0, 3, 6, 9]N5case4-FLOAT-inf-propagatestartinf, finite step显式 N5全 infinf i*step infIEEE 传播case5-FLOAT-nan-propagatestartnan显式 N5全 nannan i*step nanIEEE 传播用例 4/5 演示了一个重要边界inf/nan 属 README调用方前置约束之外的值级输入接口不做值级校验此类输入下 N 无法由ceil((end-start)/step)稳健推导故由调用方显式给定一个小 N算子按 IEEE 语义逐元素生成out[i]starti*stepFLOAT 走纯 FP32 路径inf/nan 按 IEEE 传播且不崩溃。5.3 其余 dtype 的替换方式float16 用aclDataType::ACL_FLOAT16、bfloat16 用ACL_BF16、uint8 用ACL_UINT8、int16 用ACL_INT16并同步改对应标量/输出元素的 C 类型与sizeofstart/end/step/out四者 dtype 必须一致N ceil((end-start)/step)由调用方保证算子侧不重新校验UINT8 场景start/end/step均须非负且step0、startenduint8 不可表示负值。六、源码级原理从 InferShape 到多核 Kernel 的实现链路6.1 InferShape输出固定一维、维度动态未知arange_infershape.cpp 中InferShapeArange将输出 shape 固定为一维并将dim0置为-1动态未知y_shape-SetDimNum(1); y_shape-SetDim(0, -1);这一实现与N 由调用方计算并构造 out、算子不重新计算或校验 N的契约完全一致——InferShape 阶段不读取 start/end/step 的数值只声明输出为一维动态张量。对应的单测 test_arange_infershape.cpp 明确验证了输出 1 维、dim0-1动态未知不读 start/end/step 数值这一契约并覆盖 fp32/fp16/int32 及全 dtype 场景。6.2 Tilingdtype 分派 多核 former/tail 动态切分arange_tiling.cpp 是 host 侧 tiling 的核心按三段职责拆分1决定 dtype 字节数与 TilingKeyDecideDtypeSizeAndTilingKey按输出 dtype 映射字节数int8/uint81、int16/fp16/bf162、float/int324、int648并设置 tilingkey——仅 DT_FLOAT 走 MODE_1 纯 FP32 直算路径其余 dtype 全部走 MODE_0 Cast 路径即先转入 FP32 中间域计算、再 Cast 回目标类型。这与 kernel 侧 arange.cpp 中schMode ELEMENTWISE_TPL_SCH_MODE_0/MODE_1两个分支一一对应模板参数由 arange_tiling_key.h 声明。2计算单 UB 块元素数CalcUnitNum将 UB 空间 10 等分并按 32B 块对齐且unitNum必须按 FP32 字节统一切而非随 1B/2B dtype 放大。源码注释给出了关键原因Cast 路径有 4 份 FP32 中间 buffercalc_init/step/temp/out各unitNum*sizeof(float)若 int8 下按 1B 放大 unitNum 会使 4 份 FP32 中间缓冲膨胀到约 354KB 而撑爆 184KB 的 UB因此统一按max(dtype_size, sizeof(float))计算保证全 dtype 安全。3多核 former/tail 切分CalcCoreSplitAndFillTiling将总元素数按 32B 块粒度切到平台可用核数GetCoreNum()动态获取禁止写死前formerNum个核各多分 1 个 32B 块former 段与 tail 段因负载不同各算一套 UB 子循环参数formerUnitLoops/formerTailNumvstailUnitLoops/tailTailNum写入 arange_tiling_data.h 定义的ArangeTilingData结构。小 shape 时块数小于核数则只开块数个核、至少 1 核coreNum通过context-SetBlockDim(coreNum)下发且 workspace 大小恒为 0与示例中本算子 workspaceSize 通常为 0一致。6.3 KernelArithProgression 单向量生成 多核区间解析arange.h 中实现两个 Kernel 类KernelArangeFP32 直算仅 FLOAT 使用。work_init中本核首元素叠加coreOffset*stepbaseStart start coreOffset*step随后用单条向量指令ArithProgressionfloat一次生成baseStart i*step替代了传统标量 SetValue 循环造 iota Duplicate/Mul/Add的做法是贡献说明中提到的 ArithProgression 性能优化落地。块间通过calc_temp blockStep步进unitNum*step递推避免逐元素重复计算。KernelArange_CastCast 路径FP16/BF16/INT8/UINT8/INT16/INT32/INT64 统一转入 FP32 中间域计算出口再 Cast 回目标类型。其中有两个值得注意的硬件适配细节一是float→int8/uint8硬件不支持直转必须两段式float→half(CAST_ROUND)→int8/uint8(CAST_ROUND)且 half→int8/uint8 硬件默认饱和——这正是 README 中整数越界按硬件 Cast 饱和语义处理的源码依据二是INT32 走原生整数域计算work_init_int32全程 int32 整数运算 ArithProgressionint32_t无 FP32 中转、无出口 Cast精确到 2^31规避了 FP32 2^24 精度天花板。ParseCoreParams按GetBlockIdx()将每个核分派到 former 或 tail 段参数coreOffset用 int64 计算防大 N 下 uint32 溢出ArangeCopyOutImpl提供末块 OOB 防护段长按 32B 对齐放大后末核名义coreLen可能超过真实剩余元素数故以realNum min(num, totalNum - globalOffset)兜底满 32B 对齐走DataCopy快路径否则用DataCopyPad按真实字节精确写杜绝 1B/2B 窄整型尾轴越界。6.4 多核切分的测试验证tiling 单测 test_arange_tiling.cpp 独立复算CalcUnitLoops算法并断言 host tiling 输出覆盖多核 fp32 大 shape 不可整除multicore_fp32_large_not_divisible与 former 段长度归零multicore_former_zero_iff_length_zero等场景验证了 former/tail 切分的正确性。七、贡献记录贡献者贡献方贡献算子贡献时间贡献内容forge个人贡献者Arange2026-06扩展 INT8/UINT8/INT16 数据类型动态多核 former/tail 切分ArithProgression 等性能优化结语Arange 是 CANN ops-math 实验性数学算子库中一个接口简单、实现精细的代表性算子对外它是标准的 aclnn 两段式接口调用方只需遵守四者 dtype 一致 自行计算 N两条契约即可完成调用对内它涵盖了 dtype 分派FP32 直算 / Cast 路径、动态多核 former/tail 切分、ArithProgression 向量化生成、int32 原生整数域计算与窄整型越界饱和处理等一系列 NPU 算子开发的关键实践。读者可结合 README、接口文档 与 调用示例 三份文件快速在自己的 A2/A3 环境上复现并扩展验证。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-math 算子开发指南：Arange 等差序列算子的 aclnn 接口实现与两段式调用实战
CANN ops-math 算子开发指南Arange 等差序列算子的 aclnn 接口实现与两段式调用实战【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读本文以 CANN ops-math 开源仓库中 experimental/math/arange/README.md 为核心系统讲解实验性数学算子 Arange等差序列生成的接口定义、参数约束、两段式 aclnn 调用流程与底层实现原理。该算子在昇腾 NPU 上实现与 PyTorchtorch.arange语义一致的一维等差序列生成覆盖 FLOAT 至 INT64 共 8 种数据类型。读完本文你将掌握aclnnArangeGetWorkspaceSizeaclnnArange的完整调用套路、调用方前置约束特别是 N 的计算契约以及算子 host 侧 tiling 与 kernel 侧的动态多核切分实现。一、Arange 算子功能与产品支持1.1 功能说明Arange 从start起始、以step为步长、到end结束左闭右开不含end生成一个一维等差序列张量并写入out。其中start、end、step均为 Host 侧标量aclScalarout为一维输出张量aclTensor功能与昇腾内置aclnnArange、PyTorchtorch.arange一致。计算公式如下序列元素满足$$ \text{out}_i \text{start} i \times \text{step}, \quad i 0, 1, \dots, N-1 $$输出元素个数 N左闭右开向上取整$$ N \left\lceil \frac{\text{end} - \text{start}}{\text{step}} \right\rceil $$取整口径为ceil左闭右开与昇腾内置aclnnArange/ PyTorchtorch.arange一致。关键契约out的元素个数 N 由调用方按上式计算并据此分配、构造out张量shape 为[N]算子本身不重新计算或校验 N详见下文调用方前置约束。1.2 产品支持情况产品是否支持Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas A3 训练系列产品/Atlas A3 推理系列产品√产品支持范围在源码中同样有明确体现算子原型通过 arange_def.cpp 中this-AICore().AddConfig(ascend910b).AddConfig(ascend910_93)仅向 A2ascend910b与 A3ascend910_93两个目标平台注册 kernel 变体。从源码结构看其余昇腾产品线如 Ascend 950 系列、Atlas 200I/500 A2 等不在该实验性实现的支持范围内。二、参数说明与数据类型start、end、step、out四个入参的详细说明如下参数名输入/输出/属性描述数据类型数据格式start输入Host 侧的 aclScalar取值范围的起始位置对应公式中的 startFLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16、INT32、INT64NDend输入Host 侧的 aclScalar取值范围的结束位置左闭右开不含 end对应公式中的 endFLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16、INT32、INT64NDstep输入Host 侧的 aclScalar取值的步长对应公式中的 stepFLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16、INT32、INT64NDout输出一维输出张量存放等差序列shape 为 [N]对应公式中的 outFLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16、INT32、INT64ND需要特别强调的两点dtype 一致性start、end、step、out四者的数据类型必须保持一致不做跨数据类型推导。INT32 / INT64 为兼容保留项算子原型实际注册 8 种数据类型FLOAT / FLOAT16 / BFLOAT16 / INT8 / UINT8 / INT16 / INT32 / INT64其中 INT8 / UINT8 / INT16 为必测数据类型INT32 / INT64 为兼容保留项INT32 同时是 examples 与性能对标的主用例。INT32 / INT64 同样走 FP32 中间域计算由于 FP32 尾数仅 24 位当序列值的绝对值超过 2^2416777216时存在精度损失无法精确表示该量级的整数调用方应在此约束内使用或避免对超大值域使用 INT32 / INT64。2.1 原型注册的源码佐证在 arange_def.cpp 中ARANGE_SCALAR_DTYPE_LIST宏定义了 8 种数据类型的注册列表start/end/step三个输入通过.Scalar()标记为标量输入out输出镜像同一 dtype 列表四个操作数的Format与UnknownShapeFormat均为纯 ND。原型层面即保证了四者 dtype 必须一致、格式只支持 ND的约束。三、约束说明3.1 算子约束start、end、step、out四者的数据类型必须保持一致且数据格式只支持 ND。out不支持空 Tensor要求 N ≥ 1。整数类型INT8、UINT8、INT16输出当序列值超出对应类型值域时按硬件 Cast饱和clamp语义处理例如 INT8 越界值截断到 [-128, 127]。该行为已在 NPU 上实测确认调用方应保证序列值落在目标类型值域内以获得与 CPU 标杆一致的结果。确定性算子为纯逐元素等差序列生成out[i] start i*step无 Reduce、无核间累加相同输入恒产生相同输出默认确定性实现。3.2 调用方前置约束值级由调用方保证接口不做值级校验aclnnArange接口的入参校验仅覆盖数据类型白名单 四者一致性与空指针以下值级约束属于调用方前置条件接口不做值级校验。调用方须在调用前自行保证否则行为未定义前置约束调用方须保证step ≠ 0step 非零step 符号匹配step 0 时 start endstep 0 时 start end即 (end - start) 与 step 同号N ≥ 1UINT8 非负out 为 UINT8 时start / end / step 均需为非负且需 step 0、start endUINT8 不可表示负值N 由调用方计算out 的元素个数 N ceil((end - start) / step)由调用方按该公式计算并据此分配、构造 out 张量算子不重新计算或校验 NN ≥ 1不支持空 TensorN ≤ 0 为非法输入四、两段式 aclnn 接口调用说明4.1 函数原型每个算子分为两段式接口必须先调用aclnnArangeGetWorkspaceSize接口获取计算所需 workspace 大小以及包含了算子计算流程的执行器再调用aclnnArange接口执行计算。aclnnStatus aclnnArangeGetWorkspaceSize( const aclScalar *start, const aclScalar *end, const aclScalar *step, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnArange( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)4.2 aclnnArangeGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)startaclScalar*输入Host 侧标量序列起始值对应公式中 startstep 大于 0 时需满足 start 小于 endstep 小于 0 时需满足 start 大于 end数据类型需与 end、step、out 一致FLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16ND-endaclScalar*输入Host 侧标量序列结束值左闭右开不含 end对应公式中 end取值约束同 start数据类型需与 start、step、out 一致FLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16ND-stepaclScalar*输入Host 侧标量步长对应公式中 stepstep 不等于 0数据类型需与 start、end、out 一致FLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16ND-outaclTensor*输出一维输出张量存放等差序列对应公式中 out不支持空 Tensorshape 为一维 [N]Nceil((end-start)/step)由调用方按该公式计算并构造数据类型需与 start、end、step 一致FLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16ND1workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程----说明Atlas A2 / A3 系列产品下start、end、step、out 支持 FLOAT、FLOAT16、BFLOAT16、INT8、UINT8、INT16四者数据类型须保持一致UINT8 不可表示负值UINT8 场景下 start、end、step 均需为非负且需满足 step 大于 0、start 小于 end。4.3 返回值与错误码第一段接口aclnnArangeGetWorkspaceSize完成入参校验返回aclnnStatus状态码出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001start、end、step、out 存在空指针。ACLNN_ERR_PARAM_INVALID161002start、end、step 或 out 的数据类型不在支持的范围之内。ACLNN_ERR_PARAM_INVALID161002start、end、step、out 的数据类型不一致。ACLNN_ERR_PARAM_INVALID161002step 等于 0或 step 与 (end-start) 的符号关系不满足约束。4.4 aclnnArange 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnArangeGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream五、调用示例test_aclnn_arange.cpp 全流程解析完整可运行示例位于 examples/test_aclnn_arange.cpp可通过以下命令编译并运行bash build.sh --run_example arange eager cust --vendor_namecustom --experimental该命令的详细说明可参考 build.sh 调用说明。示例覆盖 FLOAT 升序、FLOAT 负 step 降序、INT8 窄整型升序以及 FLOAT 非有限值inf / nan传播共五组用例核心流程如下。5.1 关键步骤拆解第 1 步device / stream 初始化固定写法int Init(int32_t deviceId, aclrtStream* stream) { auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return 0; }第 2 步调用方计算 N算子侧不计算、不校验 Nint64_t ComputeN(double start, double end, double step) { return static_castint64_t(std::ceil((end - start) / step)); }第 3 步创建一维连续输出 aclTensor仅分配 device 内存无需拷入初值int CreateOutTensor(int64_t n, size_t elemSize, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size static_castsize_t(n) * elemSize; auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); std::vectorint64_t shape {n}; // 一维 [N] std::vectorint64_t strides {1}; *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); CHECK_RET(*tensor ! nullptr, return -1); return 0; }第 4 步两段式调用核心流程FLOAT 路径// 4.1 构造 start/end/step 三个 Host 侧标量aclScalardtype 须四者一致 aclScalar* sStart aclCreateScalar(start, aclDataType::ACL_FLOAT); aclScalar* sEnd aclCreateScalar(end, aclDataType::ACL_FLOAT); aclScalar* sStep aclCreateScalar(step, aclDataType::ACL_FLOAT); // 4.2 构造一维输出张量 outshape[N]dtype 与标量一致 void* outDeviceAddr nullptr; aclTensor* out nullptr; auto ret CreateOutTensor(n, sizeof(float), outDeviceAddr, aclDataType::ACL_FLOAT, out); // 4.3 第一段获取 workspace 大小与执行器 uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; ret aclnnArangeGetWorkspaceSize(sStart, sEnd, sStep, out, workspaceSize, executor); // 4.4 按需申请 workspace本算子 workspaceSize 通常为 0 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 4.5 第二段执行计算 同步等待 ret aclnnArange(workspaceAddr, workspaceSize, executor, stream); ret aclrtSynchronizeStream(stream); // 4.6 拷回 host 并打印结果 std::vectorfloat result(n); ret aclrtMemcpy(result.data(), result.size() * sizeof(float), outDeviceAddr, static_castsize_t(n) * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); // 4.7 释放标量 / 张量 / device 内存 aclDestroyScalar(sStart); aclDestroyScalar(sEnd); aclDestroyScalar(sStep); aclDestroyTensor(out); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); }5.2 用例清单与预期输出用例路径/参数预期输出case1-FLOAT-ascstart0, end10, step1[0, 1, 2, ..., 9]N10case2-FLOAT-neg-stepstart5, end-5, step-2[5, 3, 1, -1, -3]N5case3-INT8-ascstart-3, end12, step3[-3, 0, 3, 6, 9]N5case4-FLOAT-inf-propagatestartinf, finite step显式 N5全 infinf i*step infIEEE 传播case5-FLOAT-nan-propagatestartnan显式 N5全 nannan i*step nanIEEE 传播用例 4/5 演示了一个重要边界inf/nan 属 README调用方前置约束之外的值级输入接口不做值级校验此类输入下 N 无法由ceil((end-start)/step)稳健推导故由调用方显式给定一个小 N算子按 IEEE 语义逐元素生成out[i]starti*stepFLOAT 走纯 FP32 路径inf/nan 按 IEEE 传播且不崩溃。5.3 其余 dtype 的替换方式float16 用aclDataType::ACL_FLOAT16、bfloat16 用ACL_BF16、uint8 用ACL_UINT8、int16 用ACL_INT16并同步改对应标量/输出元素的 C 类型与sizeofstart/end/step/out四者 dtype 必须一致N ceil((end-start)/step)由调用方保证算子侧不重新校验UINT8 场景start/end/step均须非负且step0、startenduint8 不可表示负值。六、源码级原理从 InferShape 到多核 Kernel 的实现链路6.1 InferShape输出固定一维、维度动态未知arange_infershape.cpp 中InferShapeArange将输出 shape 固定为一维并将dim0置为-1动态未知y_shape-SetDimNum(1); y_shape-SetDim(0, -1);这一实现与N 由调用方计算并构造 out、算子不重新计算或校验 N的契约完全一致——InferShape 阶段不读取 start/end/step 的数值只声明输出为一维动态张量。对应的单测 test_arange_infershape.cpp 明确验证了输出 1 维、dim0-1动态未知不读 start/end/step 数值这一契约并覆盖 fp32/fp16/int32 及全 dtype 场景。6.2 Tilingdtype 分派 多核 former/tail 动态切分arange_tiling.cpp 是 host 侧 tiling 的核心按三段职责拆分1决定 dtype 字节数与 TilingKeyDecideDtypeSizeAndTilingKey按输出 dtype 映射字节数int8/uint81、int16/fp16/bf162、float/int324、int648并设置 tilingkey——仅 DT_FLOAT 走 MODE_1 纯 FP32 直算路径其余 dtype 全部走 MODE_0 Cast 路径即先转入 FP32 中间域计算、再 Cast 回目标类型。这与 kernel 侧 arange.cpp 中schMode ELEMENTWISE_TPL_SCH_MODE_0/MODE_1两个分支一一对应模板参数由 arange_tiling_key.h 声明。2计算单 UB 块元素数CalcUnitNum将 UB 空间 10 等分并按 32B 块对齐且unitNum必须按 FP32 字节统一切而非随 1B/2B dtype 放大。源码注释给出了关键原因Cast 路径有 4 份 FP32 中间 buffercalc_init/step/temp/out各unitNum*sizeof(float)若 int8 下按 1B 放大 unitNum 会使 4 份 FP32 中间缓冲膨胀到约 354KB 而撑爆 184KB 的 UB因此统一按max(dtype_size, sizeof(float))计算保证全 dtype 安全。3多核 former/tail 切分CalcCoreSplitAndFillTiling将总元素数按 32B 块粒度切到平台可用核数GetCoreNum()动态获取禁止写死前formerNum个核各多分 1 个 32B 块former 段与 tail 段因负载不同各算一套 UB 子循环参数formerUnitLoops/formerTailNumvstailUnitLoops/tailTailNum写入 arange_tiling_data.h 定义的ArangeTilingData结构。小 shape 时块数小于核数则只开块数个核、至少 1 核coreNum通过context-SetBlockDim(coreNum)下发且 workspace 大小恒为 0与示例中本算子 workspaceSize 通常为 0一致。6.3 KernelArithProgression 单向量生成 多核区间解析arange.h 中实现两个 Kernel 类KernelArangeFP32 直算仅 FLOAT 使用。work_init中本核首元素叠加coreOffset*stepbaseStart start coreOffset*step随后用单条向量指令ArithProgressionfloat一次生成baseStart i*step替代了传统标量 SetValue 循环造 iota Duplicate/Mul/Add的做法是贡献说明中提到的 ArithProgression 性能优化落地。块间通过calc_temp blockStep步进unitNum*step递推避免逐元素重复计算。KernelArange_CastCast 路径FP16/BF16/INT8/UINT8/INT16/INT32/INT64 统一转入 FP32 中间域计算出口再 Cast 回目标类型。其中有两个值得注意的硬件适配细节一是float→int8/uint8硬件不支持直转必须两段式float→half(CAST_ROUND)→int8/uint8(CAST_ROUND)且 half→int8/uint8 硬件默认饱和——这正是 README 中整数越界按硬件 Cast 饱和语义处理的源码依据二是INT32 走原生整数域计算work_init_int32全程 int32 整数运算 ArithProgressionint32_t无 FP32 中转、无出口 Cast精确到 2^31规避了 FP32 2^24 精度天花板。ParseCoreParams按GetBlockIdx()将每个核分派到 former 或 tail 段参数coreOffset用 int64 计算防大 N 下 uint32 溢出ArangeCopyOutImpl提供末块 OOB 防护段长按 32B 对齐放大后末核名义coreLen可能超过真实剩余元素数故以realNum min(num, totalNum - globalOffset)兜底满 32B 对齐走DataCopy快路径否则用DataCopyPad按真实字节精确写杜绝 1B/2B 窄整型尾轴越界。6.4 多核切分的测试验证tiling 单测 test_arange_tiling.cpp 独立复算CalcUnitLoops算法并断言 host tiling 输出覆盖多核 fp32 大 shape 不可整除multicore_fp32_large_not_divisible与 former 段长度归零multicore_former_zero_iff_length_zero等场景验证了 former/tail 切分的正确性。七、贡献记录贡献者贡献方贡献算子贡献时间贡献内容forge个人贡献者Arange2026-06扩展 INT8/UINT8/INT16 数据类型动态多核 former/tail 切分ArithProgression 等性能优化结语Arange 是 CANN ops-math 实验性数学算子库中一个接口简单、实现精细的代表性算子对外它是标准的 aclnn 两段式接口调用方只需遵守四者 dtype 一致 自行计算 N两条契约即可完成调用对内它涵盖了 dtype 分派FP32 直算 / Cast 路径、动态多核 former/tail 切分、ArithProgression 向量化生成、int32 原生整数域计算与窄整型越界饱和处理等一系列 NPU 算子开发的关键实践。读者可结合 README、接口文档 与 调用示例 三份文件快速在自己的 A2/A3 环境上复现并扩展验证。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考