Unfold3dZeroPaddingAccKernelImpl Class — pytorch Architecture

Source Code

aten/src/ATen/native/Unfold3d.cpp lines 302–353

template <typename T>
void Unfold3dZeroPaddingAccKernelImpl(
    int64_t C,
    int64_t X_D,
    int64_t X_H,
    int64_t X_W,
    int64_t Y_D,
    int64_t Y_H,
    int64_t Y_W,
    int64_t kernel_d,
    int64_t kernel_h,
    int64_t kernel_w,
    int64_t stride_d,
    int64_t stride_h,
    int64_t stride_w,
    const T* src,
    T* dst) {
  const int64_t X_size = X_D * X_H * X_W;
  const int64_t Y_size = Y_D * Y_H * Y_W;
  const int64_t kernel_size = kernel_d * kernel_h * kernel_w;
  at::parallel_for(0, C, 0, [=](int64_t begin, int64_t end) {
    std::memset(dst + begin * X_size, 0, (end - begin) * X_size * sizeof(T));
    for (const auto c : c10::irange(begin, end)) {
      for (const auto kd : c10::irange(kernel_d)) {
        for (const auto kh : c10::irange(kernel_h)) {
          for (const auto kw : c10::irange(kernel_w)) {
            const int64_t p =
                c * kernel_size + kd * kernel_h * kernel_w + kh * kernel_w + kw;
            for (const auto yd : c10::irange(Y_D)) {
              const int64_t xd = yd * stride_d + kd;
              const T* src_ptr = src + p * Y_size + yd * Y_H * Y_W;
              T* dst_ptr = dst + c * X_size + xd * X_H * X_W + kh * X_W + kw;
              if (stride_w == 1) {
                MatAdd<T>(Y_H, Y_W, Y_W, stride_h * X_W, src_ptr, dst_ptr);
              } else {
                MatAdd<T>(
                    Y_H,
                    Y_W,
                    Y_W,
                    1,
                    stride_h * X_W,
                    stride_w,
                    src_ptr,
                    dst_ptr);
              }
            }
          }
        }
      }
    }
  });
}