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

人工智能深度学习算子库CANNAscend【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址https://gitcode.com/cann/asc-devkit点击查看免费下载导读reg_load_gather是 asc-devkit 开源仓库中基于 Ascend C C API 的 Reg 矢量计算样例演示如何将 Unified BufferUB中的非连续数据按索引离散加载到矢量数据寄存器Vector Register再写回 UB 并搬运至 GMGlobal Memory。本文以样例 README 为主体结合样例内核源码、构建配置与数据脚本完整讲解两种离散搬入场景按元素索引asc_gather与按 32B DataBlock 索引asc_gather_datablock的实现原理、接口签名、编译运行方式与结果校验方法。读完本文你将能够独立编译运行该样例并掌握在 Ascend 950PR950DT 系列产品CANN 9.2.0上使用 Reg 矢量接口完成 Gather 类非连续数据搬运的核心套路。一、样例概述与两种场景本样例通过 CMake 参数SCENARIO_NUM选择两种离散搬入场景SCENARIO_NUM离散搬入场景1使用asc_gather按元素索引离散搬入数据2使用asc_gather_datablock按 32B DataBlock 离散搬入数据两种场景的差异本质上是索引粒度的不同场景 1 的索引精确到单个元素下标场景 2 的索引精确到 32B 的 DataBlock字节偏移。对应到接口选择上场景 1 使用vector_uint16_t索引寄存器调用 asc_gatherUB 源版本场景 2 使用vector_uint32_t索引寄存器调用 asc_gather_datablock。支持的产品及 CANN 软件版本产品CANN 软件版本Ascend 950PR950DT 系列产品 CANN 9.2.0注意样例的构建配置对架构做了强校验仅允许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 // 英文说明文档整体数据流样例的执行链路可概括为四段流水GM → UBMTE2 搬运输入数据src先由asc_copy_gm2ub从 GM 搬运到 UB索引数据index同样搬运到 UBUB → 矢量数据寄存器SIMD VF 函数内入口函数通过asc_vf_call调用 SIMD VF 函数按单个矢量数据寄存器对应的元素数分块循环完成索引加载 → 离散搬入 → 寄存器写回 UBUB → GMMTE3 搬运SIMD VF 函数执行完成后由asc_copy_ub2gm将结果从 UB 搬运回 GM流水同步内核末尾调用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 长度推导分块参数元素数 / 每 VFDataBlock 数 / 每 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_castuint16_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_callgather_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_callgather_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_implfunc_ptr(args...)真正派发执行在非 AIV 核上则直接跳过从而保证多核场景下的正确性。三、场景 1按元素索引离散搬入asc_gather语义与数据规格场景 1 的语义是dst[i] src[index[i]]即按索引数组逐个元素离散搬入。数据规格如下项目名称Shape数据类型输入src[1, 1024]half输入index[1, 256]uint16_t输出dst[1, 256]half索引值表示相对于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); } }分两次加载索引使用asc_loadalign将 256 个uint16_t索引分两次每次一个 VF 长度加载到类型为vector_uint16_t的矢量数据寄存器中。asc_loadalign的声明位于 loadalign.h返回/写入vector_xxx_t寄存器支持包括uint16_t、half、uint32_t等在内的广泛数据类型离散搬入调用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寄存器写回 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。数据规格如下项目名称Shape数据类型输入src[1, 1024]half输入index[1, 16]uint32_t输出dst[1, 256]half输出为 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); } }加载索引使用asc_loadalign将 16 个有效uint32_t索引加载到类型为vector_uint32_t的矢量数据寄存器中示例中索引数组按 VF 长度进行了对齐填充UB 索引缓冲区长度为DATABLOCK_INDEX_LENGTH uint32_elements_per_vf - datablocks_per_vf保证循环中每轮asc_loadalign都能对齐读取按 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连续写回使用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, ..., 2016golden 数据按每个 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_NUM1 # 选择场景1 mkdir -p build cd build; # 创建并进入build目录 cmake -DSCENARIO_NUM$SCENARIO_NUM -DCMAKE_ASC_ARCHITECTURESdav-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_MODEsimcmake -DSCENARIO_NUM$SCENARIO_NUM -DCMAKE_ASC_RUN_MODEsim -DCMAKE_ASC_ARCHITECTURESdav-3510 ..;make -j; # NPU仿真模式切换场景或运行模式前请在build目录删除CMakeCache.txt后重新执行 CMake避免复用缓存配置。编译选项选项可选值说明SCENARIO_NUM1、2选择按元素索引或按 DataBlock 加载场景。CMAKE_ASC_RUN_MODEnpu、sim选择 NPU 运行或 NPU 仿真模式。CMAKE_ASC_ARCHITECTURESdav-3510指定 NPU 架构对应 Ascend 950PR950DT 系列产品。CMake 工程配置细节样例的 CMakeLists.txt 体现了 C API 样例的典型工程结构值得展开说明缓存参数声明CMAKE_ASC_RUN_MODE默认npuCMAKE_ASC_ARCHITECTURES默认dav-3510SCENARIO_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_ERRORCMAKE_ASC_RUN_MODE非npu/sim时FATAL_ERRORSCENARIO_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入口展示了完整的上位机调用流程ACL 初始化aclInit(nullptr)→aclrtSetDevice(device_id)→aclrtCreateStream(stream)内存分配与数据搬入分别用aclrtMallocHost分配宿主内存、aclrtMalloc分配设备内存ACL_MEM_MALLOC_HUGE_FIRST通过工具函数ReadFile读取./input/input_x.binsrc与./input/input_y.binindex再经aclrtMemcpy(..., ACL_MEMCPY_HOST_TO_DEVICE)拷贝到设备端内核启动以reg_load_gather1, 0, stream(src_device, index_device, dst_device)语法在指定 stream 上启动内核block 数 1blockDim 为 0 表示默认配置随后aclrtSynchronizeStream(stream)等待内核完成结果回读与落盘aclrtMemcpy(..., ACL_MEMCPY_DEVICE_TO_HOST)将结果拷回宿主WriteFile写入./output/output.bin资源释放依次aclrtFree/aclrtFreeHost释放设备与宿主内存aclrtDestroyStream、aclrtResetDevice、aclFinalize完成收尾。其中ReadFile/WriteFile定义在 data_utils.hReadFile会校验文件存在性、是否为普通文件、大小是否与预期字节数完全一致例如场景 1 中 src 文件必须恰为1024 * sizeof(half) 2048字节WriteFile处理EINTR中断重试保证写入完整。所有 ACL 调用均通过CHECK_ACL宏检查返回码并在失败时打印错误位置文件、行号、ret 值。七、接口声明位置速查源码级佐证为便于在仓库中继续深入将本文涉及的 C API 声明位置汇总如下接口声明位置说明asc_gatherUB 源include/c_api/reg_compute/gather/ub_gather.h按元素索引离散搬入支持vector_uint16_t/vector_uint32_t两种索引寄存器覆盖 int8/uint8/int16/uint16/int32/uint32/half/float/bfloat16/fp8 等多种数据类型asc_gather_datablockinclude/c_api/reg_compute/gather/ub_gather.h按 32B DataBlock 离散搬入使用vector_uint32_t索引寄存器返回单个完整 VF 的数据asc_loadaligninclude/c_api/reg_compute/load/loadalign.h将 UB 连续数据加载到矢量数据寄存器另有_brc、_upsample、_downsample、_unpack、_datablock_strided等变体asc_storealigninclude/c_api/reg_compute/store/storealign.h将矢量数据寄存器写回 UB另有_1st、_pack、_intlv、_postupdate等变体asc_update_mask_b16include/c_api/reg_compute/compute/reg_mask.h构造 16 位有效元素掩码同族还有asc_update_mask_b8/b32asc_get_vf_leninclude/c_api/utils/sys_var.h获取单次 VF 操作处理的字节数用于推导分块参数asc_vf_callinclude/utils/base/helpers.hSIMD VF 函数调用入口仅在 AIV 核上派发执行asc_sync_notify/wait/pipeinclude/c_api/sync/sync.hMTE2→V、V→MTE3 的流水间事件同步与全流水同步八、关键要点总结两种索引粒度元素索引asc_gatheruint16_t下标适合稀疏单元素收集DataBlock 索引asc_gather_datablockuint32_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的组合在矢量数据寄存器层面高效实现。赞分享人工智能深度学习算子库CANNAscend【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址https://gitcode.com/cann/asc-devkit点击查看免费下载相关推荐asc-devkit 中 reg_load_store_align 样例解析Ascend C C API 寄存器对齐加载与存储实战asc devkit 中 reg_load_store_align 样例解析Ascend C C API 寄存器对齐加载与存储实战 本篇技术指南围绕 asc人工智能深度学习算子库CANNAscendasc-devkit PhiloxRandom 随机数生成样例详解基于 Ascend C 高阶 API 的端到端实现asc devkit PhiloxRandom 随机数生成样例详解基于 Ascend C 高阶 API 的端到端实现 导读 本文以 asc devkit 开源人工智能深度学习算子库CANNAscendAscend C Transpose 高阶 API 实战基于 asc-devkit 的 NZ2ND 转置样例深度解析Ascend C Transpose 高阶 API 实战基于 asc devkit 的 NZ2ND 转置样例深度解析 本文以 asc devkit 仓库中的人工智能深度学习算子库CANNAscend上一篇3分钟掌握如何用Python通过手机号快速查询QQ号下一篇PotatoNV vs HCU Client华为Bootloader解锁技术方案深度评估与实践指南创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考