fix: Add schemas to conv replace

- Add support for transposed conv2d and conv3d, as well as for conv3d
- Add testing for all convolution functions, rename test accordingly

Author: gs-olive

core/lowering/lowering.cpp

@@ -124,7 +124,9 @@ void LowerGraph(std::shared_ptr<torch::jit::Graph>& g, std::vector<torch::jit::I
   passes::Conv1DToConvolution(g);
   passes::ConvTransposed1DToConvolution(g);
   passes::Conv2DToConvolution(g);
+  passes::ConvTransposed2DToConvolution(g);
   passes::Conv3DToConvolution(g);
+  passes::ConvTransposed3DToConvolution(g);
   passes::FuseAddMMBranches(g);
   passes::RemoveBNDimCheck(g);
   // torch::jit::UnrollLoops(g);

core/lowering/passes/convNd_to_convolution.cpp

@@ -82,6 +82,20 @@ void Conv2DToConvolution(std::shared_ptr<torch::jit::Graph>& graph) {
   LOG_GRAPH("Post map conv2d -> _convolution: " << *graph);
 }
 
+void ConvTransposed2DToConvolution(std::shared_ptr<torch::jit::Graph>& graph) {
+  const std::string conv_transpose2d_node_kind = "aten::conv_transpose2d";
+  const std::string convolution_pattern = R"IR(
+        graph(%x, %w, %b, %s, %p, %o, %g, %d):
+            %1 : bool = prim::Constant[value=1]()
+            %2 : bool = prim::Constant[value=1]()
+            %4 : Tensor = aten::_convolution(%x, %w, %b, %s, %p, %d, %1, %o, %g, %2, %2, %2, %2)
+            return (%4))IR";
+
+  // Schema is aten::conv_transpose2d(%x, %w, %b, %s, %p, %o, %g, %d) --> 8 inputs
+  replaceConv(graph->block(), conv_transpose2d_node_kind, convolution_pattern, 8);
+  LOG_GRAPH("Post map conv_transpose2d -> _convolution: " << *graph);
+}
+
 void Conv3DToConvolution(std::shared_ptr<torch::jit::Graph>& graph) {
   const std::string conv3d_node_kind = "aten::conv3d";
   const std::string convolution_pattern = R"IR(
@@ -96,6 +110,20 @@ void Conv3DToConvolution(std::shared_ptr<torch::jit::Graph>& graph) {
   LOG_GRAPH("Post map conv3d -> _convolution: " << *graph);
 }
 
+void ConvTransposed3DToConvolution(std::shared_ptr<torch::jit::Graph>& graph) {
+  const std::string conv_transpose3d_node_kind = "aten::conv_transpose3d";
+  const std::string convolution_pattern = R"IR(
+        graph(%x, %w, %b, %s, %p, %o, %g, %d):
+            %1 : bool = prim::Constant[value=1]()
+            %2 : bool = prim::Constant[value=1]()
+            %4 : Tensor = aten::_convolution(%x, %w, %b, %s, %p, %d, %1, %o, %g, %2, %2, %2, %2)
+            return (%4))IR";
+
+  // Schema is aten::conv_transpose3d(%x, %w, %b, %s, %p, %o, %g, %d) --> 8 inputs
+  replaceConv(graph->block(), conv_transpose3d_node_kind, convolution_pattern, 8);
+  LOG_GRAPH("Post map conv_transpose3d -> _convolution: " << *graph);
+}
+
 } // namespace passes
 } // namespace lowering
 } // namespace core

core/lowering/passes/passes.h

@@ -15,7 +15,9 @@ void NotateModuleForFallback(
 void Conv1DToConvolution(std::shared_ptr<torch::jit::Graph>& graph);
 void ConvTransposed1DToConvolution(std::shared_ptr<torch::jit::Graph>& graph);
 void Conv2DToConvolution(std::shared_ptr<torch::jit::Graph>& graph);
+void ConvTransposed2DToConvolution(std::shared_ptr<torch::jit::Graph>& graph);
 void Conv3DToConvolution(std::shared_ptr<torch::jit::Graph>& graph);
+void ConvTransposed3DToConvolution(std::shared_ptr<torch::jit::Graph>& graph);
 void FuseAddMMBranches(std::shared_ptr<torch::jit::Graph> graph);
 void LinearToAddMM(std::shared_ptr<torch::jit::Graph>& graph);
 void EliminateExceptionOrPassPattern(std::shared_ptr<torch::jit::Graph> graph);

tests/core/lowering/BUILD

@@ -28,7 +28,7 @@ cc_test(
 )
 
 lowering_test(
-    name = "test_conv1d_pass",
+    name = "test_conv_pass",
 )
 
 lowering_test(
@@ -102,7 +102,7 @@ lowering_test(
 test_suite(
     name = "lowering_tests",
     tests = [
-        ":test_conv1d_pass",
+        ":test_conv_pass",
         ":test_device_casting",
         ":test_exception_elimination_pass",
         ":test_linear_to_addmm",