企业官网建设流程全解析

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 种数据类型。读完本文，你将掌握aclnnArangeGetWorkspaceSize+aclnnArange的完整调用套路、调用方前置约束（特别是 N 的计算契约），以及算子 host 侧 tiling 与 kernel 侧的动态多核切分实现。

一、Arange 算子功能与产品支持

1.1 功能说明

Arange 从start起始、以step为步长、到end结束（左闭右开，不含end），生成一个一维等差序列张量并写入out。其中start、end、step均为 Host 侧标量（aclScalar），out为一维输出张量（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 产品支持情况

产品支持范围在源码中同样有明确体现：算子原型通过 arange_def.cpp 中this->AICore().AddConfig("ascend910b").AddConfig("ascend910_93")仅向 A2（ascend910b）与 A3（ascend910_93）两个目标平台注册 kernel 变体。从源码结构看，其余昇腾产品线（如 Ascend 950 系列、Atlas 200I/500 A2 等）不在该实验性实现的支持范围内。

二、参数说明与数据类型

start、end、step、out四个入参的详细说明如下：

需要特别强调的两点：

1. dtype 一致性：start、end、step、out四者的数据类型必须保持一致，不做跨数据类型推导。
2. INT32 / INT64 为兼容保留项：算子原型实际注册 8 种数据类型（FLOAT / FLOAT16 / BFLOAT16 / INT8 / UINT8 / INT16 / INT32 / INT64），其中 INT8 / UINT8 / INT16 为必测数据类型，INT32 / INT64 为兼容保留项（INT32 同时是 examples 与性能对标的主用例）。INT32 / INT64 同样走 FP32 中间域计算，由于 FP32 尾数仅 24 位，当序列值的绝对值超过 2^24（16777216）时存在精度损失（无法精确表示该量级的整数），调用方应在此约束内使用，或避免对超大值域使用 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接口的入参校验仅覆盖数据类型（白名单 + 四者一致性）与空指针；以下值级约束属于调用方前置条件，接口不做值级校验。调用方须在调用前自行保证，否则行为未定义：

四、两段式 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 参数说明

说明：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状态码，出现以下场景时报错：

4.4 aclnnArange 参数说明

五、调用示例：test_aclnn_arange.cpp 全流程解析

完整可运行示例位于 examples/test_aclnn_arange.cpp，可通过以下命令编译并运行：

bash build.sh --run_example arange eager cust --vendor_name=custom --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（算子侧不计算、不校验 N）

int64_t ComputeN(double start, double end, double step) { return static_cast<int64_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_cast<size_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::vector<int64_t> shape = {n}; // 一维 [N] std::vector<int64_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 侧标量（aclScalar），dtype 须四者一致 aclScalar* sStart = aclCreateScalar(&start, aclDataType::ACL_FLOAT); aclScalar* sEnd = aclCreateScalar(&end, aclDataType::ACL_FLOAT); aclScalar* sStep = aclCreateScalar(&step, aclDataType::ACL_FLOAT); // 4.2 构造一维输出张量 out（shape=[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::vector<float> result(n); ret = aclrtMemcpy(result.data(), result.size() * sizeof(float), outDeviceAddr, static_cast<size_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 用例清单与预期输出

用例 4/5 演示了一个重要边界：inf/nan 属 README"调用方前置约束"之外的值级输入（接口不做值级校验），此类输入下 N 无法由ceil((end-start)/step)稳健推导，故由调用方显式给定一个小 N；算子按 IEEE 语义逐元素生成out[i]=start+i*step（FLOAT 走纯 FP32 路径），inf/nan 按 IEEE 传播且不崩溃。

5.3 其余 dtype 的替换方式

• float16 用aclDataType::ACL_FLOAT16、bfloat16 用ACL_BF16、uint8 用ACL_UINT8、int16 用ACL_INT16，并同步改对应标量/输出元素的 C++ 类型与sizeof；
• start/end/step/out四者 dtype 必须一致；N = ceil((end-start)/step)由调用方保证，算子侧不重新校验；
• UINT8 场景start/end/step均须非负且step>0、start<end（uint8 不可表示负值）。

六、源码级原理：从 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 Tiling：dtype 分派 + 多核 former/tail 动态切分

arange_tiling.cpp 是 host 侧 tiling 的核心，按三段职责拆分：

（1）决定 dtype 字节数与 TilingKey：DecideDtypeSizeAndTilingKey按输出 dtype 映射字节数（int8/uint8=1、int16/fp16/bf16=2、float/int32=4、int64=8），并设置 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 中间 buffer（calc_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 Kernel：ArithProgression 单向量生成 + 多核区间解析

arange.h 中实现两个 Kernel 类：

• KernelArange（FP32 直算）：仅 FLOAT 使用。work_init中本核首元素叠加coreOffset*step（baseStart = start + coreOffset*step），随后用单条向量指令ArithProgression<float>一次生成baseStart + i*step，替代了传统"标量 SetValue 循环造 iota + Duplicate/Mul/Add"的做法，是贡献说明中提到的 ArithProgression 性能优化落地。块间通过calc_temp += blockStep（步进unitNum*step）递推，避免逐元素重复计算。
• KernelArange_Cast（Cast 路径）：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 整数运算 +ArithProgression<int32_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 切分的正确性。

七、贡献记录

结语

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