fix: fix the prim::Loop fallback issue by bowang007 · Pull Request #1691 · pytorch/TensorRT (original) (raw)

Description

In current fallback logic, prim::Loop is set to be run in Torch firstly by this line:

if (!conversion::OpSupported(n)) {
      // If the op is not supported by the conversion phase it should run in PyTorch
      ctx->setNodeExecutorDecision(n, NodeExecutorDecision::kUNSUPPORTED);

However, in segmentation part, there is a contradiction:

else if (n->kind() == torch::jit::prim::Loop) {
        if (!in_prog_pyt_blk_nodes.empty()) {
          finalizeNewBlock(segmented_blocks, SegmentedBlock::kTorch, in_prog_pyt_blk_nodes);
          cur_pyt_nodes_uses.clear();
        }
        if (checkLoopEvaluatable(n)) {
          in_prog_trt_blk_nodes.push_back(n);
          cur_trt_nodes_uses.insert(dependent_nodes.begin(), dependent_nodes.end());
        } else {
          auto loop_node = std::vector<torch::jit::Node*>{n};
          finalizeNewBlock(segmented_blocks, SegmentedBlock::kTorch, loop_node);
          segmented_blocks.back().do_not_merge(true);
        }
        continue;
      }

This introduces many issues and bugs that we have now.
The reason is that firstly setting LOOP running in Torch will make many nodes that produce values consumed in this LOOP block or nodes that consume the values dependent in this Block fallback.
In other words, setNonTensorConnectedNodes won't work as what we expected, and this would introduce a lot of NonTensor Input TensorRT segment after segmentation.

Fixes #1617

Type of change

Checklist: