feat(conv2d_to_convolution): A pass to map aten::conv2d to _convolution

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

Diff for: core/lowering/lowering.cpp

@@ -25,6 +25,7 @@ void LowerGraph(std::shared_ptr<torch::jit::Graph>& g) {
     torch::jit::FuseLinear(g);
     passes::RemoveDropout(g);
     passes::FuseFlattenLinear(g);
+    passes::Conv2DToConvolution(g);
     passes::UnpackAddMM(g);
     passes::UnpackLogSoftmax(g);
     //passes::RemoveDimExeception(g);

Diff for: core/lowering/passes/BUILD

@@ -6,6 +6,7 @@ cc_library(
         "passes.h",
     ],
     srcs = [
+        "conv2d_to_convolution.cpp",
         "exception_elimination.cpp",
         "fuse_flatten_linear.cpp",
         "remove_dropout.cpp",

Diff for: core/lowering/passes/conv2d_to_convolution.cpp

@@ -0,0 +1,34 @@
+#include <torch/csrc/jit/passes/subgraph_rewrite.h>
+
+#include "core/util/prelude.h"
+
+namespace trtorch {
+namespace core {
+namespace lowering {
+namespace passes {
+
+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(
+        graph(%x, %w, %b, %s, %p, %d, %g):
+            %1 : bool = prim::Constant[value=1]()
+            %2 : int[] = prim::Constant[value=[0, 0]]()
+            %3 : bool = prim::Constant[value=0]()
+            %4 : Tensor = aten::_convolution(%x, %w, %b, %s, %p, %d, %1, %2, %g, %1, %1, %3)
+            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);
+    LOG_GRAPH("Post map conv2d -> _convolution: " << *graph);
+}
+
+} // namespace passes
+} // namespace lowering
+} // namespace core
+} // namespace trtorch

Diff for: core/lowering/passes/passes.h

@@ -7,6 +7,7 @@ namespace core {
 namespace lowering {
 namespace passes {
 
+void Conv2DToConvolution(std::shared_ptr<torch::jit::Graph>& graph);
 void FuseFlattenLinear(std::shared_ptr<torch::jit::Graph>& graph);
 void RemoveDropout(std::shared_ptr<torch::jit::Graph>& graph);
 void UnpackAddMM(std::shared_ptr<torch::jit::Graph>& graph);