CANN ops-cv UpsampleTrilinear3dBackward 算子深度解析:三维三线性上采样的反向梯度计算与 aclnn 接口实战
2026/9/18 3:09:08 网站建设 项目流程

CANN ops-cv UpsampleTrilinear3dBackward 算子深度解析:三维三线性上采样的反向梯度计算与 aclnn 接口实战

【免费下载链接】ops-cv本项目是CANN提供的图像处理、目标检测相关的算子库,实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-cv

本篇技术指南以 CANN ops-cv 开源算子库中的 UpsampleTrilinear3dBackward 算子文档 为核心,系统讲解三维三线性插值(Trilinear Interpolation)反向算子的数学原理、参数约束、两段式 aclnn 调用方式,并结合该算子目录下的源码与测试用例剖析其在 NPU 上的实现机制。读者读完将能够独立完成该算子的 aclnn 接口调用、理解其梯度传播公式,并掌握其 host 侧算子定义、InferShape、Tiling 与 kernel 侧 matmul 实现的全链路结构。

产品支持情况

根据 README.md 中列出的产品支持矩阵,UpsampleTrilinear3dBackward 算子的支持情况如下:

产品是否支持
Ascend 950PR/Ascend 950DT×
Atlas A3 训练系列产品/Atlas A3 推理系列产品
Atlas A2 训练系列产品/Atlas A2 推理系列产品
Atlas 200I/500 A2 推理产品×
Atlas 推理系列产品×
Atlas 训练系列产品×

此外,在算子对应的 aclnnUpsampleTrilinear3dBackward 接口文档 中也列示了各产品的支持情况,两处文档在部分产品(如 Atlas 推理系列、Atlas 训练系列)上的表述存在差异,实际使用时建议以与自身环境匹配的版本化文档为准。从 host 侧算子注册文件 upsample_trilinear3d_backward_def.cpp 可以看到,该算子当前注册的 AICore 配置为ascend910bascend910_93两个平台。

功能说明:三维三线性上采样反向算子

UpsampleTrilinear3dBackward 是 ResizeUpsampleTrilinear 算子(即 aclnnUpsampleTrilinear3d 接口)的反向计算算子。其输入为正向算子输出的梯度gradOutput,输出为反向传播得到的输入侧梯度gradInput,shape 与正向算子的输入一致。该算子通常用于深度学习训练链路中F.interpolate(..., mode='trilinear')这类三维特征图缩放操作的自动求导。

正向核心算法逻辑

要理解反向计算,首先要弄清正向三线性插值的计算过程,其核心逻辑分为三步:

  1. 将目标图像的每一个点映射回原图,得到一个带小数点的坐标;
  2. 根据这个浮点数坐标,计算前后相邻的原始图像的点;
  3. 分别计算相邻点到对应目标点的权重,按照权重相乘累加即可得到目标点值。

缩放系数与坐标映射公式

缩放方式分为角对齐(alignCorners 为 true)与边对齐(alignCorners 为 false)两种:角对齐表示按照原始图片左上角像素中心点对齐;边对齐表示按照原始图片左上角顶点及两条边对齐。二者在缩放系数和坐标位置的计算上存在差异。对于三维插值点 $(N, C, D, H, W)$,三个维度的缩放系数定义如下:

$$ scale_d =\begin{cases} (inputSize[2]-1) / (outputSize[0]-1) & alignCorners=true \ 1 / scales_d & alignCorners=false&scales_d>0\ inputSize[2] / outputSize[0] & alignCorners=false \end{cases} $$

$$ scale_h =\begin{cases} (inputSize[3]-1) / (outputSize[1]-1) & alignCorners=true \ 1 / scales_h & alignCorners=false&scales_h>0\ inputSize[3] / outputSize[1] & alignCorners=false \end{cases} $$

$$ scale_w =\begin{cases} (inputSize[4]-1) / (outputSize[2]-1) & alignCorners=true \ 1 / scales_w & alignCorners=false&scales_w>0\ inputSize[4] / outputSize[2] & alignCorners=false \end{cases} $$

其中inputSize是正向输入在 D、H、W 维度上的大小,outputSize是正向输出在 D、H、W 维度上的大小(对应到反向场景中,分别是gradInputgradOutput的空间尺寸),scales_d/scales_h/scales_w为用户显式指定的缩放因子。

对于输出(反向场景下即 gradOutput)某个方向上的点 p(x, y, z),映射回原始图像中的点记为 q(x', y', z'),则有:

$$ x' =\begin{cases} x * scale_d & alignCorners=true \ MAX(0,{(x+0.5)*scale_d-0.5}) & alignCorners=false \end{cases} $$

$$ y' =\begin{cases} y * scale_h & alignCorners=true \ MAX(0,{(y+0.5)*scale_h-0.5}) & alignCorners=false \end{cases} $$

$$ z' =\begin{cases} z * scale_w & alignCorners=true \ MAX(0,{(z+0.5)*scale_w-0.5}) & alignCorners=false \end{cases} $$

三线性插值:八个相邻点的加权累加

得到浮点坐标后,取其整数部分与整数部分加一,得到三个方向上各自的前后相邻点索引,并计算相应权重:

$$ x_{0} =int(x'),x_{1} =int(x')+1, lambda_{0} = x_{1}-x', lambda_{1} = 1-lambda_{0} $$

$$ y_{0} =int(y'),y_{1} =int(y')+1, lambdb_{0} = y_{1}-y', lambdb_{1} = 1-lambdb_{0} $$

$$ z_{0} =int(z'),z_{1} =int(z')+1, lambdc_{0} = z_{1}-z', lambdc_{1} = 1-lambdc_{0} $$

目标点像素值由包围它的 8 个原始点按三维权重加权求和得到:

$$ {V(p_{x, y, z})} = {V(p_{x0, y0, z0})} * {lambda_{0}} * {lambdb_{0}} * {lambdc_{0}} + {V(p_{x0, y0, z1})} * {lambda_{0}} * {lambdb_{0}} * {lambdc_{1}} + {V(p_{x0, y1, z0})} * {lambda_{0}} * {lambdb_{1}} * {lambdc_{0}} + {V(p_{x0, y1, z1})} * {lambda_{0}} * {lambdb_{1}} * {lambdc_{1}} + {V(p_{x1, y0, z0})} * {lambda_{1}} * {lambdb_{0}} * {lambdc_{0}} + {V(p_{x1, y0, z1})} * {lambda_{1}} * {lambdb_{0}} * {lambdc_{1}} + {V(p_{x1, y1, z0})} * {lambda_{1}} * {lambdb_{1}} * {lambdc_{0}} + {V(p_{x1, y1, z1})} * {lambda_{1}} * {lambdb_{1}} * {lambdc_{1}} $$

反向梯度传播公式

反向计算的核心思路是:将 gradOutput 中每个点的梯度,按正向插值时的权重反哺给其参与插值的 8 个原始点。假设正向插值的输出图像 out(x, y, z) 受原图像 input(x_i, y_j, z_k) 影响,则:

$$ gradInput(x_i,y_j,z_k) += gradOutput(x,y,z) * lambda(x_i,y_j,z_k)* lambdb(x_i,y_j,z_k)* lambdc(x_i,y_j,z_k) $$

即每个原始点累积所有“以它为插值源”的目标点梯度乘以其对应的三轴权重。由于一个原始点可能被多个输出点引用,因此梯度使用+=累加语义。这一公式在 kernel 侧被进一步转化为三个方向的连续矩阵乘(matmul)流水实现,具体见后文源码剖析。

参数说明

算子层面的输入、输出与属性定义如下表(对应 README.md 的参数说明章节):

参数名输入/输出/属性描述数据类型数据格式
grad_output输入表示反向计算的梯度 Tensor,对应公式中的gradOutputFLOAT32、FLOAT16、DOUBLENCDHW
input_size属性表示输出gradInput分别在 N、C、D、H 和 W 维度上的空间大小,对应公式中的inputSize。size 为 5,且各元素均大于零。LISTINT-
output_size属性表示输入grad_output在 D、H 和 W 维度上的空间大小,对应公式中的outputSize。size 为 3,且各元素均大于零。LISTINT-
scales可选属性指定沿每个维度的缩放数组,包含 3 个元素:scale_depth、scale_height、scale_width,对应公式中的scales_dscales_hscales_w。默认值为空。LISTFLOAT-
align_corners可选属性决定是否对齐角像素点,对应公式中的alignCorners。为 true 时输入和输出张量的角像素点会被对齐,否则不对齐。默认值为 false。BOOL-
y输出表示反向计算的输出张量,对应公式中的gradInput。数据类型和数据格式与入参grad_output保持一致。shape 依赖input_sizeoutput_size/scalesFLOAT32、FLOAT16、DOUBLENCDHW

在 host 侧算子定义文件 upsample_trilinear3d_backward_def.cpp 中,上述参数被注册为:

  • grad_output为必选输入,支持DT_BF16DT_FLOAT16DT_FLOAT三种数据类型(注意此处与算子文档表格中列的 DOUBLE 并不一致,实际支持的以注册表为准);
  • output_sizeinput_size为必选属性(ListInt);
  • align_corners为可选属性,默认 false;
  • scale_dscale_hscale_w为可选属性(Float);
  • grad_input为必选输出,数据类型与输入一致。

约束说明:README 中明确该算子无额外约束。但在 aclnn 接口层面存在一系列 shape 与内存约束,详见下文「接口约束」小节。

aclnn 两段式接口

CANN 的 aclnn 算子接口采用「两段式」调用模型(参见 two_phase_api 说明):必须先调用aclnnUpsampleTrilinear3dBackwardGetWorkspaceSize获取计算所需的 workspace 大小以及封装了算子计算流程的执行器,再调用aclnnUpsampleTrilinear3dBackward执行计算。

函数原型

aclnnStatus aclnnUpsampleTrilinear3dBackwardGetWorkspaceSize( const aclTensor *gradOut, const aclIntArray *outputSize, const aclIntArray *inputSize, bool alignCorners, double scalesD, double scalesH, double scalesW, aclTensor *gradInput, uint64_t *workspaceSize, aclOpExecutor **executor)
aclnnStatus aclnnUpsampleTrilinear3dBackward( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

GetWorkspaceSize 参数说明

参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensor
gradOut(aclTensor*)输入表示反向计算的梯度 Tensor,对应公式中的gradOut不支持空 Tensor;当数据格式为 ND 时,默认按照 NCDHW 格式处理。FLOAT32、FLOAT16、BFLOAT16、DOUBLENCDHW、NDHWC、ND5
outputSize(aclIntArray*)输入表示输入gradOut在 D、H 和 W 维度上的空间大小,对应公式中的outputSizesize 为 3,且各元素均大于零。INT64---
inputSize(aclIntArray*)输入表示输出gradInput分别在 N、C、D、H 和 W 维度上的空间大小,对应公式中的inputSizesize 为 5,且最后三个元素均大于零。INT64---
alignCorners(bool)输入表示是否对齐角像素点,对应公式中的alignCorners为 true 则输入和输出张量的角像素点会对齐,否则不对齐。----
scalesD(double)输入表示输出gradInput的 depth 维度乘数,对应公式中的scales_d-----
scalesH(double)输入表示输出gradInput的 height 维度乘数,对应公式中的scales_h-----
scalesW(double)输入表示输出gradInput的 width 维度乘数,对应公式中的scales_w-----
gradInput(aclTensor*)输出表示反向计算的输出张量,对应公式中的gradInput不支持空 Tensor;shape 在 N、C、D、H、W 上的大小需与inputSize一致;数据类型、数据格式、shape 与入参gradOut保持一致。FLOAT32、FLOAT16、BFLOAT16、DOUBLENCDHW、NDHWC、ND5
workspaceSize(uint64_t*)输出返回需要在 Device 侧申请的 workspace 大小。-----
executor(aclOpExecutor**)输出返回 op 执行器,包含了算子计算流程。-----

需要注意的是,接口文档中注明:Atlas 训练系列产品上,参数gradOutgradInput的数据类型不支持 BFLOAT16

返回值与错误码

第一段接口完成入参校验,出现以下场景时报错:

返回码错误码描述
ACLNN_ERR_PARAM_NULLPTR161001传入的 gradOut、outputSize、inputSize 或 gradInput 是空指针。
ACLNN_ERR_PARAM_INVALID161002gradOut 的数据类型和数据格式不在支持的范围内;或 gradOut 与 gradInput 数据类型不一致;或 gradOut 维度不为 5 维;或 outputSize 的 size 不等于 3;或 outputSize 某元素不大于 0;或 inputSize 的 size 不等于 5;或 inputSize 最后三个元素中某个不大于 0;或 gradOut 与 inputSize 在 N、C 维度上的 size 不同;或 gradOut 在 D、H、W 维度上的 size 与 outputSize[0..2] 不一致;或 gradInput 的 shape 与 inputSize[0..4] 不一致。

这些校验逻辑在 aclnn_upsample_trilinear_3d_backward.cpp 的CheckParams中实现:依次检查空指针(CheckNotNull2In2Out)、数据类型合法性(CheckDtypeValid1Out1In)、shape 与 format(CheckUpsampleShapeCheckFormat)、输入元素合法性(CheckInputElement,包括各维 size 大于 0、gradOut/gradInput 的 shape 与 inputSize/outputSize 完全匹配)以及上边界(CheckUplimit,各维不得超过 INT32_MAX)。CheckFormat仅接受 ND、NCDHW、NDHWC 三种存储格式(源码 L35-L44)。

接口约束

  • shape 约束gradOutgradInput每个维度的取值小于等于 2^20;gradInput的 N 轴和 C 轴与gradOut保持一致;内存占用需小于 60GB。内存占用的计算公式为:

    $$ N * C * (gradOut_D * gradOut_H * gradOut_W + gradInput_D * gradInput_H * gradInput_W + gradOut_D * gradOut_H * gradInput_W + gradOut_D * gradInput_H * gradInput_W) * sizeof(dtype) < 60 * 1024 * 1024 * 1024 $$

    其中 N、C 分别为输入输出的 N 轴与 C 轴,dtype 为输入张量的数据类型。另有N * C * gradOut_D * gradOut_H < 2^31gradInput_W * gradInput_H < 2^31两条约束。

  • scales 约束:当输入 scalesD、scalesH、scalesW 取值均大于 0 时,参数 inputSize、outputSize、scales 需满足:

    $$ outputSize_D = floor(inputSize_D * scalesD), \quad outputSize_H = floor(inputSize_H * scalesH), \quad outputSize_W = floor(inputSize_W * scalesW) $$

    这一“output_size 与 scales 二选一”的约束在 InferShape 中同样有体现:在 upsample_trilinear3d_backward_infershape.cpp 中,若output_size非空且scales为空,则校验 output_size 的三个元素必须等于 grad_output 的 D、H、W 维;若output_size为空且scales非空,则校验floor(input_size[i] * scales[i])必须等于 grad_output 对应维度;两者同时为空或同时非空则直接报错。

  • 确定性计算:aclnnUpsampleTrilinear3dBackward 默认为确定性实现,多次执行结果可复现。

调用示例

仓库中提供了完整的可运行样例 test_aclnn_upsample_trilinear3d_backward.cpp,其编译与执行过程可参考 编译与运行样例。核心调用流程如下:

#include <iostream> #include <vector> #include "acl/acl.h" #include "aclnnop/aclnn_upsample_trilinear_3d_backward.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::vector<int64_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::vector<T>& hostData, const std::vector<int64_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::vector<int64_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_NCDHW, 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 == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret); // 2. 构造输入与输出,需要根据API的接口自定义构造 std::vector<int64_t> gradOutShape = {2, 2, 2, 2, 2}; std::vector<int64_t> gradInputShape = {2, 2, 1, 1, 1}; void* gradOutDeviceAddr = nullptr; void* gradInputDeviceAddr = nullptr; aclTensor* gradOut = nullptr; aclTensor* gradInput = nullptr; std::vector<float> gradOutHostData = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32}; std::vector<float> gradInputHostData = {2.0, 2, 2, 2}; std::vector<int64_t> outputSizeData = {2, 2, 2}; std::vector<int64_t> inputSizeData = {2, 2, 1, 1, 1}; bool alignCorners = false; double scalesD = 0.0; double scalesH = 0.0; double scalesW = 0.0; // 创建gradOut aclTensor ret = CreateAclTensor(gradOutHostData, gradOutShape, &gradOutDeviceAddr, aclDataType::ACL_FLOAT, &gradOut); CHECK_RET(ret == ACL_SUCCESS, return ret); // 创建gradInput aclTensor ret = CreateAclTensor(gradInputHostData, gradInputShape, &gradInputDeviceAddr, aclDataType::ACL_FLOAT, &gradInput); CHECK_RET(ret == ACL_SUCCESS, return ret); const aclIntArray* outputSize = aclCreateIntArray(outputSizeData.data(), outputSizeData.size()); CHECK_RET(outputSize != nullptr, return ACL_ERROR_INTERNAL_ERROR); const aclIntArray* inputSize = aclCreateIntArray(inputSizeData.data(), inputSizeData.size()); CHECK_RET(inputSize != nullptr, return ACL_ERROR_INTERNAL_ERROR); // 3. 调用CANN算子库API,需要修改为具体的API uint64_t workspaceSize = 0; aclOpExecutor* executor; // 调用aclnnUpsampleTrilinear3dBackward第一段接口 ret = aclnnUpsampleTrilinear3dBackwardGetWorkspaceSize(gradOut, outputSize, inputSize, alignCorners, scalesD, scalesH, scalesW, gradInput, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnUpsampleTrilinear3dBackwardGetWorkspaceSize 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); } // 调用aclnnUpsampleTrilinear3dBackward第二段接口 ret = aclnnUpsampleTrilinear3dBackward(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnUpsampleTrilinear3dBackward 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(gradInputShape); std::vector<float> 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和aclScalar,需要根据具体API的接口定义修改 aclDestroyTensor(gradOut); aclDestroyTensor(gradInput); // 7. 释放device资源,需要根据具体API的接口定义修改 aclrtFree(gradOutDeviceAddr); aclrtFree(gradInputDeviceAddr); if (workspaceSize > 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }

示例的调用流程可归纳为七步:设备与流初始化 → 构造输入输出 aclTensor 与 aclIntArray → 调用第一段接口获取 workspace 与 executor → 申请 workspace 设备内存 → 调用第二段接口执行计算 → 同步流并回拷结果 → 释放资源。示例中gradOutshape 为{2, 2, 2, 2, 2}gradInputshape 为{2, 2, 1, 1, 1},对应一次2x2x2 → 1x1x1的下采样反向,alignCorners=false且三个 scales 传 0(此时完全依据 inputSize/outputSize 计算缩放系数)。

源码级实现剖析

1. 接口层:从 NDHWC 到 NCDHW 的自动转置

在 aclnn_upsample_trilinear_3d_backward.cpp 中,第一段接口完成参数校验后会执行以下关键逻辑:

  • gradOut调用l0op::Contiguous做连续性整理;
  • scalesD > 0 && scalesH > 0 && scalesW > 0,则把三个 scale 组装进scalesList,并把outputSize置为空数组(L203-L209),即「显式指定 scales 时不再需要 outputSize」,与 InferShape 的二选一约束完全对应;同时生成castScales数组保存1/scalesD、1/scalesH、1/scalesW,对应公式中scale = 1/scales的分支;
  • 核心计算函数upsampleTrilinear3dBackwardCompute(L149-L174)中,若gradOut为 NDHWC 格式,会先按{0, 4, 1, 2, 3}置换为 NCDHW 调用底层l0op::UpsampleTrilinear3dGradNcdhw,再按{0, 2, 3, 4, 1}置换回 NDHWC 输出;ND 与 NCDHW 格式则直接计算;
  • 最后通过l0op::ViewCopy将中间结果写入用户提供的gradInput,并将整个执行图所需的 workspace 大小通过executor->GetWorkspaceSize()返回(L230-L235)。

第二段接口aclnnUpsampleTrilinear3dBackward则直接调用CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成图执行。

2. InferShape:输出 shape 推导与约束校验

upsample_trilinear3d_backward_infershape.cpp 中通过IMPL_OP_INFERSHAPE(UpsampleTrilinear3dGrad)注册 shape 推导:要求 grad_output 为 5 维、input_size 的 size 必须为 5,随后按上文所述的「output_size 与 scales 二选一」规则校验 grad_output 的 D、H、W 维,最终输出 shape 的 N、C、D、H、W 全部取自input_size

3. Tiling:分核策略与 workspace 规划

Tiling 逻辑位于 upsample_trilinear3d_backward_tiling_func.cpp,其中几个关键常量与算法值得注意:

  • SLIDE_SIZE = 16:每个核心按 16 为滑动步长切分数据;
  • MIN_SUPPORT_SCALE = 0.02f:三个方向实际缩放系数(ComputeScales)不得小于 0.02,否则 tiling 失败;
  • AreaPixelComputeScale(L228-L239)与数学公式一致:outputSize == inputSize时缩放系数为 1(该方向无需 resize),alignCorners时取(inputSize-1)/(outputSize-1),否则取scaleinputSize/outputSize
  • 三个方向(D/H/W)分别计算各自的单核 K 维度、分核数量与 matmul tiling,最终needCoreNum取三个方向的最大值(L241-L267);
  • workspace 由「W 方向中间矩阵 + H 方向中间矩阵 + 各核权重矩阵 + 32MB 预留空间」组成(GetWorkSpace),其中权重矩阵大小为SLIDE_SIZE * singleCoreKsingleCoreK = min(ceil((16+4)/scale)+4, outputSize)
  • Tiling 数据通过 upsample_trilinear3d_backward_tiling.h 中BEGIN_TILING_DATA_DEF定义的结构体传递,包含数据类型、batch、三维输入输出 shape、三轴 scale、alignCorners、各方向是否需要 resize、滑动/分核参数以及 W/H/D 三个TCubeTiling

4. Kernel:三方向 matmul 流水

kernel 入口 upsample_trilinear3d_backward.cpp 按dataType(half / float / bfloat16)实例化UpsampleTrilinear3dBackwardND<T>模板类。该类定义于 upsample_trilinear3d_backward.h,其核心设计是:

  • 将「梯度按权重散射回原图」转化为W → H → D 三个方向的连续 matmul,类内持有matmulWmatmulHmatmulD三个 Cube 侧 matmul 对象;
  • CalculateRadioTensor计算每个方向的权重矩阵(即各输出点对应的 lambda/lambdb/lambdc),坐标映射逻辑AreaPixelComputeSourceIndex(L256-L263)与公式完全对应:alignCorners 时scale * dstIndex,否则max(scale*(dst+0.5)-0.5, 0)
  • Process(L141-L162)按 W → H → D 顺序依次执行DirectionExpansion,每步之间以SyncAll同步多核,中间结果写入 GM workspace;某方向无需 resize(needResizeX == false)时跳过该方向的 matmul,并把上一阶段结果直接作为输出或下一阶段输入;
  • 在只对一个方向 resize 的简化场景下,可直接将 matmul 结果写入最终输出outTensorsGM,避免额外的中间矩阵搬运。

整体上,该 kernel 利用 Cube 算力将逐点插值权重计算转化为矩阵乘,属于「权重矩阵 × 数据矩阵」的经典算子实现范式。

测试与验证

仓库为该算子提供了多层次的测试支撑:

  • ST 级用例:在 atk_aclnnUpsampleTrilinear3dBackward.json 中配置了上百组测试场景,对应 PyTorch 的torch.nn.functional.interpolate(mode='trilinear'),覆盖 bf16/fp16 多种数据类型、align_corners两种取值、多样化的 shape(含[1,64,8,128,128]等大 shape)以及 ±inf 等边界输入值,精度标准采用double_benchmark双精度基准比对;
  • UT 级用例tests/ut/op_host/op_api/test_aclnn_upsample_trilinear_3d_backward.cpp验证接口层行为,tests/ut/op_host/test_upsample_trilinear3d_backward_tiling.cpp验证 tiling 数据生成,tests/ut/op_host/test_upsample_trilinear3d_grad_infershape.cpp验证 shape 推导,tests/ut/op_kernel/test_upsample_trilinear3d_backward.cpp配合gen_data.py/compare_data.py完成 kernel 数值比对。

小结

UpsampleTrilinear3dBackward 是 CANN ops-cv 图像算子库中三维三线性上采样(F.interpolate(..., mode='trilinear'))的反向梯度算子。本文从产品支持、数学原理(角对齐/边对齐缩放、8 邻域加权插值、梯度散射公式)、算子参数、aclnn 两段式接口与错误码、完整调用示例,到 host 侧算子定义 / InferShape / Tiling 与 kernel 侧三方向 matmul 流水做了全链路解析。开发者可直接参照 测试样例 完成调用,并依据 接口文档 中的 shape、内存与 scales 约束合理设计输入规模。

【免费下载链接】ops-cv本项目是CANN提供的图像处理、目标检测相关的算子库,实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-cv

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询