CANN ops-nn aclnnInplacePut 算子详解：两段式接口实现原地 Put/累加写入

人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载aclnnInplacePut 是 CANN ops-nn 仓库 index/scatter_nd_update 模块下提供的神经网络计算算子它把selfRef视为一维张量以index中的元素作为下标在 NPU 上原地完成「替换」或「累加」写入。本文以 docs/aclnnInplacePut.md 为主体结合仓库源码与单元测试完整讲解其产品支持情况、计算语义、两段式 API 用法、参数约束、错误码以及可编译运行的调用示例读完即可在自己的工程中直接调用该算子。功能说明与计算公式接口功能将selfRef视为一维张量把index张量中元素值作为索引。当accumulate为true时把source中的元素和selfRef对应位置上的元素做累加操作当accumulate为false时把source中的元素替换掉selfRef对应位置上的元素。计算公式如下accumulate True$$selfRef[index] selfRef[index] source$$accumulate False$$selfRef[index] source$$从源码看该接口由两段式 CANN 算子 API 构成分别位于 op_api/aclnn_put.cpp 与 op_api/aclnn_put.h。其中注释清晰地给出了完整计算流程selfRef index source accumulate | | | | Contiguous Contiguous Contiguous / \ | | / Reshape Reshape Reshape / \ | / / ScatterNdAdd(ScatterNdUpdate) | Reshape | ViewCopy | selfRef即selfRef、index、source先分别经过Contiguous转成连续张量再Reshape展平随后按accumulate标志选择底层算子——accumulatetrue时调用ScatterNdAddaccumulatefalse时调用ScatterNdUpdate见 op_api/aclnn_put.cpp计算结果Reshape回原形状后若selfRef本身非连续则通过ViewCopy把结果写回selfRef。产品支持情况本算子在 CANN ops-nn 仓库当前版本支持的产品如下Ascend 950PR/Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品支持Atlas 训练系列产品支持注意Atlas 训练系列产品、Atlas 推理系列产品上数据类型不支持 BFLOAT16、COMPLEX64、COMPLEX128。这一点与源码中GetDtypeSupportList的差异一致——op_api/aclnn_put.cpp 分别定义了ASCEND910_DTYPE_DTYPE_SUPPORT_LIST不含 BF16/COMPLEX与ASCEND910B_DTYPE_DTYPE_SUPPORT_LIST含 BF16仅在特定架构且accumulatefalse时才启用后者。两段式接口与函数原型每个算子分为两段式接口必须先调用aclnnInplacePutGetWorkspaceSize接口获取计算所需 workspace 大小以及包含了算子计算流程的执行器再调用aclnnInplacePut接口执行计算。aclnnStatus aclnnInplacePutGetWorkspaceSize( aclTensor* selfRef, const aclTensor* index, const aclTensor* source, bool accumulate, uint64_t* workspaceSize, aclOpExecutor** executor)aclnnStatus aclnnInplacePut( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)两段式接口的通用约定可参考 docs/zh/context/two_phase_api.mdworkspace 指除输入/输出外算子在 NPU 上完成计算所需的临时内存且第二段接口aclnnInplacePut(...)不可重复调用。aclnnInplacePutGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfRefaclTensor*输入/输出公式中的 selfRef数据类型和 source 一致BOOL、FLOAT、FLOAT16、BFLOAT16、DOUBLE、INT8、INT16、INT32、INT64、UINT8、COMPLEX64、COMPLEX128ND-√indexaclTensor*输入公式中的 index元素个数要求和 source 保持一致。index 中的索引数据不支持越界-INT32、INT64ND-√sourceaclTensor*输入公式中的 source数据类型和 selfRef 一致元素个数和 index 一致-和 selfRef 一致ND-√accumulatebool输入累加或更新的操作类型标志位accumulate 为 True 时为累加accumulate 为 False 时为更新----workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----上述参数约束在源码的CheckParams中逐条落实op_api/aclnn_put.cppCheckNotNull检查三个 tensor 是否为空指针CheckDtypeValid检查self是否在支持列表内、index是否属于 INT32/INT64INDEX_DTYPE_SUPPORT_LIST、self与source数据类型是否一致CheckShape检查self、index的维度不超过MAX_DIM_LEN 8并校验index与source元素个数相等GetViewShape().GetShapeSize()比较。另外源码中还对空 tensor 做了特殊处理若index为空则workspaceSize直接置 0 返回成功若selfRef为空而index非空则直接报ACLNN_ERR_PARAM_INVALID。返回值与错误码aclnnStatus返回状态码具体参见 aclnn返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 selfRef、index、source 是空指针ACLNN_ERR_PARAM_INVALID161002selfRef 和 index 的数据类型不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002selfRef 和 source 的数据类型不同ACLNN_ERR_PARAM_INVALID161002source 和 index 的元素数量不等ACLNN_ERR_PARAM_INVALID161002selfRef 是空 tensorindex 不是空 tensor这些错误场景均被 tests/ut/op_api/test_aclnn_inplace_put.cpp 中的 UT 用例覆盖例如aclnnInplacePut_input_nullptr空指针、aclnnInplacePut_self_dtype_errorself 类型不支持、aclnnInplacePut_index_dtype_errorindex 类型不支持、aclnnInplacePut_self_and_source_dtype_non_consistentself 与 source 类型不一致、aclnnInplacePut_self_shape_out_of_8维度超过 8、aclnnInplacePut_index_and_source_elements_not_same元素个数不等以及aclnnInplacePut_self_empty_tensor_and_index_not_emptyself 为空、index 非空等。aclnnInplacePut 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnInplacePutGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream返回值同样是aclnnStatus具体参见 aclnn返回码。该接口在源码中为固定写法直接调用CommonOpExecutorRun完成计算op_api/aclnn_put.cpp。约束说明确定性计算Ascend 950PR/Ascend 950DT默认非确定性支持通过aclrtSetSysParamOpt开启确定性。Atlas A3 训练系列产品 / Atlas A3 推理系列产品、Atlas A2 训练系列产品 / Atlas A2 推理系列产品、Atlas 推理系列产品、Atlas 训练系列产品默认确定性。关于确定性计算的更多说明可参考 docs/zh/context/determinism_compute.md。调用示例下面示例摘自 docs/aclnnInplacePut.md仅供参考具体编译和执行过程请参考编译与运行样例。仓库中还提供了可直接编译的完整样例 examples/test_aclnn_inplace_put.cpp该样例输入self {1..8}、乱序index {0,2,4,6,1,3,5,7}、source {10..80}以accumulatefalse演示原地更新。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_put.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手册 // 根据自己的实际device填写deviceId int32_t deviceId 0; 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的接口自定义构造 std::vectorint64_t selfShape {4, 2}; std::vectorint64_t indexShape {4, 2}; std::vectorint64_t sourceShape {4, 2}; void* selfDeviceAddr nullptr; void* indexDeviceAddr nullptr; void* sourceDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* index nullptr; aclTensor* source nullptr; std::vectorfloat selfHostData {0,0,0,0,0,0,0,0}; std::vectorint64_t indexHostData {0,1,2,3,4,5,6,7}; std::vectorfloat sourceHostData{10,10,10,10,10,10,10,10}; // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_INT32, 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); // 创建source aclTensor ret CreateAclTensor(sourceHostData, sourceShape, sourceDeviceAddr, aclDataType::ACL_INT32, source); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API需要修改为具体的API名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnInplacePut第一段接口 ret aclnnInplacePutGetWorkspaceSize(self, index, source,false, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnInplacePutGetWorkspaceSize 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); } // 调用aclnnInplacePut第二段接口 ret aclnnInplacePut(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnInplacePut 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(selfShape); std::vectorfloat resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), selfDeviceAddr, size * sizeof(float), 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(source); // 7. 释放device资源需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(indexDeviceAddr); aclrtFree(sourceDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例运行流程为初始化 device/stream → 创建三个 aclTensorself形状{4,2}全 0、index为0~7的一维展开索引、source全 10→ 调用第一段接口得到 workspaceSize 并申请内存 → 调用第二段接口执行此处accumulatefalse即替换语义→ 同步等待 → 拷回结果逐元素打印。由于selfRef按一维共 8 个元素处理index {0,...,7}恰好覆盖全部位置因此执行后self的每个元素都会更新为 10。单元测试覆盖仓库为 aclnnInplacePut 提供了完整的 UT 用例 tests/ut/op_api/test_aclnn_inplace_put.cpp覆盖维度如下基础功能aclnnInplacePut_8_8_float_nd_...等用例覆盖 FLOAT/BOOL 等多类型下accumulatetrue走 ScatterNdAdd与accumulatefalse走 ScatterNdUpdate两条路径类型与格式矩阵覆盖 INT8/FLOAT/FLOAT16 及 NCHW、HWCN、NDHWC、ND 等格式以及 1~7 维 shape测试注释中还明确了底层算子支持矩阵ScatterNdAdd 的 AICORE 支持 FLOAT/FLOAT16/INT32、AICPU 支持 INT8/UINT8ScatterNdUpdate 的 AICORE 支持 BOOL/FLOAT/FLOAT16、AICPU 支持 DOUBLE/INT8/INT16/INT32/INT64/UINT8/COMPLEX64/COMPLEX128边界与特殊场景空 tensorself 非空、index 为空、边界值如 FLOAT16 上下限 65504/-65504、非连续 tensor通过 view/storage 维度构造异常路径空指针、dtype 不支持、self 与 source 类型不一致、维度超 8、index 与 source 元素个数不等、self 空而 index 非空等均断言返回ACLNN_ERR_PARAM_INVALID/ACLNN_ERR_PARAM_NULLPTR。源码中的底层实现API 层op_api/aclnn_put.cpp 实现两段式接口aclnnInplacePutGetWorkspaceSize内部通过CREATE_EXECUTOR创建执行器、CheckParams校验参数随后用l0op::Contiguous、l0op::Reshape、l0op::ScatterNdAdd/l0op::ScatterNdUpdate、l0op::ViewCopy编排计算图最终由uniqueExecutor-GetWorkspaceSize()汇总 workspace 大小op_api/aclnn_put.cpp。底层算子ScatterNdUpdate的 l0 接口声明见 op_api/scatter_nd_update.h其算子定义含 var/indices/updates 输入、var 输出、可选属性use_locking默认 false见 op_host/scatter_nd_update_def.cpp算子内核实现位于 op_kernel/scatter_nd_update.cpp 及 arch22/arch35 子目录下的多种实现如 deterministic_simd/simt、large_index、linear_index、no_sort 等变体。融合/图下沉仓库还提供了 TensorScatterUpdate 融合 passop_graph/fusion_pass/tensor_scatter_update_fusion_pass.cpp说明该类原地更新语义在构图阶段也有对应的融合与下沉路径。相关文档两段式接口说明aclnn 返回码编译与运行样例确定性计算说明同目录姊妹算子aclnnScatterNdUpdate赞分享人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载相关推荐CANN ops-nn ForeachLog2 算子详解aclnnForeachLog2 两段式接口使用与实现原理CANN ops nn ForeachLog2 算子详解aclnnForeachLog2 两段式接口使用与实现原理 本文以 CANN ops nn 算子库中的人工智能算子库深度学习CANNAscend9大网盘直链下载助手告别限速体验全速下载的终极解决方案9大网盘直链下载助手告别限速体验全速下载的终极解决方案 还在为网盘下载速度慢而烦恼吗LinkSwift网盘直链下载助手是一款基于JavaScript开发的人工智能算子库深度学习CANNAscendCANN ops-nn 算子详解aclnnSquaredRelu 两段式接口使用指南与实现原理CANN ops nn 算子详解aclnnSquaredRelu 两段式接口使用指南与实现原理 本文以 CANN ops nn 仓库中 aclnnSquare人工智能算子库深度学习CANNAscend上一篇5个高效技巧让ComfyUI-WanVideoWrapper发挥最大价值下一篇推荐开源项目Liebling精致的Ghost博客主题创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考