Pytorch CUDA Reflect Padding 算子实现详解

CUDA 简介

• CUDA（Compute Unified Device Architecture）是由NVIDIA开发的一种并行计算平台和应用编程接口（API），允许软件开发者和软件工程师使用NVIDIA的图形处理单元（GPU）进行通用计算。自2007年推出以来，CUDA已经使得利用GPU的强大计算能力进行高性能计算（HPC）和复杂图形渲染成为可能，广泛应用于科学计算、工程、机器学习和深度学习等领域。
• CUDA 相关资料
  • 官方文档：https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html
  • 入门样例：https://cuda-tutorial.readthedocs.io/en/latest/tutorials/tutorial01/

Reflect Padding 介绍

• 反射填充是一种常见的图像边缘填充技术，用于卷积神经网络中，特别是在处理图像数据时。它通过镜像边缘像素来扩展图像的尺寸，从而使得边缘信息在卷积操作中得到更好的保留。reflect padding 样例如下图所示：
  
• Q: 反射填充与零填充在实际应用中有何不同？
  • A: 反射填充通过复制边缘像素来扩展图像，保持了图像边缘的自然连续性，而零填充则在边缘添加零值，可能会在卷积后引入人为的边缘效应。

Pytorch Reflect Padding 实现

• torch reflect padding 文档：https://pytorch.org/docs/stable/generated/torch.nn.ReflectionPad2d.html

>>> import torch.nn as nn
>>> import torch
>>> m = nn.ReflectionPad2d(2)
>>> input = torch.arange(9, dtype=torch.float).reshape(1, 1, 3, 3)
>>> input
tensor([[[[0., 1., 2.],
          [3., 4., 5.],
          [6., 7., 8.]]]])
>>> m(input)
tensor([[[[8., 7., 6., 7., 8., 7., 6.],
          [5., 4., 3., 4., 5., 4., 3.],
          [2., 1., 0., 1., 2., 1., 0.],
          [5., 4., 3., 4., 5., 4., 3.],
          [8., 7., 6., 7., 8., 7., 6.],
          [5., 4., 3., 4., 5., 4., 3.],
          [2., 1., 0., 1., 2., 1., 0.]]]])
>>> # using different paddings for different sides
>>> m = nn.ReflectionPad2d((1, 1, 2, 0))
>>> m(input)
tensor([[[[7., 6., 7., 8., 7.],
          [4., 3., 4., 5., 4.],
          [1., 0., 1., 2., 1.],
          [4., 3., 4., 5., 4.],
          [7., 6., 7., 8., 7.]]]])

CUDA Reflect Padding 代码实现理解

forward

• reflection_pad2d_out_template 实现，用于执行二维反射填充。

// 定义一个函数，用于对输入Tensor进行二维反射填充，并将结果输出到output Tensor。
void reflection_pad2d_out_template(
    Tensor &output, const Tensor &input_, IntArrayRef padding) {

  // 检查输入Tensor是否可以使用32位索引数学运算。
  TORCH_CHECK(canUse32BitIndexMath(input_),
    "input tensor must fit into 32-bit index math");

  // 初始化一些维度标识符和批次大小。
  int plane_dim = 0;
  int dim_h = 1;
  int dim_w = 2;
  int nbatch = 1;

  // 检查输入Tensor和padding参数是否合法。
  at::native::padding::check_valid_input<2>(input_, padding);

  // 如果输入Tensor是4维的，说明有批次维度，需要相应调整其他维度的索引，并更新批次大小。
  if (input_.ndimension() == 4) {
    nbatch = input_.size(0);
    plane_dim++;
    dim_h++;
    dim_w++;
  }

  // 从padding参数中提取左、右、上、下四个方向的填充大小。
  int64_t pad_l = padding[0];
  int64_t pad_r = padding[1];
  int64_t pad_t = padding[2];
  int64_t pad_b = padding[3];

  // 获取输入Tensor在不同维度上的大小。
  int nplane = input_.size(plane_dim);
  int input_h = input_.size(dim_h);
  int input_w = input_.size(dim_w);

  // 检查左右填充大小是否小于输入宽度，上下填充大小是否小于输入高度。
  TORCH_CHECK(pad_l < input_w && pad_r < input_w, ...);
  TORCH_CHECK(pad_t < input_h && pad_b < input_h, ...);

  // 计算输出Tensor的高度和宽度。
  int output_h = input_h + pad_t + pad_b;
  int output_w = input_w + pad_l + pad_r;

  // 确保计算出的输出Tensor尺寸是有效的。
  TORCH_CHECK(output_w >= 1 || output_h >= 1, ...);

  // 根据输入Tensor的维度，调整输出Tensor的尺寸。
  if (input_.ndimension() == 3) {
    output.resize_({nplane, output_h, output_w});
  } else {
    output.resize_({nbatch, nplane, output_h, output_w});
  }
  // 如果输出Tensor为空，则不执行后续操作。
  if (output.numel() == 0) {
    return;
  }

  // 确保输入Tensor是连续的，便于后续处理。
  Tensor input = input_.contiguous();

  // 计算输出平面的大小，用于配置CUDA核函数的参数。
  int64_t output_plane_size = output_h * output_w;
  dim3 block_size(output_plane_size > 256 ? 256 : output_plane_size);

  // 准备在CUDA核函数中使用的变量。
  int64_t size_y = nplane;
  int64_t size_z = nbatch;

  // 对所有数据类型执行反射填充操作
  AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(kHalf, kBFloat16,
    input.scalar_type(), "reflection_pad2d_out_template", [&] {

      // 遍历所有平面和批次进行填充
      for (int64_t block_y = 0; block_y < size_y; block_y += 65535) {
        int64_t block_y_size = std::min(size_y - block_y, static_cast<int64_t>(65535));
        for (int64_t block_z = 0; block_z < size_z; block_z += 65535) {
          int64_t block_z_size = std::min(size_z - block_z, static_cast<int64_t>(65535));

          // 计算网格大小并启动CUDA核心
          dim3 grid_size(at::ceil_div(output_plane_size, static_cast<int64_t>(256)), block_y_size, block_z_size);

          // 计算网格大小并启动CUDA核心
          // 这里使用了CUDA的核心启动语法，`<<<grid_size, block_size, 0, at::cuda::getCurrentCUDAStream()>>>`，
          // 其中grid_size和block_size是CUDA核心执行时网格和块的维度配置，这里的0表示使用默认的共享内存大小，
          // at::cuda::getCurrentCUDAStream()获取当前CUDA流，用于并行计算。
          reflection_pad2d_out_kernel<<<
            grid_size, block_size, 0, at::cuda::getCurrentCUDAStream()>>>(
              // 传递给核心函数的参数，包括输入和输出张量的数据指针，
              // 输入的宽度和高度，四个方向的填充大小，当前处理的平面和批次索引，以及平面的总数。
              input.const_data_ptr<scalar_t>(), output.mutable_data_ptr<scalar_t>(),
              input_w, input_h,
              pad_t, pad_b, pad_l, pad_r, block_y, block_z, nplane);
          // 检查CUDA核心启动后是否有错误发生
          C10_CUDA_KERNEL_LAUNCH_CHECK();
        }
      }
    }
  );
}

代码的最后部分是关键的，它展示了如何调用CUDA核心函数（reflection_pad2d_out_kernel）来实际执行反射填充操作。这个核心函数利用 CUDA 的并行计算能力，对输入张量的每个元素进行填充处理，确保在 GPU 上高效地完成操作。C10_CUDA_KERNEL_LAUNCH_CHECK() 是用于检测核心启动后是否有任何错误发生。

• reflection_pad2d_out_kernel 实现：CUDA reflect pad2d 核函数。它接收输入和输出张量的指针、输入尺寸、填充尺寸和平面偏移量，然后计算每个线程应处理的输出张量中的像素位置，并根据输入张量中相应位置的值来填充它。

template<typename scalar_t>
__global__ void reflection_pad2d_out_kernel(
    const scalar_t * input, scalar_t * output,
    int64_t input_dim_x, int64_t input_dim_y,
    int pad_t, int pad_b, int pad_l, int pad_r, int y_shift, int z_shift, int nplane) {
  // 计算当前线程负责的输出位置
  auto output_xy = threadIdx.x + blockIdx.x * blockDim.x;
  // 计算输出维度
  auto output_dim_x = input_dim_x + pad_l + pad_r;
  auto output_dim_y = input_dim_y + pad_t + pad_b;

  // 如果当前线程负责的位置在输出范围内
  if (output_xy < output_dim_x * output_dim_y) {
    // 获取输入和输出索引映射
    auto index_pair = get_index_mapping2d(
      input_dim_x, input_dim_y,
      output_dim_x, output_dim_y,
      pad_l, pad_t,
      output_xy, y_shift, z_shift, nplane);

    // 根据映射关系复制数据
    output[index_pair.second] = input[index_pair.first];
  }
}

• get_index_mapping2d 函数实现：基于输出像素位置、填充参数和偏移量，计算出反射填充后的输入和输出索引。这个函数利用了 CUDA 的内置函数 abs 来处理反射逻辑，确保输出位置正确地映射到输入张量上

// 定义一个 mapping 函数，用于计算从输出位置到输入位置的索引映射。
__device__
inline thrust::pair<int64_t, int64_t>  get_index_mapping2d(
    int64_t input_dim_x, int64_t input_dim_y,
    int64_t output_dim_x, int64_t output_dim_y,
    int64_t pad_l, int64_t pad_t,
    int64_t output_xy, int y_shift, int z_shift, int nplane) {
  
  // 计算输入和输出的偏移量，考虑了批次和通道的变化。
  auto input_offset =
    ((blockIdx.y + y_shift) + (blockIdx.z + z_shift) * nplane) * input_dim_x * input_dim_y;
  auto output_offset =
    ((blockIdx.y + y_shift) + (blockIdx.z + z_shift) * nplane) * output_dim_x * output_dim_y;

  // 根据线性索引计算输出坐标。
  auto output_x = output_xy % output_dim_x;
  auto output_y = output_xy / output_dim_x;

  // 计算输入和输出坐标的起始点。
  auto i_start_x = ::max(int64_t(0), -pad_l);
  auto i_start_y = ::max(int64_t(0), -pad_t);
  auto o_start_x = ::max(int64_t(0), pad_l);
  auto o_start_y = ::max(int64_t(0), pad_t);

  // 根据反射逻辑计算输入坐标。
  auto input_x = ::abs(output_x - pad_l)
                 - ::abs(output_x - (input_dim_x + pad_l - 1))
                 - output_x
                 + 2 * pad_l + input_dim_x - 1
                 - o_start_x + i_start_x;

  auto input_y = ::abs(output_y - pad_t)
                 - ::abs(output_y - (input_dim_y + pad_t - 1))
                 - output_y
                 + 2 * pad_t + input_dim_y - 1
                 - o_start_y + i_start_y;

  // 返回输入和输出坐标的线性索引对。
  return thrust::make_pair<int64_t, int64_t>(
    input_offset + input_y * input_dim_x + input_x,
    output_offset + output_y * output_dim_x + output_x);
}

backward

• backward 与 forward 整体实现思路接近，主要是梯度反传时逻辑与前传时需要反过来，代码实现思路基本和之前介绍的 forward 部分一致
• backward 函数入口

// 定义一个函数，用于计算二维反射填充的梯度输出。
void reflection_pad2d_backward_out_template(
    Tensor &grad_input, const Tensor &grad_output_,
    const Tensor &input, IntArrayRef padding) {

  // 如果梯度输入的元素数为0，则不执行任何操作。
  if (grad_input.numel() == 0) {
    return;
  }

  // 检查输入张量和梯度输出张量是否可以使用32位索引进行数学运算，如果不可以则抛出错误。
  TORCH_CHECK(canUse32BitIndexMath(input),
    "input tensor must fit into 32-bit index math");
  TORCH_CHECK(canUse32BitIndexMath(grad_output_),
    "output gradient tensor must fit into 32-bit index math");

  // 初始化一些维度和批次的变量，用于后续的张量尺寸计算。
  int plane_dim = 0;
  int dim_h = 1;
  int dim_w = 2;
  int nbatch = 1;

  // 如果输入张量的维度是4，说明有一个批次维度，需要相应地调整其他维度的索引，并计算批次大小。
  if (input.ndimension() == 4) {
    nbatch = input.size(0);
    plane_dim++;
    dim_h++;
    dim_w++;
  }

  // 解析padding参数，得到左、右、上、下的填充尺寸。
  int64_t pad_l = padding[0];
  int64_t pad_r = padding[1];
  int64_t pad_t = padding[2];
  int64_t pad_b = padding[3];

  // 计算输入张量在特定维度上的尺寸。
  int nplane = input.size(plane_dim);
  int input_h = input.size(dim_h);
  int input_w = input.size(dim_w);

  // 根据输入尺寸和填充尺寸计算输出尺寸。
  int output_h = input_h + pad_t + pad_b;
  int output_w  = input_w + pad_l + pad_r;

  // 检查梯度输出张量的尺寸是否与预期一致，如果不一致则抛出错误。
  TORCH_CHECK(output_w == grad_output_.size(dim_w), "grad_output width unexpected. Expected: ", output_w, ", Got: ", grad_output_.size(dim_w));
  TORCH_CHECK(output_h == grad_output_.size(dim_h), "grad_output height unexpected. Expected: ", output_h, ", Got: ", grad_output_.size(dim_h));

  // 为了保证数据的连续性，将梯度输出张量转换为连续的。
  Tensor grad_output = grad_output_.contiguous();

  // 计算输出平面的大小，用于后续的CUDA核函数配置。
  int64_t output_plane_size = output_h * output_w;
  // 配置CUDA核函数的线程块大小，取256或输出平面大小的较小者。
  dim3 block_size(output_plane_size > 256 ? 256 : output_plane_size);

  // 准备循环遍历的尺寸变量。
  int64_t size_y = nplane;
  int64_t size_z = nbatch;

  // 对输入张量的数据类型进行分派，支持多种浮点和复数类型。
  AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND2(kHalf, kBFloat16,
    input.scalar_type(), "reflection_pad2d_backward_out_template", [&] {

      // 对每个平面（通道）和批次进行循环，处理大于65535的情况。
      for (int64_t block_y = 0; block_y < size_y; block_y += 65535) {
        int64_t block_y_size = std::min(size_y - block_y, static_cast<int64_t>(65535));
        for (int64_t block_z = 0; block_z < size_z; block_z += 65535) {
          int64_t block_z_size = std::min(size_z - block_z, static_cast<int64_t>(65535));

          // 计算网格大小，用于CUDA核函数的配置。
          dim3 grid_size(at::ceil_div(output_plane_size, static_cast<int64_t>(256)), block_y_size, block_z_size);

          // 调用CUDA核函数，计算梯度输入。
          reflection_pad2d_backward_out_kernel<<<
            grid_size, block_size, 0, at::cuda::getCurrentCUDAStream()>>>(
              grad_input.mutable_data_ptr<scalar_t>(), grad_output.const_data_ptr<scalar_t>(),
              input_w, input_h,
              pad_t, pad_b, pad_l, pad_r, block_y, block_z, nplane);
          // 检查CUDA核函数的启动是否有错误。
          C10_CUDA_KERNEL_LAUNCH_CHECK();
        }
      }
    }
  );
}

• reflection_pad2d_backward_out_kernel 实现：

// 定义模板函数，用于CUDA内核，处理反射填充的梯度反向传播。
template <typename scalar_t>
__global__ void reflection_pad2d_backward_out_kernel(
    scalar_t * grad_input, // 指向梯度输入的指针，即对应前向传播输入的梯度
    const scalar_t * grad_output, // 指向梯度输出的指针，即损失函数对输出的偏导
    int64_t input_dim_x, // 输入的宽度
    int64_t input_dim_y, // 输入的高度
    int pad_t, // 顶部填充的大小
    int pad_b, // 底部填充的大小
    int pad_l, // 左侧填充的大小
    int pad_r, // 右侧填充的大小
    int y_shift, // 平面（plane）的偏移量，用于多通道数据处理
    int z_shift, // 批量的偏移量，用于批处理
    int nplane) { // 通道数或平面数
  auto output_xy = threadIdx.x + blockIdx.x * blockDim.x; // 计算当前线程处理的输出位置索引
  auto output_dim_x = input_dim_x + pad_l + pad_r; // 计算经过填充后的输出宽度
  auto output_dim_y = input_dim_y + pad_t + pad_b; // 计算经过填充后的输出高度

  // 判断当前线程负责的输出位置是否在有效范围内
  if (output_xy < output_dim_x * output_dim_y) {
    // 计算输出位置对应的输入位置索引
    auto index_pair = get_index_mapping2d(
      input_dim_x, input_dim_y,
      output_dim_x, output_dim_y,
      pad_l, pad_t,
      output_xy, y_shift, z_shift, nplane);

    // 使用原子操作累加计算梯度输入。这里的原子操作确保了多个线程更新同一位置时的正确性。
    gpuAtomicAddNoReturn(&grad_input[index_pair.first], grad_output[index_pair.second]);
  }
}

总结

• PyTorch 中的 CUDA 反射填充通过两个 CUDA 核函数实现：reflection_pad2d_out_kernel 和 reflection_pad2d_backward_out_kernel。
  这两个核函数利用了 CUDA 的并行计算能力，可以高效地执行反射填充操作。
  • 其中 reflection_pad2d_out_kernel 理解了之后 reflection_pad2d_backward_out_kernel 理解起来就比较容易了
  • 代码的核心逻辑主要是在 padding 时输入输出之间的映射关系实现部分，也即 get_index_mapping2d 函数实现需要关注下具体实现细节
• 通过这篇博客，我们简单介绍了 CUDA 和反射填充的概念和应用，提供了实际的代码实现理解和对应资源的链接，希望能帮助读者更深入地理解并利用这些技术。