refactor(dynamo): make _requires_output_allocator a pure predicate (#4369)

Signed-off-by: cehongwang <wangcehong@gmail.com>
(cherry picked from commit 7eec04f)

Author: cehongwang

py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py

@@ -53,7 +53,12 @@ def is_node_supported(
                 )
             return False
 
-        if TorchTensorRTOperatorSupport._requires_output_allocator(node):
+        settings = CONVERTERS.compilation_settings
+        if (
+            settings is not None
+            and settings.fallback_data_dependent_ops
+            and TorchTensorRTOperatorSupport._requires_output_allocator(node)
+        ):
             # data-dependent output shape needs a TRT output allocator, which some
             # runtimes cannot consume; honor the fallback and run the node in PyTorch
             if not node.is_impure():

py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py

@@ -168,11 +168,9 @@ def _dtype(n: torch.fx.Node) -> Optional[torch.dtype]:
 
     @staticmethod
     def _requires_output_allocator(node: torch.fx.Node) -> bool:
-        # A converter that needs a TRT output allocator has a data-dependent output
-        # shape; route the node to PyTorch when fallback_data_dependent_ops is set.
-        settings = CONVERTERS.compilation_settings
-        if settings is None or not settings.fallback_data_dependent_ops:
-            return False
+        # True if the converter selected for this node needs a TRT output allocator,
+        # i.e. the node has a data-dependent output shape (e.g. nonzero or boolean
+        # index). The fallback_data_dependent_ops setting is honored by the caller.
         converter_packet = CONVERTERS.get(node)
         return converter_packet is not None and converter_packet[2].get(
             "requires_output_allocator", False
@@ -192,7 +190,12 @@ def is_node_supported(
                 )
             return False
 
-        if self._requires_output_allocator(node):
+        settings = CONVERTERS.compilation_settings
+        if (
+            settings is not None
+            and settings.fallback_data_dependent_ops
+            and self._requires_output_allocator(node)
+        ):
             # data-dependent output shape needs a TRT output allocator, which some
             # runtimes cannot consume; honor the fallback and run the node in PyTorch
             if not node.is_impure():

tests/py/dynamo/models/test_fallback_data_dependent_ops.py

@@ -43,21 +43,29 @@ def forward(self, x):
     )
 
 
+def test_requires_output_allocator_is_setting_independent():
+    # _requires_output_allocator is a pure predicate: it reports whether the
+    # converter selected for the node needs a TRT output allocator (decided per node
+    # via the selected converter, not by op target), independent of any setting.
+    node = _nonzero_node()
+    assert TorchTensorRTOperatorSupport._requires_output_allocator(node) is True
+
+
 def test_output_allocator_node_falls_back_only_when_enabled():
-    # nonzero's selected converter requires an output allocator, so the partitioner
-    # marks the node unsupported only when the flag is on. This is decided per node
-    # (via the selected converter), not by op target.
+    # The node is routed to PyTorch (unsupported) only when fallback_data_dependent_ops
+    # is on; with it off the node stays on TensorRT.
     node = _nonzero_node()
+    support = TorchTensorRTOperatorSupport()
     original = CONVERTERS.compilation_settings
     try:
         CONVERTERS.set_compilation_settings(
             CompilationSettings(fallback_data_dependent_ops=False)
         )
-        assert TorchTensorRTOperatorSupport._requires_output_allocator(node) is False
+        assert support.is_node_supported({}, node) is True
         CONVERTERS.set_compilation_settings(
             CompilationSettings(fallback_data_dependent_ops=True)
         )
-        assert TorchTensorRTOperatorSupport._requires_output_allocator(node) is True
+        assert support.is_node_supported({}, node) is False
     finally:
         CONVERTERS.compilation_settings = original