chore: optimize minor code problems according to PR

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

Author: bowang007

core/compiler.cpp

@@ -17,7 +17,6 @@
 #include "torch/custom_class.h"
 
 #include "core/compiler.h"
-#include "core/util/prelude.h"
 
 #include "core/conversion/conversion.h"
 #include "core/lowering/lowering.h"
@@ -31,7 +30,8 @@ void AddEngineToGraph(
     torch::jit::script::Module mod,
     std::shared_ptr<torch::jit::Graph>& g,
     const std::string& serialized_engine,
-    int engine_id = 0) {
+    int engine_id = 0,
+    bool fallback = false) {
   auto engine_ptr =
       c10::make_intrusive<runtime::TRTEngine>(mod._ivalue()->name() + std::to_string(engine_id), serialized_engine);
   // Get required metadata about the engine out
@@ -96,7 +96,7 @@ void AddEngineToGraph(
 
   // If there are multiple output tensors from TensorRT we wrap them in a tuple
   // to return, convert to tuple only when we only have 1 segmented graph
-  if (!engine_id && unpack_node->outputs().size() > 1) {
+  if (!fallback && unpack_node->outputs().size() > 1) {
     // Creates prim::TupleConstruct(<output tensors>) using outputs of the
     // unpack node
     auto return_tuple_node = g->createTuple(unpack_node->outputs());
@@ -196,10 +196,11 @@ torch::jit::script::Module CompileGraphWithFallback(const torch::jit::script::Mo
       // segment the graph and convert segmented TensorRT block
       auto segmented_blocks = partitioning::Partition(g, convert_cfg.input_ranges, cfg.partition_info);
       if (segmented_blocks.size() == 1 && segmented_blocks[0].target() == partitioning::SegmentedBlock::kTorch) {
+        LOG_WARNING("Didn't generate any TensorRT engines, the compiler did nothing\n");
         return mod;
       }
 
-      int trt_engine_id = 1;
+      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()) {
@@ -216,7 +217,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++);
+          AddEngineToGraph(new_mod, temp_g, engine, trt_engine_id++, true);
 
           seg_block.update_graph(temp_g);
           AddSegmentedBlockToGraph(new_g, seg_block, old_to_new_g);

core/partitioning/SegmentedBlock.cpp

@@ -34,6 +34,7 @@ torch::jit::Value* SegmentedBlock::getOrAddInputForValue(torch::jit::Value* old_
     if (node->kind() == torch::jit::prim::Constant) {
       auto new_const = g_->createClone(node, {nullptr});
       g_->block()->prependNode(new_const);
+      old_to_new_[old_value] = new_const->output();
       return new_const->output();
     }
     auto new_value = g_->block()->addInput();

core/partitioning/SegmentedBlock.h

@@ -76,7 +76,7 @@ struct SegmentedBlock {
 
  private:
   SegmentedBlockTarget target_;
-  std::vector<ir::InputRange> in_shape_; // REVIEW: This should just be ir::InputRange
+  std::vector<ir::InputRange> in_shape_;
   std::vector<torch::jit::Value*> inputs_;
   std::vector<torch::jit::Value*> outputs_;
   std::vector<torch::jit::Node*> nodes_;

core/partitioning/partitioning.cpp

@@ -4,6 +4,7 @@
 #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 {
 namespace core {
@@ -203,13 +204,17 @@ void registerSegmentsOutputs(PartitionedGraph& segmented_blocks, std::shared_ptr
       }
     }
   }
-  // erase segments which still have no output
-  segmented_blocks.erase(
-      std::remove_if(
-          segmented_blocks.begin(),
-          segmented_blocks.end(),
-          [](SegmentedBlock& seg_block) { return seg_block.raw_outputs().empty(); }),
-      segmented_blocks.end());
+  std::for_each(
+      segmented_blocks.begin(),
+      segmented_blocks.end(),
+      [](SegmentedBlock& seg_block) { torch::jit::EliminateDeadCode(seg_block.g()); })
+      // erase segments which still have no output
+      segmented_blocks.erase(
+          std::remove_if(
+              segmented_blocks.begin(),
+              segmented_blocks.end(),
+              [](SegmentedBlock& seg_block) { return seg_block.raw_outputs().empty(); }),
+          segmented_blocks.end());
 
   return;
 }
@@ -225,8 +230,9 @@ std::vector<SegmentedBlock> segment_graph(std::shared_ptr<torch::jit::Graph> g,
   // segment the nodes
   std::vector<torch::jit::Node*> tensorrt_nodes, pytorch_nodes;
   for (const auto n : nodes) {
-    if (n->kind() == torch::jit::prim::Constant)
+    if (n->kind() == torch::jit::prim::Constant) {
       continue;
+    }
 
     std::string node_string(n->kind().toQualString());
     if (conversion::OpSupported(n) && !forced_fallback_operators.count(node_string)) {

core/partitioning/shape_analysis.cpp

@@ -61,9 +61,10 @@ void getSegmentsOutputByRunning(
       jit_inputs_ivalues.push_back(ivalues_maps[input].toBool());
     } else if (input->type()->kind() == torch::jit::TypeKind::ListType) {
       jit_inputs_ivalues.push_back(ivalues_maps[input].toList());
-    } else {
-      TRTORCH_CHECK(input->type()->kind() == torch::jit::TypeKind::TupleType, "Input for mini graph is not TupleType.");
+    } else if (input->type()->kind() == torch::jit::TypeKind::TupleType){
       jit_inputs_ivalues.push_back(ivalues_maps[input].toTuple());
+    } else {
+      TRTORCH_THROW_ERROR("Unable to find type for value: " << input->debugName() << " to get the ivalues.\n");
     }
   }
 

core/util/jit_util.h

@@ -33,6 +33,20 @@ inline std::vector<int64_t> toVec(c10::IntArrayRef a) {
   return arr;
 }
 
+inline c10::FunctionSchema GenerateGraphSchema(std::string method_name, std::shared_ptr<torch::jit::Graph>& g) {
+  std::vector<c10::Argument> args;
+  for (auto in : g->inputs()) {
+    args.push_back(c10::Argument(in->debugName(), in->type()));
+  }
+
+  std::vector<c10::Argument> returns;
+  for (auto out : g->outputs()) {
+    returns.push_back(c10::Argument(out->debugName(), out->type()));
+  }
+
+  return c10::FunctionSchema(method_name, method_name, args, returns);
+}
+
 } // namespace util
 } // namespace core
 } // namespace trtorch

core/util/trt_util.cpp

@@ -350,20 +350,6 @@ c10::optional<nvinfer1::DataType> toTRTDataType(caffe2::TypeMeta dtype) {
   }
 }
 
-c10::FunctionSchema GenerateGraphSchema(std::string method_name, std::shared_ptr<torch::jit::Graph>& g) {
-  std::vector<c10::Argument> args;
-  for (auto in : g->inputs()) {
-    args.push_back(c10::Argument(in->debugName(), in->type()));
-  }
-
-  std::vector<c10::Argument> returns;
-  for (auto out : g->outputs()) {
-    returns.push_back(c10::Argument(out->debugName(), out->type()));
-  }
-
-  return c10::FunctionSchema(method_name, method_name, args, returns);
-}
-
 torch::jit::Value* getOrAddInputForValue(
     torch::jit::Value* old_value,
     std::shared_ptr<torch::jit::Graph>& graph,

core/util/trt_util.h

@@ -109,7 +109,6 @@ std::string toStr(nvinfer1::Dims d);
 at::ScalarType toATenDType(nvinfer1::DataType t);
 nvinfer1::DataType toTRTDataType(at::ScalarType t);
 c10::optional<nvinfer1::DataType> toTRTDataType(caffe2::TypeMeta dtype);
-c10::FunctionSchema GenerateGraphSchema(std::string method_name, std::shared_ptr<torch::jit::Graph>& g);
 torch::jit::Value* getOrAddInputForValue(
     torch::jit::Value* old_value,
     std::shared_ptr<torch::jit::Graph>& graph,