基于 Arange 高阶 API 实现等差数列生成：CANN asc-devkit 索引算子（Index）示例深度解析

原创2026-10-03 09:28:181,667 阅读

文章标签：人工智能深度学习高性能计算CANN

基于 Arange 高阶 API 实现等差数列生成：CANN asc-devkit 索引算子（Index）示例深度解析

本篇文章围绕 CANN asc-devkit 仓库中 examples/01_simd_cpp_api/04_advanced_api/06_index 索引（Index）算子示例集，以其中的 Arange（等差数列生成）端到端示例为主体，完整讲解 Arange 高阶 API 的接口语义、算子核实现、Tiling 传递、构建运行方式，并结合仓库源码剖析其底层向量化实现原理。读完本文，你将掌握如何在昇腾 AI 处理器上通过高阶 API 快速实现“给定起始值、等差值与长度生成等差数列”的算子，并能够独立完成 NPU 运行、CPU 调试与 NPU 仿真三种模式下的编译与验证。

一、示例集定位：Index 索引类高阶 API 的典型用法

examples/01_simd_cpp_api/04_advanced_api/06_index 目录是一个面向“索引（Index）计算类算子”的示例集合。索引类算子用于按给定规则生成或变换序列下标，Arange 正是其中最典型的一员。该目录下的 README 给出了当前示例清单：

目录名	说明	支持产品
arange	基于核函数直调算子工程，介绍如何调用 Arange 高阶 API 实现 arange 单算子：给定起始值（start value）、等差值（step value）和长度，返回一个等差数列	Ascend 950PR/Ascend 950DT；Atlas A3 训练/推理系列产品；Atlas A2 训练/推理系列产品

可以看出，该示例集采用**核函数直调（kernel direct call）**工程形态：不依赖 MindSpore 等上层框架，直接在 host 侧通过 <<<>>> 语法启动核函数，适合开发者快速验证单算子功能与精度。

二、Arange 高阶 API：接口语义与支持范围

2.1 接口原型与数学公式

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)
{
    ArithProgressionImpl(dst, firstValue, diffValue, count);
}

参数说明：

参数	方向	含义
`dst`	输出	输出 LocalTensor，存放生成的等差数列
`firstValue`	输入	等差数列的起始值（第一个元素）
`diffValue`	输入	等差值，相邻元素之间的差值
`count`	输入	元素个数，即序列长度

2.2 支持的数据类型

从头文件注释可以确认 Arange 的数据类型支持范围：

通用平台：half、float、int16_t、int32_t；
Ascend 950 系列（Ascend 950PR/Ascend 950DT）：额外支持 int64_t。

同时，arithprogression.h 明确限制了 SIMT 纯线程模式：该头文件不能在开启 --enable-simt 编译标志（__NPU_COMPILER_INTERNAL_PURE_SIMT__）时使用，即 Arange 属于 SIMD（矢量）编程模型下的高阶 API。

2.3 旧头文件迁移提示

仓库中还保留了一个被标记为弃用的旧头文件 include/adv_api/index/kernel_operator_arithprogression_intf.h，其中通过 [[deprecated]] 属性提示开发者：

kernel_operator_arithprogression_intf.h is deprecated, please use arithprogression.h instead!

因此新代码应统一包含 adv_api/index/arithprogression.h。

2.4 配套的 Tiling 查询接口

若需在 host 侧 Tiling 阶段预估工作空间（workspace），可参考 include/adv_api/index/arithprogression_tiling.h 提供的两个接口：

void GetArithProgressionMaxMinTmpSize(uint32_t& maxValue, uint32_t& minValue);
void GetArangeMaxMinTmpSize(uint32_t& maxValue, uint32_t& minValue);

它们用于获取该算子所需临时空间（temporary space）的最大值与最小值区间，开发者可据此选择合适的 Tiling 参数。

三、示例功能与输入输出规格

Arange 示例位于 examples/01_simd_cpp_api/04_advanced_api/06_index/arange，其功能与规格如下：

示例功能：实现等差数列生成。给定起始值、等差值与长度，返回对应等差数列。
示例规格：采用固定 shape —— 输入 firstGm[1, 1]、diffGm[1, 1]，输出 outputGm[104, 1]，数据类型均为 float，格式为 ND。核函数名为 arange_custom。

类别	名称	shape	数据类型	格式
输入	firstGm	[1, 1]	float	ND
输入	diffGm	[1, 1]	float	ND
输出	outputGm	[104, 1]	float	ND
核函数名	arange_custom	-	-	-

说明：逻辑长度为 100 个 float 元素，因昇腾侧数据搬运按 32 字节对齐，实际输出为 104 个元素（100 向上对齐到 32 字节 = 104 × 4 字节）。

示例工程的文件结构如下：

examples/01_simd_cpp_api/04_advanced_api/06_index/arange
├── scripts
│   └── gen_data.py         // 输入数据与真值（golden）数据生成脚本
├── CMakeLists.txt          // 构建工程文件
├── data_utils.h            // 数据读写函数
├── arange.asc              // Ascend C 示例实现与调用示例
└── README.md               // 示例说明文档（英文版为 README_en.md）

四、算子核实现源码解析

核心实现位于 arange.asc，分为 Tiling 数据定义、核内实现与 host 侧调用三部分。

4.1 Tiling 数据定义与生成

示例在 optiling 命名空间中用宏定义 Tiling 结构，将输出长度 outSize 传给核函数：

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

配套的 GenerateTiling(outSize) 将 Tiling 数据序列化到 buffer，供 host 侧通过 aclrtMemcpy 拷贝到设备端 tilingDevice。

4.2 核内 KernelArange 类

核内实现封装为模板类 KernelArange<T>，遵循 Ascend C 典型的 CopyIn → Compute → CopyOut 三段式流程：

__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();
}

要点：

Init 中将输出全局内存指针绑定到 GlobalTensor<T>，并按 32 字节（BYTES_ALIGN_SIZE）对齐在 VECOUT 队列中申请输出缓冲区；
Compute 中通过 outDst.AllocTensor<T>() 申请 LocalTensor，随后直接调用高阶 API 完成计算：

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);

CopyOut 中将结果从 LocalTensor 通过 AscendC::DataCopy 搬运回全局内存，搬运长度同样按 32 字节对齐（(count_ + BLOCK_NUM - 1) / BLOCK_NUM * BLOCK_NUM，其中 BLOCK_NUM = 32 / sizeof(T)）。

4.3 核函数入口与 host 侧直调

核函数入口为 arange_custom，通过 CopyTiling 从全局内存解析 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();
}

host 侧 main 函数展示了完整的 ACL 初始化 → 内存管理 → 数据搬运 → 核函数启动 → 结果校验 → 资源释放 闭环：

aclInit(nullptr);
aclrtSetDevice(deviceId);
aclrtCreateContext(&context, deviceId);
aclrtCreateStream(&stream);
// ... aclrtMallocHost / aclrtMalloc 申请 host 与 device 内存
// ... aclrtMemcpy 将 first、diff、tiling 拷贝到设备端
arange_custom<<<numBlocks, 0, stream>>>(firstDevice, diffDevice, outputDevice, workspaceDevice, tilingDevice);
aclrtSynchronizeStream(stream);
aclrtMemcpy(outputHost, outSize, outputDevice, outSize, ACL_MEMCPY_DEVICE_TO_HOST);

其中 NUM_BLOCKS = 1（单核单 block）、COUNT = 100、WORKSPACE_SIZE = 1024。启动后调用 CompareResult 读取 ./output/golden.bin 逐元素比对，判定条件为绝对误差与相对误差均大于 EPS = 1e-5 时记为错误，全部通过则打印 test pass!。

4.4 数据工具与真值生成

data_utils.h 提供了 ReadFile/WriteFile 两个二进制文件读写函数，负责读取输入 first.bin、diff.bin 与写出 output.bin、读取 golden；
scripts/gen_data.py 使用 NumPy 生成输入与真值：

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

真值通过 np.linspace(start, stop, length) 生成，与 Arange 的递推公式 dst[i] = first + diff * i 等价。

五、Arange 底层实现原理（源码级）

Arange 的实际计算逻辑位于内部实现头文件 impl/adv_api/detail/index/arithprogression/arithprogression_common_impl.h（该文件为内部头文件，官方提示开发者应使用公开接口 arithprogression.h）。

5.1 基础序列的标量生成

对于长度不超过一个 BLOCK 的场景，直接用标量运算逐元素写入：

template <typename T>
__aicore__ inline void GetBaseArithProgression(
    const LocalTensor<T>& dstLocal, const T firstValue, const T diffValue, const int32_t count)
{
    for (int i = 0; i < count; i++) {
        dstLocal.SetValue(i, static_cast<T>(firstValue) + static_cast<T>(diffValue) * static_cast<T>(i));
    }
}

其中 half 类型有专门特化：先转为 float 计算，再在 SetValue 时转回 half，以避免半精度中间运算的精度损失。整型（int16_t/int32_t/int64_t）在计算步进 GetArithProgressionStep 时统一提升到 int64_t 乘法再截断，降低溢出风险。

5.2 长序列的向量化广播填充

当 count 较大时（超过一个 BLOCK 或一个 REPEAT），实现采用“先生成一段基础等差数列，再用向量加法指令 Adds 广播扩展”的策略：

先生成长度为 BLOCK_NUM（一个 32 字节块可容纳的元素数）的基础序列；
通过 SetFlag<HardEvent::S_V> / WaitFlag<HardEvent::S_V> 完成标量到向量流水（S→V）的同步；

若 count > REPEAT_NUM（一个 REPEAT 可容纳的元素数），先用 Adds 以 stride 1 将序列从 1 个 BLOCK 广播到 8 个 BLOCK：

SetVectorMask<T>(0, (((static_cast<uint64_t>(1)) << static_cast<uint32_t>(BLOCK_NUM)) - 1));
for (int i = 0; i < DEFAULT_BLK_NUM - 1; i++) {
    Adds<T, false>(dstLocal[(i + 1) * BLOCK_NUM], dstLocal[i * BLOCK_NUM],
                   GetArithProgressionStep<T>(diffValue, BLOCK_NUM), MASK_PLACEHOLDER, (uint16_t)1, addsParamsStride1);
}

再以 DEFAULT_REPEAT_STRIDE 步长对每个 REPEAT 块执行 Adds，逐块累加步进 diffValue * REPEAT_NUM，直至填满 count；
末尾不足一个 REPEAT 的 tail 部分，通过 SetVectorMask<T>(tail_aligned) 设置向量掩码后做最后一次部分填充，随后 ResetMask() 恢复默认掩码。

这一实现充分利用了昇腾矢量单元（Vector Pipe）的广播与累加特性：无论序列多长，基础序列只需标量生成一次，其余全部由高效的向量 Adds 指令批量完成。需要说明的是，上述分块常量（ONE_BLK_SIZE、ONE_REPEAT_BYTE_SIZE、DEFAULT_BLK_NUM、DEFAULT_REPEAT_STRIDE）随不同 NPU 架构（如 dav-2201 与 dav-3510）可能有差异，仓库在 include/adv_api/index/arithprogression.h 中按 __NPU_ARCH__ 区分了 3510/9201/5102/5162/3003/3113 等架构专用实现与通用实现两条分支。

六、构建与运行：三种模式全覆盖

6.1 前置环境准备

在示例根目录下按 CANN 开发套件的安装方式配置环境变量：

source ${install_path}/cann/set_env.sh

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

6.2 NPU 运行模式

在示例目录执行以下命令完成构建、数据生成与运行：

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

6.3 CPU 调试与 NPU 仿真模式

在 cmake 配置阶段追加 -DCMAKE_ASC_RUN_MODE 参数即可切换运行模式：

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 仿真模式

注意：切换构建模式前，需在 build 目录下清除 cmake 缓存（rm CMakeCache.txt），然后重新执行 cmake。

6.4 构建选项说明

选项	取值	说明
`CMAKE_ASC_RUN_MODE`	`npu`（默认）、`cpu`、`sim`	运行模式：NPU 运行、CPU 调试、NPU 仿真
`CMAKE_ASC_ARCHITECTURES`	`dav-2201`（默认）、`dav-3510`	NPU 架构：`dav-2201` 对应 Atlas A2/A3 训练与推理系列产品；`dav-3510` 对应 Ascend 950PR/Ascend 950DT

6.5 工程构建文件

CMakeLists.txt 展示了 Ascend C 示例工程的标准组织方式：

cmake_minimum_required(VERSION 3.16)
find_package(ASC REQUIRED)
project(kernel_samples LANGUAGES ASC CXX)

add_executable(demo
    arange.asc
)

set(CMAKE_ASC_ARCHITECTURES "dav-2201" CACHE STRING "CMAKE_ASC_ARCHITECTURES, e.g. dav-2201, dav-3510")

target_link_libraries(demo PRIVATE
    tiling_api
    register
    platform
    unified_dlog
    dl
    graph_base
)

target_compile_options(demo PRIVATE
    $<$<COMPILE_LANGUAGE:ASC>:--npu-arch=${CMAKE_ASC_ARCHITECTURES}>
)

其中 find_package(ASC REQUIRED) 引入 Ascend C 语言编译支持，arange.asc 作为源码参与编译（ASC 语言），链接库覆盖 Tiling API、注册、平台与日志等运行依赖。

6.6 执行结果验证

运行成功且精度比对通过时输出：

test pass!

若比对失败，则输出 test failed! 并在 CompareResult 中打印首个错误位置的输出值与 golden 值，便于定位问题。

七、支持产品与 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

八、小结：从示例到自研算子的迁移路径

通过本文可以提炼出一条清晰的索引类算子开发范式：

Tiling 侧：定义 Tiling 数据结构（如 outSize），并可通过 GetArangeMaxMinTmpSize 查询临时空间区间；
核内计算：在 TPipe + TQue 框架下申请 LocalTensor，一行 AscendC::Arange<T>(dstLocal, first, diff, count) 完成计算，全程无需手写逐元素循环——长序列的向量化广播填充由底层 ArithProgressionImpl 自动完成；
host 侧直调：借助 ACL 运行时接口与 <<<>>> 语法完成数据搬移、核函数启动与精度校验；
构建验证：通过 CMAKE_ASC_RUN_MODE 与 CMAKE_ASC_ARCHITECTURES 两个开关即可在 NPU 实跑、CPU 调试、NPU 仿真三种模式间切换。

如需将本示例改造成自己的算子，可替换输入输出 shape、数据类型（注意 half 的浮点精度特化与整型的溢出处理）与核函数名称，并保持 32 字节对齐与 test pass! 精度校验流程不变。

登录后查看全文

人工智能深度学习高性能计算CANN

asc-devkit

基于 Arange 高阶 API 实现等差数列生成：CANN asc-devkit 索引算子（Index）示例深度解析

基于 Arange 高阶 API 实现等差数列生成：CANN asc-devkit 索引算子（Index）示例深度解析

一、示例集定位：Index 索引类高阶 API 的典型用法

二、Arange 高阶 API：接口语义与支持范围

2.1 接口原型与数学公式

2.2 支持的数据类型

2.3 旧头文件迁移提示

2.4 配套的 Tiling 查询接口

三、示例功能与输入输出规格

四、算子核实现源码解析

4.1 Tiling 数据定义与生成

4.2 核内 KernelArange 类

4.3 核函数入口与 host 侧直调

4.4 数据工具与真值生成

五、Arange 底层实现原理（源码级）

5.1 基础序列的标量生成

5.2 长序列的向量化广播填充

六、构建与运行：三种模式全覆盖

6.1 前置环境准备

6.2 NPU 运行模式

6.3 CPU 调试与 NPU 仿真模式

6.4 构建选项说明

6.5 工程构建文件

6.6 执行结果验证

七、支持产品与 CANN 版本约束

八、小结：从示例到自研算子的迁移路径

相关推荐

热门文章