Merge branch 'batch_norm_alt' into pytorch_1.5.0

Closes: #31

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

Author: narendasan

core/compiler.cpp

@@ -71,47 +71,29 @@ void AddEngineToGraph(torch::jit::script::Module mod, std::shared_ptr<torch::jit
 
 bool CheckMethodOperatorSupport(const torch::jit::script::Module& mod,
                                 std::string method_name) {
-    auto g = mod.get_method(method_name).graph();
-    // Go through PyTorch Lowering to simplify graph and extract weight parameters
-    auto graph_and_parameters = torch::jit::LowerGraph(*g, mod._ivalue());
+    // Go through Lowering to simplify graph and extract weight parameters
+    auto graph_and_parameters = lowering::Lower(mod, method_name);
 
-    g = graph_and_parameters.first;
-
-    // Go through TRTorch Lowering to reformat graph to be conversion friendly
-    // and also segment for accelerators and executors (TRT-DLA, TRT-GPU, PYT)
-    lowering::LowerGraph(g);
-
-    auto params = graph_and_parameters.second;
-    auto named_params = conversion::get_named_params(g->inputs(), params);
+    auto g = graph_and_parameters.first;
     LOG_DEBUG(*g << "(CheckMethodOperatorSupport)\n");
 
-    // Is this necessary?
-    lowering::LowerBlock(g->block());
-
     return conversion::VerifyConverterSupportForBlock(g->block());
 }
 
 std::string ConvertGraphToTRTEngine(const torch::jit::script::Module& mod,
                                     std::string method_name,
                                     ExtraInfo cfg) {
-    auto convert_cfg = std::move(cfg.convert_info);
-
-    auto g = mod.get_method(method_name).graph();
-    // Go through PyTorch Lowering to simplify graph and extract weight parameters
-    auto graph_and_parameters = torch::jit::LowerGraph(*g, mod._ivalue());
-
-    g = graph_and_parameters.first;
 
-    // Go through TRTorch Lowering to reformat graph to be conversion friendly
-    // and also segment for accelerators and executors (TRT-DLA, TRT-GPU, PYT)
-    lowering::LowerGraph(g);
+    // Go through Lowering to simplify graph and extract weight parameters
+    auto graph_and_parameters = lowering::Lower(mod, method_name);
 
+    auto convert_cfg = std::move(cfg.convert_info);
+    auto g = graph_and_parameters.first;
     auto params = graph_and_parameters.second;
     auto named_params = conversion::get_named_params(g->inputs(), params);
+
     LOG_INFO(*g << "(CompileGraph)\n");
 
-    // Is this necessary?
-    lowering::LowerBlock(g->block());
     auto engine = ConvertBlockToEngine(g->block(), convert_cfg, named_params);
     return std::move(engine);
 }

core/lowering/lowering.cpp

@@ -1,5 +1,7 @@
-#include "torch/csrc/jit/passes/fuse_linear.h"
 #include "torch/csrc/jit/passes/dead_code_elimination.h"
+#include "torch/csrc/jit/passes/fuse_linear.h"
+#include "torch/csrc/jit/passes/lower_graph.h"
+#include "torch/csrc/jit/passes/quantization.h"
 
 #include "core/lowering/lowering.h"
 #include "core/lowering/irfusers/irfusers.h"
@@ -22,7 +24,29 @@ void LowerGraph(std::shared_ptr<torch::jit::Graph>& g) {
     //irfusers::UnpackBatchNorm(g);
     //torch::jit::EliminateDeadCode(g);
 }
-    
+
+void LowerModule(const torch::jit::script::Module& mod) {
+    torch::jit::FoldConvBatchNorm2d(mod);
+}
+
+std::pair<std::shared_ptr<torch::jit::Graph>, std::vector<at::Tensor>> Lower(const torch::jit::script::Module& mod,
+                                                                            std::string method_name) {
+    LowerModule(mod);
+    auto g = mod.get_method(method_name).graph();
+    // Go through PyTorch Lowering to simplify graph and extract weight parameters
+    auto graph_and_parameters = torch::jit::LowerGraph(*g, mod._ivalue());
+
+    g = graph_and_parameters.first;
+
+    // Go through TRTorch Lowering to reformat graph to be conversion friendly
+    // and also segment for accelerators and executors (TRT-DLA, TRT-GPU, PYT)
+    lowering::LowerGraph(g);
+    // Is this necessary?
+    lowering::LowerBlock(g->block());
+    return graph_and_parameters;
+}
+
+
 } // namespace lowering
 } // namespace core
 } // namespace trtorch

core/lowering/lowering.h

@@ -8,6 +8,9 @@ namespace lowering {
 
 void LowerBlock(torch::jit::Block* b);
 void LowerGraph(std::shared_ptr<torch::jit::Graph>& g);
+void LowerModule(const torch::jit::script::Module& mod);
+std::pair<std::shared_ptr<torch::jit::Graph>, std::vector<at::Tensor>> Lower(const torch::jit::script::Module& mod,
+                                                                             std::string method_name);
 
 } // namespace lowering
 } // namespace core