asc-devkit 中 reg_load_gather 样例详解：基于 Ascend C C API 的矢量寄存器离散搬入

• 人工智能
• 深度学习
• 算子库
• CANN
• Ascend

导读

reg_load_gather是 asc-devkit 开源仓库中基于 Ascend C C API 的 Reg 矢量计算样例，演示如何将 Unified Buffer（UB）中的非连续数据按索引离散加载到矢量数据寄存器（Vector Register），再写回 UB 并搬运至 GM（Global Memory）。本文以样例 README 为主体，结合样例内核源码、构建配置与数据脚本，完整讲解两种离散搬入场景（按元素索引asc_gather与按 32B DataBlock 索引asc_gather_datablock）的实现原理、接口签名、编译运行方式与结果校验方法。读完本文，你将能够独立编译运行该样例，并掌握在 Ascend 950PR&950DT 系列产品（CANN >= 9.2.0）上使用 Reg 矢量接口完成 Gather 类非连续数据搬运的核心套路。

一、样例概述与两种场景

本样例通过 CMake 参数SCENARIO_NUM选择两种离散搬入场景：

两种场景的差异本质上是索引粒度的不同：场景 1 的索引精确到单个元素（下标），场景 2 的索引精确到 32B 的 DataBlock（字节偏移）。对应到接口选择上，场景 1 使用vector_uint16_t索引寄存器调用 asc_gather（UB 源版本），场景 2 使用vector_uint32_t索引寄存器调用 asc_gather_datablock。

支持的产品及 CANN 软件版本

注意：样例的构建配置对架构做了强校验，仅允许dav-3510架构（对应上述产品系列），详见下文"编译选项"。

二、目录结构与数据流总览

目录结构

reg_load_gather/ ├── scripts/ │ ├── gen_data.py // 生成输入和 golden 数据 │ └── verify_result.py // 校验输出数据 ├── CMakeLists.txt // 编译配置 ├── data_utils.h // 输入输出文件工具 ├── reg_load_gather.asc // Ascend C C API 实现和调用入口 ├── README.md // 中文说明文档 └── README_en.md // 英文说明文档

整体数据流

样例的执行链路可概括为四段流水：

1. GM → UB（MTE2 搬运）：输入数据（src）先由asc_copy_gm2ub从 GM 搬运到 UB，索引数据（index）同样搬运到 UB；
2. UB → 矢量数据寄存器（SIMD VF 函数内）：入口函数通过asc_vf_call调用 SIMD VF 函数，按单个矢量数据寄存器对应的元素数分块，循环完成"索引加载 → 离散搬入 → 寄存器写回 UB"；
3. UB → GM（MTE3 搬运）：SIMD VF 函数执行完成后，由asc_copy_ub2gm将结果从 UB 搬运回 GM；
4. 流水同步：内核末尾调用asc_sync_pipe(PIPE_ALL)，确保全部流水完成。

在内核源码 reg_load_gather.asc 中，搬运与同步的实现在reg_load_gather内核函数中顺序展开：

__global__ __vector__ void reg_load_gather(__gm__ half* src, __gm__ void* index, __gm__ half* dst) { asc_init(); // 按 VF 长度推导分块参数（元素数 / 每 VF，DataBlock 数 / 每 VF，循环次数） constexpr uint32_t half_elements_per_vf = asc_get_vf_len() / sizeof(half); constexpr uint32_t uint32_elements_per_vf = asc_get_vf_len() / sizeof(uint32_t); constexpr uint32_t datablocks_per_vf = half_elements_per_vf / HALF_ELEMENTS_PER_BLOCK; constexpr uint16_t repeat_times = static_cast<uint16_t>((TOTAL_LENGTH + half_elements_per_vf - 1) / half_elements_per_vf); constexpr uint32_t datablock_index_ub_length = DATABLOCK_INDEX_LENGTH + uint32_elements_per_vf - datablocks_per_vf; __ubuf__ half src_ub[SRC_LENGTH]; __ubuf__ half dst_ub[TOTAL_LENGTH]; asc_copy_gm2ub(src_ub, src, SRC_LENGTH * sizeof(half)); if constexpr (scenarioNum == 1) { __ubuf__ uint16_t index_ub[TOTAL_LENGTH]; asc_copy_gm2ub(index_ub, reinterpret_cast<__gm__ uint16_t*>(index), TOTAL_LENGTH * sizeof(uint16_t)); asc_sync_notify(PIPE_MTE2, PIPE_V, EVENT_ID0); asc_sync_wait(PIPE_MTE2, PIPE_V, EVENT_ID0); asc_vf_call<gather_elements>(dst_ub, src_ub, index_ub, TOTAL_LENGTH, half_elements_per_vf, repeat_times); } else { __ubuf__ uint32_t index_ub[datablock_index_ub_length]; asc_copy_gm2ub(index_ub, reinterpret_cast<__gm__ uint32_t*>(index), DATABLOCK_INDEX_LENGTH * sizeof(uint32_t)); asc_sync_notify(PIPE_MTE2, PIPE_V, EVENT_ID0); asc_sync_wait(PIPE_MTE2, PIPE_V, EVENT_ID0); asc_vf_call<gather_datablocks>( dst_ub, src_ub, index_ub, TOTAL_LENGTH, half_elements_per_vf, datablocks_per_vf, repeat_times); } asc_sync_notify(PIPE_V, PIPE_MTE3, EVENT_ID0); asc_sync_wait(PIPE_V, PIPE_MTE3, EVENT_ID0); asc_copy_ub2gm(dst, dst_ub, TOTAL_LENGTH * sizeof(half)); asc_sync_pipe(PIPE_ALL); }

其中关键的流水同步接口在 sync.h 中声明：asc_sync_notify(pipe_t pipe, pipe_t tpipe, event_t id)负责发送事件通知，asc_sync_wait(pipe_t pipe, pipe_t tpipe, event_t id)负责等待事件，asc_sync_pipe(pipe_t pipe)负责同步指定流水。样例通过 MTE2→V、V→MTE3 的两级事件同步，保证"搬运完成后才执行矢量计算、矢量计算完成后才回搬"的顺序正确性。

asc_vf_call是 SIMD VF 函数的调用入口，其实现位于 helpers.h：在 AIV（矢量核）上通过asc::details::asc_vf_call_impl<func_ptr>(args...)真正派发执行，在非 AIV 核上则直接跳过，从而保证多核场景下的正确性。

三、场景 1：按元素索引离散搬入（asc_gather）

语义与数据规格

场景 1 的语义是dst[i] = src[index[i]]，即按索引数组逐个元素离散搬入。数据规格如下：

索引值表示相对于src首地址的元素下标，取值范围为[0, 1023]。

实现步骤

根据 README 与内核源码，场景 1 的 SIMD VF 函数gather_elements按以下三步循环处理：

__simd_vf__ inline void gather_elements( __ubuf__ half* dst, __ubuf__ half* src, __ubuf__ uint16_t* index, uint32_t count, uint32_t one_repeat_size, uint16_t repeat_times) { vector_half dst_reg; vector_uint16_t index_reg; vector_bool mask; for (uint16_t i = 0; i < repeat_times; ++i) { mask = asc_update_mask_b16(count); asc_loadalign(index_reg, index + i * one_repeat_size); asc_gather(dst_reg, src, index_reg, mask); asc_storealign(dst + i * one_repeat_size, dst_reg, mask); } }

1. 分两次加载索引：使用asc_loadalign将 256 个uint16_t索引分两次（每次一个 VF 长度）加载到类型为vector_uint16_t的矢量数据寄存器中。asc_loadalign的声明位于 loadalign.h，返回/写入vector_xxx_t寄存器，支持包括uint16_t、half、uint32_t等在内的广泛数据类型；
2. 离散搬入：调用asc_gather(dst_reg, src, index_reg, mask)按元素索引完成离散搬入。对应声明见 ub_gather.h，其原型为asc_gather(__ubuf__ half* src, vector_uint16_t index, vector_bool mask)，返回值写入vector_half寄存器；
3. 写回 UB：使用asc_storealign将类型为vector_half的矢量数据寄存器中的结果连续写回 UB 的dst缓冲区。对应声明见 storealign.h。

mask 的作用：asc_update_mask_b16(count)在 reg_mask.h 中声明，用于按有效元素数count构造 16 位掩码。由于 256 个元素按 VF 分块时，最后一轮可能不足一个完整 VF（例如 512B VF 下每 VF 容纳 64 个 half 元素，需要 4 轮，恰好整除；若数据量不能整除，则最后一轮用 mask 屏蔽无效通道），mask 保证离散搬入只作用于有效元素。

数据生成与 golden 逻辑

场景 1 的测试数据由 gen_data.py 生成：

def generate_element_gather_data(): src_length = 1024 output_length = 256 src = np.random.uniform(0, 1, src_length).astype(np.float16) index = np.random.randint(0, src_length, output_length).astype(np.uint16) golden = src[index] src.tofile("./input/input_x.bin") index.tofile("./input/input_y.bin") golden.tofile("./output/golden.bin")

即：src为 1024 个随机half值；index为 256 个[0, 1023]范围内的随机uint16_t下标；golden 数据直接用 NumPy 的花式索引src[index]计算得到，与内核语义dst[i] = src[index[i]]完全对应。

四、场景 2：按 DataBlock 离散搬入（asc_gather_datablock）

语义与数据规格

场景 2 的索引粒度升级为 32B DataBlock：

• 一个 DataBlock 为 32B，对应 16 个half元素（32 / sizeof(half) = 16）；
• 索引数组每个值表示相对于src首地址的字节偏移，必须为 32B 的整数倍；
• src包含 1024 个half元素，共 2048B，因此index的合法取值为0, 32, ..., 2016（共 64 个 DataBlock）。

数据规格如下：

输出为 16 个 DataBlock × 16 个元素 = 256 个half。

实现步骤

场景 2 的 SIMD VF 函数gather_datablocks同样按三步循环处理：

__simd_vf__ inline void gather_datablocks( __ubuf__ half* dst, __ubuf__ half* src, __ubuf__ uint32_t* index, uint32_t count, uint32_t one_repeat_size, uint32_t index_repeat_size, uint16_t repeat_times) { vector_half dst_reg; vector_uint32_t index_reg; vector_bool mask; for (uint16_t i = 0; i < repeat_times; ++i) { mask = asc_update_mask_b16(count); asc_loadalign(index_reg, index + i * index_repeat_size); asc_gather_datablock(dst_reg, src, index_reg, mask); asc_storealign(dst + i * one_repeat_size, dst_reg, mask); } }

1. 加载索引：使用asc_loadalign将 16 个有效uint32_t索引加载到类型为vector_uint32_t的矢量数据寄存器中（示例中索引数组按 VF 长度进行了对齐填充，UB 索引缓冲区长度为DATABLOCK_INDEX_LENGTH + uint32_elements_per_vf - datablocks_per_vf，保证循环中每轮asc_loadalign都能对齐读取）；
2. 按 DataBlock 离散搬入：调用asc_gather_datablock(dst_reg, src, index_reg, mask)加载 16 个 DataBlock。对应声明见 ub_gather.h，原型为asc_gather_datablock(__ubuf__ half* src, vector_uint32_t index, vector_bool mask)，返回vector_half；
3. 连续写回：使用asc_storealign将 256 个half结果连续写回 UB。

与场景 1 的关键差异在于：场景 1 的index是元素下标（每个索引取一个元素），场景 2 的index是字节偏移（每个索引取整个 32B DataBlock，即 16 个连续元素），因此单次asc_gather_datablock搬入的数据量是asc_gather的 16 倍。

数据生成与 golden 逻辑

场景 2 的测试数据由 gen_data.py 生成：

def generate_datablock_gather_data(): src_length = 1024 output_length = 256 elements_per_block = 32 // np.dtype(np.float16).itemsize src = np.random.uniform(0, 1, src_length).astype(np.float16) block_index = np.random.randint(0, src_length // elements_per_block, output_length // elements_per_block).astype( np.uint32 ) index = block_index * 32 golden = np.concatenate( [src[block * elements_per_block : (block + 1) * elements_per_block] for block in block_index] ) src.tofile("./input/input_x.bin") index.tofile("./input/input_y.bin") golden.tofile("./output/golden.bin")

即：先在[0, 63]范围内随机生成 16 个 DataBlock 序号，乘以 32 得到字节偏移索引（保证合法取值0, 32, ..., 2016）；golden 数据按每个 DataBlock 的 16 个连续half元素拼接而成，与内核按 DataBlock 离散搬入的语义一致。

五、编译与运行

配置环境变量

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

source ${install_path}/cann/set_env.sh

说明：${install_path}为 CANN 包安装目录；root 用户默认为/usr/local/Ascend，非 root 用户默认为${HOME}/Ascend。

运行样例

在样例根目录执行以下命令。SCENARIO_NUM可取 1 或 2，默认值为 1。

SCENARIO_NUM=1 # 选择场景1 mkdir -p build && cd build; # 创建并进入build目录 cmake -DSCENARIO_NUM=$SCENARIO_NUM -DCMAKE_ASC_ARCHITECTURES=dav-3510 ..;make -j; # 编译工程，默认NPU模式 python3 ../scripts/gen_data.py -scenarioNum=$SCENARIO_NUM # 生成测试输入数据和golden数据 ./demo # 执行样例 python3 ../scripts/verify_result.py output/output.bin output/golden.bin # 校验输出数据与golden数据

默认使用 NPU 运行模式。使用 NPU 仿真模式时，增加-DCMAKE_ASC_RUN_MODE=sim：

cmake -DSCENARIO_NUM=$SCENARIO_NUM -DCMAKE_ASC_RUN_MODE=sim -DCMAKE_ASC_ARCHITECTURES=dav-3510 ..;make -j; # NPU仿真模式

切换场景或运行模式前，请在build目录删除CMakeCache.txt后重新执行 CMake，避免复用缓存配置。

编译选项

CMake 工程配置细节

样例的 CMakeLists.txt 体现了 C API 样例的典型工程结构，值得展开说明：

• 缓存参数声明：CMAKE_ASC_RUN_MODE默认"npu"，CMAKE_ASC_ARCHITECTURES默认"dav-3510"，SCENARIO_NUM默认"1"，均以CACHE STRING形式暴露给用户；
• ASC 工具链引入：通过find_package(ASC REQUIRED)查找 ASC 编译工具链，随后project(c_api_reg_load_gather LANGUAGES ASC CXX)声明项目使用 ASC 与 C++ 语言；
• 合法性校验：CMake 层面对参数做严格校验——CMAKE_ASC_ARCHITECTURES非dav-3510时FATAL_ERROR，CMAKE_ASC_RUN_MODE非npu/sim时FATAL_ERROR，SCENARIO_NUM不在 1~2 范围时FATAL_ERROR，保证错误配置在构建阶段即被拦截；
• 编译定义与架构选项：通过add_definitions(-DSCENARIO_NUM=${SCENARIO_NUM})将场景号注入源码（内核中的constexpr uint32_t scenarioNum = SCENARIO_NUM;即来源于此），并通过target_compile_options为 ASC 语言传递--npu-arch=${CMAKE_ASC_ARCHITECTURES}架构参数。

运行结果

输出结果与 golden 数据一致时，校验脚本输出如下信息：

test pass!

校验脚本 verify_result.py 的逻辑为：以float16逐元素读取output/output.bin与output/golden.bin，先比对元素个数，再逐元素比对值；若存在不一致，会打印最多前 101 个失配元素的下标、期望值与实际值，并抛出[ERROR] result error退出码 1；全部一致则打印test pass!。

六、宿主侧调用与 ACL 资源管理

reg_load_gather.asc中同时包含内核实现与宿主侧main入口，展示了完整的上位机调用流程：

1. ACL 初始化：aclInit(nullptr)→aclrtSetDevice(device_id)→aclrtCreateStream(&stream)；
2. 内存分配与数据搬入：分别用aclrtMallocHost分配宿主内存、aclrtMalloc分配设备内存（ACL_MEM_MALLOC_HUGE_FIRST），通过工具函数ReadFile读取./input/input_x.bin（src）与./input/input_y.bin（index），再经aclrtMemcpy(..., ACL_MEMCPY_HOST_TO_DEVICE)拷贝到设备端；
3. 内核启动：以reg_load_gather<<<1, 0, stream>>>(src_device, index_device, dst_device)语法在指定 stream 上启动内核（block 数 1，blockDim 为 0 表示默认配置），随后aclrtSynchronizeStream(stream)等待内核完成；
4. 结果回读与落盘：aclrtMemcpy(..., ACL_MEMCPY_DEVICE_TO_HOST)将结果拷回宿主，WriteFile写入./output/output.bin；
5. 资源释放：依次aclrtFree/aclrtFreeHost释放设备与宿主内存，aclrtDestroyStream、aclrtResetDevice、aclFinalize完成收尾。

其中ReadFile/WriteFile定义在 data_utils.h：ReadFile会校验文件存在性、是否为普通文件、大小是否与预期字节数完全一致（例如场景 1 中 src 文件必须恰为1024 * sizeof(half) = 2048字节）；WriteFile处理EINTR中断重试，保证写入完整。所有 ACL 调用均通过CHECK_ACL宏检查返回码并在失败时打印错误位置（文件、行号、ret 值）。

七、接口声明位置速查（源码级佐证）

为便于在仓库中继续深入，将本文涉及的 C API 声明位置汇总如下：

八、关键要点总结

• 两种索引粒度：元素索引（asc_gather，uint16_t下标）适合稀疏单元素收集；DataBlock 索引（asc_gather_datablock，uint32_t字节偏移）以 32B 为最小搬入单位，适合连续块收集，单次数据吞吐是元素级 gather 的 16 倍（对 half）；
• 索引取值范围是正确性关键：场景 1 的index必须是[0, 1023]的元素下标；场景 2 的index必须是0, 32, ..., 2016的 32B 对齐字节偏移，超出范围即越界访问；
• 分块与 mask 配合：数据量按asc_get_vf_len()推导的单 VF 元素数分块循环，末尾不足一个 VF 时由asc_update_mask_b16屏蔽无效通道；
• 流水同步不可省略：MTE2 搬运完成后才允许 V 流水计算，V 流水完成后才允许 MTE3 回搬，最后以asc_sync_pipe(PIPE_ALL)收尾；
• 工程约束：样例仅支持dav-3510架构与npu/sim两种运行模式，切换场景或模式前需清理build/CMakeCache.txt，防止复用旧缓存配置导致编译结果与预期不符。

通过本样例，读者可以举一反三：凡是"从 UB 中按不连续下标/偏移收集数据到连续输出"的需求（如查表、Embedding 收集、稀疏特征聚合等），都可以参照gather_elements/gather_datablocks的结构，以asc_loadalign + asc_gather/asc_gather_datablock + asc_storealign的组合在矢量数据寄存器层面高效实现。

• 人工智能
• 深度学习
• 算子库
• CANN
• Ascend