Merge branch 'main' into aten_size_fix

Author: peri044

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,78 +8,91 @@ 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)
-            return (%4))IR";
-  std::string convolution_pattern = R"IR(
+  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);
 }
 

docs/_cpp_api/classtorch__tensorrt_1_1DataType.html

@@ -10,7 +10,7 @@
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
 
-  <title>Class DataType &mdash; Torch-TensorRT v1.4.0.dev0+b388010 documentation</title>
+  <title>Class DataType &mdash; Torch-TensorRT v1.4.0.dev0+2ea9f00 documentation</title>
 
 
 
@@ -215,7 +215,7 @@
 
 
                 <div class="version">
-                  v1.4.0.dev0+b388010
+                  v1.4.0.dev0+2ea9f00
                 </div>
 
 

docs/_cpp_api/classtorch__tensorrt_1_1Device_1_1DeviceType.html

@@ -10,7 +10,7 @@
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
 
-  <title>Class Device::DeviceType &mdash; Torch-TensorRT v1.4.0.dev0+b388010 documentation</title>
+  <title>Class Device::DeviceType &mdash; Torch-TensorRT v1.4.0.dev0+2ea9f00 documentation</title>
 
 
 
@@ -215,7 +215,7 @@
 
 
                 <div class="version">
-                  v1.4.0.dev0+b388010
+                  v1.4.0.dev0+2ea9f00
                 </div>
 
 

docs/_cpp_api/classtorch__tensorrt_1_1TensorFormat.html

@@ -10,7 +10,7 @@
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
 
-  <title>Class TensorFormat &mdash; Torch-TensorRT v1.4.0.dev0+b388010 documentation</title>
+  <title>Class TensorFormat &mdash; Torch-TensorRT v1.4.0.dev0+2ea9f00 documentation</title>
 
 
 
@@ -215,7 +215,7 @@
 
 
                 <div class="version">
-                  v1.4.0.dev0+b388010
+                  v1.4.0.dev0+2ea9f00
                 </div>
 
 

docs/_cpp_api/classtorch__tensorrt_1_1ptq_1_1Int8CacheCalibrator.html

@@ -10,7 +10,7 @@
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
 
-  <title>Template Class Int8CacheCalibrator &mdash; Torch-TensorRT v1.4.0.dev0+b388010 documentation</title>
+  <title>Template Class Int8CacheCalibrator &mdash; Torch-TensorRT v1.4.0.dev0+2ea9f00 documentation</title>
 
 
 
@@ -215,7 +215,7 @@
 
 
                 <div class="version">
-                  v1.4.0.dev0+b388010
+                  v1.4.0.dev0+2ea9f00
                 </div>
 
 

docs/_cpp_api/classtorch__tensorrt_1_1ptq_1_1Int8Calibrator.html

@@ -10,7 +10,7 @@
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
 
-  <title>Template Class Int8Calibrator &mdash; Torch-TensorRT v1.4.0.dev0+b388010 documentation</title>
+  <title>Template Class Int8Calibrator &mdash; Torch-TensorRT v1.4.0.dev0+2ea9f00 documentation</title>
 
 
 
@@ -215,7 +215,7 @@
 
 
                 <div class="version">
-                  v1.4.0.dev0+b388010
+                  v1.4.0.dev0+2ea9f00
                 </div>
 
 

docs/_cpp_api/define_macros_8h_1a18d295a837ac71add5578860b55e5502.html

@@ -10,7 +10,7 @@
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
 
-  <title>Define STR &mdash; Torch-TensorRT v1.4.0.dev0+b388010 documentation</title>
+  <title>Define STR &mdash; Torch-TensorRT v1.4.0.dev0+2ea9f00 documentation</title>
 
 
 
@@ -215,7 +215,7 @@
 
 
                 <div class="version">
-                  v1.4.0.dev0+b388010
+                  v1.4.0.dev0+2ea9f00
                 </div>
 
 

docs/_cpp_api/define_macros_8h_1a282fd3c0b1c3a215148ae372070e1268.html

@@ -10,7 +10,7 @@
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
 
-  <title>Define TORCH_TENSORRT_PATCH_VERSION &mdash; Torch-TensorRT v1.4.0.dev0+b388010 documentation</title>
+  <title>Define TORCH_TENSORRT_PATCH_VERSION &mdash; Torch-TensorRT v1.4.0.dev0+2ea9f00 documentation</title>
 
 
 
@@ -215,7 +215,7 @@
 
 
                 <div class="version">
-                  v1.4.0.dev0+b388010
+                  v1.4.0.dev0+2ea9f00
                 </div>
 
 

docs/_cpp_api/define_macros_8h_1a31398a6d4d27e28817afb0f0139e909e.html

@@ -10,7 +10,7 @@
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
 
-  <title>Define TORCH_TENSORRT_MAJOR_VERSION &mdash; Torch-TensorRT v1.4.0.dev0+b388010 documentation</title>
+  <title>Define TORCH_TENSORRT_MAJOR_VERSION &mdash; Torch-TensorRT v1.4.0.dev0+2ea9f00 documentation</title>
 
 
 
@@ -215,7 +215,7 @@
 
 
                 <div class="version">
-                  v1.4.0.dev0+b388010
+                  v1.4.0.dev0+2ea9f00
                 </div>
 
 

docs/_cpp_api/define_macros_8h_1a35703561b26b1a9d2738ad7d58b27827.html

@@ -10,7 +10,7 @@
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
 
-  <title>Define TORCH_TENSORRT_MINOR_VERSION &mdash; Torch-TensorRT v1.4.0.dev0+b388010 documentation</title>
+  <title>Define TORCH_TENSORRT_MINOR_VERSION &mdash; Torch-TensorRT v1.4.0.dev0+2ea9f00 documentation</title>
 
 
 
@@ -215,7 +215,7 @@
 
 
                 <div class="version">
-                  v1.4.0.dev0+b388010
+                  v1.4.0.dev0+2ea9f00
                 </div>