CANN ops-math 算子实战ReflectionPad3dGrad 反射填充反向算子与 aclnnReflectionPad3dBackward 两段式接口【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本文是 CANN / ops-math 仓库中 ReflectionPad3dGrad 算子的完整技术指南。该算子实现三维反射填充ReflectionPad3d的反向传播是 PyTorchnn.ReflectionPad3d在 NPU 上训练时梯度回传的关键一环。读完本文你将掌握该算子的产品支持范围、输入输出参数与约束、aclnnReflectionPad3dBackward两段式接口的完整调用流程并能结合仓库源码理解其 tiling 分派与内核实现的底层原理。一、算子定位ReflectionPad3d 的反向传播ReflectionPad3dGrad是 CANN ops-math 数学算子库中conversion数据转换类目录下的一个反向算子功能是计算aclnnReflectionPad3d接口的正向反射填充操作对应的梯度。在深度学习中反射填充常用于卷积网络的边界处理训练阶段通过本算子把上游传来的梯度gradOutput映射回未填充前的输入张量gradInput从而完成对前向输入的梯度回传。正向算子aclnnReflectionPad3d的接口文档位于 conversion/mirror_pad/docs/aclnnReflectionPad3d.md其语义为对输入 tensor 的左右上下前后六个方向做镜像式填充填充值取自边缘的反射映射而非补零。例如输入[[[[[0,1],[2,3]]],[[[4,5],[6,7]]]]]、padding 为[1,1,1,1,1,1]时输出会沿各维度把边缘值镜像复制到填充区域。本反向算子则与之对应把填充后张量的梯度按同样规则折回到原始尺寸上。二、产品支持情况不同硬件平台对该算子的支持情况如下表所示产品是否支持Ascend 950PR/Ascend 950DT×Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×需要注意的是算子接口文档 aclnnReflectionPad3dBackward.md 中的产品支持矩阵与上述算子 README 略有差异该文档标注 Ascend 950PR/950DT、Atlas 推理系列产品、Atlas 训练系列产品为支持这从侧面说明产品支持情况会随 CANN 版本演进而变化实际部署时请以当前 CANN 版本配套的算子清单为准。从算子注册配置可以印证当前的 AICore 适配范围op_host/reflection_pad3d_grad_def.cpp 中通过this-AICore().AddConfig(ascend910b)与this-AICore().AddConfig(ascend910_93)注册了ascend910b对应 Atlas A2 系列和ascend910_93对应 Atlas A3 系列两个平台的核函数配置与 README 中Atlas A2/A3 系列支持的结论一致。三、参数说明算子输入输出参数如下表前向的输入张量self在反向中仅用于形状参考不参与梯度计算参数名输入/输出/属性描述数据类型数据格式gradOutput输入反向传播的输入即正向反射填充输出张量的梯度FLOAT16、FLOAT32、DOUBLE、COMPLEX64、COMPLEX128NDself输入正向的输入张量未填充前的张量用于确定 gradInput 的形状FLOAT16、FLOAT32、DOUBLE、COMPLEX64、COMPLEX128NDpadding输入长度为 6 的数组数值依次代表左、右、上、下、前、后需要填充的值aclIntArray 数组INT64-gradInput输出反向传播的输出与 self 形状一致FLOAT16、FLOAT32、DOUBLE、COMPLEX64、COMPLEX128ND3.1 padding 的语义细节padding长度为 6顺序依次为[左, 右, 上, 下, 前, 后]。结合 tiling 源码 reflection_pad3d_grad_tiling.cpp 中GetInputInfo的解析逻辑可以确认六个值分别被命名为wPad1/wPad2宽方向左右、hPad1/hPad2高方向上下、dPad1/dPad2深方向前后params.dPad1 paddingsValue[4]; // 前 params.dPad2 paddingsValue[5]; // 后 params.hPad1 paddingsValue[6]; // 上 params.hPad2 paddingsValue[7]; // 下 params.wPad1 paddingsValue[8]; // 左 params.wPad2 paddingsValue[9]; // 右值得注意底层算子定义reflection_pad3d_grad_def.cpp中paddings输入支持DT_INT32与DT_INT64两种类型对应图中注册的 int32/int64 组合而对外 API 层aclnn接口统一以aclIntArrayINT64形式传入。3.2 数据格式与形状所有张量均使用 ND 格式支持四维batch、channel、H、W 展开为三维填充场景或五维batch、channel、D、H、W输入。op 定义文件中对输入x与输出y声明了FLOAT、FLOAT16、BF16三种实测数据类型及 ND 格式API 层在此基础上扩展支持了 DOUBLE、COMPLEX64、COMPLEX128最终以 API 层声明为准。四、约束说明维度一致性gradOutput、self、gradInput的维度需一致支持四维或五维且它们的形状需与reflection_pad3d正向传播的输出形状相互一致。即满足gradOutput形状 正向输出形状gradInput形状 self形状且正向输出形状 self 形状 padding。padding 值域限制padding 前两个数值左右需小于 self 最后一维度的数值中间两个数值上下需小于 self 倒数第二维度的数值后两个数值前后需小于 self 倒数第三维度的数值。该限制在 tiling 阶段由 reflection_pad3d_grad_tiling.cpp 中的形状校验逻辑保证outHeight height - hPad1 - hPad2等关系不满足时直接返回失败。确定性计算aclnnReflectionPad3dBackward为默认确定性实现当 gradOutput 中元素个数大于300 * 1024 * 1024时存在运行超时风险超大张量场景需评估分块或降低单次计算规模。五、调用说明两段式接口 aclnnReflectionPad3dBackward该算子通过aclnn接口调用采用 CANN 标准的两段式接口模式先调用aclnnReflectionPad3dBackwardGetWorkspaceSize获取 workspace 大小并生成执行器再调用aclnnReflectionPad3dBackward执行计算。接口声明位于 op_api/aclnn_reflection_pad3d_backward.h。5.1 第一段接口aclnnReflectionPad3dBackwardGetWorkspaceSizeaclnnStatus aclnnReflectionPad3dBackwardGetWorkspaceSize( const aclTensor* gradOutput, const aclTensor* self, const aclIntArray* padding, aclTensor* gradInput, uint64_t* workspaceSize, aclOpExecutor** executor)各参数使用说明参数名输入/输出描述使用说明数据类型数据格式维度非连续 TensorgradOutputaclTensor*输入反向传播的输入shape 需与 reflection_pad3d 正向传播的 output 一致与 self 保持一致ND与 self 一致√selfaclTensor*输入正向的输入张量-BFLOAT16、FLOAT16、FLOAT32、DOUBLE、COMPLEX64、COMPLEX128ND4/5√paddingaclIntArray*输入填充值数组长度 6前两值小于 self 最后一维、中间两值小于倒数第二维、后两值小于倒数第三维INT64---gradInputaclTensor*输出反向传播的输出-与 self 保持一致ND与 self 一致√workspaceSizeuint64_t*输出需在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出包含算子计算流程的 op 执行器-----返回值aclnnStatus具体参见 aclnn 返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001Tensor 为空指针ACLNN_ERR_PARAM_INVALID161002gradOutput、self、padding 和 gradInput 的数据类型或数据格式不在支持范围之内ACLNN_ERR_PARAM_INVALID161002gradOutput、self、padding 和 gradInput 的输入 shape 在支持范围之外ACLNN_ERR_PARAM_INVALID161002self 为空 tensor 且后 4 维存在维度大小为 0ACLNN_ERR_PARAM_INVALID161002padding 的 size 不等于 6ACLNN_ERR_PARAM_INVALID161002padding 内的数值大于等于 self 的维度ACLNN_ERR_PARAM_INVALID161002gradOutput shape 与 reflection_pad3d 正向传播的 output 不一致5.2 第二段接口aclnnReflectionPad3dBackwardaclnnStatus aclnnReflectionPad3dBackward( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream)参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口获取executor输入op 执行器包含算子计算流程stream输入指定执行任务的 Stream使用注意第二段接口不能重复调用一次GetWorkspaceSize对应的executor只能用于一次执行workspace指除输入输出外算子在 NPU 上完成计算所需的临时内存详见两段式接口说明。六、完整调用示例C仓库提供了可直接参考的完整示例 examples/test_aclnn_reflection_pad3d_grad.cpp其完整流程如下编译与运行步骤请参考编译与运行样例#include acl/acl.h #include aclnnop/aclnn_reflection_pad3d_backward.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手册 int32_t deviceId 0; // 根据自己的实际device填写deviceId aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret 0, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2.构造输入与输出需要根据API的接口定义构造 std::vectorint64_t gradOutputShape {1, 1, 4, 4, 4}; // 正向填充后的输出形状 std::vectorint64_t selfShape {1, 1, 2, 2, 2}; // 正向输入形状 std::vectorint64_t gradInputShape {1, 1, 2, 2, 2}; // 梯度输出形状与self一致 void* gradOutputDeviceAddr nullptr; void* selfDeviceAddr nullptr; void* gradInputDeviceAddr nullptr; aclTensor* gradOutput nullptr; aclTensor* self nullptr; aclIntArray* padding nullptr; aclTensor* gradInput nullptr; std::vectorfloat gradOutputHostData(64); for (int64_t i 0; i 64; i) { gradOutputHostData[i] 1; // 上游梯度全为1便于观察结果 } std::vectorfloat selfHostData {1, 2, 3, 4, 5, 6, 7, 8}; std::vectorint64_t paddingData {1, 1, 1, 1, 1, 1}; // 左、右、上、下、前、后 std::vectorfloat gradInputHostData {0, 0, 0, 0, 0, 0, 0, 0}; // 创建gradOutput、self、gradInput的aclTensor及padding的aclIntArray ret CreateAclTensor(gradOutputHostData, gradOutputShape, gradOutputDeviceAddr, aclDataType::ACL_FLOAT, gradOutput); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); padding aclCreateIntArray(paddingData.data(), 6); CHECK_RET(padding ! nullptr, LOG_PRINT(aclCreateIntArray failed.\n); return ACL_ERROR_INTERNAL_ERROR;); ret CreateAclTensor(gradInputHostData, gradInputShape, gradInputDeviceAddr, aclDataType::ACL_FLOAT, gradInput); CHECK_RET(ret ACL_SUCCESS, return ret); // 3.调用CANN算子库API两段式 uint64_t workspaceSize 0; aclOpExecutor* executor; // 第一段接口计算workspace大小并生成执行器 ret aclnnReflectionPad3dBackwardGetWorkspaceSize(gradOutput, self, padding, gradInput, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnReflectionPad3dBackwardGetWorkspaceSize 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;); } // 第二段接口执行计算 ret aclnnReflectionPad3dBackward(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnReflectionPad3dBackward 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侧 auto size GetShapeSize(gradInputShape); std::vectorfloat resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), gradInputDeviceAddr, 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 aclDestroyTensor(gradOutput); aclDestroyTensor(self); aclDestroyIntArray(padding); aclDestroyTensor(gradInput); // 7.释放device资源 aclrtFree(gradOutputDeviceAddr); aclrtFree(selfDeviceAddr); aclrtFree(gradInputDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例中self形状为[1,1,2,2,2]、padding 全为 1对应正向输出形状[1,1,4,4,4]满足gradOutput 形状与正向输出一致的约束。梯度全为 1 时反射填充会把每个原始元素的梯度累加其被镜像复制的次数。七、源码级原理剖析7.1 算子定义与平台配置reflection_pad3d_grad_def.cpp 通过OP_ADD(ReflectionPad3dGrad)注册算子声明输入x、paddings与输出yy带.InitValue(0)初始化并为ascend910b、ascend910_93两个平台添加 AICore 配置。对应的二进制 kernel 清单见 op_host/config/ascend910_93/reflection_pad3d_grad_binary.json覆盖 float32/float16/bfloat16 与 int32/int64 padding 的六种组合平台简化 key 配置见 op_host/config/ascend910_93/reflection_pad3d_grad_simplified_key.ini。7.2 Tiling 策略按形状分派多种计算模式Tiling 是算子高性能的关键。从 reflection_pad3d_grad_tiling.cpp 可以看到Host 侧 tiling 逻辑会根据输入维度、数据类型与可用的 UBUnified Buffer空间选择不同的 tilingKey核数分配GetUsedCore以batch × channel为基本任务单元当前后方向无填充时再乘以 depth在可用 AIV 核数内均分任务并计算尾块tailNCUB 切分SplitUb扣除 tiling 数据与预留的 16KB 后将可用 UB 按 5 份float32或 8 份float16/bfloat16因涉及升 float32 计算切分得到ubFactorElement并按 256 字节对齐模式选择FillTilingKey根据对齐后的alignHeight × alignWidth与ubFactorElement的关系决定走小图Small、扁平Flat、中图Mid仅 float32或大图Big路径workspace 计算float 类型按max(alignHeight, alignWidth) × 32 × blockNum × 4计算float16/bfloat16 则按alignHeight × alignWidth × blockNum × 4升 float32 中转所需计算并叠加系统 workspace原子累加SetNeedAtomic(true)——反射填充中一个原始元素可能被多个填充位置引用梯度回填需要原子加来保证累加正确性。从 tiling 常量定义可以归纳出各模式的 tilingKey 编码float32 为 100/101/102/103Small/Mid/Flat/Bigfloat16 为 200/202/203bfloat16 为 300/302/303。7.3 Kernel 实现按 tilingKey 多路分发内核入口 op_kernel/reflection_pad3d_grad.cpp 是一个典型的按 tilingKey 分发的调度函数根据TILING_KEY_IS(...)宏把任务路由到对应模板类与处理函数ReflectionPad3dGradfloat处理 float32支持SmallProcess、MidProcess、FlatProcess、BigProcess四种路径ReflectionPad3dGradF16half处理 float16支持 Small/Flat/Big 三种路径ReflectionPad3dGradF16bfloat16_t处理 bfloat16支持 Small/Flat/Big 三种路径。以 op_kernel/reflection_pad3d_grad_small.h 中的SmallProcess为例实现采用流水式三段结构CopyInSmall经inQueueX队列把 GM 数据搬入 UB→ComputeSmall计算梯度并做转置→CopyOutSmall写回 GMisAtomicAdd true走原子加。值得注意float16/bfloat16 路径在ComputeSmall中会先将数据Cast到 float32 做梯度计算ComputeSmallBasicfloat再Cast回原类型输出以换取更高的计算精度这也是 tiling 阶段为其预留更大 workspace 的原因。7.4 测试验证仓库为该算子提供了多层测试保障Tiling 单元测试tests/ut/op_host/test_reflection_pad3d_grad_tiling.cpp覆盖 float32 小图key100、中图key101、float16 扁平图key202、bfloat16 小/扁/大图key300/302/303、int32/int64 padding 兼容等成功场景以及输入维度非 5、padding 长度错误、输出形状不匹配等失败场景并断言具体的 workspace 大小可直接作为 tiling 行为的行为基线算子 API L2 测试tests/ut/op_api/test_aclnn_reflection_pad3d_backward_l2.cppSTSystem Test用例用例定义见 tests/st/aclnnReflectionPad3dBackward/atk_aclnnReflectionPad3dBackward.json覆盖 4/5 维、bf16/fp16/fp32 及非对称 padding 的 200 组形状执行脚本 tests/st/aclnnReflectionPad3dBackward/executor_aclnnReflectionPad3dBackward.py 中通过torch.ops.aten.reflection_pad3d_backward作为参考实现进行数值比对验证算子与 PyTorch 语义一致。八、总结ReflectionPad3dGrad是 CANN ops-math 中实现 3D 反射填充反向传播的标准算子采用Host 侧 tiling 决策 Device 侧多模式内核的架构tiling 阶段依据形状与数据类型在 Small/Mid/Flat/Big 及 float16/bfloat16 升精度路径间智能分派kernel 阶段通过原子累加保证反射填充梯度回填的正确性。调用侧只需遵循GetWorkspaceSize → 申请 workspace → 执行 → 同步 → 取结果的标准两段式流程即可在 Atlas A2/A3 系列产品上完成梯度计算可直接对接 PyTorchnn.ReflectionPad3d的训练反向过程。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-math 算子实战：ReflectionPad3dGrad 反射填充反向算子与 aclnnReflectionPad3dBackward 两段式接口
CANN ops-math 算子实战ReflectionPad3dGrad 反射填充反向算子与 aclnnReflectionPad3dBackward 两段式接口【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本文是 CANN / ops-math 仓库中 ReflectionPad3dGrad 算子的完整技术指南。该算子实现三维反射填充ReflectionPad3d的反向传播是 PyTorchnn.ReflectionPad3d在 NPU 上训练时梯度回传的关键一环。读完本文你将掌握该算子的产品支持范围、输入输出参数与约束、aclnnReflectionPad3dBackward两段式接口的完整调用流程并能结合仓库源码理解其 tiling 分派与内核实现的底层原理。一、算子定位ReflectionPad3d 的反向传播ReflectionPad3dGrad是 CANN ops-math 数学算子库中conversion数据转换类目录下的一个反向算子功能是计算aclnnReflectionPad3d接口的正向反射填充操作对应的梯度。在深度学习中反射填充常用于卷积网络的边界处理训练阶段通过本算子把上游传来的梯度gradOutput映射回未填充前的输入张量gradInput从而完成对前向输入的梯度回传。正向算子aclnnReflectionPad3d的接口文档位于 conversion/mirror_pad/docs/aclnnReflectionPad3d.md其语义为对输入 tensor 的左右上下前后六个方向做镜像式填充填充值取自边缘的反射映射而非补零。例如输入[[[[[0,1],[2,3]]],[[[4,5],[6,7]]]]]、padding 为[1,1,1,1,1,1]时输出会沿各维度把边缘值镜像复制到填充区域。本反向算子则与之对应把填充后张量的梯度按同样规则折回到原始尺寸上。二、产品支持情况不同硬件平台对该算子的支持情况如下表所示产品是否支持Ascend 950PR/Ascend 950DT×Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×需要注意的是算子接口文档 aclnnReflectionPad3dBackward.md 中的产品支持矩阵与上述算子 README 略有差异该文档标注 Ascend 950PR/950DT、Atlas 推理系列产品、Atlas 训练系列产品为支持这从侧面说明产品支持情况会随 CANN 版本演进而变化实际部署时请以当前 CANN 版本配套的算子清单为准。从算子注册配置可以印证当前的 AICore 适配范围op_host/reflection_pad3d_grad_def.cpp 中通过this-AICore().AddConfig(ascend910b)与this-AICore().AddConfig(ascend910_93)注册了ascend910b对应 Atlas A2 系列和ascend910_93对应 Atlas A3 系列两个平台的核函数配置与 README 中Atlas A2/A3 系列支持的结论一致。三、参数说明算子输入输出参数如下表前向的输入张量self在反向中仅用于形状参考不参与梯度计算参数名输入/输出/属性描述数据类型数据格式gradOutput输入反向传播的输入即正向反射填充输出张量的梯度FLOAT16、FLOAT32、DOUBLE、COMPLEX64、COMPLEX128NDself输入正向的输入张量未填充前的张量用于确定 gradInput 的形状FLOAT16、FLOAT32、DOUBLE、COMPLEX64、COMPLEX128NDpadding输入长度为 6 的数组数值依次代表左、右、上、下、前、后需要填充的值aclIntArray 数组INT64-gradInput输出反向传播的输出与 self 形状一致FLOAT16、FLOAT32、DOUBLE、COMPLEX64、COMPLEX128ND3.1 padding 的语义细节padding长度为 6顺序依次为[左, 右, 上, 下, 前, 后]。结合 tiling 源码 reflection_pad3d_grad_tiling.cpp 中GetInputInfo的解析逻辑可以确认六个值分别被命名为wPad1/wPad2宽方向左右、hPad1/hPad2高方向上下、dPad1/dPad2深方向前后params.dPad1 paddingsValue[4]; // 前 params.dPad2 paddingsValue[5]; // 后 params.hPad1 paddingsValue[6]; // 上 params.hPad2 paddingsValue[7]; // 下 params.wPad1 paddingsValue[8]; // 左 params.wPad2 paddingsValue[9]; // 右值得注意底层算子定义reflection_pad3d_grad_def.cpp中paddings输入支持DT_INT32与DT_INT64两种类型对应图中注册的 int32/int64 组合而对外 API 层aclnn接口统一以aclIntArrayINT64形式传入。3.2 数据格式与形状所有张量均使用 ND 格式支持四维batch、channel、H、W 展开为三维填充场景或五维batch、channel、D、H、W输入。op 定义文件中对输入x与输出y声明了FLOAT、FLOAT16、BF16三种实测数据类型及 ND 格式API 层在此基础上扩展支持了 DOUBLE、COMPLEX64、COMPLEX128最终以 API 层声明为准。四、约束说明维度一致性gradOutput、self、gradInput的维度需一致支持四维或五维且它们的形状需与reflection_pad3d正向传播的输出形状相互一致。即满足gradOutput形状 正向输出形状gradInput形状 self形状且正向输出形状 self 形状 padding。padding 值域限制padding 前两个数值左右需小于 self 最后一维度的数值中间两个数值上下需小于 self 倒数第二维度的数值后两个数值前后需小于 self 倒数第三维度的数值。该限制在 tiling 阶段由 reflection_pad3d_grad_tiling.cpp 中的形状校验逻辑保证outHeight height - hPad1 - hPad2等关系不满足时直接返回失败。确定性计算aclnnReflectionPad3dBackward为默认确定性实现当 gradOutput 中元素个数大于300 * 1024 * 1024时存在运行超时风险超大张量场景需评估分块或降低单次计算规模。五、调用说明两段式接口 aclnnReflectionPad3dBackward该算子通过aclnn接口调用采用 CANN 标准的两段式接口模式先调用aclnnReflectionPad3dBackwardGetWorkspaceSize获取 workspace 大小并生成执行器再调用aclnnReflectionPad3dBackward执行计算。接口声明位于 op_api/aclnn_reflection_pad3d_backward.h。5.1 第一段接口aclnnReflectionPad3dBackwardGetWorkspaceSizeaclnnStatus aclnnReflectionPad3dBackwardGetWorkspaceSize( const aclTensor* gradOutput, const aclTensor* self, const aclIntArray* padding, aclTensor* gradInput, uint64_t* workspaceSize, aclOpExecutor** executor)各参数使用说明参数名输入/输出描述使用说明数据类型数据格式维度非连续 TensorgradOutputaclTensor*输入反向传播的输入shape 需与 reflection_pad3d 正向传播的 output 一致与 self 保持一致ND与 self 一致√selfaclTensor*输入正向的输入张量-BFLOAT16、FLOAT16、FLOAT32、DOUBLE、COMPLEX64、COMPLEX128ND4/5√paddingaclIntArray*输入填充值数组长度 6前两值小于 self 最后一维、中间两值小于倒数第二维、后两值小于倒数第三维INT64---gradInputaclTensor*输出反向传播的输出-与 self 保持一致ND与 self 一致√workspaceSizeuint64_t*输出需在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出包含算子计算流程的 op 执行器-----返回值aclnnStatus具体参见 aclnn 返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001Tensor 为空指针ACLNN_ERR_PARAM_INVALID161002gradOutput、self、padding 和 gradInput 的数据类型或数据格式不在支持范围之内ACLNN_ERR_PARAM_INVALID161002gradOutput、self、padding 和 gradInput 的输入 shape 在支持范围之外ACLNN_ERR_PARAM_INVALID161002self 为空 tensor 且后 4 维存在维度大小为 0ACLNN_ERR_PARAM_INVALID161002padding 的 size 不等于 6ACLNN_ERR_PARAM_INVALID161002padding 内的数值大于等于 self 的维度ACLNN_ERR_PARAM_INVALID161002gradOutput shape 与 reflection_pad3d 正向传播的 output 不一致5.2 第二段接口aclnnReflectionPad3dBackwardaclnnStatus aclnnReflectionPad3dBackward( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream)参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口获取executor输入op 执行器包含算子计算流程stream输入指定执行任务的 Stream使用注意第二段接口不能重复调用一次GetWorkspaceSize对应的executor只能用于一次执行workspace指除输入输出外算子在 NPU 上完成计算所需的临时内存详见两段式接口说明。六、完整调用示例C仓库提供了可直接参考的完整示例 examples/test_aclnn_reflection_pad3d_grad.cpp其完整流程如下编译与运行步骤请参考编译与运行样例#include acl/acl.h #include aclnnop/aclnn_reflection_pad3d_backward.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手册 int32_t deviceId 0; // 根据自己的实际device填写deviceId aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret 0, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2.构造输入与输出需要根据API的接口定义构造 std::vectorint64_t gradOutputShape {1, 1, 4, 4, 4}; // 正向填充后的输出形状 std::vectorint64_t selfShape {1, 1, 2, 2, 2}; // 正向输入形状 std::vectorint64_t gradInputShape {1, 1, 2, 2, 2}; // 梯度输出形状与self一致 void* gradOutputDeviceAddr nullptr; void* selfDeviceAddr nullptr; void* gradInputDeviceAddr nullptr; aclTensor* gradOutput nullptr; aclTensor* self nullptr; aclIntArray* padding nullptr; aclTensor* gradInput nullptr; std::vectorfloat gradOutputHostData(64); for (int64_t i 0; i 64; i) { gradOutputHostData[i] 1; // 上游梯度全为1便于观察结果 } std::vectorfloat selfHostData {1, 2, 3, 4, 5, 6, 7, 8}; std::vectorint64_t paddingData {1, 1, 1, 1, 1, 1}; // 左、右、上、下、前、后 std::vectorfloat gradInputHostData {0, 0, 0, 0, 0, 0, 0, 0}; // 创建gradOutput、self、gradInput的aclTensor及padding的aclIntArray ret CreateAclTensor(gradOutputHostData, gradOutputShape, gradOutputDeviceAddr, aclDataType::ACL_FLOAT, gradOutput); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); padding aclCreateIntArray(paddingData.data(), 6); CHECK_RET(padding ! nullptr, LOG_PRINT(aclCreateIntArray failed.\n); return ACL_ERROR_INTERNAL_ERROR;); ret CreateAclTensor(gradInputHostData, gradInputShape, gradInputDeviceAddr, aclDataType::ACL_FLOAT, gradInput); CHECK_RET(ret ACL_SUCCESS, return ret); // 3.调用CANN算子库API两段式 uint64_t workspaceSize 0; aclOpExecutor* executor; // 第一段接口计算workspace大小并生成执行器 ret aclnnReflectionPad3dBackwardGetWorkspaceSize(gradOutput, self, padding, gradInput, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnReflectionPad3dBackwardGetWorkspaceSize 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;); } // 第二段接口执行计算 ret aclnnReflectionPad3dBackward(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnReflectionPad3dBackward 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侧 auto size GetShapeSize(gradInputShape); std::vectorfloat resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), gradInputDeviceAddr, 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 aclDestroyTensor(gradOutput); aclDestroyTensor(self); aclDestroyIntArray(padding); aclDestroyTensor(gradInput); // 7.释放device资源 aclrtFree(gradOutputDeviceAddr); aclrtFree(selfDeviceAddr); aclrtFree(gradInputDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例中self形状为[1,1,2,2,2]、padding 全为 1对应正向输出形状[1,1,4,4,4]满足gradOutput 形状与正向输出一致的约束。梯度全为 1 时反射填充会把每个原始元素的梯度累加其被镜像复制的次数。七、源码级原理剖析7.1 算子定义与平台配置reflection_pad3d_grad_def.cpp 通过OP_ADD(ReflectionPad3dGrad)注册算子声明输入x、paddings与输出yy带.InitValue(0)初始化并为ascend910b、ascend910_93两个平台添加 AICore 配置。对应的二进制 kernel 清单见 op_host/config/ascend910_93/reflection_pad3d_grad_binary.json覆盖 float32/float16/bfloat16 与 int32/int64 padding 的六种组合平台简化 key 配置见 op_host/config/ascend910_93/reflection_pad3d_grad_simplified_key.ini。7.2 Tiling 策略按形状分派多种计算模式Tiling 是算子高性能的关键。从 reflection_pad3d_grad_tiling.cpp 可以看到Host 侧 tiling 逻辑会根据输入维度、数据类型与可用的 UBUnified Buffer空间选择不同的 tilingKey核数分配GetUsedCore以batch × channel为基本任务单元当前后方向无填充时再乘以 depth在可用 AIV 核数内均分任务并计算尾块tailNCUB 切分SplitUb扣除 tiling 数据与预留的 16KB 后将可用 UB 按 5 份float32或 8 份float16/bfloat16因涉及升 float32 计算切分得到ubFactorElement并按 256 字节对齐模式选择FillTilingKey根据对齐后的alignHeight × alignWidth与ubFactorElement的关系决定走小图Small、扁平Flat、中图Mid仅 float32或大图Big路径workspace 计算float 类型按max(alignHeight, alignWidth) × 32 × blockNum × 4计算float16/bfloat16 则按alignHeight × alignWidth × blockNum × 4升 float32 中转所需计算并叠加系统 workspace原子累加SetNeedAtomic(true)——反射填充中一个原始元素可能被多个填充位置引用梯度回填需要原子加来保证累加正确性。从 tiling 常量定义可以归纳出各模式的 tilingKey 编码float32 为 100/101/102/103Small/Mid/Flat/Bigfloat16 为 200/202/203bfloat16 为 300/302/303。7.3 Kernel 实现按 tilingKey 多路分发内核入口 op_kernel/reflection_pad3d_grad.cpp 是一个典型的按 tilingKey 分发的调度函数根据TILING_KEY_IS(...)宏把任务路由到对应模板类与处理函数ReflectionPad3dGradfloat处理 float32支持SmallProcess、MidProcess、FlatProcess、BigProcess四种路径ReflectionPad3dGradF16half处理 float16支持 Small/Flat/Big 三种路径ReflectionPad3dGradF16bfloat16_t处理 bfloat16支持 Small/Flat/Big 三种路径。以 op_kernel/reflection_pad3d_grad_small.h 中的SmallProcess为例实现采用流水式三段结构CopyInSmall经inQueueX队列把 GM 数据搬入 UB→ComputeSmall计算梯度并做转置→CopyOutSmall写回 GMisAtomicAdd true走原子加。值得注意float16/bfloat16 路径在ComputeSmall中会先将数据Cast到 float32 做梯度计算ComputeSmallBasicfloat再Cast回原类型输出以换取更高的计算精度这也是 tiling 阶段为其预留更大 workspace 的原因。7.4 测试验证仓库为该算子提供了多层测试保障Tiling 单元测试tests/ut/op_host/test_reflection_pad3d_grad_tiling.cpp覆盖 float32 小图key100、中图key101、float16 扁平图key202、bfloat16 小/扁/大图key300/302/303、int32/int64 padding 兼容等成功场景以及输入维度非 5、padding 长度错误、输出形状不匹配等失败场景并断言具体的 workspace 大小可直接作为 tiling 行为的行为基线算子 API L2 测试tests/ut/op_api/test_aclnn_reflection_pad3d_backward_l2.cppSTSystem Test用例用例定义见 tests/st/aclnnReflectionPad3dBackward/atk_aclnnReflectionPad3dBackward.json覆盖 4/5 维、bf16/fp16/fp32 及非对称 padding 的 200 组形状执行脚本 tests/st/aclnnReflectionPad3dBackward/executor_aclnnReflectionPad3dBackward.py 中通过torch.ops.aten.reflection_pad3d_backward作为参考实现进行数值比对验证算子与 PyTorch 语义一致。八、总结ReflectionPad3dGrad是 CANN ops-math 中实现 3D 反射填充反向传播的标准算子采用Host 侧 tiling 决策 Device 侧多模式内核的架构tiling 阶段依据形状与数据类型在 Small/Mid/Flat/Big 及 float16/bfloat16 升精度路径间智能分派kernel 阶段通过原子累加保证反射填充梯度回填的正确性。调用侧只需遵循GetWorkspaceSize → 申请 workspace → 执行 → 同步 → 取结果的标准两段式流程即可在 Atlas A2/A3 系列产品上完成梯度计算可直接对接 PyTorchnn.ReflectionPad3d的训练反向过程。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考