昇腾 Ascend C Arange 索引计算算子实战：基于 asc-devkit 高阶 API 的等差数列生成与端到端实现解析

导读

本文围绕 asc-devkit 仓库中 索引计算算子样例集 的核心样例——Arange 算子展开，系统讲解如何基于 Kernel 直调工程，在昇腾 AI 处理器上使用 Arange 高阶 API 生成等差数列。读完本文，你将掌握：Arange 接口的语义与底层实现原理、Kernel 直调样例的工程结构（含 tiling 传递与 host 侧调用链）、完整的编译运行流程（NPU 运行 / CPU 调试 / NPU 仿真三种模式），以及精度校验方法，可直接在 Atlas A2/A3 训练推理系列产品与 Ascend 950 系列上复现该样例。

概述：索引计算算子与 Arange 高阶 API

索引计算（Index Computation）是算子开发中负责生成索引序列、坐标网格等基础数据的算子类别，其计算结果通常作为后续 gather、scatter、坐标变换等算子的输入。asc-devkit 的 索引计算算子样例集 专门收录了这类算子的典型用法，并给出端到端实现。

其中 Arange 样例基于 Kernel 直调算子工程，演示了如何使用 Arange 高阶 API 实现"给定起始值、等差值和长度，返回一个等差数列"的单算子功能。该功能是昇腾 Ascend C 语言（SIMD 向量编程模型）中最基础、最具代表性的索引生成场景之一，也是理解高阶 API 如何封装底层向量/标量指令的绝佳入口。

样例规格与支持的产品

输入输出规格

Arange 样例采用固定 shape 设计，核心规格如下表所示：

项目	名称	shape	数据类型	format
输入	firstGm	[1，1]	float	ND
输入	diffGm	[1，1]	float	ND
输出	outputGm	[104，1]	float	ND
核函数名	arange_custom	-	-	-

其中 firstGm 为等差数列的起始值，diffGm 为等差值（步长），输出 outputGm 为长度 104 的等差数列。值得说明的是：样例逻辑上生成 100 个 float 元素，但由于昇腾向量搬运要求 32 字节对齐，实际按 (100 + 8 - 1) / 8 * 8 = 104 向上对齐，这也是输出 shape 为 [104，1] 的原因（详见下文"数据生成与精度校验"）。

支持的产品及 CANN 版本

产品	CANN 软件版本
Ascend 950PR / Ascend 950DT	>= CANN 9.1.0
Atlas A3 训练系列产品 / Atlas A3 推理系列产品	>= CANN 9.0.0
Atlas A2 训练系列产品 / Atlas A2 推理系列产品	>= CANN 9.0.0

目录结构与工程文件

样例完整源码位于 examples/01_simd_cpp_api/04_advanced_api/06_index/arange 目录，结构如下：

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

各文件职责明确：

• arange.asc：样例核心文件，同时包含设备侧核函数实现（arange_custom）与 host 侧调用、校验逻辑；
• CMakeLists.txt：基于 find_package(ASC REQUIRED) 的编译工程，通过 add_executable(demo arange.asc) 生成 demo 可执行程序；
• gen_data.py：使用 NumPy 生成输入 first.bin、diff.bin 与真值 golden.bin；
• data_utils.h：提供二进制文件的 ReadFile / WriteFile 读写封装，供 host 侧搬运数据与落盘使用。

核心实现解析：从核函数到 host 调用链

设备侧：KernelArange 类与 Arange 高阶 API 调用

核函数主体定义在 arange.asc 的 KernelArange<T> 模板类中，采用标准的 Init → Process（CopyIn → Compute → CopyOut）三段式结构：

template <typename T>
class KernelArange {
public:
    __aicore__ inline void Init(
        __gm__ uint8_t* dstGm, float firstValue, float diffValue, uint32_t count, AscendC::TPipe* pipeIn)
    {
        pipe = pipeIn;
        firstValue_ = firstValue;
        diffValue_ = diffValue;
        count_ = count;
        dst_global.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(dstGm), count_);
        pipe->InitBuffer(outDst, 1, (sizeof(T) * count_ + BYTES_ALIGN_SIZE - 1) / BYTES_ALIGN_SIZE * BYTES_ALIGN_SIZE);
    }
    __aicore__ inline void Process()
    {
        CopyIn();
        Compute();
        CopyOut();
    }
    // ...
    __aicore__ inline void Compute()
    {
        AscendC::LocalTensor<T> dstLocal = outDst.AllocTensor<T>();
        // dstLocal: 输出 Tensor，存储生成的等差数列
        // 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);
    }
    __aicore__ inline void CopyOut()
    {
        AscendC::LocalTensor<T> dstLocal = outDst.DeQue<T>();
        const int32_t BLOCK_NUM = BYTES_ALIGN_SIZE / sizeof(T);
        AscendC::DataCopy(dst_global, dstLocal, (count_ + BLOCK_NUM - 1) / BLOCK_NUM * BLOCK_NUM);
        outDst.FreeTensor(dstLocal);
    }
    // ...
    AscendC::TPipe* pipe;
    AscendC::TQue<AscendC::TPosition::VECOUT, 1> outDst;
    AscendC::GlobalTensor<T> dst_global;
};

几个关键点值得展开：

1. 队列管理：通过 TPipe::InitBuffer 为 VECOUT 队列 outDst 申请片上缓冲区，缓冲区大小按 32 字节（BYTES_ALIGN_SIZE）对齐计算；AllocTensor / EnQue / DeQue / FreeTensor 构成完整的"申请—入队—出队—释放"流水线，与昇腾向量单元的执行节奏配合，实现片上存储复用。
2. Compute 即核心计算：仅一行 AscendC::Arange<T>(dstLocal, firstValue, diffValue, count_) 即完成整个等差数列生成，这正是高阶 API 的价值——把复杂的底层指令序列封装为语义清晰的单接口调用。
3. CopyOut 与对齐搬运：DataCopy 从 LocalTensor 搬回 GlobalTensor 时，搬运块数同样按 32 字节对齐向上取整（BLOCK_NUM = 32 / sizeof(T)），与 Init 阶段的缓冲区对齐策略保持一致。

核函数入口与 tiling 传递

核函数入口 arange_custom 通过 __global__ __vector__ 声明为 SIMD 向量核函数，其签名遵循 Kernel 直调的标准五参形式（输入、输出、workspace、tiling）：

__global__ __vector__ void arange_custom(
    __gm__ uint8_t* firstGm, __gm__ uint8_t* diffGm, __gm__ uint8_t* dstGm, __gm__ uint8_t* workspace,
    __gm__ uint8_t* tiling)
{
    AscendC::TPipe pipe;
    MyCustomKernel::VecTiling tilingData;
    CopyTiling(&tilingData, tiling);
    auto first = reinterpret_cast<__gm__ float*>(firstGm);
    auto diff = reinterpret_cast<__gm__ float*>(diffGm);
    MyCustomKernel::KernelArange<float> op;
    op.Init(dstGm, *first, *diff, tilingData.outSize, &pipe);
    op.Process();
}

tiling（即算子切分/尺寸信息）的生成与传递链路清晰：

• host 侧生成：使用 optiling 命名空间下的 tiling 框架（BEGIN_TILING_DATA_DEF / TILING_DATA_FIELD_DEF(uint32_t, outSize) / REGISTER_TILING_DATA_CLASS）定义 ArithprogressionCustomTilingData，GenerateTiling(outSize) 将 tiling 数据序列化到缓冲区；
• 设备侧接收：CopyTiling 将 tiling GM 内存按 uint32_t 逐个拷入核函数局部结构体 VecTiling，取出 outSize 作为等差数列长度传入 KernelArange::Init。

值得一提的是，Arange 算子无需任何临时空间。这一点可从仓库源码得到印证：tiling 接口 arithprogression_tiling.h 声明了 GetArangeMaxMinTmpSize(uint32_t& maxValue, uint32_t& minValue)，而其实现 arithprogression_tiling_impl.cpp 中 maxValue 与 minValue 均置 0。因此样例中 workspace 仅按 1024 字节申请、实际不参与计算。

host 侧：ACL 流程与内核调用符

host 侧主函数 arange.asc 完整演示了 Kernel 直调工程的 ACL（Ascend Computing Language）运行时流程：

aclInit(nullptr);
aclrtContext context;
aclrtSetDevice(deviceId);
aclrtCreateContext(&context, deviceId);
aclrtCreateStream(&stream);
// 申请 host/device 内存（aclrtMallocHost / aclrtMalloc）
// 读入输入：ReadFile("./input/first.bin", ...) 与 ReadFile("./input/diff.bin", ...)
// 搬运：aclrtMemcpy(...ACL_MEMCPY_HOST_TO_DEVICE)
// 生成并搬运 tiling：aclrtMemcpy(tilingDevice, ..., GenerateTiling(count_aligned), ...)
arange_custom<<<numBlocks, 0, stream>>>(firstDevice, diffDevice, outputDevice, workspaceDevice, tilingDevice);
aclrtSynchronizeStream(stream);
aclrtMemcpy(outputHost, outSize, outputDevice, outSize, ACL_MEMCPY_DEVICE_TO_HOST);
WriteFile("./output/output.bin", outputHost, outSize);
bool goldenResult = CompareResult(outputHost, outSize);
// ... 释放资源：aclrtFree / aclrtFreeHost / aclrtDestroyStream / aclrtDestroyContext / aclrtResetDevice / aclFinalize

调用链要点：

1. 使用内核调用符 <<<numBlocks, 0, stream>>> 以单 block（NUM_BLOCKS = 1）方式启动核函数；
2. 输入起始值、等差值分别从 first.bin、diff.bin 以标量形式读入并拷入设备端（scalarSize = sizeof(float)）；
3. 输出按对齐后的 outSize（(sizeof(float) * 100 + 32 - 1) / 32 * 32 字节）申请、回拷并落盘 output.bin；
4. 结束时按对称顺序销毁 stream、context 并复位设备，保证资源正确释放。

底层原理：Arange 接口语义与寄存器级实现

接口定义与语义

Arange 高阶 API 的公开接口定义在 include/adv_api/index/arithprogression.h，其数学语义为：

dst[i+1] = dst[i] + diffValue

即从 firstValue 开始，每相邻两个元素相差 diffValue，共 count 个元素。接口签名与约束如下：

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 为输出 LocalTensor（必须是输出侧 pipe，接口标注 __out_pipe__(V, S)）；firstValue 为序列首元素；diffValue 为步长；count 为序列长度；
• 数据类型：支持 half、float、int16_t、int32_t；在 Ascend 950 系列上额外支持 int64_t；
• 编译期保护：头文件通过 __NPU_COMPILER_INTERNAL_PURE_SIMT__ 宏检查，禁止在 --enable-simt 纯 SIMT 编译模式下使用该向量接口。

实现层面的三项检查

底层实现 arithprogression_3510_impl.h（3510 系列架构，同时覆盖 9201/5102/5162/3003/3113 等架构，其余架构走 arithprogression_common_impl.h）在调用前做了三项校验：

ASCENDC_ASSERT((dstLocal.GetSize() >= count), { KERNEL_LOG(KERNEL_ERROR, "dst length must equal with Arange length"); });
ASCENDC_ASSERT((static_cast<float>(diffValue) >= static_cast<float>(0)), { KERNEL_LOG(KERNEL_ERROR, "diff value must bigger than 0"); });
static_assert(SupportType<T, int16_t, int32_t, half, float, int64_t>(),
              "current data type is not supported on current device!");

即：输出缓冲区容量必须不小于 count、等差值必须非负、数据类型必须受支持，否则会在编译期或运行期报错。

寄存器级计算流程

从源码结构看，Arange 在 3510 系列上的实现分为"基序列生成"与"分段累加写出"两个阶段：

• 基序列生成（GetBaseArithProgression）：先用 Reg::Arange 将寄存器初始化为 0,1,2,… 的等差索引，再 Reg::Muls 乘以 diffValue、Reg::Adds 加上 firstValue，从而得到首段序列 firstValue, firstValue+diffValue, …；
• 分段写出与步进（VfCallArithProgression）：按单次 repeat 可承载的元素数（ONE_REPEAT_BYTE_SIZE / sizeof(T)，8 字节类型翻倍）切分 repeatTimes 轮，每轮用 Reg::UpdateMask 更新掩码、Reg::StoreAlign 将当前段序列写入片上 dstLocalAddr，随后 Reg::Add 累加 stepReg = repeatLen * diffValue，迭代推进到下一段。

这种"寄存器索引 + 标量运算 + 循环写出"的实现方式，避免了逐元素标量循环，充分发挥了向量单元并行能力，是理解高阶 API 性能特性的关键视角。

编译与运行：三种模式全流程

1. 配置环境变量

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

source ${install_path}/cann/set_env.sh

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

2. 样例执行（NPU 模式）

在样例根目录（examples/01_simd_cpp_api/04_advanced_api/06_index/arange）下依次执行：

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 A2 推理系列产品和 Atlas A3 训练系列产品/Atlas A3 推理系列产品，dav-3510 对应 Ascend 950PR/Ascend 950DT

工程文件 CMakeLists.txt 中通过 set(CMAKE_ASC_ARCHITECTURES "dav-2201" CACHE STRING ...) 设置默认架构，并以 --npu-arch=${CMAKE_ASC_ARCHITECTURES} 编译选项传入 ASC 编译器；同时链接了 tiling_api、register、platform、unified_dlog、graph_base 等 Ascend C 运行所依赖的库。

5. 执行结果

执行成功后输出：

test pass!

表示设备侧输出与 golden 真值精度对比通过。

数据生成与精度校验

测试数据生成

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)
# 写出 ./input/first.bin、./input/diff.bin、./input/length.bin、./output/golden.bin

这里的关键细节是：逻辑长度 100 的 float 序列按 32 字节对齐后实际生成 104 个元素，golden 真值同样用 104 长度生成，与设备侧对齐逻辑保持一致，避免搬运末尾的对齐填充字节干扰精度比对。

精度对比

host 侧 CompareResult（arange.asc）从 ./output/golden.bin 读入真值，逐元素对比设备输出：

constexpr float EPS = 1e-5;
for (int i = 0; i < outSize / sizeof(float); i++) {
    float a = (reinterpret_cast<const float*>(outputData))[i];
    float b = (reinterpret_cast<const float*>(goldenData))[i];
    float ae = std::abs(a - b);
    float re = ae / abs(b);
    if (ae > EPS && re > EPS) {
        // 记录错误并打印差异
    }
}

判定条件同时考察绝对误差与相对误差（均需超过 EPS = 1e-5 才算失败），全部通过则打印 CompareResult golden.bin success! 并最终由主函数输出 test pass!。

小结与延伸

本文从索引计算算子样例集出发，以 Arange 样例为主线，完整梳理了：

1. 接口语义：Arange(dst, firstValue, diffValue, count) 生成 dst[i+1] = dst[i] + diffValue 的等差数列，支持 half/float/int16_t/int32_t（Ascend 950 额外支持 int64_t）；
2. 工程结构：Kernel 直调工程的设备侧（KernelArange 类 + tiling 结构 + __global__ __vector__ 核函数）与 host 侧（ACL 运行时 + 内核调用符 + 精度校验）的完整调用链；
3. 底层原理：寄存器级"基序列生成 + 分段累加写出"实现，以及 Arange 无需临时空间（GetArangeMaxMinTmpSize 返回 0）的 tiling 设计；
4. 实操流程：环境变量配置、NPU/CPU 调试/NPU 仿真三种模式的编译运行命令、编译选项表与对齐/精度细节。

后续若需将本样例扩展到动态 shape 场景，可参照仓库中其他高阶 API 样例（如 examples/01_simd_cpp_api/04_advanced_api 下的 math、reduce 等模块）的 tiling 切分方式；若需在纯 SIMT 编程模型中生成索引，则应改用 SIMT 侧的索引接口（如 include/basic_api/reg_compute/kernel_reg_compute_vec_arange_intf.h 所声明的寄存器级 arange 能力）而非本向量接口。建议结合 docs/zh/quick_start.md 完成环境准备后，亲自动手跑通该样例，体会高阶 API 从"一行调用"到"底层指令序列"的完整映射过程。