chore: refactor code structures according to PR

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

Author: bowang007

core/compiler.cpp

@@ -12,6 +12,7 @@
 #include "torch/csrc/jit/frontend/function_schema_parser.h"
 #include "torch/csrc/jit/ir/ir.h"
 #include "torch/csrc/jit/passes/graph_fuser.h"
+#include "torch/csrc/jit/passes/loop_unrolling.h"
 #include "torch/csrc/jit/passes/lower_graph.h"
 #include "torch/csrc/jit/passes/pass_manager.h"
 #include "torch/custom_class.h"
@@ -30,10 +31,9 @@ void AddEngineToGraph(
     torch::jit::script::Module mod,
     std::shared_ptr<torch::jit::Graph>& g,
     const std::string& serialized_engine,
-    int engine_id = 0,
+    std::string engine_id = "",
     bool fallback = false) {
-  auto engine_ptr =
-      c10::make_intrusive<runtime::TRTEngine>(mod._ivalue()->name() + std::to_string(engine_id), serialized_engine);
+  auto engine_ptr = c10::make_intrusive<runtime::TRTEngine>(mod._ivalue()->name() + engine_id, serialized_engine);
   // Get required metadata about the engine out
   auto num_io = engine_ptr->num_io;
   auto name = engine_ptr->name;
@@ -200,14 +200,17 @@ torch::jit::script::Module CompileGraphWithFallback(const torch::jit::script::Mo
         return mod;
       }
 
-      int trt_engine_id = 0;
       std::unordered_map<torch::jit::Value*, torch::jit::Value*> old_to_new_g;
       // add global graph's input to old_to_new_g mapping
       for (auto input : g->inputs()) {
         util::getOrAddInputForValue(input, new_g, old_to_new_g);
       }
       for (auto& seg_block : segmented_blocks) {
-        LOG_INFO(*g << "(MiniGraphInSegmentedBlock)\n");
+        std::string cur_block_target =
+            seg_block.target() == partitioning::SegmentedBlock::kTensorRT ? "TensorRT" : "Torch";
+        LOG_INFO(*g << "(MiniGraphIn" << cur_block_target << "Block\n");
+        std::ostringstream trt_engine_id;
+        trt_engine_id << reinterpret_cast<const int*>(&seg_block);
         if (seg_block.target() == partitioning::SegmentedBlock::kTensorRT) {
           std::vector<ir::InputRange> input_ranges;
           for (auto& shape : seg_block.in_shape()) {
@@ -217,7 +220,7 @@ torch::jit::script::Module CompileGraphWithFallback(const torch::jit::script::Mo
           convert_cfg.input_ranges = input_ranges;
           auto engine = conversion::ConvertBlockToEngine(seg_block.block(), convert_cfg, named_params);
           auto temp_g = std::make_shared<torch::jit::Graph>();
-          AddEngineToGraph(new_mod, temp_g, engine, trt_engine_id++, true);
+          AddEngineToGraph(new_mod, temp_g, engine, trt_engine_id.str(), true);
 
           seg_block.update_graph(temp_g);
           AddSegmentedBlockToGraph(new_g, seg_block, old_to_new_g);

core/partitioning/partitioning.cpp

@@ -3,7 +3,6 @@
 #include <queue>
 #include "core/conversion/conversion.h"
 #include "core/partitioning/shape_analysis.h"
-#include "torch/csrc/jit/passes/constant_pooling.h"
 #include "torch/csrc/jit/passes/dead_code_elimination.h"
 
 namespace trtorch {
@@ -30,7 +29,7 @@ bool isAllNodesSupported(const std::vector<torch::jit::Node*>& nodes) {
   return true;
 }
 
-bool containNonTensorInputs(torch::jit::Node* n, const std::unordered_set<torch::jit::Value*>& target_inputs) {
+bool containTargetInputs(torch::jit::Node* n, const std::unordered_set<torch::jit::Value*>& target_inputs) {
   for (auto input : n->inputs()) {
     if (!isTensorOrTensorList(input) && target_inputs.count(input)) {
       return true;
@@ -94,7 +93,7 @@ std::vector<SegmentedBlock> injectNodesForNonTensorInputs(SegmentedBlock& seg_bl
     bool prev_non_tensor_outputs = false;
     for (auto n : seg_block.raw_nodes()) {
       // it's a kTorch block if it uses the nonTensor input and the nonTensor input is produced in kTorch block
-      if (containNonTensorInputs(n, nontensor_inputs_set) || prev_non_tensor_outputs) {
+      if (containTargetInputs(n, nontensor_inputs_set) || prev_non_tensor_outputs) {
         if (!tensorrt_nodes.empty()) {
           new_seg_blocks.emplace_back(SegmentedBlock::kTensorRT, tensorrt_nodes);
           tensorrt_nodes.clear();
@@ -278,18 +277,8 @@ std::vector<SegmentedBlock> Partition(
   // register input/output torch::jit::Value for segmented graphs
   registerSegmentsOutputs(segmented_blocks, g);
 
-  // store the mapping from lowering graph torch::jit::Value => torch::jit::IValue that we get by running segments
-  std::unordered_map<torch::jit::Value*, torch::jit::IValue> ivalues_maps;
-  std::vector<torch::jit::IValue> random_inputs = generateRandomInputs(input_ranges);
-  for (size_t i = 0; i < g->inputs().size(); ++i) {
-    ivalues_maps[g->inputs()[i]] = random_inputs[i];
-  }
-
-  // register every segment's input shape, and it's running output IValues
-  for (auto& seg_block : segmented_blocks) {
-    torch::jit::ConstantPooling(seg_block.g());
-    getSegmentsOutputByRunning(seg_block, ivalues_maps);
-  }
+  // run shape analysis on each segmented block
+  runShapeAnalysis(segmented_blocks, input_ranges, g);
 
   return segmented_blocks;
 }

core/partitioning/shape_analysis.cpp

@@ -1,6 +1,7 @@
 #include "core/partitioning/shape_analysis.h"
 #include "core/util/prelude.h"
 #include "torch/csrc/jit/api/module.h"
+#include "torch/csrc/jit/passes/constant_pooling.h"
 
 namespace trtorch {
 namespace core {
@@ -97,6 +98,25 @@ void getSegmentsOutputByRunning(
   seg_block.register_inshape(input_shape);
 }
 
+void runShapeAnalysis(
+    std::vector<SegmentedBlock>& segmented_blocks,
+    std::vector<ir::InputRange>& input_ranges,
+    std::shared_ptr<torch::jit::Graph> g) {
+  // store the mapping from lowering graph torch::jit::Value => torch::jit::IValue that we get by running segments
+  std::unordered_map<torch::jit::Value*, torch::jit::IValue> ivalues_maps;
+  std::vector<torch::jit::IValue> random_inputs = generateRandomInputs(input_ranges);
+  for (size_t i = 0; i < g->inputs().size(); ++i) {
+    ivalues_maps[g->inputs()[i]] = random_inputs[i];
+  }
+
+  // register every segment's input shape, and it's running output IValues
+  for (auto& seg_block : segmented_blocks) {
+    torch::jit::ConstantPooling(seg_block.g());
+    getSegmentsOutputByRunning(seg_block, ivalues_maps);
+  }
+  return;
+}
+
 } // namespace partitioning
 } // namespace core
 } // namespace trtorch

core/partitioning/shape_analysis.h

@@ -1,15 +1,17 @@
 #include "core/ir/ir.h"
 #include "core/partitioning/SegmentedBlock.h"
+#include "torch/csrc/jit/ir/ir.h"
 
 namespace trtorch {
 namespace core {
 namespace partitioning {
 
 std::vector<torch::jit::IValue> generateRandomInputs(std::vector<ir::InputRange>& input_ranges);
 
-void getSegmentsOutputByRunning(
-    SegmentedBlock& seg_block,
-    std::unordered_map<torch::jit::Value*, torch::jit::IValue>& ivalues_maps);
+void runShapeAnalysis(
+    std::vector<SegmentedBlock>& segmented_blocks,
+    std::vector<ir::InputRange>& input_ranges,
+    std::shared_ptr<torch::jit::Graph> g);
 
 } // namespace partitioning
 } // namespace core