add channels last support for thnn_conv2d (non-dilated)

ghstack-source-id: 0584ef4b7b499004aa55c1d62db34c584ee60aab
Pull Request resolved: pytorch#49582

aten/src/ATen/native/Convolution.cpp

@@ -570,8 +570,9 @@ static at::Tensor subtensor(at::Tensor& tensor, int dim, int groups, int g) {
   if (!tensor.defined()) {
     return at::Tensor();
   }
+  auto memory_format = tensor.suggest_memory_format();
   int64_t n = tensor.sizes()[dim] / groups;
-  return tensor.narrow(dim, n * g, n).contiguous();
+  return tensor.narrow(dim, n * g, n).contiguous(memory_format);
 }
 
 
@@ -979,12 +980,20 @@ at::Tensor _convolution(
           params.stride,
           params.padding);
   } else if (input.device().is_cpu() || input.is_cuda()) {
+    bool is_channels_last_supported = !params.transposed && (input.ndimension() == 4) &&
+        !params.use_nnpack(input, weight) && input.device().is_cpu() &&
+        !params.is_dilated();
+    if (is_channels_last_supported) {
+      auto memory_format = input.suggest_memory_format();
+      input = input.contiguous(memory_format);
+    } else {
+      input = input.contiguous();
+    }
     if (params.groups == 1) {
       output = at::_convolution_nogroup(
-          input.contiguous(), weight, bias, params.stride, params.padding, params.dilation, params.transposed, params.output_padding);
+          input, weight, bias, params.stride, params.padding, params.dilation, params.transposed, params.output_padding);
     } else {
       std::vector<Tensor> outputs(params.groups);
-      input = input.contiguous();
       for (int g = 0; g < params.groups; ++g) {
         auto input_g = subtensor(input, 1, params.groups, g);
         auto weight_g = subtensor(weight, 0, params.groups, g);