refactor: extract stitching phase out of compiler.cpp

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

Author: bowang007

core/compiler.cpp

@@ -11,7 +11,6 @@
 
 #include "torch/csrc/jit/frontend/function_schema_parser.h"
 #include "torch/csrc/jit/ir/ir.h"
-#include "torch/csrc/jit/ir/ir_views.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"
@@ -128,193 +127,54 @@ bool CheckMethodOperatorSupport(const torch::jit::script::Module& mod, std::stri
   return conversion::VerifyConverterSupportForBlock(g->block());
 }
 
-void AddSegmentedBlockToGraph(
-    std::shared_ptr<torch::jit::Graph>& g,
-    partitioning::SegmentedBlock& seg,
-    std::unordered_map<torch::jit::Value*, torch::jit::Value*>& old_to_new_g) {
-  // old_to_new_g contains: original global graph value => new global graph value,
-  // mini_to_new_g: mini graph value -> new graph value
-  std::unordered_map<torch::jit::Value*, torch::jit::Value*> mini_to_new_g;
-  size_t input_idx = 0;
-  if (seg.target() == partitioning::SegmentedBlock::kTensorRT && g->inputs().size() > 0) {
-    if (g->inputs()[0]->type()->str().find("__torch__") == std::string::npos) {
-      auto self = g->insertInput(0, "self_1");
-      self->setType(seg.inputs()[0]->type());
-    }
-    mini_to_new_g[seg.inputs()[input_idx++]] = g->inputs()[0];
-  }
-
-  for (auto& raw_input : seg.raw_inputs()) {
-    if (old_to_new_g.count(raw_input)) {
-      mini_to_new_g[seg.inputs()[input_idx++]] = old_to_new_g[raw_input];
-    }
-  }
-
-  for (const auto n : seg.nodes()) {
-    util::cloneNode(n, g, mini_to_new_g);
-  }
-
-  // original graph value => new global graph value
-  for (size_t i = 0; i < seg.raw_outputs().size(); ++i) {
-    old_to_new_g[seg.raw_outputs()[i]] = mini_to_new_g[seg.outputs()[i]];
-  }
-  size_t offset = seg.target() == partitioning::SegmentedBlock::kTensorRT ? 1 : 0;
-  for (size_t i = 0; i < seg.raw_inputs().size(); ++i) {
-    if (!old_to_new_g.count(seg.raw_inputs()[i])) {
-      old_to_new_g[seg.raw_inputs()[i]] = mini_to_new_g[seg.inputs()[i + offset]];
-    }
-  }
-
-  return;
-}
-
-typedef std::pair<std::shared_ptr<torch::jit::Graph>, std::unordered_map<torch::jit::Value*, torch::jit::Value*>>
-    GraphAndMapping;
-
-void AddIfBlockToGraph(
-    std::shared_ptr<torch::jit::Graph>& new_g,
-    torch::jit::Node* if_node,
-    const std::vector<GraphAndMapping>& graph_and_mappings,
-    std::unordered_map<torch::jit::Value*, torch::jit::Value*>& old_to_new_g) {
-  torch::jit::IfView if_view(if_node);
-
-  // create a new if node in new_g and add corresponding inputs
-  auto new_if = new_g->insertNode(new_g->create(torch::jit::prim::If, {}, 0));
-  new_if->addInput(util::getOrAddInputForValue(if_view.cond(), new_g, old_to_new_g));
-
-  // iterate over all blocks and add them to new created prim::If
-  for (auto graph_and_mapping : graph_and_mappings) {
-    auto new_if_block = new_if->addBlock();
-    auto cur_block_graph = graph_and_mapping.first;
-    auto cur_block_mapping = graph_and_mapping.second;
-    std::unordered_map<torch::jit::Value*, torch::jit::Value*> block_graph_to_new_g;
-    for (auto& i : cur_block_mapping) {
-      // for every pair in then_mapping, old_value => mini graph value, if old_value also appears in old_to_new_g, then
-      // it's mini graph's input
-      if (old_to_new_g.count(i.first)) {
-        block_graph_to_new_g[i.second] = old_to_new_g[i.first];
-      }
-    }
-
-    auto env = [&](torch::jit::Value* v) { return util::getOrAddInputForValue(v, new_g, block_graph_to_new_g); };
-    new_if_block->cloneFrom(cur_block_graph->block(), env);
-    if (cur_block_graph->inputs().size() &&
-        cur_block_graph->inputs()[0]->type()->str().find("__torch__") != std::string::npos) {
-      if (new_g->inputs()[0]->type()->str().find("__torch__") == std::string::npos) {
-        auto self = new_g->insertInput(0, "self_1");
-        self->setType(cur_block_graph->inputs()[0]->type());
-      }
-      block_graph_to_new_g[cur_block_graph->inputs()[0]] = new_g->inputs()[0];
-    }
-    for (int i = cur_block_graph->inputs().size() - 1; i >= 0; --i) {
-      new_if_block->inputs()[i]->replaceAllUsesWith(block_graph_to_new_g[cur_block_graph->inputs()[i]]);
-      new_if_block->eraseInput(i);
-    }
-  }
-  for (auto ov : if_view.outputs()) {
-    auto no = new_if->addOutput();
-    old_to_new_g[ov] = no;
-    no->copyMetadata(ov);
-  }
-  return;
-}
-
-GraphAndMapping ConstructFallbackGraph_(
+partitioning::GraphAndMapping BuildHybridGraph(
     torch::jit::script::Module& new_mod,
     torch::jit::Block* block,
-    partitioning::PartitioningCtx* partitioning_ctx,
-    conversion::ConversionInfo convert_info,
+    CompileSpec cfg,
     ir::StaticParams static_params,
-    std::unordered_map<const torch::jit::Value*, torch::jit::IValue> example_tensor_map) {
-  auto new_g = std::make_shared<torch::jit::Graph>();
+    ir::CollectionTypeMap first_use_types) {
+  auto convert_info = cfg.convert_info;
+  auto partitioning_info = cfg.partitioning_info;
+
+  auto partitioning_ctx = partitioning::PartitioningCtx(block, partitioning_info);
+  auto collection_input_ivalues_map =
+      partitioning::generateRandomInputs(partitioning_info.collection_input_spec_map, first_use_types);
 
-  auto segmented_blocks = partitioning::Partition(partitioning_ctx, block, example_tensor_map);
+  partitioning::Partition(&partitioning_ctx, collection_input_ivalues_map);
 
-  // the mapping from lowering graph => fallback global graph
-  std::unordered_map<torch::jit::Value*, torch::jit::Value*> old_to_new_g;
-  for (auto input : block->inputs()) {
-    util::getOrAddInputForValue(input, new_g, old_to_new_g);
-  }
+  for (auto &partitioned_block : partitioning_ctx.partitioned_blocks) {
+    partitioning::PartitionedGraph& segmented_blocks = partitioned_block.second;
 
-  for (auto& seg_block : segmented_blocks) {
-    LOG_INFO("Block segment:" << seg_block);
-    std::ostringstream trt_engine_id;
-    trt_engine_id << reinterpret_cast<const int*>(&seg_block);
-
-    if (seg_block.target() == partitioning::SegmentedBlock::kTensorRT) {
-      auto shapes = seg_block.in_shapes();
-      auto types = seg_block.in_types();
-      std::vector<ir::Input> inputs;
-      for (size_t i = 0; i < shapes.size(); i++) {
-        auto in = ir::Input(shapes[i]);
-        in.dtype = util::ScalarTypeToTRTDataType(types[i]);
-        inputs.push_back(in);
-      }
-      // update the input ranges for each segments
-      convert_info.inputs = ir::associate_specs_with_inputs(seg_block.g(), inputs, static_params);
-
-      // TODO mapping Inputs Ivalue to flatten one here
-      auto engine = conversion::ConvertBlockToEngine(seg_block.block(), convert_info, static_params);
-      auto temp_g = std::make_shared<torch::jit::Graph>();
-      auto device_spec = convert_info.engine_settings.device;
-      auto cuda_device = runtime::CudaDevice(device_spec.gpu_id, device_spec.device_type);
-      AddEngineToGraph(new_mod, temp_g, engine, cuda_device, trt_engine_id.str(), true);
-
-      seg_block.update_graph(temp_g);
-      AddSegmentedBlockToGraph(new_g, seg_block, old_to_new_g);
-    } else {
-      if (seg_block.raw_nodes()[0]->kind() == torch::jit::prim::If) {
-        auto if_node = seg_block.raw_nodes()[0];
-
-        // convert the 2 blocks in prim::if and get the converted graph with mappings
-        std::vector<GraphAndMapping> graph_and_mappings;
-        for (auto cur_block : if_node->blocks()) {
-          graph_and_mappings.push_back(ConstructFallbackGraph_(
-              new_mod, cur_block, partitioning_ctx, convert_info, static_params, example_tensor_map));
+    for (auto& seg_block : segmented_blocks) {
+      LOG_INFO("Block segment:" << seg_block);
+      std::ostringstream trt_engine_id;
+      trt_engine_id << reinterpret_cast<const int*>(&seg_block);
+
+      if (seg_block.target() == partitioning::SegmentedBlock::kTensorRT) {
+        auto shapes = seg_block.in_shapes();
+        auto types = seg_block.in_types();
+        std::vector<ir::Input> inputs;
+        for (size_t i = 0; i < shapes.size(); i++) {
+          auto in = ir::Input(shapes[i]);
+          in.dtype = util::ScalarTypeToTRTDataType(types[i]);
+          inputs.push_back(in);
         }
-        AddIfBlockToGraph(new_g, if_node, graph_and_mappings, old_to_new_g);
+        // update the input ranges for each segments
+        convert_info.inputs = ir::associate_specs_with_inputs(seg_block.g(), inputs, static_params);
 
-      } else {
-        AddSegmentedBlockToGraph(new_g, seg_block, old_to_new_g);
-      }
-    }
-  }
+        // TODO mapping Inputs Ivalue to flatten one here
+        auto engine = conversion::ConvertBlockToEngine(seg_block.block(), convert_info, static_params);
+        auto temp_g = std::make_shared<torch::jit::Graph>();
+        auto device_spec = convert_info.engine_settings.device;
+        auto cuda_device = runtime::CudaDevice(device_spec.gpu_id, device_spec.device_type);
+        AddEngineToGraph(new_mod, temp_g, engine, cuda_device, trt_engine_id.str(), true);
 
-  if (block->outputs().size() > 1) {
-    std::vector<torch::jit::Value*> fallback_graph_vector;
-    for (auto& output : block->outputs()) {
-      if (old_to_new_g.count(output)) {
-        fallback_graph_vector.push_back(old_to_new_g[output]);
+        seg_block.update_graph(temp_g);
       }
     }
-    torch::jit::ArrayRef<torch::jit::Value*> fallback_graph_outputs(fallback_graph_vector);
-    auto return_tuple_node = new_g->createTuple(fallback_graph_outputs);
-    new_g->block()->appendNode(return_tuple_node);
-    // Set the output as the produced tuple
-    new_g->registerOutput(return_tuple_node->outputs()[0]);
-  } else {
-    if (block->outputs().size() && old_to_new_g.count(block->outputs()[0])) {
-      new_g->registerOutput(old_to_new_g[block->outputs()[0]]);
-    }
   }
-  return {new_g, old_to_new_g};
-}
-
-GraphAndMapping ConstructFallbackGraph(
-    torch::jit::script::Module& new_mod,
-    torch::jit::Block* block,
-    CompileSpec cfg,
-    ir::StaticParams static_params,
-    ir::CollectionTypeMap first_use_types) {
-  auto convert_info = cfg.convert_info;
-  auto partitioning_info = cfg.partitioning_info;
-
-  auto partitioning_ctx = partitioning::PartitioningCtx(block, partitioning_info);
-  auto collection_input_ivalues_map =
-      partitioning::generateRandomInputs(partitioning_info.collection_input_spec_map, first_use_types);
 
-  return ConstructFallbackGraph_(
-      new_mod, block, &partitioning_ctx, convert_info, static_params, collection_input_ivalues_map);
+  return partitioning::Stitch(&partitioning_ctx, block);
 }
 
 void MapInputsAndDetermineDTypes(
@@ -451,7 +311,7 @@ torch::jit::Module CompileGraph(const torch::jit::Module& mod, CompileSpec cfg)
           (!(cfg.lower_info.forced_fallback_modules.size() == 0 &&
              cfg.partitioning_info.forced_fallback_operators.size() == 0 && isBlockConvertible) ||
            outputIsCollection)) {
-        auto graph_and_mapping = ConstructFallbackGraph(new_mod, g->block(), cfg, static_params, first_use_types);
+        auto graph_and_mapping = BuildHybridGraph(new_mod, g->block(), cfg, static_params, first_use_types);
         new_g = graph_and_mapping.first;
         // renaming the input name of graph after fallback to ensure pytorch deserialize it correctly
         for (size_t i = 0; i < new_g->inputs().size(); ++i) {

core/partitioning/BUILD

@@ -15,6 +15,7 @@ cc_library(
     srcs = [
         "partitioning.cpp",
         "shape_analysis.cpp",
+        "stitching.cpp"
     ],
     hdrs = [
         "partitioning.h",