Hide convolution_same from jit IR

peterbell10 · peterbell10 · commit 28ff4634f205 · 2020-09-06T21:10:56.000+01:00
diff --git a/aten/src/ATen/native/Convolution.cpp b/aten/src/ATen/native/Convolution.cpp
@@ -644,43 +644,43 @@ static Tensor convolution_same(
                          dilation, false, output_padding, groups);
 }
 
-at::Tensor conv1d(
+Tensor _convolution_mode(
     const Tensor& input, const Tensor& weight, const Tensor& bias,
     IntArrayRef stride, std::string padding, IntArrayRef dilation,
     int64_t groups) {
   if (padding == "same") {
-    return at::native::convolution_same(input, weight, bias, stride, dilation, groups);
+    return at::native::convolution_same(
+        input, weight, bias, stride, dilation, groups);
   } else if (padding == "valid") {
     const int64_t padding_[] = {0};
-    return at::native::conv1d(input, weight, bias, stride, padding_, dilation, groups);
+    return at::native::convolution(
+        input, weight, bias, stride, padding_, dilation, false, padding_, groups);
   }
   TORCH_CHECK(false, "Invalid padding mode '", padding, "'");
 }
 
+at::Tensor conv1d(
+    const Tensor& input, const Tensor& weight, const Tensor& bias,
+    IntArrayRef stride, std::string padding, IntArrayRef dilation,
+    int64_t groups) {
+  return at::_convolution_mode(
+      input, weight, bias, stride, std::move(padding), dilation, groups);
+}
+
 at::Tensor conv2d(
     const Tensor& input, const Tensor& weight, const Tensor& bias,
     IntArrayRef stride, std::string padding, IntArrayRef dilation,
     int64_t groups) {
-  if (padding == "same") {
-    return at::native::convolution_same(input, weight, bias, stride, dilation, groups);
-  } else if (padding == "valid") {
-    const int64_t padding_[] = {0, 0};
-    return at::native::conv2d(input, weight, bias, stride, padding_, dilation, groups);
-  }
-  TORCH_CHECK(false, "Invalid padding mode '", padding, "'");
+  return at::_convolution_mode(
+      input, weight, bias, stride, std::move(padding), dilation, groups);
 }
 
 at::Tensor conv3d(
     const Tensor& input, const Tensor& weight, const Tensor& bias,
     IntArrayRef stride, std::string padding, IntArrayRef dilation,
     int64_t groups) {
-  if (padding == "same") {
-    return at::native::convolution_same(input, weight, bias, stride, dilation, groups);
-  } else if (padding == "valid") {
-    const int64_t padding_[] = {0, 0, 0};
-    return at::native::conv3d(input, weight, bias, stride, padding_, dilation, groups);
-  }
-  TORCH_CHECK(false, "Invalid padding mode '", padding, "'");
+  return at::_convolution_mode(
+      input, weight, bias, stride, std::move(padding), dilation, groups);
 }
 
 at::Tensor conv_transpose1d(
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -938,6 +938,9 @@
 - func: _convolution.deprecated(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, bool transposed, int[] output_padding, int groups, bool benchmark, bool deterministic, bool cudnn_enabled) -> Tensor
   use_c10_dispatcher: full
 
+- func: _convolution_mode(Tensor input, Tensor weight, Tensor? bias, int[] stride, str padding, int[] dilation, int groups) -> Tensor
+  use_c10_dispatcher: full
+
 - func: _convolution_nogroup(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, bool transposed, int[] output_padding) -> Tensor
   use_c10_dispatcher: full
 
