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

@@ -1,4 +1,5 @@
 #include <torch/csrc/jit/passes/subgraph_rewrite.h>
+#include "torch/csrc/jit/ir/irparser.h"
 
 #include "core/util/prelude.h"
 
@@ -7,81 +8,122 @@ namespace core {
 namespace lowering {
 namespace passes {
 
-void Conv1DToConvolution(std::shared_ptr<torch::jit::Graph>& graph) {
-  std::string conv1d_pattern = R"IR(
-        graph(%x, %w, %b, %s, %p, %d, %g):
-            %4 : Tensor = aten::conv1d(%x, %w, %b, %s, %p, %d, %g)
-            return (%4))IR";
+void replaceConv(
+    torch::jit::Block* block,
+    const std::string& node_kind,
+    const std::string& unwrapped_conv,
+    const size_t num_input_args) {
+  // Iterate through nodes in block, seaching for aten::conv*
+  for (auto it = block->nodes().begin(); it != block->nodes().end(); it++) {
+    auto n = *it;
+
+    // Recursively explore nested blocks, such as those arising from prim::If
+    for (auto nested_block : n->blocks()) {
+      replaceConv(nested_block, node_kind, unwrapped_conv, num_input_args);
+    }
+
+    // If node matches desired kind and number of input arguments, replace it
+    if ((n->kind().toQualString() == node_kind) && (n->inputs().size() == num_input_args)) {
+      // Establish insert point within block
+      torch::jit::WithInsertPoint guard(*it);
+
+      // Initialize new fused subgraph from IR code provided
+      auto fused_g = std::make_shared<torch::jit::Graph>();
+      torch::jit::parseIR(unwrapped_conv, fused_g.get());
+
+      // Insert subgraph in place of aten::conv*, replacing inputs and outputs accordingly
+      torch::jit::Value* new_output = insertGraph(*it->owningGraph(), *fused_g, it->inputs()).at(0);
+      new_output->setType(it->output()->type());
+      it->output()->replaceAllUsesWith(new_output);
+      it.destroyCurrent();
+    }
+  }
+}
 
-  std::string convolution_pattern = R"IR(
+void Conv1DToConvolution(std::shared_ptr<torch::jit::Graph>& graph) {
+  const std::string conv1d_node_kind = "aten::conv1d";
+  const std::string convolution_pattern = R"IR(
         graph(%x, %w, %b, %s, %p, %d, %g):
             %1 : bool = prim::Constant[value=0]()
             %2 : int[] = prim::Constant[value=[0]]()
             %4 : Tensor = aten::_convolution(%x, %w, %b, %s, %p, %d, %1, %2, %g, %1, %1, %1, %1)
             return (%4))IR";
 
-  torch::jit::SubgraphRewriter map_conv1d_to_convolution;
-  map_conv1d_to_convolution.RegisterRewritePattern(conv1d_pattern, convolution_pattern);
-  map_conv1d_to_convolution.runOnGraph(graph);
+  // Schema is aten::conv1d(%x, %w, %b, %s, %p, %d, %g) --> 7 inputs
+  replaceConv(graph->block(), conv1d_node_kind, convolution_pattern, 7);
   LOG_GRAPH("Post map conv1d -> _convolution: " << *graph);
 }
 
 void ConvTransposed1DToConvolution(std::shared_ptr<torch::jit::Graph>& graph) {
-  std::string conv_transpose1d_pattern = R"IR(
-        graph(%x, %w, %b, %s, %p, %o, %g, %d):
-            %4 : Tensor = aten::conv_transpose1d(%x, %w, %b, %s, %p, %o, %g, %d)
-            return (%4))IR";
-  std::string convolution_pattern = R"IR(
+  const std::string conv_transpose1d_node_kind = "aten::conv_transpose1d";
+  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";
 
-  torch::jit::SubgraphRewriter map_conv_transpose1d_to_convolution;
-  map_conv_transpose1d_to_convolution.RegisterRewritePattern(conv_transpose1d_pattern, convolution_pattern);
-  map_conv_transpose1d_to_convolution.runOnGraph(graph);
+  // Schema is aten::conv_transpose1d(%x, %w, %b, %s, %p, %o, %g, %d) --> 8 inputs
+  replaceConv(graph->block(), conv_transpose1d_node_kind, convolution_pattern, 8);
   LOG_GRAPH("Post map conv_transpose1d -> _convolution: " << *graph);
 }
 
 void Conv2DToConvolution(std::shared_ptr<torch::jit::Graph>& graph) {
-  std::string conv2d_pattern = R"IR(
-        graph(%x, %w, %b, %s, %p, %d, %g):
-            %4 : Tensor = aten::conv2d(%x, %w, %b, %s, %p, %d, %g)
-            return (%4))IR";
-  std::string convolution_pattern = R"IR(
+  const std::string conv2d_node_kind = "aten::conv2d";
+  const std::string convolution_pattern = R"IR(
         graph(%x, %w, %b, %s, %p, %d, %g):
             %1 : bool = prim::Constant[value=0]()
             %2 : int[] = prim::Constant[value=[0, 0]]()
             %4 : Tensor = aten::_convolution(%x, %w, %b, %s, %p, %d, %1, %2, %g, %1, %1, %1, %1)
             return (%4))IR";
 
-  // replace matmul + add pattern to linear
-  torch::jit::SubgraphRewriter map_conv2d_to_convolution;
-  map_conv2d_to_convolution.RegisterRewritePattern(conv2d_pattern, convolution_pattern);
-  map_conv2d_to_convolution.runOnGraph(graph);
+  // Schema is aten::conv2d(%x, %w, %b, %s, %p, %d, %g) --> 7 inputs
+  replaceConv(graph->block(), conv2d_node_kind, convolution_pattern, 7);
   LOG_GRAPH("Post map conv2d -> _convolution: " << *graph);
 }
 
-void Conv3DToConvolution(std::shared_ptr<torch::jit::Graph>& graph) {
-  std::string conv3d_pattern = R"IR(
-        graph(%x, %w, %b, %s, %p, %d, %g):
-            %4 : Tensor = aten::conv3d(%x, %w, %b, %s, %p, %d, %g)
+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";
-  std::string convolution_pattern = R"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(
         graph(%x, %w, %b, %s, %p, %d, %g):
             %1 : bool = prim::Constant[value=0]()
             %2 : int[] = prim::Constant[value=[0, 0, 0]]()
             %4 : Tensor = aten::_convolution(%x, %w, %b, %s, %p, %d, %1, %2, %g, %1, %1, %1, %1)
             return (%4))IR";
 
-  // replace matmul + add pattern to linear
-  torch::jit::SubgraphRewriter map_conv3d_to_convolution;
-  map_conv3d_to_convolution.RegisterRewritePattern(conv3d_pattern, convolution_pattern);
-  map_conv3d_to_convolution.runOnGraph(graph);
+  // Schema is aten::conv3d(%x, %w, %b, %s, %p, %d, %g) --> 7 inputs
+  replaceConv(graph->block(), conv3d_node_kind, convolution_pattern, 7);
   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",