CANN ops-nn ScatterElementsV2 算子深度解析：基于 aclnnScatter/aclnnInplaceScatter 的 Tensor 散落更新实战指南

CANN ops-nn ScatterElementsV2 算子深度解析基于 aclnnScatter/aclnnInplaceScatter 的 Tensor 散落更新实战指南【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nnScatterElementsV2 是 CANN ops-nn 算子库中负责“按索引将源张量 src 的值替换、累加或累乘到目标张量 self”的散落Scatter类算子是索引类算子家族index/目录中应用最广泛的基础组件之一。本文以 scatter_elements_v2/README.md 为骨架结合仓库内算子定义、Host 侧 Tiling、Kernel 侧实现与单元测试源码系统讲解 ScatterElementsV2 的功能语义、参数规范、两段式 aclnn 调用方式与底层实现原理读完可直接上手编写、编译并运行基于 aclnnScatter / aclnnInplaceScatter 的 NPU 加速程序。产品支持情况ScatterElementsV2 算子在不同产品形态上的支持情况如下以当前仓库 README 声明为准产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×从算子注册源码 scatter_elements_v2_def.cpp 可以看到该算子的 AICore 配置分三档注册ascend910b、ascend910_93走同一套配置ascend310p走单独的scatter_elements_v2_apt配置ascend950与ascend350共用注册配置且ExtendCfgInfo(opFile.value, scatter_elements_v2_apt)指向专用于高算力产品的 APT 实现这从源码层面印证了不同产品的差异化适配路径。功能说明算子功能将 tensorsrc中的值按指定的轴dim与方向以及对应的位置关系index逐个替换none、累加add或累乘mul至 tensorself中。对于一个 3D tensorself会按照如下规则进行更新self[index[i][j][k]][j][k] src[i][j][k] # 如果 dim 0 reduction 1累加 self[i][index[i][j][k]][k] * src[i][j][k] # 如果 dim 1 reduction 2累乘 self[i][j][index[i][j][k]] src[i][j][k] # 如果 dim 2 reduction 0替换也就是说index中每个元素index[i][j][k]表示在dim轴上src[i][j][k]应写入self的目标位置其余轴的坐标由(i, j, k)原样保留。在计算时需要满足以下要求self、index和src的维度数量必须相同。对于每一个维度d有index.size(d) src.size(d)的限制。对于每一个维度d如果d ! dim有index.size(d) self.size(d)的限制。dim的值必须在[-self的维度数量, self的维度数量-1]之间。self的维度数应小于等于 8。index中对应维度dim的值索引值大小必须在[0, self.size(dim)-1]之间。上述约束在 aclnn 接口的 Host 侧校验中均有对应实现见 aclnn_scatter.cpp 中的CheckShape、CheckTensorDim、CheckDimRange等函数CheckTensorDim校验三个张量维度数一致CheckShape逐维比较index.shape src.shape、除dim外index.shape self.shape并调用OP_CHECK_MAX_DIM限制不超过 8 维CheckDimRange将dim限制在[-N, N-1]。参数说明下表为 ScatterElementsV2 算子的核心参数源自 README 参数表结合接口头文件 aclnn_scatter.h 补充说明参数名输入/输出/属性描述数据类型数据格式self输入公式中的selfDevice 侧 aclTensor数据类型需与 src 一致UINT8、INT8、INT16、INT32、INT64、BOOL、FLOAT16、FLOAT32、DOUBLE、COMPLEX64、COMPLEX128、BFLOAT16NDdim输入用来 scatter 的维度数据类型为 INT64范围[-N, N-1]int64_t-index输入公式中的indexDevice 侧 aclTensor指定 src 写入 self 的位置INT32、INT64NDsrc输入公式中的srcDevice 侧 aclTensor数据类型需与 self 一致UINT8、INT8、INT16、INT32、INT64、BOOL、FLOAT16、FLOAT32、DOUBLE、COMPLEX64、COMPLEX128、BFLOAT16NDreduction输入Host 侧的字符串选择应用的 reduction 操作string-out输出公式中的输出shape 与数据类型需与 self 一致FLOAT、FLOAT16、DOUBLE、BFLOAT16、INT8、INT16、INT32、INT64、UINT8、BOOLND关于reduction与数值取值的对应关系aclnn_scatter.cpp 中定义了常量REDUCE_NONE 0替换、REDUCE_ADD 1累加、REDUCE_MUL 2累乘并通过GetReduceStr将整数映射为字符串none/add/mul。此外接口层还扩展支持REDUCE_MAX 3、REDUCE_MIN 4、REDUCE_MEAN 5映射到max/min/mean但对应分支会打印Maximum/Minimum/Mean mode is experimental!警告日志属于实验性能力正式业务请优先使用 0/1/2 三档。需要说明的是README 参数表中的self数据类型集合与算子注册文件 scatter_elements_v2_def.cpp 中按产品分档的数据类型集合并不完全一致例如ascend910b/ascend910_93档位注册了DT_FLOAT/DT_FLOAT16/DT_INT32/DT_UINT8/DT_INT8/DT_BF16/DT_BOOL的组合而ascend950/ascend350档位支持DT_INT64/DT_INT32/DT_INT16/DT_UINT8/DT_INT8/DT_FLOAT/DT_FLOAT16/DT_BF16。实际可用类型以具体产品形态为准调用前建议核对目标产品的算子配置。约束说明约束说明无额外约束README 原文。确定性说明源自同目录文档 aclnnScatteraclnnInplaceScatter.mdAscend 950PR/Ascend 950DT默认非确定性支持通过aclrtCtxSetSysParamOpt开启确定性但当 reduction 为累乘场景时开启确定性也不生效依旧为非确定性。Atlas A3 训练/推理系列产品、Atlas A2 训练/推理系列产品默认确定性。调用说明aclnn 两段式接口ScatterElementsV2 算子通过 aclnn 接口对外暴露调用入口统一为aclnnScatter及其 inplace 变体aclnnInplaceScatter两者功能相同区别在于aclnnScatter需要新建输出张量对象存储计算结果aclnnInplaceScatter无需新建输出张量直接在输入张量selfRef的内存中存储计算结果。调用方式调用样例说明aclnn 调用test_aclnn_scatter.cpp通过 aclnnScatteraclnnInplaceScatter 接口方式调用 ScatterElementsV2 算子每个算子采用“两段式接口”设计必须依次调用两段接口才能完成计算第一段接口获取 workspace 大小与执行器。入参校验空指针、数据类型、数据格式、shape、dim 范围等均在此阶段完成aclnnStatus aclnnScatterGetWorkspaceSize( const aclTensor* self, int64_t dim, const aclTensor* index, const aclTensor* src, int64_t reduce, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)第二段接口执行计算。使用第一段返回的 workspace 与 executor 在指定 stream 上异步执行aclnnStatus aclnnScatter( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream)inplace 变体两段接口原型如下aclnnStatus aclnnInplaceScatterGetWorkspaceSize( aclTensor* selfRef, int64_t dim, const aclTensor* index, const aclTensor* src, int64_t reduce, uint64_t* workspaceSize, aclOpExecutor** executor) aclnnStatus aclnnInplaceScatter( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)第一段接口返回码说明aclnnScatterGetWorkspaceSize完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001self、index、src 或 out 存在空指针ACLNN_ERR_PARAM_INVALID161002self、index、src 或 out 的数据类型不在支持范围内self、src、out 的数据类型不一致self、index、src 的维度数不一致self 与 out 的 shape 不一致dim 不在[-N, N-1]之间self 的维度数超过 8shape 不满足index.size(d) src.size(d)且d ! dim时index.size(d) self.size(d)的限制aclnnInplaceScatterGetWorkspaceSize的校验逻辑相同只是将 self 换成 selfRef同时省略了 self 与 out shape 一致性检查因为二者共用同一内存。无转置快速路径在进入通用执行路径前接口会先调用 scatter_elements.h 中声明的l0op::SupportNoTranspose判断是否满足“无转置”快速路径条件若满足则直接走ExecScatterNoTranspose跳过dim轴搬移从源码注释看该路径对self与out内存地址不同的场景会自动插入TensorMove拷贝避免破坏输入数据。调用示例完整可编译代码以下示例源自 aclnnScatteraclnnInplaceScatter.md 的调用示例完整展示了 aclnnScatter 从初始化到结果回收的七个步骤。示例参数为self形状{3, 4}、index形状{2, 3}、src形状{2, 3}、dim 1、reduce 1累加#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_scatter.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法资源初始化 auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return 0; } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请device侧内存 auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); // 计算连续tensor的strides std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1.固定写法device/stream初始化参考acl API手册 int32_t deviceId 0; // 根据自己的实际device填写deviceId aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出需要根据API的接口自定义构造 int64_t dim 1; int64_t reduce 1; std::vectorint64_t selfShape {3, 4}; std::vectorint64_t indexShape {2, 3}; std::vectorint64_t srcShape {2, 3}; std::vectorint64_t outShape {3, 4}; void* selfDeviceAddr nullptr; void* indexDeviceAddr nullptr; void* srcDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* index nullptr; aclTensor* src nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}; std::vectorint64_t indexHostData {0, 0, 2, 1, 0, 2}; std::vectorfloat srcHostData {-1, -2, -3, -4, -5, -6}; std::vectorfloat outHostData {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建index aclTensor ret CreateAclTensor(indexHostData, indexShape, indexDeviceAddr, aclDataType::ACL_INT64, index); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建src aclTensor ret CreateAclTensor(srcHostData, srcShape, srcDeviceAddr, aclDataType::ACL_FLOAT, src); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建out aclTensor ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API需要修改为具体的API名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnScatter第一段接口 ret aclnnScatterGetWorkspaceSize(self, dim, index, src, reduce, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnScatterGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } // 调用aclnnScatter第二段接口 ret aclnnScatter(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnScatter failed. ERROR: %d\n, ret); return ret); // 4.固定写法同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 获取输出的值将device侧内存上的结果拷贝至host侧需要根据具体API的接口定义修改 auto size GetShapeSize(outShape); std::vectorfloat resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i size; i) { LOG_PRINT(result[%ld] is: %f\n, i, resultData[i]); } // 6. 释放aclTensor和aclScalar需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyTensor(index); aclDestroyTensor(src); aclDestroyTensor(out); // 7. 释放device资源需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(indexDeviceAddr); aclrtFree(srcDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }inplace 版本示例整体结构与上述代码一致差异仅在以下几点不需要创建out张量selfRef直接作为输入/输出使用第一段接口换为aclnnInplaceScatterGetWorkspaceSize(selfRef, dim, index, src, reduce, workspaceSize, executor)第二段接口换为aclnnInplaceScatter(workspaceAddr, workspaceSize, executor, stream)结果直接从selfRefDeviceAddr拷回 Host 侧查看。算子执行流程与底层实现解析1. 算子定义与 shape 推导Host 侧算子通过 scatter_elements_v2_def.cpp 中的OP_ADD(ScatterElementsV2)注册输入为var、indices、updates输出为var属性axis默认 0、reduction默认none、include_self默认true并显式开启DynamicCompileStaticFlag、DynamicFormatFlag、DynamicRankSupportFlag、DynamicShapeSupportFlag即该算子支持动态 shape、动态 rank 与动态格式编译。shape 推导在 scatter_elements_infershape.cpp 中实现逻辑非常简单直接InferShape4ScatterElementsV2将输出 shape 直接赋值为输入var的 shape*var_out_shape *var_in_shapeInferDataType4ScatterElementsV2将输出数据类型继承自var即输出与目标张量同 shape 同类型。2. aclnn 入口参数校验与执行路由接口层 aclnn_scatter.cpp 是理解整个算子执行流程的最佳入口其核心流程可概括为入参检查CheckNotNull空指针、CheckDtypeValiddtype 支持列表与 self/src/out 一致性、格式检查仅支持 NDNZ 格式打印警告。特殊场景短路index为空时直接执行ViewCopy(self → out)返回即“无索引则原样输出 self”。0 维张量处理InitializeTensor先通过l0op::Contiguous转为连续张量0 维张量通过UnsqueezeNd扩为 1 维计算输出侧再SqueezeNd还原。BF16 降级非 910B/Regbase 平台将 BF16 强转为 FP32 计算AICore 侧不支持 BF16 精度的平台。special 路由expandFlag当 index 的 strides 呈现[..., 1, 0]的广播/扩展形态且规模满足条件时走DoScatterAddWithSorted基于 scatter_add_with_sorted 与Sort的确定性累加路径该路径还会依据ACL_OPT_DETERMINISTIC系统参数决定是否启用 sorted 分支Regbase 平台 2D 场景下可路由到 AICore 的ScatterAdd广播场景IsRouteToUpdate命中且满足约束时路由到ScatterUpdate否则进入通用ScatterElements路径。转置统一当dim不在最后一维时构造 permutation 将dim轴交换到末维调用l0op::Transpose对 self/index/src 统一转置计算完成后对输出再转置回来从而让 Kernel 侧只需处理“最后一维 scatter”这一种布局。线性索引化910B 系列将 index 通过l0op::LinearIndex转换为全局线性索引同时处理负数索引再交由底层 scatter 实现。3. Kernel 侧实现Kernel 入口为 scatter_elements_v2.cpp 中的scatter_elements_v2通过TILING_KEY宏根据(valueType, indexType)组合如 110 对应 floatint、220 对应 halflong 等实例化模板。内部根据IsCacheOpTiling判定结果分流命中 Cache 算子 TilingcoreNums/xDim0/xDim1/indicesDim0/...任一非 0时走 scatter_elements_v2_low_memory 目录下的ExecTransposeAndScatterElements低内存实现含transpose_tile_forward/backward、transpose_batch_forward/backward、init_gather_offset、scatter_elements_cache等模块否则走传统KernelScatterElementsV2的ProcessSmall/ProcessScatter分支。同时 arch35 目录提供了scatter_elements_v2_deterministic.h、scatter_elements_v2_multi.h、scatter_elements_with_sorted.h三个变体头文件分别对应确定性模式、多核模式与 sorted 模式与文档中“Ascend 950 默认非确定性、可通过系统参数开启确定性”的说明相互印证。4. Tiling 与多架构适配Tiling 计算在 arch22/scatter_elements_v2_tiling.cpp 与 arch35 目录下分架构实现scatter_elements_v2_asc_tiling.cpp、scatter_elements_v2_base_tiling.cpp不同产品通过独立的 binary 配置op_host/config 下的ascend310p、ascend350、ascend910_93、ascend910b、ascend950五个子目录挂载对应的算子二进制。测试与验证仓库为该算子提供了完整的单测与系统测试覆盖aclnn 接口 UTtests/ut/op_api/test_aclnn_scatter.cpp 使用OP_API_UT框架分别对reduction 0/1/2/3/4五种模式调用TestGetWorkspaceSize校验成功路径并用test_run_invalid覆盖非法场景预期返回ACLNN_ERR_PARAM_INVALID。值得关注的是空 tensor 用例selfindexsrc全空、仅 index 空、仅 src 空被断言为合法而self空或src空此时 index 非空被断言为非法与接口实现中“index 为空直接拷贝输出”的短路逻辑一致。Host 侧 UTtests/ut/op_host覆盖 infershape 单测test_scatter_elements_v2_infershape.cpp以及 arch22/arch35 两套 Tiling 单测。Kernel 侧 UTtests/ut/op_kernel 通过gen_data.py/gen_tiling.py生成输入与 Tiling 数据运行 test_scatter_elements_v2.cpp 验证 Kernel 正确性。系统测试 STtests/st下按接口分组aclnnScatter、aclnnScatterAdd、aclnnScatterReduce、aclnnScatterValue每组的 JSON 描述用例参数、executor_*.py负责组装与执行另有arch35目录下的自适应adaptive/Warp 累加专项用例test_scatter_elements_v2_adaptive_add.py、test_scatter_elements_v2_warp_add.py覆盖确定性模式下的多核累加正确性。总结ScatterElementsV2 是 CANN ops-nn 索引类算子中语义最基础、覆盖产品最广的 scatter 实现通过 aclnnScatter 系列两段式接口用户只需依次完成“GetWorkspaceSize 获取 workspace 与 executor”和“执行计算”两步即可在 NPU 上完成替换/累加/累乘三种模式的按索引散落更新。从源码看其工程实现充分体现了 CANN 算子库的典型分层设计——aclnn 入口层负责校验与路由、Host 侧负责 shape 推导与 Tiling、Kernel 侧通过模板特化与低内存/确定性/多核等变体适配不同算力平台。理解本算子后可以同类推广到index/目录下的 scatter_add、scatter_nd、scatter_update 等系列算子快速掌握 CANN 索引类算子的通用开发与调优方法。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考