slow_conv3d_update_output_frame Class — pytorch Architecture

Source Code

aten/src/ATen/native/ConvolutionMM3d.cpp lines 288–335

template <typename scalar_t>
void slow_conv3d_update_output_frame(
    TensorAccessor<const scalar_t, 4> input,
    TensorAccessor<scalar_t, 4> output,
    TensorAccessor<const scalar_t, 2> weight,
    bool has_bias,
    TensorAccessor<const scalar_t, 2> finput,
    int64_t kernel_depth,
    int64_t kernel_height,
    int64_t kernel_width,
    int64_t stride_depth,
    int64_t stride_height,
    int64_t stride_width,
    int64_t pad_depth,
    int64_t pad_height,
    int64_t pad_width,
    int64_t n_input_plane,
    int64_t groups,
    int64_t input_depth,
    int64_t input_height,
    int64_t input_width,
    int64_t n_output_plane,
    int64_t output_depth,
    int64_t output_height,
    int64_t output_width) {
  const int beta = has_bias ? 1 : 0;

  // Compute out = weight * input
  // Note gemm expects fortran order, so all 3 matrices are transposed.
  // Swapping argument order cancels this, since C == AB <=> T(C) == T(B)T(A)
  const int64_t m = output_depth * output_height * output_width;
  const int64_t n = (n_output_plane / groups);
  const int64_t k = (n_input_plane / groups) * kernel_depth * kernel_height * kernel_width;

  const int64_t lda = m;
  const int64_t ldb = k;
  const int64_t ldc = m;

  at::native::cpublas::gemm_batched_with_stride(
      TransposeType::NoTranspose,
      TransposeType::NoTranspose,
      groups, m, n, k,
      static_cast<scalar_t>(1),
      finput.data(), lda, finput.stride(0) * k,
      weight.data(), ldb, weight.stride(0) * n,
      static_cast<scalar_t>(beta),
      output.data(), ldc, output.stride(0) * n);
}