Convolution.cpp

#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
#include <ATen/core/Tensor.h>
#include <ATen/Config.h>
#include <ATen/Parallel.h>
#include <ATen/TensorOperators.h>
#include <ATen/native/ConvolutionMM3d.h>
#include <ATen/native/ConvUtils.h>
#include <ATen/native/Pool.h>
#include <ATen/native/cpu/DepthwiseConvKernel.h>
#include <ATen/native/utils/ParamUtils.h>
#include <ATen/native/xnnpack/Engine.h>
#include <c10/util/accumulate.h>
#include <c10/util/irange.h>
#include <c10/macros/Macros.h>
#include <limits>
#include <utility>

#ifndef AT_PER_OPERATOR_HEADERS
#include <ATen/Functions.h>
#else
#include <ATen/ops/permute.h>
#endif

#if AT_NNPACK_ENABLED()
#include <nnpack.h>
#endif

#if AT_MKLDNN_ENABLED()
#include <ATen/native/mkldnn/Utils.h>
#endif

#ifndef AT_PER_OPERATOR_HEADERS
#include <ATen/Functions.h>
#include <ATen/NativeFunctions.h>
#else
#include <ATen/ops/_conv_depthwise2d.h>
#include <ATen/ops/_convolution.h>
#include <ATen/ops/_convolution_double_backward_native.h>
#include <ATen/ops/_convolution_mode.h>
#include <ATen/ops/_convolution_mode_native.h>
#include <ATen/ops/_convolution_native.h>
#include <ATen/ops/_mps_convolution.h>
#include <ATen/ops/_mps_convolution_transpose.h>
#include <ATen/ops/_nnpack_available.h>
#include <ATen/ops/_nnpack_spatial_convolution.h>
#include <ATen/ops/_slow_conv2d_backward.h>
#include <ATen/ops/_unsafe_view.h>
#include <ATen/ops/cat.h>
#include <ATen/ops/constant_pad_nd.h>
#include <ATen/ops/conv1d_native.h>
#include <ATen/ops/conv2d_native.h>
#include <ATen/ops/conv3d_native.h>
#include <ATen/ops/conv_depthwise3d.h>
#include <ATen/ops/conv_transpose1d_native.h>
#include <ATen/ops/conv_transpose2d_native.h>
#include <ATen/ops/conv_transpose3d_native.h>
#include <ATen/ops/convolution.h>
#include <ATen/ops/convolution_backward_native.h>
#include <ATen/ops/convolution_backward_overrideable.h>
#include <ATen/ops/convolution_backward_overrideable_native.h>
#include <ATen/ops/convolution_native.h>
#include <ATen/ops/convolution_overrideable.h>
#include <ATen/ops/convolution_overrideable_native.h>
#include <ATen/ops/cudnn_convolution.h>
#include <ATen/ops/cudnn_convolution_transpose.h>
#include <ATen/ops/empty.h>
#include <ATen/ops/empty_like.h>
#include <ATen/ops/empty_native.h>
#include <ATen/ops/miopen_convolution.h>
#include <ATen/ops/miopen_convolution_transpose.h>
#include <ATen/ops/miopen_depthwise_convolution.h>
#include <ATen/ops/mkldnn_convolution.h>
#include <ATen/ops/mps_convolution_backward.h>
#include <ATen/ops/mps_convolution_transpose_backward.h>
#include <ATen/ops/slow_conv3d.h>
#include <ATen/ops/slow_conv_dilated2d.h>
#include <ATen/ops/slow_conv_dilated3d.h>
#include <ATen/ops/slow_conv_transpose2d.h>
#include <ATen/ops/slow_conv_transpose3d.h>
#include <ATen/ops/thnn_conv2d.h>
#include <ATen/ops/view_as_real.h>
#include <ATen/ops/zeros.h>
#include <ATen/ops/zeros_like.h>
#endif

constexpr int MIOPEN_DIM_MAX = 5;

namespace at { namespace native {

// Check workload to activate fast depthwise FP16 cudnn conv kernels
template <typename T>
bool check_cudnn_depthwise_workload(const at::Tensor& input, int stride) {
  auto w = at::symint::size<T>(input, 3);  // same as h
  auto ch = at::symint::size<T>(input, 1);
  auto bs = at::symint::size<T>(input, 0);
  if (stride==1) {
    if (w >= 7) {
      // All batch sizes and nb_channels
      if (w >= 112) {
        return true;
      }

      // large nb_channels
      if (ch >= 1024) {
        // NOLINTNEXTLINE(bugprone-branch-clone,cppcoreguidelines-avoid-magic-numbers)
        if (w >= 56) {
          return true;
        } else if (bs >= 32) {
          return true;
        }
      }

      // batch_size specific
      if (bs >= 128) {
        // NOLINTNEXTLINE(bugprone-branch-clone,cppcoreguidelines-avoid-magic-numbers)
        if (ch >= 512) {
          return true;
        } else if (ch >= 64) {
          if (w >= 14) {
            return true;
          }
        } else if ((ch >= 32) && (w >=28)) {
          return true;
        }
      } else if (bs >= 64) {
        // NOLINTNEXTLINE(bugprone-branch-clone,cppcoreguidelines-avoid-magic-numbers)
        if ((ch >= 256) && (w >= 14)) {
          return true;
        } else if ((ch >= 32) && (w >= 28)) {
          return true;
        }
      } else if (bs >= 32) {
        // NOLINTNEXTLINE(bugprone-branch-clone,cppcoreguidelines-avoid-magic-numbers)
        if ((ch >= 256) && (w >= 14)) {
          return true;
        } else if ((ch >= 128) && (w >= 28)) {
          return true;
        } else if ((ch >= 32) && (w >= 56)) {
          return true;
        }
      } else if (bs >= 16) {
        if ((ch >= 1024) && (w >= 14)) {
          return true;
        }
        // NOLINTNEXTLINE(bugprone-branch-clone,cppcoreguidelines-avoid-magic-numbers)
        if ((ch >= 256) && (w >= 28)) {
          return true;
        } else if ((ch >= 32) && (w >= 56)) {
          return true;
        }
      } else if (bs >= 8) {
        // NOLINTNEXTLINE(bugprone-branch-clone,cppcoreguidelines-avoid-magic-numbers)
        if ((ch >= 512) && (w >= 28)) {
          return true;
        } else if ((ch >= 64) && (w >= 56)) {
          return true;
        }
      }
    }
  } else if (stride==2) {
    if (ch < 256) {
      return false;
    }

    if (w >= 7) {
      if (bs >= 128) {
        // NOLINTNEXTLINE(bugprone-branch-clone,cppcoreguidelines-avoid-magic-numbers)
        if (ch >= 1024) {
          return true;
        } else if ((ch >= 512) && (w >= 14)) {
          return true;
        } else if (w >= 28) {
          return true;
        }
      } else if (bs >= 64) {
        // NOLINTNEXTLINE(bugprone-branch-clone,cppcoreguidelines-avoid-magic-numbers)
        if ((ch >= 512) && (w >= 14)) {
          return true;
        } else if (w >= 28) {
          return true;
        }
      } else if (bs >= 32) {
        // NOLINTNEXTLINE(bugprone-branch-clone,cppcoreguidelines-avoid-magic-numbers)
        if ((ch >= 1024) && (w >= 14)) {
          return true;
        } else if (w >= 28) {
          return true;
        }
      } else if (bs >= 16) {
        // NOLINTNEXTLINE(bugprone-branch-clone,cppcoreguidelines-avoid-magic-numbers)
        if ((ch >= 512) && (w >= 28)) {
          return true;
        } else if (w >= 56) {
          return true;
        }
      } else if (bs >= 8) {
        // NOLINTNEXTLINE(bugprone-branch-clone,cppcoreguidelines-avoid-magic-numbers)
        if ((ch >= 1024) && (w >= 28)) {
          return true;
        } else if (w >= 56) {
          return true;
        }
      } else if (bs >= 1) {
        if ((ch >= 512) && (w >=112)) {
          return true;
        }
      }
    }
  }
  return false;
}

// simplified version for cudnn 8.2 and above
template <typename T>
bool check_cudnn_depthwise_workload_with_filter(const at::Tensor& input, int stride, const at::Tensor& weight) {
  // 1D conv
  if(at::symint::size<T>(input, 2) == 1 && stride == 1){
    return true;
  }

  // 2d conv
  // only square filters
  if (at::symint::size<T>(weight, 2) != at::symint::size<T>(weight, 3)) return false;
  auto filter = at::symint::size<T>(weight, 3);
  // only 1/3/5 filter
  if (filter != 1 && filter != 3 && filter != 5) return false;
  // we don't enforce square input but only check width to reduce heuristic space
  if (at::symint::size<T>(input, 3) < 7) return false; // min width 7
  auto w = at::symint::size<T>(input, 3);
  // only 1/2 stride, use cudnn for all stride 1
  if (stride == 1) return true;
  if (stride != 2) return false;

  auto ch = at::symint::size<T>(input, 1);
  auto bs = at::symint::size<T>(input, 0);
  // special case since bs1 show good perf in lots of cases
  if (bs == 1) {
    if (filter == 1 && w <= 28) return true;
    if (filter == 3 || filter == 5) return true;
  } else {
    if (filter == 1 && bs <= 16 && ch >= 128 && w <= 7) return true;
    if (filter == 3 || filter == 5) {
      if ((ch >= 512) || (ch >= 256 && w >= 28)) return true;
    }
  }
  return false;
}


#if defined(C10_MOBILE)
static bool xnnpack_use_convolution2d(
    const Tensor& input,
    const Tensor& weight,
    const at::OptionalIntArrayRef bias_sizes_opt,
    const IntArrayRef padding,
    const IntArrayRef stride,
    const IntArrayRef dilation,
    const int64_t groups,
    const bool transposed) {
  return xnnpack::use_convolution2d(input, weight, bias_sizes_opt, padding, stride, dilation, groups, transposed);
}

static bool xnnpack_use_convolution2d(
    const Tensor& input,
    const Tensor& weight,
    const at::OptionalSymIntArrayRef bias_sizes_opt,
    const SymIntArrayRef padding,
    const IntArrayRef stride,
    const IntArrayRef dilation,
    const int64_t groups,
    const bool transposed) {
  // Never use xnnpack for symbolic tracing
  return false;
}
#endif

// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
// This struct is templated so that we can run backend selection in a dynamic
// shapes context; all of the real kernel selection in eager mode runs with
// int64_t
template <typename T>
struct ConvParams {
  std::vector<int64_t> stride;
  std::vector<T> padding;
  std::vector<int64_t> dilation;
  bool transposed;
  std::vector<T> output_padding;
  int groups;
  bool benchmark;
  bool deterministic;
  bool cudnn_enabled;
  bool allow_tf32;

  bool is_strided() const {
    bool is_strided = false;
    for (auto s : stride) {
      is_strided |= (s != 1);
    }
    return is_strided;
  }

  bool is_dilated() const {
    bool is_dilated = false;
    for (auto d : dilation) {
      is_dilated |= (d != 1);
    }
    return is_dilated;
  }

  bool is_padded() const {
    bool is_padded = false;
    for (auto p : padding) {
      is_padded |= (p != 0);
    }
    return is_padded;
  }

  bool is_output_padding_neg() const {
    bool is_non_neg = false;
    for (const auto& p : output_padding) {
      is_non_neg |= (p < 0);
    }
    return is_non_neg;
  }

  bool is_output_padding_big() const {
    bool is_big = false;
    for (auto i: c10::irange(output_padding.size())) {
      is_big |= (output_padding[i] >= stride[i]);
    }
    return is_big;
  }

  bool is_padding_neg() const {
    bool is_non_neg = false;
    for (const auto& p : padding) {
      is_non_neg |= (p < 0);
    }
    return is_non_neg;
  }

  bool is_stride_nonpos() const {
    bool is_nonpos = false;
    for (auto s : stride) {
      is_nonpos |= (s <= 0);
    }
    return is_nonpos;
  }

  void view1d_as_2d() {
    if (stride.size() == 1) {
      stride.insert(stride.begin(), 1);
      padding.insert(padding.begin(), 0);
      dilation.insert(dilation.begin(), 1);
      output_padding.insert(output_padding.begin(), 0);
    }
  }

  bool use_cpu_depthwise3x3_winograd(const at::Tensor& input, const at::Tensor& weight, const c10::optional<at::Tensor>& bias) const {
#if defined(__ARM_NEON__)
    // Currently only 3x3 depthwise convolutions on tensors of float are supported.
    return (input.ndimension() == 4) &&
           (at::symint::size<T>(input, 1) == groups) &&
           (weight.ndimension() == 4 ) &&
           (at::symint::size<T>(weight, 0) % at::symint::size<T>(input, 1) == 0) &&
           (at::symint::size<T>(weight, 1) == 1) &&
           (at::symint::size<T>(weight, 2) == 3) &&
           (at::symint::size<T>(weight, 3) == 3) &&
           (input.device().is_cpu()) &&
           (input.scalar_type() == at::kFloat) &&
           input.is_contiguous() &&
           (weight.device().is_cpu()) &&
           (weight.scalar_type() == at::kFloat) &&
           weight.is_contiguous() &&
           (!bias.has_value() || bias->is_contiguous()) &&
           !is_strided() &&
           !is_dilated() &&
           !transposed;
#else
    return false;
#endif
  }

  bool needs_64bit_indexing_no_split(const at::Tensor& input, const at::Tensor& weight) const {
    constexpr int64_t int_max = std::numeric_limits<int>::max();
    auto numel_input = at::symint::numel<T>(input);
    // empty input
    if (numel_input == 0) {
      return false;
    }
    // input size can not be reduced to the range of int by splitting the batch dim
    auto n = at::symint::size<T>(input, 0);
    if (numel_input / n > int_max) {
      return true;
    }
    // output size can not be reduced to the range of int by splitting the batch dim
    T outsize = 1;
    if (transposed) {
      auto o = conv_input_size(at::symint::sizes<T>(input), at::symint::sizes<T>(weight), padding, output_padding, stride, dilation, groups);
      outsize = c10::multiply_integers(o.begin() + 1, o.end());
    } else {
      auto o = conv_output_size(at::symint::sizes<T>(input), at::symint::sizes<T>(weight), padding, stride, dilation);
      outsize = c10::multiply_integers(o.begin() + 1, o.end());
    }
    return outsize > int_max;
  }

  bool use_cudnn(const at::Tensor& input, const at::Tensor& weight) const {
  // Note [Mobile check segfaults]
  // cudnn and miopen are guaranteed not to be on mobile, and T102591915 / T110194934 suggest
  // that maybe the compiledWithCuDNN() check sometimes segfaults (though I can't imagine how)
#if !defined(C10_MOBILE)
    if (needs_64bit_indexing_no_split(input, weight)) {
      return false;
    }
    if (!detail::getCUDAHooks().compiledWithCuDNN()) {
      return false;
    }
    if (!input.is_cuda() || !cudnn_enabled) {
      return false;
    }
    if (input.scalar_type() == at::kBFloat16 || weight.scalar_type() == at::kBFloat16) {
      if (!(detail::getCUDAHooks().supportsBFloat16ConvolutionWithCuDNNv8() && at::native::cudnnv8_enabled_check_debug())) {
        return false;
      }
    }
    if (cudnn_conv_suggest_memory_format(input, weight) == at::MemoryFormat::Contiguous) {
      // bypass dilation checks for channels_last convolution
      if (deterministic && is_dilated()) {
        // cudnn doesn't support deterministic dilated convolution fully yet
        return false;
      }
      if (is_dilated()) {
        return detail::getCUDAHooks().supportsDilatedConvolutionWithCuDNN() && !is_output_padding_big();
      }
    }
    return !is_output_padding_big();
#else
    return false;
#endif
  }

  // Use cudnn for FP16 depthwise convolutions
  bool use_cudnn_depthwise(const at::Tensor& input, const at::Tensor& weight) const  {
    if (cudnn_conv_suggest_memory_format(input, weight) != at::MemoryFormat::Contiguous && use_cudnn(input, weight)) {
      // always use cudnn_depthwise for channels_last format
      return true;
    }
    if (detail::getCUDAHooks().supportsDepthwiseConvolutionWithCuDNN()) {
      long cudnn_version = detail::getCUDAHooks().versionCuDNN();
      if (cudnn_version >= 8200) {
        bool kernel_cond =  (use_cudnn(input, weight) &&
                             input.scalar_type() == kHalf && // only for FP16
                             weight.scalar_type() == kHalf &&
                             is_depthwise(input, weight) &&
                             input.ndimension() == 4 &&   // TODO: 5-D contiguous depthwise is not supported yet, need benchmarks
                             !is_dilated() && // no dilation supported
                             (stride[0] == stride[1] || at::symint::size<T>(input, 2) == 1) && // square or 1d
                             at::symint::size<T>(input, 1) >= 32); // min 32 channels supported)
        if (kernel_cond) {
          return check_cudnn_depthwise_workload_with_filter<T>(input, stride[1], weight);
        }
      }
      // keep (7600 <= cudnn < 8200) code unchanged
      bool kernel_cond =  (cudnn_version >= 7600 &&
                           use_cudnn(input, weight) &&
                           input.scalar_type() == kHalf && // only for FP16
                           weight.scalar_type() == kHalf &&
                           is_depthwise(input, weight) &&
                           input.ndimension() == 4 &&   // TODO: 5-D contiguous depthwise is not supported yet, need benchmarks
                           at::symint::size<T>(weight, 2) == at::symint::size<T>(weight, 3) && // only square kernels
                           at::symint::size<T>(input, 2) >= 7 && // min width/height 7
                           !is_dilated() && // no dilation supported
                           stride[0] == stride[1] && // equal strides
                           ((at::symint::size<T>(weight, 3) == 3) || (at::symint::size<T>(weight, 3) == 1)) &&
                           at::symint::size<T>(input, 1) >= 32); // min 32 channels supported)
      if (kernel_cond) {
        return check_cudnn_depthwise_workload<T>(input, stride[0]);
      } else {
        return false;
      }
    } else {
      return false;
    }
  }

  bool use_miopen(const at::Tensor& input, const at::Tensor& weight, bool bias_defined) const  {
    if (needs_64bit_indexing_no_split(input, weight)) {
      return false;
    }
    return ((input.scalar_type() == at::kFloat) || (input.scalar_type() == at::kHalf) || (input.scalar_type() == at::kBFloat16))
           && detail::getCUDAHooks().compiledWithMIOpen()
           && input.is_cuda()
           && input.dim() <= MIOPEN_DIM_MAX
           && !(groups > 1 && is_dilated()) // MIOpen currently does not support dilation with groups of size > 1
           && !(input.scalar_type() == at::kBFloat16 && bias_defined) // MIOpen currently doesn't support bias with bfloat16
           && cudnn_enabled
           ;
  }
  bool use_mkldnn(const at::Tensor& input, const at::Tensor& weight) const  {
#if AT_MKLDNN_ENABLED()
    if (!at::globalContext().userEnabledMkldnn()) {
      return false;
    }
    if (transposed && is_output_padding_big()) {
      return false;
    }
    if (transposed && groups > 1 && at::symint::size<T>(input, 1) == groups) {
      return false;
    }
    if (input.device().is_cpu() && input.scalar_type() == kBFloat16 && mkldnn_bf16_device_check()) {
      return true;
    }
    return (input.is_mkldnn()) || // input is mkldnn Tensor
      (input.device().is_cpu() &&
       input.scalar_type() == kFloat && // only on CPU Float Tensors
       // For 1x1 filters, MKLDNN is faster than THNN when multi-threaded,
       // but THNN is faster when single-threaded.
       (is_strided() || is_dilated() || at::symint::size<T>(input, 0) >= 16 ||
        at::symint::size<T>(weight, -1) != 1 || at::symint::size<T>(weight, -2) != 1 || at::get_num_threads() > 1) &&
       (groups > 1
        || (at::symint::size<T>(weight, -1) > 3 && at::symint::size<T>(weight, -2) > 3)
        || at::symint::size<T>(input, 0) > 1
        || at::symint::size<T>(input, 0)*at::symint::size<T>(input, 1)*at::symint::size<T>(input, 2)*at::symint::size<T>(input, 3) > 20480) // for some case, native is faster
        );

#endif
    return false;
  }
  bool use_nnpack(const at::Tensor& input, const at::Tensor& weight) const  {
#if AT_NNPACK_ENABLED()
    return at::_nnpack_available() &&
           input.device().is_cpu() &&
           input.scalar_type() == kFloat && // only on CPU Float Tensors
           !is_dilated() && // or dilation
           !transposed &&   // or transposed tensors
           input.ndimension() == 4 && // must be in NCHW format
           weight.ndimension() == 4 &&
           (at::symint::size<T>(weight, 2) < 17) && (at::symint::size<T>(weight, 3) < 17) // NNPACK only supports kernels up to 16x16
#if !defined(C10_MOBILE)
           && at::symint::size<T>(input, 0) >= 16 // ensure large enough batch size to ensure perf, tuneable
#endif
       ;
#endif
    return false;
  }
  bool use_xnnpack(const at::Tensor& input, const at::Tensor& weight,
                   const at::OptionalArrayRef<T> bias_sizes_opt) const {
#if defined(C10_MOBILE)
    if (!transposed) {
      // NB: for the call here, it MATTERS that we are templated. If you
      // untemplate this to always use SymInt, the function
      // xnnpack_use_convolution2d will always return false
      return (at::symint::size<T>(input, 1) == groups) &&
              xnnpack_use_convolution2d(
                  input,
                  weight,
                  bias_sizes_opt,
                  padding,
                  stride,
                  dilation,
                  groups,
                  transposed);
    }
#endif
    return false;
  }

  bool use_mps(const at::Tensor& input, const at::Tensor& weight) const {
    // These checks need to be expanded. Currently we have very limited set of
    // checks for MPS.
#ifdef USE_MPS
    if (needs_64bit_indexing_no_split(input, weight)) {
      return false;
    }
    if (!input.is_mps()) {
      return false;
    }
    return true;
#else
    return false;
#endif
  }

  // We currently only have depthwise support for the case where groups ==
  // nInputPlane and nInputPlane == nOutputPlane (the latter due to the lack of
  // a depthwise multiplier)
  bool is_depthwise(const at::Tensor& input, const at::Tensor& weight) const  {
    return input.is_cuda() &&
           !transposed &&
           (input.ndimension() == 4 || input.ndimension() == 5) &&
           at::symint::size<T>(input, 1) == groups &&
           groups > 1 && // no point if there is only a single group
           at::symint::size<T>(weight, 0) % at::symint::size<T>(input, 1) == 0; // output channels must be a multiple of input channels
  }
};

DEFINE_DISPATCH(conv_depthwise2d_backward_stub);
DEFINE_DISPATCH(conv_depthwise3d_backward_stub);
DEFINE_DISPATCH(cudnn_convolution_backward_stub);
DEFINE_DISPATCH(cudnn_convolution_transpose_backward_stub);
DEFINE_DISPATCH(slow_conv_transpose3d_backward_stub);
DEFINE_DISPATCH(convolution_depthwise3x3_winograd_stub);
DEFINE_DISPATCH(miopen_convolution_backward_stub);
DEFINE_DISPATCH(miopen_convolution_transpose_backward_stub);
DEFINE_DISPATCH(miopen_depthwise_convolution_backward_stub);
DEFINE_DISPATCH(mkldnn_convolution_backward_stub);
DEFINE_DISPATCH(mkldnn_convolution_transpose_stub);
DEFINE_DISPATCH(mkldnn_convolution_transpose_backward_stub);
DEFINE_DISPATCH(slow_conv_dilated2d_backward_stub);
DEFINE_DISPATCH(slow_conv_dilated3d_backward_stub);
DEFINE_DISPATCH(slow_conv_transpose2d_backward_stub);
REGISTER_NO_CPU_DISPATCH(conv_depthwise2d_backward_stub);
REGISTER_NO_CPU_DISPATCH(conv_depthwise3d_backward_stub);
REGISTER_NO_CPU_DISPATCH(cudnn_convolution_backward_stub);
REGISTER_NO_CPU_DISPATCH(cudnn_convolution_transpose_backward_stub);
REGISTER_NO_CPU_DISPATCH(miopen_convolution_backward_stub);
REGISTER_NO_CPU_DISPATCH(miopen_convolution_transpose_backward_stub);
REGISTER_NO_CPU_DISPATCH(miopen_depthwise_convolution_backward_stub);

template <typename T>
std::ostream& operator<<(std::ostream & out, const ConvParams<T>& params) {
  out << "ConvParams {"
      << "  stride = " << IntArrayRef{params.stride}
      << "  padding = " << ArrayRef<T>{params.padding}
      << "  dilation = " << IntArrayRef{params.dilation}
      << "  transposed = " << params.transposed
      << "  output_padding = " << ArrayRef<T>{params.output_padding}
      << "  groups = " << params.groups
      << "  benchmark = " << params.benchmark
      << "  deterministic = " << params.deterministic
      << "  cudnn_enabled = " << params.cudnn_enabled
      << "  allow_tf32 = " << params.allow_tf32
      << "}";
  return out;
}

template <typename T>
static void check_shape_forward(const at::Tensor& input,
                                const c10::ArrayRef<T>& weight_sizes, const at::Tensor& bias,
                                const ConvParams<T>& params) {
  int64_t k = input.ndimension();
  int64_t weight_dim = weight_sizes.size();
  int64_t groups = params.groups;
  const auto& padding = params.padding;
  const auto& dilation = params.dilation;
  bool transposed = params.transposed;

  TORCH_CHECK(!params.is_padding_neg(), "negative padding is not supported");
  TORCH_CHECK(!params.is_output_padding_neg(), "negative output_padding is not supported");
  TORCH_CHECK(!params.is_stride_nonpos(), "non-positive stride is not supported");

  TORCH_CHECK(weight_dim == k,
           "Expected ", weight_dim, "-dimensional input for ", weight_dim,
           "-dimensional weight ", weight_sizes, ", but got ", k, "-dimensional input of size ",
           at::symint::sizes<T>(input), " instead");
  TORCH_CHECK(weight_sizes[0] >= groups,
           "Given groups=", groups, ", expected weight to be at least ", groups,
           " at dimension 0, but got weight of size ", weight_sizes, " instead");
  TORCH_CHECK(weight_sizes[0] % groups == 0,
           "Given groups=", groups, ", expected weight to be divisible by ",
           groups, " at dimension 0, but got weight of size [", weight_sizes,
           "] instead");

  if (!transposed) {
    std::vector<T> input_shape;
    std::vector<T> kernel_shape;
    bool kernel_size_correct = true;

    TORCH_CHECK(at::symint::size<T>(input, 1) == (weight_sizes[1] * groups),
                "Given groups=", groups, ", weight of size ", weight_sizes,
                ", expected input", at::symint::sizes<T>(input), " to have ",
                (weight_sizes[1] * groups), " channels, but got ", at::symint::size<T>(input, 1),
                " channels instead");

    TORCH_CHECK(!bias.defined() || (bias.ndimension() == 1 && at::symint::size<T>(bias, 0) == weight_sizes[0]),
             "Given weight of size ", weight_sizes,
             ", expected bias to be 1-dimensional with ", weight_sizes[0], " elements",
             ", but got bias of size ", at::symint::sizes<T>(bias), " instead");

    for (const auto i : c10::irange(2, k)) {
      input_shape.push_back(at::symint::size<T>(input, i) + 2 * padding[i-2]);
      // log new kernel size considering dilation
      kernel_shape.push_back(dilation[i-2] * (weight_sizes[i]-1) + 1);
      if (input_shape.back() < kernel_shape.back()) {
        kernel_size_correct = false;
      }
    }

    TORCH_CHECK(input_shape.size() == kernel_shape.size(), "Inconsistent shape between Input and Kernel");

    if (!kernel_size_correct) {
      // If kernel size is incorrect
      std::ostringstream input_ss;
      std::ostringstream kernel_ss;
      std::string separator = "";

      for (int i = 0, len = input_shape.size(); i < len; ++i) {
        input_ss << separator << input_shape[i];
        kernel_ss << separator << kernel_shape[i];
        separator = " x ";
      }

      AT_ERROR("Calculated padded input size per channel: (", input_ss.str(), "). "
               "Kernel size: (", kernel_ss.str(), "). Kernel size can't be greater than actual input size");
    }
  } else { // transposed
    TORCH_CHECK(at::symint::size<T>(input, 1) == weight_sizes[0],
             "Given transposed=", transposed, ", weight of size ", weight_sizes,
             ", expected input", at::symint::sizes<T>(input), " to have ", weight_sizes[0],
             " channels, but got ", at::symint::size<T>(input, 1), " channels instead");
    TORCH_CHECK(!bias.defined() || (bias.ndimension() == 1 && at::symint::size<T>(bias, 0) == weight_sizes[1] * groups),
             "Given transposed=", transposed, ", weight of size ", weight_sizes,
             ", expected bias to be 1-dimensional with ", weight_sizes[1] * groups, " elements",
             ", but got bias of size ", at::symint::sizes<T>(bias), " instead");
  }
}

template <typename T>
static void check_shape_backward(
    const at::Tensor& input,
    const c10::ArrayRef<T>& weight_sizes,
    const ConvParams<T>& params) {
  check_shape_forward<T>(input, weight_sizes, /*bias=*/ Tensor(), params);
}

// Given an input tensor and an expected number of spatial dimensions, checks that the
// input is a valid shape and returns the batched form of the input.
//
// Args:
//     input (Tensor): Input tensor
//     num_spatial_dims (int): Number of spatial dimensions expected for the input
//     func_name (string): Function name to produce a nice error message for invalid input
//
// Returns a std::tuple containing:
//     batched_input (Tensor): Input with a batch dimension
//     is_batched (bool): Indicates whether the original input was already batched
static std::tuple<Tensor, bool> batchify(
    const Tensor& input,
    const int64_t num_spatial_dims,
    const std::string& func_name) {
  const auto dim_count_no_batch = num_spatial_dims + 1;
  const auto dim_count_batch = dim_count_no_batch + 1;
  const auto is_batched = (input.dim() == dim_count_batch);
  TORCH_CHECK(input.dim() == dim_count_no_batch || is_batched,
      "Expected ", dim_count_no_batch, "D (unbatched) or ", dim_count_batch,
      "D (batched) input to ", func_name, ", but got input of size: ", input.sizes());
  return std::make_tuple(is_batched ? input : input.unsqueeze(0), is_batched);
}

static void check_input_same_type_as_parameters(
    const Tensor& input,
    const Tensor& weight,
    const Tensor& bias) {
  TORCH_CHECK(input.options().type_equal(weight.options()),
      "Input type (", input.toString(), ") and weight type (", weight.toString(),
      ") should be the same");
  TORCH_CHECK(!bias.defined() || (input.options().type_equal(bias.options())),
      "Input type (", input.toString(), ") and bias type (", bias.toString(),
      ") should be the same");
}

static void check_input_same_type_as_parameters(
    const Tensor& input,
    const Tensor& weight) {
  check_input_same_type_as_parameters(input, weight, /*bias=*/ Tensor());
}

static void check_input_same_type_as_parameters(
    const Tensor& input,
    const Tensor& weight,
    const Tensor& bias,
    const ConvBackend backend) {
  if (backend == ConvBackend::Mkldnn || backend == ConvBackend::MkldnnTranspose) {
    TORCH_CHECK(input.options().type_equal(weight.options())
        || (input.is_mkldnn() && weight.device().is_cpu() && weight.scalar_type() == kFloat),
        "Input type (", input.toString(), ") and weight type (", weight.toString(),
        ") should be the same or input should be a MKLDNN tensor and weight is a dense tensor");
    TORCH_CHECK(!bias.defined() || (input.options().type_equal(bias.options()))
        || (input.is_mkldnn() && bias.device().is_cpu() && bias.scalar_type() == kFloat),
        "Input type (", input.toString(), ") and bias type (", bias.toString(),
        ") should be the same or input should be a MKLDNN tensor and bias is a dense tensor");
  } else {
    check_input_same_type_as_parameters(input, weight, bias);
  }
}

static auto view4d(const at::Tensor& tensor) -> at::Tensor {
  TORCH_CHECK(tensor.ndimension() == 3,
           "expected 3D tensor, got tensor with ", tensor.ndimension(),
           " dimensions instead");
  return tensor.unsqueeze(2);
}

static auto view3d(const at::Tensor& tensor) -> at::Tensor {
  TORCH_CHECK(tensor.ndimension() == 4,
           "expected 4D tensor, got tensor with ", tensor.ndimension(),
           " dimensions instead");
  return tensor.squeeze(2);
}

static at::Tensor subtensor(at::Tensor& tensor, int dim, int groups, int g) {
  if (!tensor.defined()) {
    return at::Tensor();
  }
  const auto memory_format = tensor.suggest_memory_format();
  int64_t n = tensor.sizes()[dim] / groups;
  return tensor.narrow(dim, n * g, n).contiguous(memory_format);
}

namespace {

std::pair<Tensor, Tensor> complex_to_real(const Tensor& inp) {
  auto inp_view_as_complex = at::view_as_real(inp);
  auto dim_i = inp_view_as_complex.dim() - 1;
  auto i_r = inp_view_as_complex.select(dim_i, 0);
  auto i_i = inp_view_as_complex.select(dim_i, 1);
  return std::make_pair(i_r, i_i);
}

at::Tensor complex_convolution(
    const Tensor& input,
    const Tensor& weight,
    const Tensor& bias,
    IntArrayRef stride,
    IntArrayRef padding,
    IntArrayRef dilation,
    bool transposed,
    IntArrayRef output_padding,
    int64_t groups) {
  check_input_same_type_as_parameters(input, weight, bias);
  Tensor i_r, i_i, w_r, w_i;
  std::tie(i_r, i_i) = complex_to_real(input.resolve_conj());
  std::tie(w_r, w_i) = complex_to_real(weight.resolve_conj());

  // [NOTE] Complex Convolution
  // conv(W, x, b) = conv(Wr, xr, br) - conv(Wi, xi, 0) + i(conv(Wi, xr, bi) + conv(Wr, xi, 0))
  // where W, x and b are all complex inputs.
  // With Gauss Trick:
  // a = conv(Wr, xr, br),
  // b = conv(Wi, xi, 0),
  // c = conv(Wr + Wi, xr + xi, bi + br)
  // conv(W, x, b) = a - b + i(c - a - b)
  Tensor a, b, c;
  if (!bias.defined()) {
    a = at::convolution(i_r, w_r, bias, stride, padding, dilation, transposed, output_padding, groups);
    b = at::convolution(i_i, w_i, bias, stride, padding, dilation, transposed, output_padding, groups);
    c = at::convolution(i_r + i_i, w_r + w_i, bias, stride, padding, dilation, transposed, output_padding, groups);
  } else {
    Tensor b_r, b_i;
    std::tie(b_r, b_i) = complex_to_real(bias.resolve_conj());
    a = at::convolution(i_r, w_r, b_r, stride, padding, dilation, transposed, output_padding, groups);
    b = at::convolution(i_i, w_i, Tensor(), stride, padding, dilation, transposed, output_padding, groups);
    c = at::convolution(i_r + i_i, w_r + w_i, b_r + b_i, stride, padding, dilation, transposed, output_padding, groups);
  }

  auto i = c10::Scalar(c10::complex<double>(0, 1));
  return a - b + i * (c - a - b);
}

at::Tensor complex_convolution_mode(
    const at::Tensor& input,
    const at::Tensor& weight,
    const c10::optional<at::Tensor>& bias_opt,
    at::IntArrayRef stride,
    c10::string_view padding,
    at::IntArrayRef dilation,
    int64_t groups) {
  auto bias = bias_opt.value_or(Tensor());
  check_input_same_type_as_parameters(input, weight, bias);
  Tensor i_r, i_i, w_r, w_i;
  std::tie(i_r, i_i) = complex_to_real(input.resolve_conj());
  std::tie(w_r, w_i) = complex_to_real(weight.resolve_conj());

  // See [NOTE] Complex Convolution
  Tensor a, b, c;
  if (!bias.defined()) {
    a = at::_convolution_mode(i_r, w_r, bias, stride, padding, dilation, groups);
    b = at::_convolution_mode(i_i, w_i, bias, stride, padding, dilation, groups);
    c = at::_convolution_mode(i_r + i_i, w_r + w_i, bias, stride, padding, dilation, groups);
  } else {
    Tensor b_r, b_i;
    std::tie(b_r, b_i) = complex_to_real(bias.resolve_conj());
    a = at::_convolution_mode(i_r, w_r, b_r, stride, padding, dilation, groups);
    b = at::_convolution_mode(i_i, w_i, Tensor(), stride, padding, dilation, groups);
    c = at::_convolution_mode(i_r + i_i, w_r + w_i, b_r + b_i, stride, padding, dilation, groups);
  }

  auto i = c10::Scalar(c10::complex<double>(0, 1));
  return a - b + i * (c - a - b);
}

} // namespace

at::Tensor conv1d(
    const Tensor& input_, const Tensor& weight, const c10::optional<Tensor>& bias_opt,
    IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation, int64_t groups) {
  // See [Note: hacky wrapper removal for optional tensor]
  c10::MaybeOwned<Tensor> bias_maybe_owned = at::borrow_from_optional_tensor(bias_opt);
  const Tensor& bias = *bias_maybe_owned;

  Tensor input;
  bool is_batched;
  std::tie(input, is_batched) = batchify(input_, /*num_spatial_dims=*/ 1, "conv1d");
  Tensor output;
  if (at::isComplexType(input_.scalar_type())) {
    output = complex_convolution(input, weight, bias, stride, padding, dilation, false, {0}, groups);
  } else {
    output = at::convolution(input, weight, bias, stride, padding, dilation, false, {0}, groups);
  }
  return is_batched ? std::move(output) : output.squeeze(0);
}

at::Tensor conv2d(
    const Tensor& input_, const Tensor& weight, const c10::optional<Tensor>& bias_opt,
    IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation, int64_t groups) {
  // See [Note: hacky wrapper removal for optional tensor]
  c10::MaybeOwned<Tensor> bias_maybe_owned = at::borrow_from_optional_tensor(bias_opt);
  const Tensor& bias = *bias_maybe_owned;

  TORCH_CHECK(
    !bias.defined() || bias.dtype() == input_.dtype(),
    "Input type (",
    input_.dtype().name(),
    ") and bias type (",
    bias.dtype().name(),
    ") should be the same");

  Tensor input;
  bool is_batched;
  std::tie(input, is_batched) = batchify(input_, /*num_spatial_dims=*/ 2, "conv2d");
  Tensor output;
  if (at::isComplexType(input_.scalar_type())) {
    output = complex_convolution(input, weight, bias, stride, padding, dilation, false, {{0, 0}}, groups);
  } else {
    output = at::convolution(input, weight, bias, stride, padding, dilation, false, {{0, 0}}, groups);
  }
  return is_batched ? std::move(output) : output.squeeze(0);
}

at::Tensor conv3d(
    const Tensor& input_, const Tensor& weight, const c10::optional<Tensor>& bias_opt,
    IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation, int64_t groups) {
  // See [Note: hacky wrapper removal for optional tensor]
  c10::MaybeOwned<Tensor> bias_maybe_owned = at::borrow_from_optional_tensor(bias_opt);
  const Tensor& bias = *bias_maybe_owned;

  Tensor input;
  bool is_batched;
  std::tie(input, is_batched) = batchify(input_, /*num_spatial_dims=*/ 3, "conv3d");
  Tensor output;
  if (at::isComplexType(input_.scalar_type())) {
    output = complex_convolution(input, weight, bias, stride, padding, dilation, false, {{0, 0, 0}}, groups);
  } else {
    output = at::convolution(input, weight, bias, stride, padding, dilation, false, {{0, 0, 0}}, groups);
  }
  return is_batched ? std::move(output) : output.squeeze(0);
}


static Tensor convolution_same(
    const Tensor &input, const Tensor &weight, const Tensor &bias,
    IntArrayRef stride, IntArrayRef dilation, int64_t groups) {

  auto k = weight.dim();
  TORCH_CHECK(k > 2, "weight should have at least three dimensions");
  auto dim = static_cast<size_t>(k - 2);
  auto weight_sizes = weight.sym_sizes();
  auto input_sizes = input.sym_sizes();
  TORCH_CHECK(k == input.dim(),
              "Expected ", k, "-dimensional input for ",
              k, "-dimensional weight", weight_sizes, ", but got ",
              input.dim(), "-dimensional input of size ",
              input.sizes(), " instead");
  TORCH_CHECK(stride.size() == dim || stride.size() == 1U,
              "stride cannot broadcast to ", dim, " dimensions");
  TORCH_CHECK(dilation.size() == dim || dilation.size() == 1U,
              "dilation cannot broadcast to ", dim, " dimensions");
  for (auto i: c10::irange(stride.size())) {
    TORCH_CHECK(stride[i] == 1, "padding='same' is not supported for strided convolutions");
  }

  // Calculate the correct padding
  SymDimVector padding_l, padding_r;
  bool symmetric_padding = true;
  for (auto i: c10::irange(dim)) {
    auto s = stride.size() == 1 ? stride[0] : stride[i];
    auto d = dilation.size() == 1 ? dilation[0] : dilation[i];
    auto pad = pooling_same_mode_padding_lr(
        input_sizes[i + 2], weight_sizes[i + 2], s, d);
    padding_l.push_back(pad.first);
    padding_r.push_back(pad.second);
    if (pad.first != pad.second) {
      symmetric_padding = false;
    }
  }

  if (symmetric_padding) {
    // All backends handle symmetric padding natively
    SymDimVector output_padding(static_cast<size_t>(dim));
    return at::convolution_symint(input, weight, bias, stride, padding_l, dilation,
                               false, output_padding, groups);
  }

  TORCH_WARN_ONCE("Using padding='same' with even kernel lengths and odd dilation may"
                  " require a zero-padded copy of the input be created");
  SmallVector<c10::SymInt, kDimVectorStaticSize * 2> pad_nd(static_cast<size_t>(2 * dim));
  for (auto i: c10::irange(dim)) {
    // Apply padding by the difference, leaving only a symmetric padding
    auto delta_pad = padding_r[i] - padding_l[i];
    auto pad_idx = 2 * (dim - 1 - i);  // F.pad goes from last dim to first
    if (delta_pad > 0) {
      pad_nd[pad_idx + 1] = delta_pad;
    } else {
      pad_nd[pad_idx] = delta_pad;
      padding_l[i] = padding_r[i];
    }
  }
  auto padded_input = at::constant_pad_nd_symint(input, pad_nd, 0);
  SymDimVector output_padding(static_cast<size_t>(dim));
  return at::convolution_symint(padded_input, weight, bias, stride, padding_l,
                                dilation, false, output_padding, groups);
}

Tensor _convolution_mode(
    const Tensor& input, const Tensor& weight, const c10::optional<Tensor>& bias_opt,
    IntArrayRef stride, c10::string_view padding, IntArrayRef dilation,
    int64_t groups) {
  // See [Note: hacky wrapper removal for optional tensor]
  c10::MaybeOwned<Tensor> bias_maybe_owned = at::borrow_from_optional_tensor(bias_opt);
  const Tensor& bias = *bias_maybe_owned;

  if (padding == "same") {
    return at::native::convolution_same(
        input, weight, bias, stride, dilation, groups);
  } else if (padding == "valid") {
    // NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
    const int64_t padding_[] = {0};
    return at::convolution(
        input, weight, bias, stride, padding_, dilation, false, padding_, groups);
  }
  TORCH_CHECK(false, "Invalid padding string: '", padding, "'");
}

at::Tensor conv1d(
    const Tensor& input_, const Tensor& weight, const c10::optional<Tensor>& bias,
    IntArrayRef stride, c10::string_view padding, IntArrayRef dilation,
    int64_t groups) {
  Tensor input;
  bool is_batched;
  std::tie(input, is_batched) = batchify(input_, /*num_spatial_dims=*/ 1, "conv1d");
  Tensor output;
  if (at::isComplexType(input_.scalar_type())) {
    output = complex_convolution_mode(input, weight, bias, stride, std::move(padding), dilation, groups);
  } else {
    output = at::_convolution_mode(input, weight, bias, stride, std::move(padding), dilation, groups);
  }
  return is_batched ? std::move(output) : output.squeeze(0);
}

at::Tensor conv2d(
    const Tensor& input_, const Tensor& weight, const c10::optional<Tensor>& bias,
    IntArrayRef stride, c10::string_view padding, IntArrayRef dilation,
    int64_t groups) {
  Tensor input;
  bool is_batched;
  std::tie(input, is_batched) = batchify(input_, /*num_spatial_dims=*/ 2, "conv2d");
  Tensor output;
  if (at::isComplexType(input_.scalar_type())) {
    output = complex_convolution_mode(input, weight, bias, stride, std::move(padding), dilation, groups);
  } else {
    output = at::_convolution_mode(input, weight, bias, stride, std::move(padding), dilation, groups);
  }
  return is_batched ? std::move(output) : output.squeeze(0);
}

at::Tensor conv3d(
    const Tensor& input_, const Tensor& weight, const c10::optional<Tensor>& bias,
    IntArrayRef stride, c10::string_view padding, IntArrayRef dilation,
    int64_t groups) {
  Tensor input;
  bool is_batched;
  std::tie(input, is_batched) = batchify(input_, /*num_spatial_dims=*/ 3, "conv3d");
  Tensor output;
  if (at::isComplexType(input_.scalar_type())) {
    output = complex_convolution_mode(input, weight, bias, stride, std::move(padding), dilation, groups);
  } else {
    output = at::_convolution_mode(input, weight, bias, stride, std::move(padding), dilation, groups);
  }
  return is_batched ? std::move(output) : output.squeeze(0);
}

at::Tensor conv_transpose1d(
    const Tensor& input_, const Tensor& weight, const c10::optional<Tensor>& bias_opt,
    IntArrayRef stride, IntArrayRef padding, IntArrayRef output_padding, int64_t groups, IntArrayRef dilation) {
  // See [Note: hacky wrapper removal for optional tensor]
  c10::MaybeOwned<Tensor> bias_maybe_owned = at::borrow_from_optional_tensor(bias_opt);
  const Tensor& bias = *bias_maybe_owned;

  Tensor input;
  bool is_batched;
  std::tie(input, is_batched) = batchify(input_, /*num_spatial_dims=*/ 1, "conv_transpose1d");
  Tensor output;
  if (at::isComplexType(input_.scalar_type())) {
    output = complex_convolution(
      input, weight, bias, stride, padding, dilation, true, output_padding, groups);
  } else {
    output = at::convolution(
      input, weight, bias, stride, padding, dilation, true, output_padding, groups);
  }
  return is_batched ? std::move(output) : output.squeeze(0);
}

at::Tensor conv_transpose2d(
    const Tensor& input_, const Tensor& weight, const c10::optional<Tensor>& bias_opt,
    IntArrayRef stride, IntArrayRef padding, IntArrayRef output_padding, int64_t groups, IntArrayRef dilation) {
  // See [Note: hacky wrapper removal for optional tensor]
  c10::MaybeOwned<Tensor> bias_maybe_owned = at::borrow_from_optional_tensor(bias_opt);
  const Tensor& bias = *bias_maybe_owned;

  Tensor input;
  bool is_batched;
  std::tie(input, is_batched) = batchify(input_, /*num_spatial_dims=*/ 2, "conv_transpose2d");
  Tensor output;
  if (at::isComplexType(input_.scalar_type())) {
    output = complex_convolution(
      input, weight, bias, stride, padding, dilation, true, output_padding, groups);
  } else {
    output = at::convolution(
      input, weight, bias, stride, padding, dilation, true, output_padding, groups);
  }
  return is_batched ? std::move(output) : output.squeeze(0);
}

at::Tensor conv_transpose3d(
    const Tensor& input_, const Tensor& weight, const c10::optional<Tensor>& bias_opt,
    IntArrayRef stride, IntArrayRef padding, IntArrayRef output_padding, int64_t groups, IntArrayRef dilation) {
  // See [Note: hacky wrapper removal for optional tensor]
  c10::MaybeOwned<Tensor> bias_maybe_owned = at::borrow_from_optional_tensor(bias_opt);
  const Tensor& bias = *bias_maybe_owned;

  Tensor input;
  bool is_batched;
  std::tie(input, is_batched) = batchify(input_, /*num_spatial_dims=*/ 3, "conv_transpose3d");
  Tensor output;
  if (at::isComplexType(input_.scalar_type())) {
    output = complex_convolution(
      input, weight, bias, stride, padding, dilation, true, output_padding, groups);
  } else {
    output = at::convolution(
      input, weight, bias, stride, padding, dilation, true, output_padding, groups);
  }
  return is_batched ? std::move(output) : output.squeeze(0);
}

at::Tensor convolution(
    const Tensor& input, const Tensor& weight, const c10::optional<Tensor>& bias_opt,
    IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation,
    bool transposed, IntArrayRef output_padding, int64_t groups) {
  // See [Note: hacky wrapper removal for optional tensor]
  c10::MaybeOwned<Tensor> bias_maybe_owned = at::borrow_from_optional_tensor(bias_opt);
  const Tensor& bias = *bias_maybe_owned;

  auto& ctx = at::globalContext();
  // See Note [Enabling Deterministic Operations]
  bool deterministic = ctx.deterministicCuDNN() || ctx.deterministicAlgorithms();
  return at::_convolution(input, weight, bias, stride, padding, dilation,
                          transposed, output_padding, groups,
                          ctx.benchmarkCuDNN(), deterministic, ctx.userEnabledCuDNN(), ctx.allowTF32CuDNN());
}

at::Tensor convolution_overrideable(
    const Tensor& input, const Tensor& weight, const c10::optional<Tensor>& bias_opt,
    IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation,
    bool transposed, IntArrayRef output_padding, int64_t groups) {
  TORCH_CHECK_NOT_IMPLEMENTED(false, "convolution_overrideable not implemented. You are likely triggering this with tensor backend other than CPU/CUDA/MKLDNN, if this is intended, please use TORCH_LIBRARY_IMPL to override this function ");
}

// Function to select the convolution backend based on the inputs and params.
// This overload is used within the convolution internals but not exposed to python.
// NB: The forward pass provides a bias tensor while the backward pass provides
// a bool indicating whether the bias is defined. This is done to save memory by
// avoiding saving the full bias tensor for backward.
template <typename T>
ConvBackend _select_conv_backend(
    const Tensor& input,
    const Tensor& weight,
    const c10::optional<Tensor>& bias,
    const at::OptionalArrayRef<T> bias_sizes_opt,
    const bool need_backward,
    const ConvParams<T>& params) {

  // don't send empty inputs through backends
  if (at::symint::size<T>(input, 0) == 0 || at::symint::size<T>(input, 1) == 0) {
    return input.is_mkldnn() ? ConvBackend::MkldnnEmpty : ConvBackend::Empty;
  } else if (at::symint::numel<T>(input) == 0) {
    TORCH_CHECK(false, "Only zero batch or zero channel inputs are supported, but got input shape: ", at::symint::sizes<T>(input));
  }

  if (params.is_depthwise(input, weight)) {
    if (params.use_cudnn_depthwise(input, weight)) {
      return ConvBackend::Cudnn;
    } else if (params.use_miopen(input, weight, bias_sizes_opt.has_value())) {
      return ConvBackend::MiopenDepthwise;
    } else {
      if (input.ndimension() == 4) {
        return ConvBackend::CudaDepthwise2d;
      } else if (input.ndimension() == 5) {
        return ConvBackend::CudaDepthwise3d;
      } else {
        // unsupported
      }
    }
  } else if (params.use_cudnn(input, weight)) {
    if (params.transposed) {
      return ConvBackend::CudnnTranspose;
    } else {
      return ConvBackend::Cudnn;
    }
  } else if (params.use_miopen(input, weight, bias_sizes_opt.has_value())) {
    if (params.transposed) {
      return ConvBackend::MiopenTranspose;
    } else {
      return ConvBackend::Miopen;
    }
  } else if (params.use_mkldnn(input, weight)) {
    if (params.transposed) {
      return ConvBackend::MkldnnTranspose;
    } else {
      return ConvBackend::Mkldnn;
    }
  } else if (!need_backward && params.use_xnnpack(input, weight, bias_sizes_opt)) {
    // Using prepacked conv is preferred, but XNNPACK is still the fastest
    // option for NHWC.
    return ConvBackend::Xnnpack2d;
  // 3x3 depthwith convolutions implementation is inference only
  } else if (!need_backward && params.use_cpu_depthwise3x3_winograd(input, weight, bias)) {
    return ConvBackend::Winograd3x3Depthwise;
  } else if (
      !params.transposed && (input.ndimension() == 5) &&
      (input.device().is_cpu()) &&
      !params.is_dilated()) {
    // fast path for grouped conv3d
    return ConvBackend::Slow3d;
  } else if (input.device().is_cpu() || input.is_cuda()) {
    // backends without support for groups
    if (params.transposed) {
      if (input.ndimension() == 4) {
        return ConvBackend::SlowTranspose2d;
      } else if (input.ndimension() == 5) {
        return ConvBackend::SlowTranspose3d;
      } else {
        // unsupported
      }
    } else {  /* Not transposed */
      if (input.ndimension() == 4) {
        if (params.is_dilated()) {
          return ConvBackend::SlowDilated2d;
        } else {  /* dim == 4, non-dilated */
          if (params.use_nnpack(input, weight)) {
            return ConvBackend::NnpackSpatial;
          } else {
            /* CPU implementation has specialized MM kernels
               for non-dilated case here */
            return ConvBackend::Slow2d;
          }
        }
      } else if (input.ndimension() == 5 && (input.is_cuda() || params.is_dilated())) {
        return ConvBackend::SlowDilated3d;
      } else if (input.ndimension() == 5) { /* dim == 5, CPU, non-dilated */
        /* CPU implementation has specialized MM kernels
           for non-dilated case here */
        return ConvBackend::Slow3d;
      } else {
        // unsupported
      }
    }
  } else if (params.use_mps(input, weight)) {
    if (params.transposed) {
      return ConvBackend::MpsTranspose;
    } else {
      return ConvBackend::Mps;
    }
  } else {
    // Only reach here when input is backend with out-of-source implementation.
    return ConvBackend::Overrideable;
  }

  // Error out if no suitable backend was found.
  AT_ERROR("unsupported ConvNd parameters");
}

// Selects a backend for convolution based on the inputs and params.
ConvBackend select_conv_backend(
    const Tensor& input_r, const Tensor& weight_r, const c10::optional<Tensor>& bias_opt,
    IntArrayRef stride_, SymIntArrayRef padding_, IntArrayRef dilation_,
    bool transposed_, SymIntArrayRef output_padding_, int64_t groups_, const at::OptionalSymIntArrayRef bias_sizes_opt) {
  c10::MaybeOwned<Tensor> bias_maybe_owned = at::borrow_from_optional_tensor(bias_opt);
  const Tensor& bias = *bias_maybe_owned;

  auto& ctx = at::globalContext();
  auto k = weight_r.ndimension();
  int64_t dim = k - 2;
  ConvParams<c10::SymInt> params;
  params.stride = expand_param_if_needed(stride_, "stride", dim);
  params.padding = expand_param_if_needed(padding_, "padding", dim);
  params.dilation = expand_param_if_needed(dilation_, "dilation", dim);
  params.transposed = transposed_;
  params.output_padding = expand_param_if_needed(output_padding_, "output_padding", dim);
  params.groups = groups_;
  params.benchmark = ctx.benchmarkCuDNN();
  params.deterministic = ctx.deterministicCuDNN() || ctx.deterministicAlgorithms();
  params.cudnn_enabled = ctx.userEnabledCuDNN();
  params.allow_tf32 = ctx.allowTF32CuDNN();

  auto input = input_r;
  auto weight = weight_r;
  check_shape_forward(input, weight.sym_sizes(), bias, params);

  // Expand 1d -> 2d.
  // This is only done for backends that don't natively support 1d spatial input.
  if (k == 3 && !input.is_mkldnn() && !input.is_xpu()) {
    // avoid accidentally going through NHWC for permuted 3d input.
    input = input.contiguous();
    params.view1d_as_2d();
    input = view4d(input);
    weight = view4d(weight);
  }

  auto bias_sizes = bias.defined() ? c10::optional<SymIntArrayRef>(bias.sym_sizes()) : bias_sizes_opt;
  bool need_backward = GradMode::is_enabled() &&
      (input.requires_grad() || weight.requires_grad() || (bias.defined() && bias.requires_grad()));
  return _select_conv_backend(input, weight, bias, bias_sizes, need_backward, params);
}

// For BC reasons, have a copy that does not require bias_opt
static ConvBackend select_conv_backend(
    const Tensor& input,
    const Tensor& weight,
    const at::OptionalIntArrayRef bias_sizes_opt,
    const bool need_backward,
    const ConvParams<int64_t>& params) {
  return _select_conv_backend(input, weight, {}, bias_sizes_opt, need_backward, params);
}

static at::Tensor _convolution_nogroup_backend(
    const Tensor& input,
    const Tensor& weight,
    const Tensor& bias,
    const ConvBackend backend,
    const ConvParams<int64_t>& params) {
  auto kernel_size = weight.sizes().slice(2);
  switch(backend) {
    case ConvBackend::NnpackSpatial:
#if AT_NNPACK_ENABLED()
      return at::_nnpack_spatial_convolution(input, weight, bias, params.padding, params.stride);
#else
      TORCH_INTERNAL_ASSERT(false, "NnpackSpatial backend was selected in PyTorch compiled without nnpack support");
#endif
    case ConvBackend::Slow2d:
      return at::thnn_conv2d(input, weight, kernel_size, bias, params.stride, params.padding);
    case ConvBackend::SlowDilated2d:
      return at::slow_conv_dilated2d(
          input, weight, kernel_size, bias, params.stride, params.padding, params.dilation);
    case ConvBackend::SlowDilated3d:
      return at::slow_conv_dilated3d(
          input, weight, kernel_size, bias, params.stride, params.padding, params.dilation);
    case ConvBackend::SlowTranspose2d:
      return at::slow_conv_transpose2d(
          input, weight, kernel_size, bias, params.stride, params.padding, params.output_padding, params.dilation);
    case ConvBackend::SlowTranspose3d:
      return at::slow_conv_transpose3d(
          input, weight, kernel_size, bias, params.stride, params.padding, params.output_padding, params.dilation);
    default:
      TORCH_CHECK(false, "Unsupported conv nogroup backend encountered");
  }
}

static inline std::vector<int64_t> calc_output_size(
    const Tensor& input,
    const Tensor& weight,
    const ConvParams<int64_t>& params) {
  std::vector<int64_t> output_size = params.transposed ?
    conv_input_size(input.sizes(), weight.sizes(), params.padding, params.output_padding,
        params.stride, params.dilation, params.groups) :
    conv_output_size(input.sizes(), weight.sizes(), params.padding, params.stride, params.dilation);

  // Handle empty # of channels.
  if (input.size(1) == 0) {
    output_size[input_channels_dim] = 0;
  }
  return output_size;
}

static inline at::MemoryFormat determine_backend_memory_format(
    const Tensor& input,
    const Tensor& weight,
    const ConvBackend backend) {
  at::MemoryFormat backend_memory_format = at::MemoryFormat::Contiguous;
  auto k = weight.ndimension();
#if !defined(C10_MOBILE)
  // See Note [Mobile check segfaults]
  switch(backend) {
    case ConvBackend::Cudnn:
    case ConvBackend::CudnnTranspose:
      if (detail::getCUDAHooks().compiledWithCuDNN()) {
        backend_memory_format = cudnn_conv_suggest_memory_format(input, weight);
      }
      break;
    case ConvBackend::Miopen:
    case ConvBackend::MiopenDepthwise:
    case ConvBackend::MiopenTranspose:
      if (detail::getCUDAHooks().compiledWithMIOpen() && miopen_conv_use_channels_last(input, weight)) {
        TORCH_INTERNAL_ASSERT((k == 4 || k == 5),
            "Expected 4D or 5D input for miopen memory format selection in determine_backend_memory_format()");
        backend_memory_format = (k == 5) ? at::MemoryFormat::Contiguous /*at::MemoryFormat::ChannelsLast3d*/ : at::MemoryFormat::ChannelsLast;
      }
      break;
    case ConvBackend::Mkldnn:
    case ConvBackend::MkldnnTranspose:
      if (mkldnn_conv_use_channels_last(input, weight)) {
        backend_memory_format = (k == 5) ? at::MemoryFormat::ChannelsLast3d : at::MemoryFormat::ChannelsLast;
      }
      break;
    case ConvBackend::Slow2d:
    case ConvBackend::SlowDilated2d:
    case ConvBackend::SlowTranspose2d:
      if (thnn_conv_use_channels_last(input, weight)) {
        backend_memory_format = at::MemoryFormat::ChannelsLast;
      }
      break;
    default:
      backend_memory_format = at::MemoryFormat::Contiguous;
  }
#endif
  return backend_memory_format;
}

at::MemoryFormat _determine_backend_memory_format(
    const Tensor& input,
    const Tensor& weight,
    const ConvBackend backend)  {
  return determine_backend_memory_format(input, weight, backend);
}

at::Tensor _convolution(
    const Tensor& input_r, const Tensor& weight_r, const c10::optional<Tensor>& bias_r_opt,
    IntArrayRef stride_, IntArrayRef padding_, IntArrayRef dilation_,
    bool transposed_, IntArrayRef output_padding_, int64_t groups_,
    bool benchmark, bool deterministic, bool cudnn_enabled, bool allow_tf32) {
  // See [Note: hacky wrapper removal for optional tensor]
  c10::MaybeOwned<Tensor> bias_r_maybe_owned = at::borrow_from_optional_tensor(bias_r_opt);
  const Tensor& bias_r = *bias_r_maybe_owned;

  auto input = input_r;
  auto weight = weight_r;
  auto bias = bias_r;
  auto k = weight.ndimension();
  c10::IntArrayRef weight_sizes = weight.sizes();
  int64_t dim = k - 2;

  TORCH_CHECK(dim > 0, "weight should have at least three dimensions");
  TORCH_CHECK(groups_ > 0, "non-positive groups is not supported");

  ConvParams<int64_t> params;
  params.stride = expand_param_if_needed(stride_, "stride", dim);
  params.padding = expand_param_if_needed(padding_, "padding", dim);
  params.dilation = expand_param_if_needed(dilation_, "dilation", dim);
  params.transposed = transposed_;
  params.output_padding = expand_param_if_needed(output_padding_, "output_padding", dim);
  params.groups = groups_;
  params.benchmark = benchmark;
  params.deterministic = deterministic;
  params.cudnn_enabled = cudnn_enabled;
  params.allow_tf32 = allow_tf32;

  check_shape_forward(input, weight_sizes, bias, params);

  // Expand 1d -> 2d.
  // This is only done for backends that don't natively support 1d spatial input.
  if (k == 3 && !input.is_mkldnn() && !input.is_xpu()) {
    // avoid accidentally going through NHWC for permuted 3d input.
    input = input.contiguous();
    params.view1d_as_2d();
    input = view4d(input);
    weight = view4d(weight);
  }

  // Select appropriate backend to use.
  auto bias_sizes_opt = bias.defined() ? c10::optional<IntArrayRef>(bias.sizes()) : c10::nullopt;
  bool need_backward = GradMode::is_enabled() &&
      (input.requires_grad() || weight.requires_grad() || (bias.defined() && bias.requires_grad()));
  ConvBackend backend = _select_conv_backend(input, weight, bias, c10::OptionalIntArrayRef(bias_sizes_opt), need_backward, params);
  at::MemoryFormat backend_memory_format = determine_backend_memory_format(input, weight, backend);

  // Call the backend.
  Tensor output;
  auto kernel_size = weight.sizes().slice(2);
  switch (backend) {
    case ConvBackend::CudaDepthwise2d:
      output = at::_conv_depthwise2d(input.contiguous(), weight, kernel_size, bias,
          params.stride, params.padding, params.dilation);
      break;
    case ConvBackend::CudaDepthwise3d:
      output = at::conv_depthwise3d(input.contiguous(), weight, kernel_size, bias,
          params.stride, params.padding, params.dilation);
      break;
    case ConvBackend::Cudnn:
      check_input_same_type_as_parameters(input, weight, bias);
      output = at::cudnn_convolution(
          input.contiguous(backend_memory_format), weight, params.padding, params.stride,
          params.dilation, params.groups, params.benchmark, params.deterministic, params.allow_tf32);
      if (bias.defined()) {
        output.add_(reshape_bias(input.dim(), bias));
      }
      break;
    case ConvBackend::CudnnTranspose:
      check_input_same_type_as_parameters(input, weight, bias);
      output = at::cudnn_convolution_transpose(
          input.contiguous(backend_memory_format), weight, params.padding, params.output_padding,
          params.stride, params.dilation, params.groups, params.benchmark, params.deterministic, params.allow_tf32);
      if (bias.defined()) {
        output.add_(reshape_bias(input.dim(), bias));
      }
      break;
    case ConvBackend::Empty:
    {
      Tensor weight_view;
      // Use permute and clone to avoid at::_unsafe_view(weight, -1) failure for non-contiguous cases where
      // view size is not compatible with input tensor's size and stride.
      if(weight.is_contiguous()) {
        weight_view = at::_unsafe_view(weight, -1);
      } else if (weight.is_contiguous(at::MemoryFormat::ChannelsLast)) {
        weight_view = at::_unsafe_view(at::permute(weight, {0, 2, 3, 1}), -1);
      } else if (weight.is_contiguous(at::MemoryFormat::ChannelsLast3d)) {
        weight_view = at::_unsafe_view(at::permute(weight, {0, 2, 3, 4, 1}), -1);
      } else {
        weight_view = at::_unsafe_view(weight.clone(at::MemoryFormat::Contiguous), -1);
      }

      output = (input.size(1) == 0) ? (input.view(-1) * weight_view) : (input * weight_view[0]);
      if (bias.defined()) {
        output.add_(bias[0]);
      }
      output = output.view(calc_output_size(input, weight, params));
      break;
    }
    case ConvBackend::Miopen:
      check_input_same_type_as_parameters(input, weight, bias);
      output = at::miopen_convolution(
          input.contiguous(backend_memory_format), weight, bias, params.padding, params.stride,
          params.dilation, params.groups, params.benchmark, params.deterministic);
      break;
    case ConvBackend::MiopenDepthwise:
      output = at::miopen_depthwise_convolution(
          input.contiguous(backend_memory_format), weight, bias, params.padding, params.stride,
          params.dilation, params.groups, params.benchmark, params.deterministic);
      break;
    case ConvBackend::MiopenTranspose:
      check_input_same_type_as_parameters(input, weight, bias);
      output = at::miopen_convolution_transpose(
          input.contiguous(backend_memory_format), weight, bias, params.padding, params.output_padding,
          params.stride, params.dilation, params.groups, params.benchmark, params.deterministic);
      break;
    case ConvBackend::Mkldnn:
#if AT_MKLDNN_ENABLED()
      check_input_same_type_as_parameters(input, weight, bias, backend);
      if (!input.is_mkldnn()) {
        // need to ensure contiguous for non-mkldnn tensors
        input = input.contiguous(backend_memory_format);
        weight = weight.contiguous(backend_memory_format);
        bias = bias.defined() ? bias.contiguous() : bias;
      }
      output = at::mkldnn_convolution(
          input, weight, bias, params.padding, params.stride, params.dilation, params.groups);
#else
      TORCH_INTERNAL_ASSERT(false, "Mkldnn backend was selected in PyTorch compiled without mkldnn support");
#endif
      break;
    case ConvBackend::MkldnnTranspose:
#if AT_MKLDNN_ENABLED()
      check_input_same_type_as_parameters(input, weight, bias, backend);
      if (!input.is_mkldnn()) {
        // need to ensure contiguous for non-mkldnn tensors
        input = input.contiguous(backend_memory_format);
        weight = weight.contiguous(backend_memory_format);
        bias = bias.defined() ? bias.contiguous() : bias;
      }
      output = mkldnn_convolution_transpose_stub(input.device().type(),
          input, weight, bias, params.padding, params.output_padding, params.stride, params.dilation, params.groups);
#else
      TORCH_INTERNAL_ASSERT(false, "Mkldnn backend was selected in PyTorch compiled without mkldnn support");
#endif
      break;
    case ConvBackend::MkldnnEmpty:
#if AT_MKLDNN_ENABLED()
      output = empty_mkldnn(
          calc_output_size(input, weight, params), optTypeMetaToScalarType(input.options().dtype_opt()),
          input.options().layout_opt(), input.options().device_opt(), input.options().pinned_memory_opt());
#else
      TORCH_INTERNAL_ASSERT(false, "Mkldnn backend was selected in PyTorch compiled without mkldnn support");
#endif
      break;
    case ConvBackend::Overrideable:
      output = at::convolution_overrideable(
          input, weight, bias, params.stride, params.padding, params.dilation, params.transposed,
          params.output_padding, params.groups);
      break;
    case ConvBackend::Slow3d:
      output = at::slow_conv3d(input, weight, kernel_size, bias, params.stride, params.padding);
      break;
    case ConvBackend::Winograd3x3Depthwise:
      output = convolution_depthwise3x3_winograd_stub(
          input.device().type(), input, weight, bias, params.stride, params.padding, params.groups);
      break;
    case ConvBackend::Xnnpack2d:
      output = xnnpack::convolution2d(
          input, weight, bias, params.padding, params.stride, params.dilation, params.groups);
      break;
    // Handle backends that don't natively support groups > 1.
    case ConvBackend::NnpackSpatial:
    case ConvBackend::Slow2d:
    case ConvBackend::SlowDilated2d:
    case ConvBackend::SlowDilated3d:
    case ConvBackend::SlowTranspose2d:
    case ConvBackend::SlowTranspose3d:
      input = input.contiguous(backend_memory_format);
      weight = weight.contiguous(backend_memory_format);
      if (params.groups == 1) {
        output = _convolution_nogroup_backend(input, weight, bias, backend, params);
      } else {
        std::vector<Tensor> outputs(params.groups);
        for (const auto g : c10::irange(params.groups)) {
          auto input_g = subtensor(input, 1, params.groups, g);
          auto weight_g = subtensor(weight, 0, params.groups, g);
          auto bias_g = subtensor(bias, 0, params.groups, g);
          outputs[g] = _convolution_nogroup_backend(input_g, weight_g, bias_g, backend, params);
        }
        output = at::cat(outputs, 1);
      }
      break;
    case ConvBackend::Mps:
#ifdef USE_MPS
      TORCH_CHECK(input.options().type_equal(weight.options()),
               "Input type (", input.toString(), ") and weight type (", weight.toString(),
               ") should be the same");
      TORCH_CHECK(!bias.defined() || (input.options().type_equal(bias.options())),
               "Input type (", input.toString(), ") and bias type (", bias.toString(),
               ") should be the same");

      output = at::_mps_convolution(input.contiguous(), weight, bias.defined() ? bias.contiguous() : bias,
                                     params.padding, params.stride, params.dilation,
                                     params.groups);
#else
      TORCH_INTERNAL_ASSERT(false, "MPS backend was selected in PyTorch without support");
#endif
      break;
    case ConvBackend::MpsTranspose:
#ifdef USE_MPS
      TORCH_CHECK(input.options().type_equal(weight.options()),
               "Input type (", input.toString(), ") and weight type (", weight.toString(),
               ") should be the same");
      TORCH_CHECK(!bias.defined() || (input.options().type_equal(bias.options())),
               "Input type (", input.toString(), ") and bias type (", bias.toString(),
               ") should be the same");
      output = at::_mps_convolution_transpose(
          input.contiguous(backend_memory_format), weight,
          params.padding, params.output_padding,
          params.stride, params.dilation, params.groups);
      if (bias.defined()) {
        output.add_(reshape_bias(input.dim(), bias));
      }
#else
      TORCH_INTERNAL_ASSERT(false, "MPS backend was selected in PyTorch without support");
#endif
      break;
  }

  if (k == 3 && !input.is_mkldnn() && !input.is_xpu()) {
    output = view3d(output);
  }

  return output;
}

at::Tensor _convolution(
    const Tensor& input_r, const Tensor& weight_r, const c10::optional<Tensor>& bias_r_opt,
    IntArrayRef stride_, IntArrayRef padding_, IntArrayRef dilation_,
    bool transposed_, IntArrayRef output_padding_, int64_t groups_,
    bool benchmark, bool deterministic, bool cudnn_enabled)
{
  // See [Note: hacky wrapper removal for optional tensor]
  c10::MaybeOwned<Tensor> bias_r_maybe_owned = at::borrow_from_optional_tensor(bias_r_opt);
  const Tensor& bias_r = *bias_r_maybe_owned;

  return at::_convolution(input_r, weight_r, bias_r, stride_, padding_, dilation_, transposed_, output_padding_, groups_, benchmark, deterministic, cudnn_enabled, at::globalContext().allowTF32CuDNN());
}

std::tuple<Tensor, Tensor, Tensor> convolution_backward_overrideable(
        const Tensor& grad_output, const Tensor& input, const Tensor& weight,
        IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation,
        bool transposed, IntArrayRef output_padding, int64_t groups, std::array<bool, 3> output_mask) {
   TORCH_CHECK_NOT_IMPLEMENTED(false, "convolution_backward_overrideable: You are likely triggering this with tensor backend other than CPU/CUDA/MKLDNN, if this is intended, please use TORCH_LIBRARY_IMPL to override this function ");
  return std::tuple<Tensor, Tensor, Tensor>(
          at::empty_like(input, LEGACY_CONTIGUOUS_MEMORY_FORMAT),
          at::empty_like(weight, LEGACY_CONTIGUOUS_MEMORY_FORMAT),
          at::empty({}));
}

static Tensor subvariable(const Tensor& var, int dim, int groups, int g) {
  int64_t n = var.sizes()[dim] / groups;
  auto result = var.narrow(dim, n * g, n);
  return result;
}

std::tuple<Tensor,Tensor,Tensor> _convolution_double_backward( const c10::optional<Tensor>& ggI_opt, const c10::optional<Tensor>& ggW_r_opt, const c10::optional<Tensor>& ggb_opt,
    const Tensor& gO_r, const Tensor& weight_r, const Tensor& input,
    IntArrayRef stride_, IntArrayRef padding_, IntArrayRef dilation_,
    bool transposed_, IntArrayRef output_padding_, int64_t groups_,
    std::array<bool, 3> output_mask) {
  // See [Note: hacky wrapper removal for optional tensor]
  c10::MaybeOwned<Tensor> ggI_maybe_owned = at::borrow_from_optional_tensor(ggI_opt);
  const Tensor& ggI = *ggI_maybe_owned;
  const Tensor& ggW_r = c10::value_or_else(ggW_r_opt, [] {return Tensor();});
  const Tensor& ggb = c10::value_or_else(ggb_opt, [] {return Tensor();});


  auto ggW = ggW_r;
  auto gO = gO_r;
  auto weight = weight_r;

  int64_t dim = weight.ndimension() - 2;
  ConvParams<int64_t> params;
  params.stride = expand_param_if_needed(stride_, "stride", dim);
  params.padding = expand_param_if_needed(padding_, "padding", dim);
  params.dilation = expand_param_if_needed(dilation_, "dilation", dim);
  params.transposed = transposed_;
  params.output_padding = expand_param_if_needed(output_padding_, "output_padding", dim);
  // TODO: hacky way of inferring the groups number for grouped Conv3D
  // See: https://github.com/pytorch/pytorch/pull/36355
  if (!params.transposed && input.dim() > 4) {
    // Avoid undefined behavior when num channels == 0; params are unused for that case.
    // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
    params.groups = (weight.size(1) > 0) ? input.size(1) / weight.size(1) : -1;
  } else {
    params.groups = groups_;
  }

  // Compute ggO = conv(ggI, w) + conv(i, ggW) + ggb
  Tensor ggO;
  if (input.numel() != 0) {
    if (ggI.defined()) {
      if (weight.is_cuda()) {
        weight = weight.contiguous();
      }
      ggO = at::convolution(ggI, weight, Tensor(), params.stride, params.padding, params.dilation, params.transposed, params.output_padding, params.groups);
    }

    if (ggW.defined()) {
      if (ggW.is_cuda()) {
        ggW = ggW.contiguous();
      }
      auto ggW_term = at::convolution(input, ggW, Tensor(), params.stride, params.padding, params.dilation, params.transposed, params.output_padding, params.groups);
      if (ggO.defined()) {
        ggO = ggO + ggW_term;
      } else {
        ggO = ggW_term;
      }
    }
  }

  if (ggb.defined()) {
    // View as (1, ggb.size(0), 1, 1...)

    // Expand
    std::vector<int64_t> new_size(gO.ndimension(), 1);
    new_size[1] = ggb.sizes()[0];
    auto ggb_contiguous = ggb.contiguous();
    auto ggb_view = ggb_contiguous.view(new_size);

    // Expand
    auto ggb_expanded = ggb_view.expand(gO.sizes());

    if (ggO.defined()) {
      ggO = ggO + ggb_expanded;
    } else {
      ggO = ggb_expanded;
    }
  }

  // Compute gW = conv(ggI, gO)
  Tensor gW;
  if (ggI.defined()) {

    // Modified params with correct padding
    ConvParams<int64_t> gw_conv_params(params);

    // Disable groups as they are handled separately
    auto groups = gw_conv_params.groups;
    gw_conv_params.groups = 1;
    std::swap(gw_conv_params.dilation, gw_conv_params.stride);

    // Transpose gO and ggI to accumulate over batch
    auto gOt = gO.transpose(0, 1);
    auto ggIt = ggI.transpose(0, 1);

    Tensor gWt;
    // Compute conv
    if (input.numel() != 0) {
      if (groups == 1) {

        if (gOt.is_cuda()) {
          gOt = gOt.contiguous();
        }
        // Compute conv
        if (params.transposed) {
          gw_conv_params.transposed = false;
          gWt = at::convolution(gOt, ggIt, Tensor(), gw_conv_params.stride, gw_conv_params.padding, gw_conv_params.dilation, gw_conv_params.transposed, gw_conv_params.output_padding, gw_conv_params.groups);
        } else {
          gWt = at::convolution(ggIt, gOt, Tensor(), gw_conv_params.stride, gw_conv_params.padding, gw_conv_params.dilation, gw_conv_params.transposed, gw_conv_params.output_padding, gw_conv_params.groups);
        }
      } else {
        std::vector<Tensor> gWt_list(groups);
        for (const auto g : c10::irange(groups)) {
          auto ggIt_g = subvariable(ggIt, 0, groups, g);
          auto gOt_g = subvariable(gOt, 0, groups, g);
          if (gOt_g.is_cuda()) {
            gOt_g = gOt_g.contiguous();
          }

          // Compute conv
          if (params.transposed) {
            gw_conv_params.transposed = false;
            gWt_list[g] = at::convolution(gOt_g, ggIt_g, Tensor(), gw_conv_params.stride, gw_conv_params.padding, gw_conv_params.dilation, gw_conv_params.transposed, gw_conv_params.output_padding, gw_conv_params.groups);
          } else {
            gWt_list[g] = at::convolution(ggIt_g, gOt_g, Tensor(), gw_conv_params.stride, gw_conv_params.padding, gw_conv_params.dilation, gw_conv_params.transposed, gw_conv_params.output_padding, gw_conv_params.groups);
          }
        }

        gWt = at::cat(gWt_list, 1);
      }

      // Transpose gW to match chan_in and chan_out
      gW = gWt.transpose(0, 1);

      // narrow gW to only relevant portion
      // we do it this way instead of narrowing the input itself because
      // the ConvForward kernels don't support asymmetric padding.
      auto gW_size = gW.sizes();
      auto w_size = weight.sizes();
      for (const auto i : c10::irange(2, gW_size.size())) {
        if (gW_size[i] > w_size[i]) {
            gW = gW.narrow(i, 0, w_size[i]);
            gW_size = gW.sizes();
        }
      }
    }
  }

  // Compute gI = convT(gO, ggW) if !transposed
  //         gI = conv(gO, ggw)  if transposed
  Tensor gI;
  if (input.numel() != 0) {
    if (ggW.defined()) {
      ConvParams<int64_t> gi_conv_params(params);
      gi_conv_params.transposed = !params.transposed;

      if (params.transposed) {
        if (gO.is_cuda()) {
          gO = gO.contiguous();
        }
        gI = at::convolution(gO, ggW, Tensor(), gi_conv_params.stride, gi_conv_params.padding, gi_conv_params.dilation, gi_conv_params.transposed, gi_conv_params.output_padding, gi_conv_params.groups);

        // narrow gI to only relevant portion
        // we do it this way because negative output_padding is not supported
        // TODO: figure out if we can narrow gO and save some compute,
        // rather than narrowing the computed gI
        auto gI_size = gI.sizes();
        auto i_size = input.sizes();
        for (const auto i : c10::irange(2, gI_size.size())) {
          if (gI_size[i] > i_size[i]) {
            gI = gI.narrow(i, 0, i_size[i]);
            gI_size = gI.sizes();
          }
        }
      } else {
        // calculate output_padding
        // TODO: figure out why this needs to be computed...
        auto kernel_size = weight.sizes().slice(2);
        auto input_shape = input.sizes().slice(2);
        auto grad_output_shape = gO.sizes().slice(2);

        for (const auto i : c10::irange(kernel_size.size())) {
          // Check if whole input has been used or not
          auto expected_input_shape = (kernel_size[i] - 1) * gi_conv_params.dilation[i]
            - 2 * gi_conv_params.padding[i]
            + (gi_conv_params.stride[i] * (grad_output_shape[i] - 1) + 1);
          if (expected_input_shape != input_shape[i]) {
            gi_conv_params.output_padding[i] = input_shape[i] - expected_input_shape;
          }
        }

        if (gO.is_cuda()) {
          gO = gO.contiguous();
        }

        gI = at::convolution(gO, ggW, Tensor(), gi_conv_params.stride, gi_conv_params.padding, gi_conv_params.dilation, gi_conv_params.transposed, gi_conv_params.output_padding, gi_conv_params.groups);
      }
    }
  }

  return std::tuple<Tensor,Tensor,Tensor>{ggO, gI, gW};
}

static std::tuple<at::Tensor, at::Tensor, at::Tensor> _convolution_backward_nogroup_backend(
    const Tensor& grad_output,
    const Tensor& input,
    const Tensor& weight,
    const std::array<bool, 3> output_mask,
    const ConvBackend backend,
    const ConvParams<int64_t>& params) {
  auto kernel_size = weight.sizes().slice(2);
  switch(backend) {
    case ConvBackend::Slow2d:
      return at::_slow_conv2d_backward(
        grad_output, input, weight, kernel_size, params.stride, params.padding, output_mask);
    // NB: nnpack backward does not support strided convolutions; use slow impl instead
    case ConvBackend::NnpackSpatial:
    case ConvBackend::SlowDilated2d:
      return slow_conv_dilated2d_backward_stub(
        input.device().type(),
        grad_output, input, weight, kernel_size, params.stride, params.padding, params.dilation, output_mask);
    case ConvBackend::SlowDilated3d:
      return slow_conv_dilated3d_backward_stub(
        input.device().type(),
        grad_output, input, weight, kernel_size, params.stride, params.padding, params.dilation, output_mask);
    case ConvBackend::SlowTranspose2d:
      return slow_conv_transpose2d_backward_stub(
        input.device().type(), grad_output, input, weight, kernel_size, params.stride, params.padding,
        params.output_padding, params.dilation, output_mask);
    case ConvBackend::SlowTranspose3d:
      return slow_conv_transpose3d_backward_stub(
        input.device().type(), grad_output, input, weight, kernel_size, params.stride, params.padding,
        params.output_padding, params.dilation, output_mask);
    default:
      TORCH_CHECK(false, "Unsupported conv nogroup backend encountered");
  }
}

// Backward pass for convolution. Computes gradients for input, weight, and bias depending on the
// output_mask setting. This function supports 1D, 2D, or 3D spatial convolution and currently requires
// a single batch dimension to be present.
//
// Args:
//   grad_output_: tensor of shape (N, C_out, L_out), (N, C_out, H_out, W_out), or (N, C_out, D_out, H_out, W_out)
//   input_: tensor of shape (N, C_in, L_in), (N, C_in, H_in, W_in), or (N, C_in, D_in, H_in, W_in)
//   weight_: tensor of shape (C_out, C_in // groups, *kernel_size); dimension of kernel_size must match the number
//       of input spatial dimensions
//   bias_sizes_opt: if specified, indicates that a bias was used in the forward pass and contains the shape
//       of the bias. While the bias shape can be computed from other inputs, it is provided to this function for
//       ease of use. The bias shape is (weight.shape[0]) for normal convolution and (weight.shape[1] * groups)
//       for transposed convolution.
//   stride: single value or an array with dimension matching the number of input spatial dimensions
//   padding: single value or an array with dimension matching the number of input spatial dimensions
//   dilation: single value or an array with dimension matching the number of input spatial dimensions
//   transposed: boolean indicating whether the convolution is transposed
//   output_padding: single value or dimension == number of input spatial dimensions; only supported when
//       transposed is true
//   groups: number of groups for grouped convolution
//   output_mask: 3-dim boolean array specifying which gradients to compute in input, weight, bias order
std::tuple<Tensor, Tensor, Tensor> convolution_backward(
    const Tensor& grad_output_, const Tensor& input_, const Tensor& weight_,
    const at::OptionalIntArrayRef bias_sizes_opt,
    IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation, bool transposed, IntArrayRef output_padding,
    int64_t groups, std::array<bool, 3> output_mask) {
  auto grad_output = grad_output_;
  auto input = input_;
  auto weight = weight_;

  auto k = weight.ndimension();
  int64_t dim = k - 2;

  TORCH_CHECK(dim > 0, "weight should have at least three dimensions");

  auto& ctx = at::globalContext();
  ConvParams<int64_t> params;
  params.stride = expand_param_if_needed(stride, "stride", dim);
  params.padding = expand_param_if_needed(padding, "padding", dim);
  params.dilation = expand_param_if_needed(dilation, "dilation", dim);
  params.transposed = transposed;
  params.output_padding = expand_param_if_needed(output_padding, "output_padding", dim);
  params.groups = groups;
  params.benchmark = ctx.benchmarkCuDNN();
  params.deterministic = ctx.deterministicCuDNN() || ctx.deterministicAlgorithms();
  params.cudnn_enabled = ctx.userEnabledCuDNN();
  params.allow_tf32 = ctx.allowTF32CuDNN();

  // Validate inputs.
  check_shape_backward(input, weight.sizes(), params);
  TORCH_CHECK(input.dim() == grad_output.dim(),
      "Expected input and grad_output to have the same number of dimensions, but got: ",
      input.dim(), " and ", grad_output.dim());

  // output_padding is only supported for transposed convolutions
  if (!params.transposed) {
    for (auto pad : params.output_padding) {
      TORCH_CHECK(pad == 0, "output_padding is not supported for non-transposed convolutions; got: ",
        params.output_padding);
    }
  }

  // Expand 1d -> 2d.
  // This is only done for backends that don't natively support 1d spatial input.
  if (k == 3 && !input.is_mkldnn() && !input.is_xpu()) {
    // avoid accidentally going through NHWC for permuted 3d input.
    input = input.contiguous();
    params.view1d_as_2d();
    grad_output = view4d(grad_output);
    input = view4d(input);
    weight = view4d(weight);
  }

  // Select appropriate backend to use.
  ConvBackend backend = select_conv_backend(input, weight, bias_sizes_opt, /*need_backward=*/ true, params);
  at::MemoryFormat backend_memory_format = determine_backend_memory_format(input, weight, backend);

  // Call the backend.
  Tensor backend_grad_input, backend_grad_weight, backend_grad_bias;
  auto kernel_size = weight.sizes().slice(2);
  switch(backend) {
    case ConvBackend::CudaDepthwise2d:
    {
      std::array<bool, 2> input_weight_output_mask = {output_mask[0], output_mask[1]};
      std::tie(backend_grad_input, backend_grad_weight) =
        conv_depthwise2d_backward_stub(input.device().type(), grad_output, input,
          weight, kernel_size, params.stride, params.padding, params.dilation, input_weight_output_mask);
      break;
    }
    case ConvBackend::CudaDepthwise3d:
      TORCH_CHECK(input.ndimension() == 5);
      std::tie(backend_grad_input, backend_grad_weight, backend_grad_bias) =
        conv_depthwise3d_backward_stub(
          input.device().type(), grad_output, input, weight, kernel_size, params.stride,
          params.padding, params.dilation, output_mask);
      break;
    case ConvBackend::Cudnn:
    {
      check_input_same_type_as_parameters(input, weight);
      std::array<bool, 2> input_weight_output_mask = {output_mask[0], output_mask[1]};
      std::tie(backend_grad_input, backend_grad_weight) = cudnn_convolution_backward_stub(
          input.device().type(),
          // Only make input contiguous when it is necessary for the backwards computation
          output_mask[1] ? input.contiguous(backend_memory_format) : input,
          grad_output, weight, params.padding, params.stride,
          params.dilation, params.groups, params.benchmark, params.deterministic, params.allow_tf32,
          input_weight_output_mask);
      break;
    }
    case ConvBackend::Mps:
    {
#ifdef USE_MPS
      check_input_same_type_as_parameters(input, weight);
      std::tie(backend_grad_input, backend_grad_weight, backend_grad_bias) =
        at::mps_convolution_backward(input, grad_output, weight, params.padding,
          params.stride, params.dilation, params.groups, output_mask);
#else
      TORCH_INTERNAL_ASSERT(false, "MPS backend was selected in PyTorch without support");
#endif
      break;
    }
    case ConvBackend::MpsTranspose:
    {
#ifdef USE_MPS
      check_input_same_type_as_parameters(input, weight);
      std::array<bool, 2> input_weight_output_mask = {output_mask[0], output_mask[1]};
      std::tie(backend_grad_input, backend_grad_weight) = at::mps_convolution_transpose_backward(
        // Only make input contiguous when it is necessary for the backwards computation
        output_mask[1] ? input.contiguous(backend_memory_format) : input,
        grad_output, weight, params.padding, params.output_padding,
        params.stride, params.dilation, params.groups, input_weight_output_mask);
#else
      TORCH_INTERNAL_ASSERT(false, "MPS backend was selected in PyTorch without support");
#endif
      break;
    }
    case ConvBackend::CudnnTranspose:
    {
      check_input_same_type_as_parameters(input, weight);
      std::array<bool, 2> input_weight_output_mask = {output_mask[0], output_mask[1]};
      std::tie(backend_grad_input, backend_grad_weight) = cudnn_convolution_transpose_backward_stub(
        input.device().type(),
        // Only make input contiguous when it is necessary for the backwards computation
        output_mask[1] ? input.contiguous(backend_memory_format) : input,
        grad_output, weight, params.padding, params.output_padding,
        params.stride, params.dilation, params.groups, params.benchmark, params.deterministic, params.allow_tf32,
        input_weight_output_mask);
      break;
    }
    case ConvBackend::Empty:
      if (output_mask[0]) {
        backend_grad_input = at::zeros_like(input);
      }
      if (output_mask[1]) {
        backend_grad_weight = at::zeros_like(weight);
      }
      if (output_mask[2]) {
        backend_grad_bias = at::zeros(*bias_sizes_opt, weight.options());
      }
      break;
    case ConvBackend::MkldnnEmpty:
#if AT_MKLDNN_ENABLED()
      if (output_mask[0]) {
        if (input.is_mkldnn()) {
          backend_grad_input = empty_mkldnn(input.sizes(), optTypeMetaToScalarType(input.options().dtype_opt()),
              input.options().layout_opt(), input.options().device_opt(), input.options().pinned_memory_opt());
          backend_grad_input.zero_();
        } else {
          backend_grad_input = at::zeros_like(input);
        }
      }
      if (output_mask[1]) {
        // mkldnn weight is not supported during training by the mkldnn backend
        backend_grad_weight = at::zeros_like(weight);
      }
      if (output_mask[2]) {
        // mkldnn bias is not supported during training by the mkldnn backend
        backend_grad_bias = at::zeros(*bias_sizes_opt, weight.options());
      }
#else
      TORCH_INTERNAL_ASSERT(false, "Mkldnn backend was selected in PyTorch compiled without mkldnn support");
#endif
      break;
    case ConvBackend::Miopen:
      check_input_same_type_as_parameters(input, weight);
      std::tie(backend_grad_input, backend_grad_weight, backend_grad_bias) =
        miopen_convolution_backward_stub(
          input.device().type(),
          input.contiguous(backend_memory_format), grad_output, weight, params.padding, params.stride,
          params.dilation, params.groups, params.benchmark, params.deterministic, output_mask);
      break;
    case ConvBackend::MiopenDepthwise:
      std::tie(backend_grad_input, backend_grad_weight, backend_grad_bias) =
          miopen_depthwise_convolution_backward_stub(
            input.device().type(),
            input.contiguous(backend_memory_format), grad_output, weight, params.padding, params.stride,
            params.dilation, params.groups, params.benchmark, params.deterministic, output_mask);
      break;
    case ConvBackend::MiopenTranspose:
      check_input_same_type_as_parameters(input, weight);
      std::tie(backend_grad_input, backend_grad_weight, backend_grad_bias) =
        miopen_convolution_transpose_backward_stub(
          input.device().type(),
          input.contiguous(backend_memory_format), grad_output, weight, params.padding, params.output_padding,
          params.stride, params.dilation, params.groups, params.benchmark, params.deterministic, output_mask);
      break;
    case ConvBackend::Mkldnn:
      TORCH_CHECK(!weight.is_mkldnn(),
          "The MKLDNN backend does not support weight as an MKLDNN tensor during training");
      if (!input.is_mkldnn()) {
        input = input.contiguous(backend_memory_format);
        weight = weight.contiguous(backend_memory_format);
      }
      std::tie(backend_grad_input, backend_grad_weight, backend_grad_bias) =
        mkldnn_convolution_backward_stub(input.device().type(), input, grad_output, weight, params.padding,
          params.stride, params.dilation, params.groups, output_mask);
      break;
    case ConvBackend::MkldnnTranspose:
      TORCH_CHECK(!weight.is_mkldnn(),
          "The MKLDNN backend does not support weight as an MKLDNN tensor during training");
      if (!input.is_mkldnn()) {
        input = input.contiguous(backend_memory_format);
        weight = weight.contiguous(backend_memory_format);
      }
      std::tie(backend_grad_input, backend_grad_weight, backend_grad_bias) =
        mkldnn_convolution_transpose_backward_stub(input.device().type(), input, grad_output, weight, params.padding,
        params.output_padding, params.stride, params.dilation, params.groups, output_mask);
      break;
    case ConvBackend::Overrideable:
      // Only reach here when input is backend with out-of-source implementation.
      std::tie(backend_grad_input, backend_grad_weight, backend_grad_bias) =
        at::convolution_backward_overrideable(grad_output, input, weight, params.stride, params.padding,
          params.dilation, params.transposed, params.output_padding, params.groups, output_mask);
      break;
    case ConvBackend::Slow3d:
      // Note that no CUDA implementation of this kernel exists currently.
      std::tie(backend_grad_input, backend_grad_weight, backend_grad_bias) =
        slow_conv3d_backward_cpu(
            grad_output, input, weight, kernel_size,
            params.stride, params.padding, output_mask);
      break;
    // Handle backends that don't natively support groups > 1.
    case ConvBackend::NnpackSpatial:
    case ConvBackend::Slow2d:
    case ConvBackend::SlowDilated2d:
    case ConvBackend::SlowDilated3d:
    case ConvBackend::SlowTranspose2d:
    case ConvBackend::SlowTranspose3d:
    {
      input = input.contiguous(backend_memory_format);
      weight = weight.contiguous(backend_memory_format);
      if (params.groups == 1) {
        std::tie(backend_grad_input, backend_grad_weight, backend_grad_bias) =
          _convolution_backward_nogroup_backend(
            grad_output, input, weight, output_mask, backend, params);
      } else {
        std::vector<Tensor> backend_grad_inputs(params.groups);
        std::vector<Tensor> backend_grad_weights(params.groups);
        std::vector<Tensor> backend_grad_biases(params.groups);
        for (int g = 0; g < params.groups; ++g) {
          auto grad_output_g = subtensor(grad_output, 1, params.groups, g);
          auto input_g = subtensor(input, 1, params.groups, g);
          auto weight_g = subtensor(weight, 0, params.groups, g);
          std::tie(backend_grad_inputs[g], backend_grad_weights[g], backend_grad_biases[g]) =
            _convolution_backward_nogroup_backend(
              grad_output_g, input_g, weight_g, output_mask, backend, params);
        }
        if (output_mask[0]) {
          backend_grad_input = at::cat(backend_grad_inputs, 1);
        }
        if (output_mask[1]) {
          backend_grad_weight = at::cat(backend_grad_weights, 0);
        }
        if (output_mask[2]) {
          backend_grad_bias = at::cat(backend_grad_biases, 0);
        }
      }
      break;
    }
    // Backward is not supported for these backends.
    case ConvBackend::Winograd3x3Depthwise:
      TORCH_CHECK(false, "Backward is not supported for depthwise 3x3 winograd");
      break;
    case ConvBackend::Xnnpack2d:
      TORCH_CHECK(false, "Backward is not supported for xnnpack");
      break;
  }

  // Convert 2D inputs back to 1D for backends that don't natively support 1D
  // spatial inputs.
  if (output_mask[0]) {
    if (k == 3 && !input.is_mkldnn() && !input.is_xpu()) {
      backend_grad_input = view3d(backend_grad_input);
    }
  }
  if (output_mask[1]) {
    if (k == 3 && !input.is_mkldnn() && !input.is_xpu()) {
      backend_grad_weight = view3d(backend_grad_weight);
    }
  }
  if (output_mask[2]) {
    if (!backend_grad_bias.defined()) {
      // Calculate bias gradients outside of the backend for those that don't support it.
      backend_grad_bias = grad_output.sum((dim == 3) ? IntArrayRef{0, 2, 3, 4} : IntArrayRef{0, 2, 3});
    }
  }

  return std::make_tuple(backend_grad_input, backend_grad_weight, backend_grad_bias);
}

}} // at::native