asc-devkit 中 Arange 高阶 API 实战：基于 SIMD C++ 算子实现等差数列生成

导读

本文以 asc-devkit 开源仓库中的 Arange 示例 为完整骨架，讲解如何使用昇腾 Ascend C 提供的 Arange 高阶 API 在片上（UB）存储中生成指定起始值、公差与长度的等差数列，并通过 DataCopy 将结果搬运至 Global Memory。读者将掌握该示例的目录组织、输入输出规格、Kernel 与 Tiling 的编码方式、三种运行模式（NPU 运行 / CPU 调试 / NPU 仿真）的构建命令，以及 Arange 接口的参数约束、底层实现原理与测试验证路径，可直接基于该样例模板迁移到自己的算子开发场景。

示例概述

该示例的完整代码位于 examples/01_simd_cpp_api/04_advanced_api/06_index/arange，核心功能为：给定起始值（firstValue）、等差值（diffValue）和长度（count），生成一个等差数列。

具体实现路径为：将起始值与等差值作为标量参数输入，调用 AscendC::Arange<T> 高阶 API，在片上存储（Unified Buffer）中生成指定长度的等差数列，最后将结果通过 DataCopy 搬运至外部存储（Global Memory）。

对应的 API 文档位于 docs/zh/api/SIMD-API/adv_api/index_compute/Arange.md，其中明确指出该接口的功能为：给定起始值、等差值和长度，返回一个等差数列。

支持的产品与 CANN 软件版本

产品	CANN 软件版本
Ascend 950PR&950DT 系列产品	>= CANN 9.1.0
Atlas A3 系列产品	>= CANN 9.0.0
Atlas A2 系列产品	>= CANN 9.0.0

对应到编译选项 CMAKE_ASC_ARCHITECTURES：dav-2201 对应 Atlas A2 与 Atlas A3 系列产品，dav-3510 对应 Ascend 950PR&950DT 系列产品（详见下文"编译选项说明"）。

目录结构

├── arange
│   ├── scripts
│   │   └── gen_data.py         // 输入数据和真值数据（golden）生成脚本
│   ├── CMakeLists.txt          // 编译工程文件
│   ├── data_utils.h            // 数据读入/写出函数（ReadFile / WriteFile）
│   ├── arange.asc              // Ascend C 样例实现 & 调用样例（Kernel + Tiling + main）
│   └── README.md / README_en.md // 样例说明文档

其中 arange.asc 是核心文件，它在一个文件中同时承载了 Tiling 定义、Kernel 实现、host 端设备管理与内核调用逻辑；gen_data.py 负责生成输入数据与真值数据；data_utils.h 提供二进制文件的读写工具函数。

样例输入输出规格

项	内容
样例类型（OpType）	arange
输入 firstGm	shape [1, 1]，数据类型 float，格式 ND
输入 diffGm	shape [1, 1]，数据类型 float，格式 ND
输出 outputGm	shape [104, 1]，数据类型 float，格式 ND
核函数名	arange_custom

说明：本样例为固定 shape 实现，输入为 firstGm<a href="https://link.gitcode.com/i/85c46a52bed70e59bf3e58a637bd55f7" target="_blank">1, 1]、diffGm[1, 1]，输出为 outputGm[104, 1]。输出长度为 104 而非 100，是因为 100 个 float 元素（400 字节）需按 32 字节向上对齐：(100 + 8 - 1) / 8 * 8 = 104。这一对齐逻辑同时体现在 [arange.asc 与 gen_data.py 中。

Arange 接口声明与算法实现原理

接口原型

Arange 是位于 AscendC 命名空间下的模板函数，对外接口声明在 include/adv_api/index/arithprogression.h：

template <typename T>
__aicore__ inline __in_pipe__(S) __out_pipe__(V, S) void Arange(
    const LocalTensor<T>& dst, const T firstValue, const T diffValue, const int32_t count)

其数学含义为：dst[i+1] = dst[i] + diffValue，即相邻元素之差恒定为 diffValue。

参数说明

参数	输入/输出	描述
dst	输出	目的操作数，类型为 LocalTensor<T>，支持 TPosition 为 VECIN/VECCALC/VECOUT；其大小应大于等于 count * sizeof(T)
firstValue	输入	等差数列的首个元素值
diffValue	输入	等差数列的公差，应大于等于 0
count	输入	等差数列的长度，要求 count > 0

支持的数据类型与格式约束

根据接口头文件与 API 文档：

• 格式约束：当前仅支持 ND 格式输入，不支持其他格式；
• Ascend 950PR&950DT 系列产品：支持 int16_t、half、int32_t、float、int64_t；
• Atlas A2 / Atlas A3 系列产品：支持 int16_t、half、int32_t、float。

在实现层的 static_assert（见 arithprogression_3510_impl.h）中同样固化了对 int16_t, int32_t, half, float, int64_t 的类型检查。

底层算法流程（Vector 流水）

从 docs/zh/api/SIMD-API/adv_api/index_compute/Arange.md 的算法说明可以了解到，Arange 的整体计算均在 Vector 单元上完成，采用"先构造基础段、再逐步倍增扩充"的策略，其核心步骤如下：

1. 长度 8 以内：按 firstValue 与 diffValue 的值，使用 SetValue 实现基础等差数列的构造，扩充长度最大为 8；若目标长度小于 8，算法结束；
2. 长度 8～64：对第 1 步的结果使用 Adds（标量加法）进行扩充，最多循环 7 次扩充至 64；若目标长度小于 64，算法结束；
3. 长度 64 以上：对第 2 步的结果继续使用 Adds 循环扩充，直至达到目标长度。

这一"倍增扩充"思路可以在 arithprogression_common_impl.h 的 ArithProgressionImpl 实现中得到源码级印证：当 count > BLOCK_NUM（BLOCK_NUM = 32 / sizeof(T)）时，先由 GetBaseArithProgression 构造一个 BLOCK_NUM 长度的基础等差数列，再通过 SetVectorMask + Adds 将其广播扩充到 DEFAULT_BLK_NUM 个块（1 block → 8 blocks），最后以 REPEAT_NUM 为步长逐段填充，尾部不足一段时用掩码（tail_aligned）处理；当长度小于 BLOCK_NUM 时，则直接通过标量运算（SetValue）逐元素生成。

对于 Ascend 950 等 3510 架构，arithprogression_3510_impl.h 提供了基于寄存器（Reg）的向量化实现 VfCallArithProgression：先通过 Reg::Arange + Reg::Muls + Reg::Adds 生成基础序列，再用 Reg::StoreAlign 与 Reg::Add 循环写回 UB，repeatTimes 根据 count 与单次重复可写字节数（ONE_REPEAT_BYTE_SIZE）计算。

Arange 算法框图

参数校验与临时空间

• 参数校验：CPU 调试模式（ASCENDC_CPU_DEBUG）下会启用 arithprogression_check_common.h 中的 CheckArithProgressionParamsClass，对 count > 0、diffValue >= 0、数据类型为 half/float/int16_t/int32_t、dst 位于 VECIN/VECOUT/VECCALC 等条件进行断言校验；
• 临时空间：Arange 无需临时空间，GetArangeMaxMinTmpSize 将最大/最小临时空间均返回 0（见 arithprogression_tiling_impl.cpp），tiling 参数中只需传递输出长度。

样例代码深度解析（arange.asc）

Tiling 定义与生成

arange.asc 使用 Tiling 宏定义输出长度字段，并实现 tiling 数据的生成与序列化：

namespace optiling {
BEGIN_TILING_DATA_DEF(ArithprogressionCustomTilingData)
TILING_DATA_FIELD_DEF(uint32_t, outSize);
END_TILING_DATA_DEF;

REGISTER_TILING_DATA_CLASS(ArithprogressionCustom, ArithprogressionCustomTilingData)
} // namespace optiling

void ComputeTiling(const uint32_t outSize, optiling::ArithprogressionCustomTilingData& tiling)
{
    tiling.set_outSize(outSize);
}

随后 GenerateTiling 将 tiling 数据序列化到缓冲区中，host 端通过 aclrtMemcpy 将其拷贝到设备侧，作为 arange_custom<<<>>> 的 tiling 参数传入。

Kernel 实现（KernelArange 类）

核心计算类 KernelArange<T> 实现了标准的 Init → Process（CopyIn → Compute → CopyOut） 三段式流程：

• Init：接收 Global Memory 输出地址、起始值、公差与长度，通过 SetGlobalBuffer 绑定全局输出 Tensor，并按 32 字节对齐（BYTES_ALIGN_SIZE = 32）初始化输出队列 outDst 的片上缓冲区；
• Compute：调用 AscendC::Arange<T>(dstLocal, static_cast<T>(firstValue_), static_cast<T>(diffValue_), count_) 生成等差数列，这是整个示例的核心一行：

  AscendC::Arange<T>(dstLocal, static_cast<T>(firstValue_), static_cast<T>(diffValue_), count_);
  outDst.EnQue<T>(dstLocal);
  

• CopyOut：DataCopy(dst_global, dstLocal, ...) 将 UB 中的结果按对齐后的块数搬运回 Global Memory。

核函数入口 arange_custom 使用 __global__ __vector__ 声明（Vector 核函数），从 tiling 参数中恢复 VecTiling{outSize}，读取 Global Memory 中的 first 与 diff 标量，实例化 KernelArange<float> 并执行 Init + Process。

Host 端调用

main 函数遵循标准的 ACL 流程：

1. aclInit / aclrtSetDevice / aclrtCreateContext / aclrtCreateStream 完成运行环境初始化；
2. aclrtMallocHost 与 aclrtMalloc 分别申请 host/device 内存（ACL_MEM_MALLOC_HUGE_FIRST）；
3. 通过 ReadFile 读取 ./input/first.bin、./input/diff.bin，经 aclrtMemcpy（ACL_MEMCPY_HOST_TO_DEVICE）拷入设备；
4. 使用内核调用符 <<<numBlocks, 0, stream>>> 启动核函数 arange_custom；
5. aclrtSynchronizeStream 同步后，将结果拷回 host，WriteFile 写出 ./output/output.bin；
6. 与 golden 数据对比并释放资源。

其中 numBlocks = 1，count = 100，workspaceSize = 1024 * sizeof(float)。

测试数据生成脚本

gen_data.py 使用 numpy 生成输入与真值：

start = np.float32(1.234)
difference = np.float32(5.678)
length = np.uint32((100 + 8 - 1) // 8 * 8)   # 100 个 float32 按 32 字节对齐 -> 104 个元素
stop = start + difference * (length - 1)
golden = np.linspace(start, stop, length, dtype=np.float32)

start.tofile("./input/first.bin")
difference.tofile("./input/diff.bin")
length.tofile("./input/length.bin")
golden.tofile("./output/golden.bin")

即起始值 1.234、公差 5.678、对齐后长度 104，golden 数据由 np.linspace 生成并写入 ./output/golden.bin，供核函数运行后的精度对比使用。

构建与运行

在示例根目录（examples/01_simd_cpp_api/04_advanced_api/06_index/arange）下按以下步骤操作。

1. 配置环境变量

请根据当前环境上 CANN 开发套件包的安装方式配置环境变量：

source ${install_path}/cann/set_env.sh

说明：${install_path} 为 CANN 包安装目录，未指定安装目录时默认安装至 /usr/local/Ascend 下。

2. NPU 模式运行

mkdir -p build && cd build;      # 创建并进入 build 目录
cmake .. -DCMAKE_ASC_ARCHITECTURES=dav-2201;make -j;   # 编译工程
python3 ../scripts/gen_data.py   # 生成测试输入数据
./demo                           # 执行编译生成的可执行程序

3. CPU 调试 / NPU 仿真模式

使用 CPU 调试或 NPU 仿真模式时，添加 -DCMAKE_ASC_RUN_MODE=cpu 或 -DCMAKE_ASC_RUN_MODE=sim 参数即可：

cmake -DCMAKE_ASC_RUN_MODE=cpu -DCMAKE_ASC_ARCHITECTURES=dav-2201 ..;make -j; # CPU 调试模式
cmake -DCMAKE_ASC_RUN_MODE=sim -DCMAKE_ASC_ARCHITECTURES=dav-2201 ..;make -j; # NPU 仿真模式

注意：切换编译模式前需清理 cmake 缓存，可在 build 目录下执行 rm CMakeCache.txt 后重新 cmake。

4. 编译选项说明

选项	可选值	说明
CMAKE_ASC_RUN_MODE	npu（默认）、cpu、sim	运行模式：NPU 运行、CPU 调试、NPU 仿真
CMAKE_ASC_ARCHITECTURES	dav-2201（默认）、dav-3510	NPU 架构：dav-2201 对应 Atlas A2 系列产品和 Atlas A3 系列产品，dav-3510 对应 Ascend 950PR&950DT 系列产品

对应的 CMakeLists.txt 中，工程通过 find_package(ASC REQUIRED) 引入 ASC 编译工具链，project(... LANGUAGES ASC CXX) 声明 Ascend C 语言，add_executable(demo arange.asc) 编译生成可执行文件，并通过 target_compile_options 将 --npu-arch=${CMAKE_ASC_ARCHITECTURES} 传入 ASC 编译器；链接库包括 tiling_api、register、platform、unified_dlog、dl、graph_base。

5. 执行结果

运行完成后输出如下，说明精度对比成功：

test pass!

精度对比逻辑位于 arange.asc 的 CompareResult：逐元素比较输出与 golden.bin，绝对误差（ae）与相对误差（re）均不超过 EPS = 1e-5 时判定通过，并打印 test pass!。

测试用例验证

仓库在 tests/api/adv_api/index/arithprogression/test_operator_arange.cpp 中提供了对 Arange 高阶 API 的单元测试。测试覆盖三种典型的 count 规模（15、64、260，对应长度小于 8、8～64、大于 64 三个算法分支），并验证：

• Kernel 侧通过 TQue<TPosition::VECOUT, 1> 分配输出队列，Arange<dataType>(dstLocal, ...) 生成序列后经 DataCopy 搬出；
• 输出缓冲区在测试前后保持为 0x00，即 Arange 只写入了目标区间、不产生越界写。

这为上述"倍增扩充"三段式算法流程（SetValue 构造基础段 → Adds 扩充至 64 → 循环扩充至目标长度）提供了测试层面的佐证。

小结

本示例是理解 asc-devkit 中高阶索引类 API 的入门范本：它以最精简的形态展示了从 Tiling 定义、Kernel 三段式编程、Arange 单行调用、DataCopy 搬出，到 ACL host 编排与 golden 对比的完整闭环。在此基础上，开发者可以：

• 更换 T 模板参数（half、int16_t、int32_t 等）适配不同数据类型；
• 通过修改 count 与对齐逻辑适配不同输出长度；
• 参考 examples/01_simd_cpp_api/04_advanced_api/06_index/arange 的目录与工程结构，将 Arange 应用于 mask、索引构造等需要预先生成等差数列作为索引/步长序列的算子场景（如配合 Gather、Scatter 或 CreateVecIndex 类操作）。