From 9eacf7c5e1105456bc3320bf62c76371c812ed27 Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Wed, 9 Sep 2026 21:03:50 -0700 Subject: [PATCH] fix: count unsupported operators with side effects ## Problem The compiler can report a model as fully supported even when a random or mutating operation must run in PyTorch. It then accepts `require_full_compilation=True` instead of rejecting the model. The support record excludes nodes with side effects. That excludes graph inputs and outputs as intended, but also excludes real operations such as random number generation. ## Change Record support by node kind instead of side effects. Keep graph inputs and outputs out of the counts, but include executable operations. Apply the same correction to the fast, global, and hierarchical partitioners, which divide the model into TensorRT and PyTorch sections. Their full-compilation checks and support reports then see the refused operations. ## Tests Passed 12/12 focused tests across the three partitioners. Six tests fail without this change: rejection of an unsupported operation with side effects and its presence in the support record, for each partitioner. Controls for supported models and unsupported operations without side effects pass before and after. These tests call the partitioners directly. They do not build or execute TensorRT engines. Tested on Linux x86_64 with Python 3.12 and TensorRT 11.2. Windows, aarch64, TensorRT-RTX, concurrent compilation, and models with many mutating operations were not tested. --- .../partitioning/_adjacency_partitioner.py | 15 +- .../partitioning/_global_partitioner.py | 16 +- .../partitioning/_hierarchical_partitioner.py | 4 +- .../test_000_full_support_detection.py | 161 ++++++++++++++++++ 4 files changed, 182 insertions(+), 14 deletions(-) create mode 100644 tests/py/dynamo/partitioning/test_000_full_support_detection.py diff --git a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py index 0dfb74bfc6..5b176e2dda 100644 --- a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py +++ b/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,7 +75,7 @@ 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 ) @@ -80,7 +83,7 @@ def is_node_supported( 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,7 +117,7 @@ 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: @@ -122,7 +125,7 @@ def is_node_supported( 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: diff --git a/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py index 41b775bba4..48e60331ff 100644 --- a/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py +++ b/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,7 +350,7 @@ 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: @@ -354,7 +358,7 @@ def is_node_supported( 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: diff --git a/py/torch_tensorrt/dynamo/partitioning/_hierarchical_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_hierarchical_partitioner.py index 85861adabd..113e61cbb9 100644 --- a/py/torch_tensorrt/dynamo/partitioning/_hierarchical_partitioner.py +++ b/py/torch_tensorrt/dynamo/partitioning/_hierarchical_partitioner.py @@ -83,7 +83,7 @@ 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: @@ -91,7 +91,7 @@ def is_node_supported( 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: diff --git a/tests/py/dynamo/partitioning/test_000_full_support_detection.py b/tests/py/dynamo/partitioning/test_000_full_support_detection.py new file mode 100644 index 0000000000..71d60a525e --- /dev/null +++ b/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()