fix: count unsupported operators with side effects

py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py

@@ -59,7 +59,10 @@ def is_node_supported(
                 self._non_target_device_cache,
             )
         ):
-            if not node.is_impure():
+            # Record by node kind, not by is_impure(). is_impure() is also true for
+            # random and mutating operators, so it dropped exactly the refusals this
+            # dict exists to report.
+            if node.op in CALLABLE_NODE_OPS:
                 self.unsupported_operators[node_name] = (
                     self.unsupported_operators.get(node_name, 0) + 1
                 )
@@ -72,15 +75,15 @@ def is_node_supported(
 
         if TorchTensorRTOperatorSupport._exceeds_max_tensor_rank(node):
             # Keep unrepresentable tensors entirely in the Torch partition.
-            if not node.is_impure():
+            if node.op in CALLABLE_NODE_OPS:
                 self.unsupported_operators[node_name] = (
                     self.unsupported_operators.get(node_name, 0) + 1
                 )
             return False
 
         if TorchTensorRTOperatorSupport._has_complex_dtype(node):
             # Complex-dtype tensors are not supported by TensorRT; force PyTorch fallback
-            if not node.is_impure():
+            if node.op in CALLABLE_NODE_OPS:
                 self.unsupported_operators[node_name] = (
                     self.unsupported_operators.get(node_name, 0) + 1
                 )
@@ -102,7 +105,7 @@ def is_node_supported(
         ):
             # 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():
+            if node.op in CALLABLE_NODE_OPS:
                 self.unsupported_operators[node_name] = (
                     self.unsupported_operators.get(node_name, 0) + 1
                 )
@@ -114,15 +117,15 @@ def is_node_supported(
             and node.target not in self.torch_executed_ops
         ):
             # If node is a proper, supported computational node, store the operator
-            if not node.is_impure() and node.op != "get_attr":
+            if node.op in CALLABLE_NODE_OPS:
                 if node_name not in self.supported_operators:
                     self.supported_operators[node_name] = 1
                 else:
                     self.supported_operators[node_name] += 1
 
             return True
         else:
-            if not node.is_impure():
+            if node.op in CALLABLE_NODE_OPS:
                 if node_name not in self.unsupported_operators:
                     self.unsupported_operators[node_name] = 1
                 else:

py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py

@@ -10,6 +10,7 @@
 from torch.fx.node import Target
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.passes.operator_support import OperatorSupport, SupportDict
+from torch.fx.passes.tools_common import CALLABLE_NODE_OPS
 from torch.utils._pytree import tree_flatten
 from torch_tensorrt._utils import trt_rtx_targets_turing
 from torch_tensorrt.dynamo._defaults import (
@@ -290,7 +291,10 @@ def is_node_supported(
             to_torch_device(settings.device),
             self._non_target_device_cache,
         ):
-            if not node.is_impure():
+            # Record by node kind, not by is_impure(). is_impure() is also true for
+            # random and mutating operators, so it dropped exactly the refusals this
+            # dict exists to report.
+            if node.op in CALLABLE_NODE_OPS:
                 self.unsupported_operators[node_name] = (
                     self.unsupported_operators.get(node_name, 0) + 1
                 )
@@ -303,7 +307,7 @@ def is_node_supported(
 
         if self._exceeds_max_tensor_rank(node):
             # TensorRT network inputs and outputs are limited by trt.Dims.
-            if not node.is_impure():
+            if node.op in CALLABLE_NODE_OPS:
                 self.unsupported_operators[node_name] = (
                     self.unsupported_operators.get(node_name, 0) + 1
                 )
@@ -312,7 +316,7 @@ def is_node_supported(
         if self._has_complex_dtype(node):
             # Complex-dtype tensors are not supported by TensorRT; force PyTorch fallback
             # so the graph breaks around the complex cluster inserted by complex_graph_detection.
-            if not node.is_impure():
+            if node.op in CALLABLE_NODE_OPS:
                 self.unsupported_operators[node_name] = (
                     self.unsupported_operators.get(node_name, 0) + 1
                 )
@@ -334,7 +338,7 @@ def is_node_supported(
         ):
             # 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():
+            if node.op in CALLABLE_NODE_OPS:
                 self.unsupported_operators[node_name] = (
                     self.unsupported_operators.get(node_name, 0) + 1
                 )
@@ -346,15 +350,15 @@ def is_node_supported(
             and node.target not in self.torch_executed_ops
         ):
             # If node is a proper, supported computational node, store the operator
-            if not node.is_impure() and node.op != "get_attr":
+            if node.op in CALLABLE_NODE_OPS:
                 if node_name not in self.supported_operators:
                     self.supported_operators[node_name] = 1
                 else:
                     self.supported_operators[node_name] += 1
 
             return True
         else:
-            if not node.is_impure():
+            if node.op in CALLABLE_NODE_OPS:
                 if node_name not in self.unsupported_operators:
                     self.unsupported_operators[node_name] = 1
                 else:

py/torch_tensorrt/dynamo/partitioning/_hierarchical_partitioner.py

@@ -83,15 +83,15 @@ def is_node_supported(
                 and node.target not in self.torch_executed_ops
             ):
                 # If node is a proper, supported computational node, store the operator
-                if not node.is_impure() and node.op != "get_attr":
+                if node.op in CALLABLE_NODE_OPS:
                     if node_name not in self.supported_operators:
                         self.supported_operators[f"{backend_name}_{node_name}"] = 1
                     else:
                         self.supported_operators[f"{backend_name}_{node_name}"] += 1
 
                 return True, backend_name
             else:
-                if i == len(self.backend_priority) - 1 and not node.is_impure():
+                if i == len(self.backend_priority) - 1 and node.op in CALLABLE_NODE_OPS:
                     if node_name not in self.unsupported_operators:
                         self.unsupported_operators[node_name] = 1
                     else:

tests/py/dynamo/partitioning/test_000_full_support_detection.py

@@ -0,0 +1,161 @@
+import copy
+
+import torch
+from parameterized import parameterized
+from torch.testing._internal.common_utils import TestCase, run_tests
+from torch_tensorrt.dynamo import partitioning
+from torch_tensorrt.dynamo.lowering import (
+    get_decompositions,
+    post_lowering,
+    pre_export_lowering,
+)
+
+PARTITIONERS = [
+    ("fast", partitioning.fast_partition),
+    ("global", partitioning.global_partition),
+    ("hierarchical", partitioning.hierarchical_adjacency_partition),
+]
+
+
+class TestFullSupportDetection(TestCase):
+    """A refused operator must never read as fully supported.
+
+    The operator support classes record each refusal so that require_full_compilation,
+    the dry run report and the support overview can all see it. That record used to be
+    gated on node.is_impure() being false, which is also true for random and mutating
+    operators, so a refused RNG operator was never recorded anywhere.
+    """
+
+    @staticmethod
+    def _lower(module, args):
+        exported = torch.export.export(module.eval().cuda(), args)
+        lowered = exported.run_decompositions(get_decompositions(False))
+        return post_lowering(pre_export_lowering(lowered).module())
+
+    @staticmethod
+    def _six_linear_layers():
+        return torch.nn.ModuleList([torch.nn.Linear(64, 64) for _ in range(6)])
+
+    @classmethod
+    def _impure_refusal_module(cls):
+        class WithImpureRefusal(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = cls._six_linear_layers()
+
+            def forward(self, x):
+                out = x
+                for index, layer in enumerate(self.layers):
+                    out = torch.relu(layer(out))
+                    if index == 2:
+                        # No converter, and impure, so the refusal went unrecorded.
+                        out = out + torch.normal(
+                            0.0,
+                            1.0,
+                            size=out.shape,
+                            device=out.device,
+                            dtype=out.dtype,
+                        )
+                return out
+
+        return WithImpureRefusal()
+
+    @classmethod
+    def _fully_supported_module(cls):
+        class FullySupported(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = cls._six_linear_layers()
+
+            def forward(self, x):
+                out = x
+                for layer in self.layers:
+                    out = torch.relu(layer(out))
+                return out
+
+        return FullySupported()
+
+    @staticmethod
+    def _partition(partition_fn, graph_module, **kwargs):
+        if partition_fn is partitioning.hierarchical_adjacency_partition:
+            kwargs["backend_priority"] = ["tensorrt"]
+        # Both partitioners mutate the module they are given, so hand each a copy.
+        return partition_fn(copy.deepcopy(graph_module), **kwargs)
+
+    @parameterized.expand(PARTITIONERS)
+    def test_impure_refusal_is_not_fully_supported(self, _, partition_fn):
+        graph_module = self._lower(
+            self._impure_refusal_module(), (torch.randn(8, 64, device="cuda"),)
+        )
+        with self.assertRaisesRegex(AssertionError, "not fully supported"):
+            self._partition(
+                partition_fn,
+                graph_module,
+                min_block_size=1,
+                require_full_compilation=True,
+            )
+
+    @parameterized.expand(PARTITIONERS)
+    def test_impure_refusal_is_recorded(self, _, partition_fn):
+        """The dry run report and the support overview read this dictionary, so an
+        unrecorded refusal makes both of them claim every node was supported."""
+        graph_module = self._lower(
+            self._impure_refusal_module(), (torch.randn(8, 64, device="cuda"),)
+        )
+        _, support = self._partition(partition_fn, graph_module, min_block_size=1)
+        self.assertTrue(
+            support.unsupported_operators,
+            "the refused operator was not recorded, so the support overview will "
+            "report that every node is supported",
+        )
+
+    @parameterized.expand(PARTITIONERS)
+    def test_pure_refusal_is_not_fully_supported(self, _, partition_fn):
+        """A refused pure operator was already caught. Keep it that way."""
+        graph_module = self._lower(
+            self._fully_supported_module(), (torch.randn(8, 64, device="cuda"),)
+        )
+        with self.assertRaisesRegex(AssertionError, "not fully supported"):
+            self._partition(
+                partition_fn,
+                graph_module,
+                min_block_size=1,
+                require_full_compilation=True,
+                torch_executed_ops={"torch.ops.aten.relu.default"},
+            )
+
+    @parameterized.expand(PARTITIONERS)
+    def test_fully_supported_module_is_accepted(self, name, partition_fn):
+        """Guards against over correction. This passes before the change too, so it does
+        not prove the fix; it proves the fix did not start rejecting good graphs."""
+        graph_module = self._lower(
+            self._fully_supported_module(), (torch.randn(8, 64, device="cuda"),)
+        )
+        partitioned, support = self._partition(
+            partition_fn,
+            graph_module,
+            min_block_size=1,
+            require_full_compilation=True,
+        )
+        self.assertFalse(support.unsupported_operators)
+        blocks = [child for child, _ in partitioned.named_children()]
+        self.assertTrue(
+            any("_run_on_acc" in block for block in blocks),
+            f"expected an accelerated block, got {blocks}",
+        )
+        # The global partitioner leaves a refused node inline in the parent rather than
+        # naming a torch block, so assert on what is left outside the block instead.
+        remaining = [
+            node.name
+            for node in partitioned.graph.nodes
+            if node.op == "call_function" and "_run_on_acc" not in node.name
+        ]
+        self.assertEqual(
+            remaining,
+            [],
+            f"expected no operator left outside the engine, got {remaining}",
+        )
+
+
+if __name__ == "__main__":
+    run_tests()