Fix: fix the bug that uninitialized tensor cannot be found

This PR fixes the bug that uninitialized tensor cannot be found.
Specifically, uninitialized value is copied between subgraphs
previously, this introduces unnecessary complexity and might incur
bugs. In this PR, this uninitialized value is handled like constants:
Create one uninitialized value for each subgraph in need.

Author: bowang007

core/partitioning/partitioning.cpp

@@ -70,7 +70,7 @@ void setExplicitFallbackNodes(PartitioningCtx* ctx, torch::jit::Block* block) {
   const auto to_compile_sym = c10::Symbol::attr("to_compile");
 
   for (const auto n : nodes) {
-    if (n->kind() == torch::jit::prim::Constant) {
+    if (isConstantOrUninitialized(n)) {
       continue;
     }
 
@@ -107,7 +107,7 @@ void setNonTensorConnectedNodes(PartitioningCtx* ctx, std::vector<torch::jit::No
     q.pop();
     // for every node that produces this fallback node's NonTensor input, they should fallback too
     for (auto input : cur_node->inputs()) {
-      if (!isTensor(input) && input->node()->kind() != torch::jit::prim::Constant &&
+      if (!isTensor(input) && !isConstantOrUninitialized(input->node()) &&
           ctx->shouldNodeRunInTensorRT(input->node())) {
         ctx->setNodeExecutorDecision(input->node(), NodeExecutorDecision::kNON_TENSOR);
         q.push(input->node());
@@ -118,7 +118,7 @@ void setNonTensorConnectedNodes(PartitioningCtx* ctx, std::vector<torch::jit::No
       if (!isTensor(output)) {
         for (auto use : output->uses()) {
           auto node = use.user;
-          if (node->kind() != torch::jit::prim::Constant && ctx->shouldNodeRunInTensorRT(node)) {
+          if (isConstantOrUninitialized(node) && ctx->shouldNodeRunInTensorRT(node)) {
             ctx->setNodeExecutorDecision(node, NodeExecutorDecision::kNON_TENSOR);
             q.push(node);
           }
@@ -128,11 +128,13 @@ void setNonTensorConnectedNodes(PartitioningCtx* ctx, std::vector<torch::jit::No
   }
 }
 
-std::set<torch::jit::Node*> getDependentNodes(torch::jit::Node* n) {
-  std::set<torch::jit::Node*> dependent_nodes;
+std::set<torch::jit::Node*> getUserNodes(torch::jit::Node* n) {
+  std::set<torch::jit::Node*> user_nodes;
   for (auto val : n->outputs()) {
     for (auto use : val->uses()) {
-      dependent_nodes.insert(use.user);
+      if (use.user->owningBlock()->owningNode())
+        user_nodes.insert(use.user->owningBlock()->owningNode());
+      user_nodes.insert(use.user);
     }
   }
   if (const auto* schema = n->maybeSchema()) {
@@ -142,13 +144,13 @@ std::set<torch::jit::Node*> getDependentNodes(torch::jit::Node* n) {
         for (auto use : n->inputs()[i]->uses()) {
           torch::jit::Node* use_node = use.user;
           if (use_node->isAfter(n)) {
-            dependent_nodes.insert(use_node);
+            user_nodes.insert(use_node);
           }
         }
       }
     }
   }
-  return dependent_nodes;
+  return user_nodes;
 }
 
 // Sub-function that traverses the entire block and check if TensorRT node sequence satisfy min_block_size
@@ -158,14 +160,14 @@ std::vector<torch::jit::Node*> traverseNodesForMinBlockSize(PartitioningCtx* ctx
   std::unordered_set<torch::jit::Node*> cur_trt_nodes_uses;
   std::vector<torch::jit::Node*> min_block_fallback_nodes;
   for (const auto n : nodes) {
-    if (n->kind() == torch::jit::prim::Constant) {
+    if (isConstantOrUninitialized(n)) {
       continue;
     }
 
     // check if current node fallback or not
     if (!ctx->shouldNodeRunInTorch(n)) {
       cur_trt_nodes.push_back(n);
-      auto dependent_nodes = getDependentNodes(n);
+      auto dependent_nodes = getUserNodes(n);
       cur_trt_nodes_uses.insert(dependent_nodes.begin(), dependent_nodes.end());
     } else {
       if (cur_trt_nodes_uses.count(n)) {
@@ -250,7 +252,7 @@ std::vector<torch::jit::Node*> getDependencyNodes(
     auto cur_val = q.front();
     q.pop();
     auto node = cur_val->node();
-    if (node->kind() != torch::jit::prim::Constant && !visited.count(node)) {
+    if (!isConstantOrUninitialized(node) && !visited.count(node)) {
       visited.insert(node);
       auto modifying_nodes = findModifyingNodes(cur_val, seg_block_nodes);
       stk.insert(stk.end(), modifying_nodes.rbegin(), modifying_nodes.rend());
@@ -454,10 +456,10 @@ void segmentGraph(PartitioningCtx* ctx, torch::jit::Block* block) {
   std::unordered_set<torch::jit::Node*> cur_pyt_nodes_uses;
   for (const auto n : nodes) {
     // Skip constant nodes as they are resources for both kinds of modules
-    if (n->kind() == torch::jit::prim::Constant) {
+    if (isConstantOrUninitialized(n)) {
       continue;
     }
-    auto dependent_nodes = getDependentNodes(n);
+    auto dependent_nodes = getUserNodes(n);
     // the outputs of trt subgraph shouldn't be collections
     if (ctx->shouldNodeRunInTensorRT(n)) {
       in_prog_trt_blk_nodes.push_back(n);

core/partitioning/partitioningctx/PartitioningCtx.cpp

@@ -15,13 +15,15 @@ PartitioningCtx::PartitioningCtx(torch::jit::Block* b, PartitioningInfo info)
 }
 
 void PartitioningCtx::_load_nodes_into_decision_map(torch::jit::Block* b) {
-  if (b->owningNode() && b->owningNode()->kind() == torch::jit::prim::Loop)
+  // won't load nodes if these nodes are in prim::loop or if these nodes are 2-level nested
+  if (b->owningNode() &&
+      (b->owningNode()->kind() == torch::jit::prim::Loop || b->owningNode()->owningBlock()->owningNode()))
     return;
 
   original_blocks.push_back(b);
 
   for (const auto n : b->nodes()) {
-    if (n->kind() == torch::jit::prim::Constant) {
+    if (isConstantOrUninitialized(n)) {
       continue;
     }
     node_executor_decision_map[n] = NodeExecutorDecision::kUNKNOWN;

core/partitioning/partitioningctx/PartitioningCtx.h

@@ -71,6 +71,10 @@ struct PartitioningCtx {
 
 std::ostream& operator<<(std::ostream& os, const PartitioningCtx& s);
 
+inline bool isConstantOrUninitialized(torch::jit::Node* n) {
+  return n->kind() == torch::jit::prim::Constant || n->kind() == torch::jit::prim::Uninitialized;
+}
+
 } // namespace partitioning
 } // namespace core
 } // namespace torch_tensorrt

core/partitioning/segmentedblock/SegmentedBlock.cpp

@@ -46,6 +46,12 @@ torch::jit::Value* SegmentedBlock::getOrAddInputForValue(torch::jit::Value* old_
       old_to_new_[old_value] = new_const->output();
       return new_const->output();
     }
+    if (node->kind() == torch::jit::prim::Uninitialized) {
+      auto new_uninitialized = g_->createUninitialized(old_value->type());
+      g_->block()->prependNode(new_uninitialized);
+      old_to_new_[old_value] = new_uninitialized->output();
+      return new_uninitialized->output();
+    }
     auto new_value = g_->block()->addInput();
     // every time when we addInput, we push back the corresponding lowering graph torch::jit::Value to our raw_inputs
     inputs_.push_back(old_value);

tests/core/partitioning/test_segmentation.cpp

@@ -331,6 +331,44 @@ TEST(Partitioning, SegmentModelWithDependencyAwareness) {
       checkSegmentedBlockNodesMapping(ctx.partitioned_blocks.begin()->second, g, {{0, 2, 4}, {1, 3, 5}, {6, 7}}));
 }
 
+TEST(Partitioning, ContainUninitializedValueCorrectly) {
+  auto g = std::make_shared<torch::jit::Graph>();
+  auto x = g->insertInput(0, "x");
+  auto none_const_val = g->insertConstant(torch::jit::IValue());
+  auto ivalue_1 = g->insertConstant(torch::jit::IValue(1));
+  auto ivalue_2 = g->insertConstant(torch::jit::IValue(2));
+
+  auto uninitialized_node = g->createUninitialized(torch::jit::BoolType::get());
+  g->appendNode(uninitialized_node);
+
+  auto x_dim = g->create(torch::jit::aten::dim, {x}, 1);
+  g->appendNode(x_dim);
+  x_dim->output()->setType(torch::jit::IntType::get());
+
+  auto eq1 = g->create(torch::jit::aten::eq, {ivalue_1, x_dim->output()}, 1);
+  g->appendNode(eq1);
+  eq1->output()->setType(torch::jit::BoolType::get());
+
+  torch::jit::IValue except("EXCEPTION");
+  auto exception_val = g->insertConstant(except);
+  auto if_node = g->create(torch::jit::prim::If, {eq1->output()}, 1);
+  auto if_block_0 = if_node->addBlock();
+  auto exception_node = g->create(torch::jit::prim::RaiseException, {exception_val, none_const_val}, 0);
+  if_block_0->appendNode(exception_node);
+  if_block_0->registerOutput(uninitialized_node->output());
+
+  auto if_block_1 = if_node->addBlock();
+  if_block_1->registerOutput(eq1->output());
+
+  g->insertNode(if_node);
+
+  PartitioningInfo partitioning_info;
+  partitioning_info.enabled = true;
+  PartitioningCtx ctx(g->block(), partitioning_info);
+  segmentGraph(&ctx, g->block());
+  ASSERT_TRUE(checkSegmentedBlockNumber(ctx.partitioned_blocks.begin()->second, SegmentedBlock::kTorch, 2));
+}
+
 } // namespace tests
 } // namespace partitioning
 } // namespace core