CANN ops-nn GatherElementsV3 算子详解在 NPU 上按维度聚集张量元素【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn导读GatherElementsV3 是 CANN ops-nn 神经网络算子库experimental/index/gather_elements_v3中提供的数据聚集算子它接收一个输入张量x和一个索引张量index沿指定维度dim按索引位置取出元素构成输出张量y是 embedding 检索、注意力掩码、排序重排等场景中的基础算子。本文以该算子目录下的 README 为骨架结合其算子注册、形状推导、Tiling 计算与 NPU 内核实现的完整源码链路讲解 GatherElementsV3 的数学定义、参数约束、aclnn 调用方式与底层工作原理帮助你掌握在 Atlas 系列产品上使用与理解该算子的完整方法。产品支持情况根据 README 的说明GatherElementsV3 的产品支持情况如下产品是否支持Atlas A2 训练系列产品 / Atlas 800I A2 推理产品√在源码层面这一支持关系体现在 gather_elements_v3_def.cpp 中通过this-AICore().AddConfig(ascend910b, aicoreConfig)为ascend910b平台即 Atlas A2 系列对应的 AI Core 架构注册了算子的 AICore 配置并在 gather_elements_v3_binary.json 中按不同数据类型声明了对应的算子二进制文件GatherElementsV3_fp32、GatherElementsV3_fp16、GatherElementsV3_bf16、GatherElementsV3_int32。功能说明与数学定义算子功能GatherElementsV3 的功能是对输入张量x中指定的维度dim进行数据聚集。索引张量index的每个位置给出一个在dim维度上的取值用于从x中取出对应元素最终输出张量y的形状与index相同。计算公式给定张量 $x$、维度 $d$ 和索引张量 $index$定义 $n$ 是 $x$ 的维度数$i_d$ 表示维度 $d$ 的索引$index_{i_d}$ 表示索引张量 $index$ 在维度 $d$ 上的第 $i_d$ 个索引值。沿指定维度 $d$ 的 gather 功能可以用如下数学公式表示$$ gather(X,index,d){i_0,i_1,\cdots,i{d-1},i_{d1},\cdots,i_{n-1}} x_{i_0,i_1,\cdots,i_{d-1},index_{i_d},i_{d1},\cdots,i_{n-1}} $$通俗理解输出张量在某个位置上的取值等于输入张量在除dim维外其余下标与输出位置一致、dim维下标取自index对应位置处的元素。计算示例假设输入张量 $x\begin{bmatrix}1 2 3\ 4 5 6\ 7 8 9\end{bmatrix}$索引张量 $index\begin{bmatrix}0 2\ 1 0\end{bmatrix}$$dim 0$那么输出张量 $y\begin{bmatrix}1 8\ 4 2\end{bmatrix}$具体计算过程如下$$ \begin{aligned} y_{0,0}x_{index_{0,0}, 0}x_{0,0}1 \ y_{0,1}x_{index_{0,1}, 1}x_{2,1}8 \ y_{1,0}x_{index_{1,0}, 0}x_{1,0}4 \ y_{1,1}x_{index_{1,1}, 1}x_{0,1}2 \end{aligned} $$可以看到当dim0时输出元素由index在该位置提供的行号 输出位置自身的列号从x中定位取值。同理若dim1则行号取自输出位置自身、列号取自index对应位置的值。参数说明GatherElementsV3 共包含 4 个参数其中 3 个张量参数、1 个标量属性汇总如下参数名输入/输出/属性描述数据类型数据格式x输入公式中的 x即被聚集的源张量FLOAT、FLOAT16、BFLOAT16、INT32NDindex输入公式中的 index提供各位置在 dim 维上的取值INT32NDy输出公式中的 y即聚集结果张量FLOAT、FLOAT16、BFLOAT16、INT32NDdim可选属性公式中的 d指定聚集发生的维度默认值为 0Int-上述定义与算子注册文件 gather_elements_v3_def.cpp 完全一致输入x支持DT_BF16 / DT_FLOAT16 / DT_FLOAT / DT_INT32四种类型输入index固定为DT_INT32输出y与x类型保持一致属性dim通过this-Attr(dim).Int(0)声明且默认值为 0所有张量均使用FORMAT_ND格式。相应地gather_elements_v3_binary.json 中为 float32、float16、bfloat16、int32 四种数据类型分别注册了独立的算子二进制形状均以-2动态 rank声明。关于输出形状数学定义要求输出形状与index一致。这一点在形状推导实现 gather_elements_v3_infershape.cpp 中体现InferShapeGatherElementsV3取context-GetInputShape(1)即index的形状并直接赋值给输出y*yShape *xShape从而保证运行时框架能够正确推导出输出张量的形状。约束说明根据 README本算子无额外约束说明。不过从实现细节可以推断出两条使用上的注意点dim取值必须小于x的维度数Tiling 阶段在 gather_elements_v3_tiling.cpp 中显式校验dim xShape.GetDimNum()并返回GRAPH_FAILED因此调用时传入的dim越界会直接导致算子执行失败index支持负索引内核在 gather_elements_v3.h 中对indexVal 0的情况执行indexVal xGatherDim_将负索引换算为合法的非负下标与 PyTorch 等框架的语义一致换算后仍然越界的索引由硬件内存访问保护兜底。调用说明基于 aclnn 的两阶段接口README 给出了一种调用方式——aclnn 调用样例代码位于 test_aclnn_gather_elements_v3.cpp。该样例完整展示了 CANN 算子的标准调用范式整体流程可分为以下步骤1. 环境初始化固定写法int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); // 内部依次调用 aclInit、aclrtSetDevice、aclrtCreateStreamInit函数依次执行aclInit(nullptr)初始化 ACL 运行时、aclrtSetDevice(deviceId)绑定设备、aclrtCreateStream(stream)创建任务流是所有 CANN 应用启动时的固定前置步骤。2. 构造输入输出 Tensor样例以selfShape {4, 2}、indexShape {4, 2}、outShape {4, 2}为例std::vectorfloat selfHostData {0, 1, 2, 3, 4, 5, 6, 7}; std::vectorint32_t indexHostData {0, 1, 2, 3, 3, 2, 1, 0}; int32_t dim 0;CreateAclTensor辅助函数完成三件事aclrtMalloc申请 device 侧内存、aclrtMemcpy将 host 数据拷贝到 device、aclCreateTensor基于 shape/strides/ND 格式创建aclTensor句柄。其中x使用ACL_FLOAT、index使用ACL_INT32与参数表中声明的数据类型一一对应。3. 两阶段算子调用aclnn 接口采用先算 workspace 再执行的两阶段设计uint64_t workspaceSize 0; aclOpExecutor* executor; // 第一段计算 workspace 大小并获取执行器 ret aclnnGatherElementsV3GetWorkspaceSize(self, index, dim, out, workspaceSize, executor); // 根据计算出的 workspaceSize 申请 device 内存 if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 第二段真正下发算子任务 ret aclnnGatherElementsV3(workspaceAddr, workspaceSize, executor, stream);第一段接口aclnnGatherElementsV3GetWorkspaceSize入参依次为self即 x、index、dim、out输出 workspace 大小与执行器句柄第二段接口aclnnGatherElementsV3负责将任务提交到指定 stream 执行。workspace 是算子运行时需要的临时内存由调用方按接口计算出的尺寸自行申请与释放。4. 同步、取回结果与资源释放ret aclrtSynchronizeStream(stream); // 同步等待任务执行结束 ret aclrtMemcpy(resultData.data(), ..., outDeviceAddr, ..., ACL_MEMCPY_DEVICE_TO_HOST); // 取回结果 // 释放aclDestroyTensor ×3、aclrtFree ×4含 workspace、aclrtDestroyStream、aclrtResetDevice、aclFinalize样例最终会逐元素打印result[i]。以示例数据x为 0~7 的 4×2 张量、dim0可以验证y[i][j] x[index[i][j]][j]即输出应为{0, 5, 4, 7, 6, 5, 2, 1}。源码级实现原理算子定义注册OpDefgather_elements_v3_def.cpp 中通过OP_ADD(GatherElementsV3)完成算子注册除了声明输入/输出/属性外还通过OpAICoreConfig指定了平台相关能力DynamicCompileStaticFlag(true)支持编译期静态信息参与动态编译DynamicRankSupportFlag(true)支持动态 rank即x/index的维度数可在运行期变化DynamicShapeSupportFlag(true)支持动态 shapePrecisionReduceFlag(true)允许在满足精度要求的前提下进行精度降低优化。这些标志与二进制配置文件中shape: [-2]的动态形状声明相呼应共同支撑算子的动态 shape 能力。Tiling 计算Host 侧Tiling 是算子性能的关键其职责是把大张量切分成适合 AI Core 上 UBUnified Buffer的 Tile。核心逻辑在 gather_elements_v3_tiling.cpp主要步骤包括平台信息获取通过PlatformAscendC获取 UB 大小ubSize与可用核数coreNum维度展平flatten以dim为界把x与index的 shape 分别拆成Pre / Gather / Post三段乘积例如xPre xShape[0..dim-1]、xGather xShape[dim]、xPost xShape[dim1..]。展平后一次聚集行的跨度即为xPostDim核数规划usedCores按总行数 idxPre × idxGather计算若总行数小于核数则只启用对应数量的核避免空转Tile 大小计算预留RESERVED_UB 1024字节、按双缓冲折半/2得到可用 UB再结合每元素需要sizeof(int32_t) typeLength字节index 与 x 各一份计算单次可处理的最大元素数并按 32 字节对齐约束向下取整最后不超过idxPost结果写入将上述参数写入GatherElementsV3TilingData结构体并通过context-SetBlockDim(usedCores)设定内核启动的核数。Tiling 数据结构定义在 gather_elements_v3_tiling_data.h共 8 个uint32_t字段xPreDim/xGatherDim/xPostDim/idxPreDim/idxGatherDim/idxPostDim/usedCores/tileSize由 Host 侧填充、Kernel 侧读取。内核实现Device 侧内核入口 gather_elements_v3.cpp 是一个带模板参数schMode调度模式的__global__ __aicore__函数模板参数通过 gather_elements_v3_tiling_key.h 中ASCENDC_TPL_ARGS_DECL声明的 0/1 两种调度模式实例化。入口函数注册并读取 Tiling 数据后实例化NsGatherElementsV3::GatherElementsV3DTYPE_X并依次调用Init与Process。算子类的核心实现在 gather_elements_v3.h要点包括流水结构使用TPipe配合TQueTPosition::VECIN, BUFFER_NUMindex 输入队列与TQueTPosition::VECOUT, BUFFER_NUMy 输出队列BUFFER_NUM 2实现双缓冲CopyIn → Compute → CopyOut三级流水重叠数据搬入CopyIn通过DataCopyPad按字节数拷贝 index 片段到局部内存并对齐到ALIGN_BYTES 32字节核心计算Compute逐元素处理先做负索引修正indexVal 0时indexVal xGatherDim_再按展平后的偏移公式xRealOffset xBase indexVal * xPostDim_ (postStart i)从全局内存xGm_取数写入yLocal——这正是 README 数学公式的代码化表达任务切分Process按rowId coreId_; rowId totalRows; rowId coreNum_的方式在多个核间按行轮转每行内部再按tileSize分片遍历idxPostDim兼顾了多核并行与 UB 容量限制。构建集成算子目录的 CMakeLists.txt 采用通用的子目录聚合模式默认遍历并添加各子目录op_host、op_kernel等当未开启ENABLE_TEST且未开启BENCHMARK时排除tests目录避免测试代码进入发布构建。整体算子库通过仓库根目录 CMakeLists.txt 统一组织编译。总结GatherElementsV3 是 CANN ops-nn 中一个接口简单、实现讲究的数据聚集算子对外它通过aclnnGatherElementsV3GetWorkspaceSizeaclnnGatherElementsV3两阶段接口向用户提供与 PyTorchgather语义一致的按维度取数能力对内它借助算子注册、动态 shape 支持、基于 UB 容量的 Tiling 切分、多核按行轮转以及双缓冲流水将索引寻址的访存密集型计算高效地映射到 Atlas A2 系列产品的 AI Core 上。理解其参数约束、负索引语义与两阶段调用范式即可在自研模型中安全、高效地使用该算子也可为阅读同类 aclnn 算子的源码实现提供参考范本。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考