fix: make sure that prim::if is in raw_nodes()[0] in dependency analysis

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

Author: bowang007

core/compiler.cpp

@@ -196,14 +196,15 @@ void AddIfBlockToGraph(
   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 => then value, if old_value also appears in old_to_new_g, then it's
-      // then graph's input
+      // 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];
       }
@@ -214,7 +215,7 @@ void AddIfBlockToGraph(
     if (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(loop_graph->inputs()[0]->type());
+        self->setType(cur_block_graph->inputs()[0]->type());
       }
       block_graph_to_new_g[cur_block_graph->inputs()[0]] = new_g->inputs()[0];
     }

core/partitioning/partitioning.cpp

@@ -84,8 +84,14 @@ std::vector<SegmentedBlock> injectNodesForNonTensorInputs(SegmentedBlock& seg_bl
   // if current block is kTorch or current block is TensorRT and all dependent nodes are also supported, construct only
   // one new block
   if (seg_block.target() == SegmentedBlock::kTorch || isAllNodesSupported(dependency_nodes)) {
-    dependency_nodes.insert(dependency_nodes.end(), seg_block.raw_nodes().begin(), seg_block.raw_nodes().end());
-    new_seg_blocks.emplace_back(seg_block.target(), dependency_nodes);
+    // if current node is prim::If, just ensure that we have all required input in kTorch
+    if (seg_block.raw_nodes()[0]->kind() == torch::jit::prim::If) {
+      new_seg_blocks.emplace_back(seg_block.target(), dependency_nodes);
+      new_seg_blocks.push_back(seg_block);
+    } else {
+      dependency_nodes.insert(dependency_nodes.end(), seg_block.raw_nodes().begin(), seg_block.raw_nodes().end());
+      new_seg_blocks.emplace_back(seg_block.target(), dependency_nodes);
+    }
   } else {
     // if current block is kTensorRT but the dependency nodes contain unsupported node, then we have to segment again
     std::unordered_set<torch::jit::Value*> nontensor_inputs_set(nontensor_inputs.begin(), nontensor_inputs.end());
@@ -141,8 +147,9 @@ void resolveNonTensorInputs(PartitionedGraph& segmented_blocks) {
     if (segmented_blocks[use_info.produce_id].target() == SegmentedBlock::kTensorRT && !use_info.torch_use_id.empty()) {
       int first_torch_id = use_info.torch_use_id.front();
       if (!updated_segments.count(first_torch_id)) {
-        auto new_torch_block = injectNodesForNonTensorInputs(segmented_blocks[first_torch_id]).front();
-        segmented_blocks[first_torch_id] = new_torch_block;
+        auto to_inject_blocks = injectNodesForNonTensorInputs(segmented_blocks[first_torch_id]);
+        segmented_blocks.erase(segmented_blocks.begin() + first_torch_id);
+        segmented_blocks.insert(segmented_blocks.begin() + first_torch_id, to_inject_blocks.begin(), to_inject_blocks.end());
         updated_segments.insert(first_torch_id);
       }
     } else {