fix: change the shouldNodeRunInTorch logic

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

Author: bowang007

core/partitioning/partitioning.cpp

@@ -2,7 +2,6 @@
 #include <queue>
 #include "core/conversion/conversion.h"
 #include "core/conversion/evaluators/evaluators.h"
-#include "core/partitioning/partitioningctx/PartitioningCtx.h"
 #include "torch/csrc/jit/passes/constant_pooling.h"
 #include "torch/csrc/jit/passes/dead_code_elimination.h"
 
@@ -102,8 +101,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 && node->kind() != torch::jit::prim::Return &&
-              ctx->shouldNodeRunInTensorRT(node)) {
+          if (node->kind() != torch::jit::prim::Constant && ctx->shouldNodeRunInTensorRT(node)) {
             ctx->setNodeExecutorDecision(node, NodeExecutorDecision::kNON_TENSOR);
             q.push(node);
           }
@@ -454,7 +452,6 @@ void Partition(PartitioningCtx* ctx, ExampleIValues& example_tensor_map) {
     LOG_DEBUG("Registering input/output torch::jit::Value for segmented graphs");
     RegisterSegmentsOutputs(ctx, block);
 
-
     // run shape analysis on each segmented block
     RunShapeAnalysis(ctx, block, example_tensor_map);
   }

core/partitioning/partitioning.h

@@ -7,8 +7,6 @@
 
 #include "core/ir/ir.h"
 #include "core/partitioning/partitioningctx/PartitioningCtx.h"
-#include "core/partitioning/partitioninginfo/PartitioningInfo.h"
-#include "core/partitioning/segmentedblock/SegmentedBlock.h"
 #include "core/util/prelude.h"
 
 namespace torch_tensorrt {

core/partitioning/partitioningctx/PartitioningCtx.cpp

@@ -43,14 +43,30 @@ void PartitioningCtx::setNodeExecutorDecision(torch::jit::Node* n, NodeExecutorD
 
 bool PartitioningCtx::shouldNodeRunInTorch(torch::jit::Node* n) {
   auto iter = node_executor_decision_map.find(n);
-  if (iter == node_executor_decision_map.end()) {
-    LOG_ERROR("No info about node " << *n << " execution decision status.");
+  auto decision = NodeExecutorDecision::kUNKNOWN;
+
+  if (iter != node_executor_decision_map.end()) {
+    decision = iter->second;
+  }
+  if (decision == NodeExecutorDecision::kCONVERT || decision == NodeExecutorDecision::kUNKNOWN) {
+    return false;
+  } else {
+    return true;
   }
-  return iter->second != NodeExecutorDecision::kCONVERT;
 }
 
 bool PartitioningCtx::shouldNodeRunInTensorRT(torch::jit::Node* n) {
-  return !shouldNodeRunInTorch(n);
+  auto iter = node_executor_decision_map.find(n);
+  auto decision = NodeExecutorDecision::kUNKNOWN;
+  if (iter != node_executor_decision_map.end()) {
+    decision = iter->second;
+  }
+
+  if (decision == NodeExecutorDecision::kCONVERT) {
+    return true;
+  } else {
+    return false;
+  }
 }
 
 std::vector<torch::jit::Node*> PartitioningCtx::getNodesRunInTorch() {