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()