CANN ops-math aclnnReflectionPad3d 算子详解:3D 反射填充的两段式接口、参数约束与实现原理
CANN ops-math aclnnReflectionPad3d 算子详解3D 反射填充的两段式接口、参数约束与实现原理【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本篇技术指南围绕 CANN ops-math 数学算子库中的aclnnReflectionPad3d接口展开系统讲解 3D 反射填充Reflection Padding算子在 Ascend 硬件上的使用方法包括产品支持情况、两段式 API 的调用范式与函数原型、self/padding/out三个核心入参的完整约束、常见错误码及排查方向并结合本仓库源码剖析接口内部的参数校验、连续化处理、PadV3/MirrorPad 底层算子分发与 AiCore/AiCpu 选择逻辑。读完本文你将能够独立完成aclnnReflectionPad3d的 host 侧编码、编译运行与结果校验并理解该接口与底层MirrorPad算子之间的实现关系。一、产品支持情况aclnnReflectionPad3d面向不同昇腾产品的能力支持矩阵如下与 conversion/mirror_pad/README.md 中 MirrorPad 算子的产品支持情况一致产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品支持从源码看平台差异集中体现在数据类型支持列表上见后文“平台差异与数据类型限制”小节。二、功能说明什么是 3D 反射填充aclnnReflectionPad3d的接口功能为3D 反射填充以输入 tensor 各维度边界的“镜像”为来源对 tensor 的最后三维进行边界扩充镜像时不包含边界元素本身即 REFLECT 模式区别于会包含边界本身的 SYMMETRIC 对称填充。官方示例如下输入tensor([[[[[0,1], [2,3]], [[4,5], [6,7]]]]]) padding([1,1,1,1,1,1]) 输出为 ([[[[[7,6,7,6], [5,4,5,4], [7,6,7,6], [5,4,5,4]], [[3,2,3,2], [1,0,1,0], [3,2,3,2], [1,0,1,0]], [[7,6,7,6], [5,4,5,4], [7,6,7,6], [5,4,5,4]], [[3,2,3,2], [1,0,1,0], [3,2,3,2], [1,0,1,0]]]]])该示例中输入self的 shape 为(1,1,2,2,2)padding六个值均为 1最后一维左右各补 1、倒数第二维上下各补 1、倒数第三维前后各补 1输出 shape 为(1,1,4,4,4)。观察输出可以发现每一维的扩充值都来自该维边界元素的反射镜像且边界元素本身不会被复制到填充区例如最后一维[0,1]补成[1,0,1,0]0 和 1 各自成为对方的镜像来源。三、两段式接口与函数原型aclnnReflectionPad3d采用 CANN aclnn 接口通用的两段式调用范式详见 docs/zh/context/two_phase_api.md先调用第一段接口aclnnReflectionPad3dGetWorkspaceSize完成入参校验、构建执行器aclOpExecutor并返回计算所需 workspace 大小再调用第二段接口aclnnReflectionPad3d传入第一段返回的 workspace 与 executor真正在指定 stream 上执行计算。函数原型如下aclnnStatus aclnnReflectionPad3dGetWorkspaceSize( const aclTensor *self, const aclIntArray *padding, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnReflectionPad3d( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)两段接口的声明位于 conversion/mirror_pad/op_api/aclnn_reflection_pad3d.h其中第一段接口的注释明确说明self支持非连续 Tensor、ND 格式、四维或五维、在最后三维做 padpadding为 INT64、长度 6out维度计算规则与self一致并支持非连续 Tensor。四、aclnnReflectionPad3dGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入待填充的原输入数据。维度支持四维或五维在最后三维做 pad。BOOL、INT8、UINT8、INT16、UINT16、FLOAT16、BFLOAT16、INT32、UINT32、FLOAT32、INT64、UINT64、DOUBLE、COMPLEX64、COMPLEX128、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN、FLOAT8_E8M0ND4-5√paddingaclIntArray*输入输入中需要填充的大小。长度为6数值依次代表左右上下前后需要填充的值。padding 前两个数值需小于 self 最后一维度的数值中间两个数值需小于 self 倒数第二维度的数值后两个数值需小于 self 倒数第三维度的数值。INT64ND-√outaclTensor*输出填充后的输出结果。维度与 self 一致out 倒数第三维度的数值等于 self 倒数第三维度的数值加 padding 后两个值out 倒数第二维度的数值等于 self 倒数第二维度的数值加 padding 中间两个值out 最后一维度的数值等于 self 最后一维度的数值加 padding 前两个值。BOOL、INT8、UINT8、INT16、UINT16、FLOAT16、BFLOAT16、INT32、UINT32、FLOAT32、INT64、UINT64、DOUBLE、COMPLEX64、COMPLEX128、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN、FLOAT8_E8M0ND4-5√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小。-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程。-----平台差异与数据类型限制上表中的数据类型为接口全量声明但不同产品存在差异与 aclnn_reflection_pad3d.cpp 中按平台定义的三个 dtype 支持列表一一对应Atlas A3 训练/推理系列产品、Atlas A2 训练/推理系列产品数据类型不支持UINT16、UINT32、UINT64、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN、FLOAT8_E8M0即支持 BOOL、INT8、UINT8、INT16、FLOAT16、BFLOAT16、INT32、FLOAT32、INT64、UINT64→除外、DOUBLE、COMPLEX64、COMPLEX128 等其余类型Atlas 推理系列产品、Atlas 训练系列产品数据类型不支持BFLOAT16、UINT16、UINT32、UINT64、COMPLEX64、COMPLEX128、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN、FLOAT8_E8M0对应源码中ASCEND910_DTYPE_DTYPE_SUPPORT_LIST的 9 种类型。此外源码要求self与out的数据类型必须一致CheckDtypeValid中通过OP_CHECK_DTYPE_NOT_MATCH校验且二者 view format 必须一致CheckFormat。返回值aclnnStatus返回状态码具体参见 docs/zh/context/aclnn_return_code.md。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001Tensor 为空指针。ACLNN_ERR_PARAM_INVALID161002self、padding 和 out 的数据类型或数据格式不在支持的范围之内。self、padding 和 out 的输入 shape 在支持范围之外。五维 self 为空 tensor 且存在非 batch size 维度的大小为 0。四维 self 不支持为空 tensor。padding 的数值大于等于 self 对应维度的值。out 后三维度的值不等于 self 后三维度的值加对应 padding。out 的 shape 与实际输出 shape 不匹配。这些校验在 aclnn_reflection_pad3d.cpp 的CheckParams中以固定顺序执行先检查空指针CheckNotNull再检查数据类型CheckDtypeValid然后检查数据格式CheckFormat最后检查 shape 与 padding 数值CheckShape见 L85-L120。其中CheckShape除维度与长度约束外还会校验padding六个数分别小于self对应维度以及out最后三维恰好等于self对应维加对应 padding——这与参数表中的使用说明完全一致。五、aclnnReflectionPad3d 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnReflectionPad3dGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。返回值aclnnStatus返回状态码具体参见 docs/zh/context/aclnn_return_code.md。第二段接口本身不做参数校验仅调用框架能力执行计算。在 aclnn_reflection_pad3d.cpp 中其实现即为一行CommonOpExecutorRun(workspace, workspaceSize, executor, stream)。六、约束说明使用aclnnReflectionPad3d时需注意以下约束确定性计算aclnnReflectionPad3d默认确定性实现多次运行结果一致无需额外配置。超时风险如果计算量过大可能会导致算子执行超时aicore error 类型报错errorStr 为timeout or trap error典型场景为最后 2 轴合轴小于 16、而前面的轴合轴超大。空 tensor 规则五维self允许为空 tensor但除 batch size 维度第 0 维外其余维度大小不能为 0四维self不支持为空 tensor。源码中对应逻辑见 aclnn_reflection_pad3d.cpp当self或out为空时直接置workspaceSize 0四维输入报ACLNN_ERR_PARAM_INVALID五维输入仅在非 batch 维度为 0 时报错。七、调用示例示例代码如下亦可直接参考仓库中的 test_aclnn_reflection_pad_3d.cpp具体编译和执行过程请参考 docs/zh/context/compile_and_run_sample.md。#include acl/acl.h #include aclnnop/aclnn_reflection_pad3d.h #include iostream #include vector #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 shape_size 1; for (auto i : shape) { shape_size * i; } return shape_size; } 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根据自己的需要处理 CHECK_RET(ret 0, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2.构造输入与输出需要根据API的接口定义构造 std::vectorint64_t selfShape {1, 1, 2, 2, 2}; std::vectorint64_t outShape {1, 1, 4, 4, 4}; void* selfDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclIntArray* padding nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0, 1, 2, 3, 4, 5, 6, 7}; std::vectorint64_t paddingData {1, 1, 1, 1, 1, 1}; std::vectorfloat outHostData(GetShapeSize(outShape), 0); // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建padding aclIntArray padding aclCreateIntArray(paddingData.data(), 6); CHECK_RET(padding ! nullptr, 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; // 调用aclnnReflectionPad3d第一段接口 ret aclnnReflectionPad3dGetWorkspaceSize(self, padding, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnReflectionPad3dGetWorkspaceSize 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;); } // 调用aclnnReflectionPad3d第二段接口 ret aclnnReflectionPad3d(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnReflectionPad3d 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(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); aclDestroyIntArray(padding); aclDestroyTensor(out); // 7.释放device资源需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例要点拆解资源初始化固定写法aclInit → aclrtSetDevice → aclrtCreateStream对应示例中的Init函数Tensor 构造CreateAclTensor完成 device 内存申请aclrtMalloc、host→device 数据拷贝aclrtMemcpy、连续 tensor strides 计算最后用aclCreateTensor创建aclTensorpadding通过aclCreateIntArray(paddingData.data(), 6)创建长度为 6两段式调用先aclnnReflectionPad3dGetWorkspaceSize拿到workspaceSize与executor按需aclrtMalloc申请 workspace再aclnnReflectionPad3d执行结果回收aclrtSynchronizeStream同步后将结果从 device 拷回 host 并打印最后依次释放 tensor、device 内存、stream 并aclFinalize。上述代码中selfShape {1,1,2,2,2}、padding {1,1,1,1,1,1}、outShape {1,1,4,4,4}与第二节官方示例完全对应可直接运行验证输出。八、接口内部实现流程解析aclnnReflectionPad3dGetWorkspaceSize的实现aclnn_reflection_pad3d.cpp清晰展示了该类接口的典型内部流水线创建执行器CREATE_EXECUTOR()创建aclOpExecutor唯一实例失败返回ACLNN_ERR_INNER_CREATE_EXECUTOR参数校验按“空指针 → 数据类型 → 数据格式 → shape/padding”的顺序执行CheckParams空 tensor 处理self或out为空时直接返回workspaceSize 0见第六节约束输入连续化若self非连续先通过l0op::Contiguous转为连续 tensor非连续 Tensor 的底层处理可参考 docs/zh/context/non_contiguous_tensor.md底层计算按数据类型走两条路径详见 reflection_pad_common.h复数路径COMPLEX64/COMPLEX128ProcessPadV3四维输入先UnsqueezeNd增维到五维调用l0op::PadV3(..., reflect, ...)完成反射填充再SqueezeNd还原维度其他类型路径ProcessMirrorPad将长度 6 的 padding 通过GetPaddingTensor按“后三维、每维一对”重排并Reshape为[dim, 2]调用l0op::MirrorPad(..., REFLECT, ...)输出 shape 比对CheckShapeAndScalarSame(padResult, out)校验实际计算结果 shape 与传入的out一致输出连续化回写若out是非连续 tensor通过l0op::ViewCopy把计算得到的连续结果写回out的视图返回 workspace*workspaceSize uniqueExecutor-GetWorkspaceSize()并把 executor 释放给调用方。底层 MirrorPad 的 AiCore/AiCpu 分发l0op::MirrorPadmirrorpad.cpp先执行INFER_SHAPE推导输出 shape然后依据平台 dtype 支持列表决定计算载体数据类型命中 AiCore 列表则走MirrorPadAiCoreADD_TO_LAUNCHER_LIST_AICORE否则走MirrorPadAiCpuADD_TO_LAUNCHER_LIST_AICPU属性携带Tpaddings与mode。AiCore 支持列表同样分平台定义如 910B 支持 FLOAT16/FLOAT/INT32/INT16/INT64/BF16RegBase 支持更多类型。算子定义与 shape 推导算子侧定义见 mirror_pad_def.cppMirrorPad注册输入x、paddingsValueDepend(OPTIONAL)即 shape 推导依赖其数值、输出ymode属性为 REQUIRED 且默认值REFLECTAICore 配置声明了动态编译、动态 rank、动态 shape 支持并为ascend950、ascend350挂接了mirror_pad_apt内核实现。shape 推导在 mirror_pad_infershape.cpp 中实现当输入为 unknown rankIsUnknownRank时输出同样设为 unknown rank否则复用 pad_v3 公共的InferShapeForPadWithPaddingTensor依据 paddings 数值计算输出 shape。九、测试与验证仓库为aclnnReflectionPad3d提供了完整的 UT 与 ST 用例可用于验证接口行为单元测试tests/ut/op_api/test_aclnn_reflection_pad3d.cpp覆盖正常用例如{1,1,2,2,2}FLOAT16 输入 padding 全 1 的GetWorkspaceSize调用返回ACL_SUCCESS、空 tensor 用例首维为 0 的五维输入、self/padding/out空指针场景期望ACLNN_ERR_PARAM_NULLPTR等异常分支ST 用例tests/st/aclnnReflectionPad3d/executor_aclnnReflectionPad3d.py在 CPU 侧以torch.nn.ReflectionPad3d作为 golden 参考实现生成期望输出FLOAT16 输入会先转 FLOAT32 计算再转回与 NPU 端执行结果比对数据驱动用例定义见同目录的atk_aclnnReflectionPad3d.json。十、延伸阅读aclnnReflectionPad3d是 mirror_pad 算子模块在 3D 场景的 aclnn 封装该模块还提供 1D/2D 版本接口且三者共用同一套底层实现reflection_pad_common.h中ProcessMirrorPad/ProcessPadV3按输入维度泛化处理conversion/mirror_pad/README.mdMirrorPad 算子整体说明含 REFLECT/SYMMETRIC 两种模式的差异conversion/mirror_pad/docs/aclnnReflectionPad1d.md最后一维反射填充的 1D 接口文档conversion/mirror_pad/docs/aclnnReflectionPad2d.md最后两维反射填充的 2D 接口文档docs/zh/context/two_phase_api.mdaclnn 两段式接口通用范式docs/zh/context/aclnn_return_code.mdaclnn 返回码定义docs/zh/context/compile_and_run_sample.md样例编译与运行指引conversion/pad_v3/op_api/padv3.h复数路径所复用的 PadV3 底层算子头文件由reflection_pad_common.h直接引用。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考