refactor: refactor RewriteInputsWithParams() to a lowering pass

Signed-off-by: Bo Wang <[email protected]>

Author: bowang007

core/compiler.cpp

@@ -10,7 +10,6 @@
 #include "ATen/core/jit_type.h"
 
 #include "torch/csrc/jit/frontend/function_schema_parser.h"
-#include "torch/csrc/jit/ir/constants.h"
 #include "torch/csrc/jit/ir/ir.h"
 #include "torch/csrc/jit/ir/ir_views.h"
 #include "torch/csrc/jit/passes/graph_fuser.h"
@@ -29,31 +28,6 @@
 namespace torch_tensorrt {
 namespace core {
 
-void RewriteInputsWithParams(std::shared_ptr<torch::jit::Graph> g, std::vector<torch::jit::IValue> params) {
-  auto input_size = g->inputs().size();
-  auto param_it = params.rbegin();
-  for (int i = input_size - 1; i >= 0; --i) {
-    if (g->inputs()[i]->type() != c10::TensorType::get() &&
-        g->inputs()[i]->type()->kind() != torch::jit::TypeKind::TupleType &&
-        g->inputs()[i]->type()->kind() != torch::jit::TypeKind::ListType && param_it != params.rend()) {
-      auto val = *param_it;
-      if (val.isTensor()) {
-        at::Tensor val_tensor = val.toTensor();
-        if (val_tensor.requires_grad()) {
-          val_tensor.set_requires_grad(false);
-          val = val_tensor;
-        }
-      }
-      auto new_constant = torch::jit::tryInsertConstant(*g, val);
-      ++param_it;
-      if (new_constant) {
-        g->inputs()[i]->replaceAllUsesWith(*new_constant);
-        g->eraseInput(i);
-      }
-    }
-  }
-}
-
 void AddEngineToGraph(
     torch::jit::script::Module mod,
     std::shared_ptr<torch::jit::Graph>& g,
@@ -460,9 +434,6 @@ torch::jit::Module CompileGraph(const torch::jit::Module& mod, CompileSpec cfg)
           (!(cfg.lower_info.forced_fallback_modules.size() == 0 &&
              cfg.partition_info.forced_fallback_operators.size() == 0 && isBlockConvertible) ||
            outputIsCollection)) {
-        if (!static_params.empty()) {
-          RewriteInputsWithParams(g, params);
-        }
         std::unordered_map<torch::jit::Node*, int> fallback_nodes;
         auto collection_input_ivalues_map =
             partitioning::generateRandomInputs(cfg.convert_info.collection_input_spec_map, first_use_types);

core/lowering/lowering.cpp

@@ -25,7 +25,7 @@ void LowerBlock(torch::jit::Block* b) {
   DropUnusedNodes(b);
 }
 
-void LowerGraph(std::shared_ptr<torch::jit::Graph>& g, LowerInfo lower_info) {
+void LowerGraph(std::shared_ptr<torch::jit::Graph>& g, std::vector<torch::jit::IValue>& params, LowerInfo lower_info) {
   torch::jit::EliminateRedundantGuards(g);
   torch::jit::RemoveListMutation(g);
   torch::jit::RemoveTensorMutation(g);
@@ -66,6 +66,7 @@ void LowerGraph(std::shared_ptr<torch::jit::Graph>& g, LowerInfo lower_info) {
   passes::SiluToSigmoidMultipication(g);
   passes::RemoveSingleUse0DTensors(g);
   passes::RemoveUnnecessaryCasts(g);
+  passes::RewriteInputsWithParams(g, params);
   LOG_GRAPH(*g);
 }
 
@@ -99,7 +100,7 @@ std::pair<std::shared_ptr<torch::jit::Graph>, std::vector<torch::jit::IValue>> L
   // In quantization aware trained (QAT) models, weights are passed through quantize and
   // dequantize nodes which should not be folded. So unfreeze_module is set to True for QAT models.
   LOG_GRAPH("Torch-TensorRT.TorchScript Graph Lowering");
-  lowering::LowerGraph(graph_and_ivalues.first, lower_info);
+  lowering::LowerGraph(graph_and_ivalues.first, graph_and_ivalues.second, lower_info);
 
   // Is this necessary?
   // lowering::LowerBlock(g->block());

core/lowering/passes/BUILD

@@ -27,6 +27,7 @@ cc_library(
         "remove_dropout.cpp",
         "remove_nops.cpp",
         "remove_unnecessary_casts.cpp",
+        "rewrite_inputs_with_params.cpp",
         "silu_to_sigmoid_multiplication.cpp",
         "unpack_addmm.cpp",
         "unpack_batch_norm.cpp",

core/lowering/passes/CMakeLists.txt

@@ -22,6 +22,7 @@ target_sources(${lib_name}
             "${CMAKE_CURRENT_SOURCE_DIR}/unpack_std.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/unpack_var.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/view_to_reshape.cpp"
+            "${CMAKE_CURRENT_SOURCE_DIR}/rewrite_inputs_with_params.cpp"
 )
 
 set(HEADER_FILES

core/lowering/passes/passes.h

@@ -38,6 +38,7 @@ void UnpackVar(std::shared_ptr<torch::jit::Graph>& graph);
 void AliasOperators(std::shared_ptr<torch::jit::Graph>& graph);
 void SiluToSigmoidMultipication(std::shared_ptr<torch::jit::Graph>& graph);
 void UnpackHardSwish(std::shared_ptr<torch::jit::Graph>& graph);
+void RewriteInputsWithParams(std::shared_ptr<torch::jit::Graph>& g, std::vector<torch::jit::IValue>& params);
 
 } // namespace passes
 } // namespace lowering

core/lowering/passes/rewrite_inputs_with_params.cpp

@@ -0,0 +1,41 @@
+#include "torch/csrc/jit/ir/constants.h"
+#include "core/util/prelude.h"
+
+
+namespace torch_tensorrt {
+namespace core {
+namespace lowering {
+namespace passes {
+
+
+void RewriteInputsWithParams(std::shared_ptr<torch::jit::Graph>& g, std::vector<torch::jit::IValue>& params) {
+  auto input_size = g->inputs().size();
+  auto param_it = params.rbegin();
+  for (int i = input_size - 1; i >= 0; --i) {
+    if (g->inputs()[i]->type() != c10::TensorType::get() &&
+        g->inputs()[i]->type()->kind() != torch::jit::TypeKind::TupleType &&
+        g->inputs()[i]->type()->kind() != torch::jit::TypeKind::ListType && param_it != params.rend()) {
+      auto val = *param_it;
+      if (val.isTensor()) {
+        at::Tensor val_tensor = val.toTensor();
+        if (val_tensor.requires_grad()) {
+          val_tensor.set_requires_grad(false);
+          val = val_tensor;
+        }
+      }
+      auto new_constant = torch::jit::tryInsertConstant(*g, val);
+      ++param_it;
+      if (new_constant) {
+        g->inputs()[i]->replaceAllUsesWith(*new_constant);
+        g->eraseInput(i);
+        // erase an iterator, should be safe
+        params.erase(param_it.base());
+      }
+    }
+  }
+}
+
+} // namespace passes
+} // namespace lowering
+} // namespace core
+} // namespace torch_tensorrt