diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index a5122a8524..943ec68c28 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -1440,7 +1440,20 @@ def preserve_module_specs( cpu_memory_budget=settings.cpu_memory_budget, ) - dryrun_tracker.unsupported_ops = supported_ops.unsupported_operators + dryrun_tracker.unsupported_ops = supported_ops.fallback_operators + + if supported_ops.fallback_operators: + named = "; ".join( + f"{node_name} + Operator Count: {count} " + f"(Reasons: {', '.join(sorted(supported_ops.fallback_reasons[node_name]))})" + for node_name, count in sorted(supported_ops.fallback_operators.items()) + ) + logger.info( + "%d operator(s) will run in PyTorch: %s. " + "Compile with dryrun=True for the full report.", + len(supported_ops.fallback_operators), + named, + ) # The global partitioner leaves non-TRT nodes as-is if not settings.use_fast_partitioner: diff --git a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py index 0dfb74bfc6..f287a1d1cf 100644 --- a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py +++ b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: BSD-3-Clause import logging -from typing import Collection, Dict, List, Optional, Tuple +from typing import Collection, Dict, List, Optional, Set, Tuple import torch import torch.fx.passes.operator_support as ops @@ -42,9 +42,22 @@ def __init__(self, torch_executed_ops: Collection[Target] = set()) -> None: # Initialize sets of supported/unsupported operators self.supported_operators: Dict[str, int] = {} self.unsupported_operators: Dict[str, int] = {} + # Keep impure refusals out of the counters used to decide full support. + self.fallback_operators: Dict[str, int] = {} + self.fallback_reasons: Dict[str, Set[str]] = {} self.torch_executed_ops = torch_executed_ops self._non_target_device_cache: Dict[torch.fx.Node, bool] = {} + def _record_fallback( + self, node: torch.fx.Node, node_name: str, reason: str + ) -> None: + # Structural nodes must not make a fully supported graph report fallback. + if node.op in CALLABLE_NODE_OPS: + self.fallback_operators[node_name] = ( + self.fallback_operators.get(node_name, 0) + 1 + ) + self.fallback_reasons.setdefault(node_name, set()).add(reason) + def is_node_supported( self, submodules: Dict[str, torch.nn.Module], node: torch.fx.Node ) -> bool: @@ -68,6 +81,7 @@ def is_node_supported( "non-target device region", node_name, ) + self._record_fallback(node, node_name, "explicit non-target device region") return False if TorchTensorRTOperatorSupport._exceeds_max_tensor_rank(node): @@ -76,6 +90,7 @@ def is_node_supported( self.unsupported_operators[node_name] = ( self.unsupported_operators.get(node_name, 0) + 1 ) + self._record_fallback(node, node_name, "tensor rank exceeds TensorRT limit") return False if TorchTensorRTOperatorSupport._has_complex_dtype(node): @@ -84,6 +99,7 @@ def is_node_supported( self.unsupported_operators[node_name] = ( self.unsupported_operators.get(node_name, 0) + 1 ) + self._record_fallback(node, node_name, "complex tensor dtype") return False if TorchTensorRTOperatorSupport._has_bf16_on_turing(node, settings): @@ -106,6 +122,11 @@ def is_node_supported( self.unsupported_operators[node_name] = ( self.unsupported_operators.get(node_name, 0) + 1 ) + self._record_fallback( + node, + node_name, + "data-dependent output shape (fallback_data_dependent_ops=True)", + ) return False if ( @@ -128,6 +149,16 @@ def is_node_supported( else: self.unsupported_operators[node_name] += 1 + self._record_fallback( + node, + node_name, + ( + "excluded by torch_executed_ops" + if node_name in self.torch_executed_ops + or node.target in self.torch_executed_ops + else "no validated TensorRT converter" + ), + ) return False def print_support_overview(self, num_trt_blocks: Optional[int] = None) -> None: diff --git a/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py index 41b775bba4..781ed9d1db 100644 --- a/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py +++ b/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: BSD-3-Clause import logging -from typing import Collection, Dict, List, Mapping, Optional, Sequence, Tuple +from typing import Collection, Dict, List, Mapping, Optional, Sequence, Set, Tuple import tensorrt as trt import torch @@ -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 ( @@ -149,6 +150,9 @@ def __init__( # Initialize sets of supported/unsupported operators self.supported_operators: Dict[str, int] = {} self.unsupported_operators: Dict[str, int] = {} + # Keep impure refusals out of the counters used to decide full support. + self.fallback_operators: Dict[str, int] = {} + self.fallback_reasons: Dict[str, Set[str]] = {} self.torch_executed_ops: Collection[Target] = torch_executed_ops self._non_target_device_cache: Dict[torch.fx.Node, bool] = {} @@ -279,6 +283,16 @@ def _requires_output_allocator(node: torch.fx.Node) -> bool: "requires_output_allocator", False ) + def _record_fallback( + self, node: torch.fx.Node, node_name: str, reason: str + ) -> None: + # Structural nodes must not make a fully supported graph report fallback. + if node.op in CALLABLE_NODE_OPS: + self.fallback_operators[node_name] = ( + self.fallback_operators.get(node_name, 0) + 1 + ) + self.fallback_reasons.setdefault(node_name, set()).add(reason) + def is_node_supported( self, submodules: Mapping[str, torch.nn.Module], node: torch.fx.Node ) -> bool: @@ -299,6 +313,7 @@ def is_node_supported( "non-target device region", node_name, ) + self._record_fallback(node, node_name, "explicit non-target device region") return False if self._exceeds_max_tensor_rank(node): @@ -307,6 +322,7 @@ def is_node_supported( self.unsupported_operators[node_name] = ( self.unsupported_operators.get(node_name, 0) + 1 ) + self._record_fallback(node, node_name, "tensor rank exceeds TensorRT limit") return False if self._has_complex_dtype(node): @@ -316,6 +332,7 @@ def is_node_supported( self.unsupported_operators[node_name] = ( self.unsupported_operators.get(node_name, 0) + 1 ) + self._record_fallback(node, node_name, "complex tensor dtype") return False if self._has_bf16_on_turing(node, settings): @@ -338,6 +355,11 @@ def is_node_supported( self.unsupported_operators[node_name] = ( self.unsupported_operators.get(node_name, 0) + 1 ) + self._record_fallback( + node, + node_name, + "data-dependent output shape (fallback_data_dependent_ops=True)", + ) return False if ( @@ -360,6 +382,16 @@ def is_node_supported( else: self.unsupported_operators[node_name] += 1 + self._record_fallback( + node, + node_name, + ( + "excluded by torch_executed_ops" + if node_name in self.torch_executed_ops + or node.target in self.torch_executed_ops + else "no validated TensorRT converter" + ), + ) return False def print_support_overview( diff --git a/tests/py/dynamo/partitioning/test_000_fallback_reasons.py b/tests/py/dynamo/partitioning/test_000_fallback_reasons.py new file mode 100644 index 0000000000..e83e69da37 --- /dev/null +++ b/tests/py/dynamo/partitioning/test_000_fallback_reasons.py @@ -0,0 +1,166 @@ +from unittest.mock import patch + +import tensorrt as trt +import torch +from parameterized import parameterized +from torch.testing._internal.common_utils import TestCase, run_tests +from torch_tensorrt.dynamo._settings import CompilationSettings +from torch_tensorrt.dynamo.conversion._ConverterRegistry import ( + DYNAMO_CONVERTERS, + ConverterRegistry, + ConverterSupport, +) +from torch_tensorrt.dynamo.partitioning._adjacency_partitioner import OpSupportTester +from torch_tensorrt.dynamo.partitioning._global_partitioner import ( + TorchTensorRTOperatorSupport, +) + +SUPPORT_CLASSES = [ + ("fast", OpSupportTester), + ("global", TorchTensorRTOperatorSupport), +] + + +class TestFallbackReasons(TestCase): + def setUp(self): + super().setUp() + for name, value in ( + ("compilation_settings", CompilationSettings()), + ("disallowed_targets", set()), + ): + patcher = patch.object(DYNAMO_CONVERTERS, name, value) + patcher.start() + self.addCleanup(patcher.stop) + + @staticmethod + def _node(target, input_value, output_value): + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = input_value + node = graph.call_function(target, (x,)) + node.meta["val"] = output_value + graph.output(node) + return node + + @parameterized.expand(SUPPORT_CLASSES) + def test_refusal_reasons(self, _, support_class): + x = torch.empty(2, device="cuda") + high_rank = torch.empty((1,) * (trt.Dims.MAX_DIMS + 1), device="cuda") + complex_value = torch.empty(2, dtype=torch.complex64, device="cuda") + cases = [ + ( + torch.ops.aten._to_copy.default, + x, + x.cpu(), + "explicit non-target device region", + ), + ( + torch.ops.aten.clone.default, + high_rank, + high_rank, + "tensor rank exceeds TensorRT limit", + ), + ( + torch.ops.aten.clone.default, + complex_value, + complex_value, + "complex tensor dtype", + ), + ( + torch.ops.aten.nonzero.default, + x, + torch.empty(1, 1, dtype=torch.int64, device="cuda"), + "data-dependent output shape (fallback_data_dependent_ops=True)", + ), + (torch.ops.aten.rand_like.default, x, x, "no validated TensorRT converter"), + ] + DYNAMO_CONVERTERS.compilation_settings.fallback_data_dependent_ops = True + for target, input_value, output_value, reason in cases: + with self.subTest(reason=reason): + support = support_class() + node = self._node(target, input_value, output_value) + name = ConverterRegistry.qualified_name_or_str(target) + if target == torch.ops.aten.nonzero.default: + self.assertTrue( + DYNAMO_CONVERTERS[node][2]["requires_output_allocator"] + ) + self.assertFalse(support.is_node_supported({}, node)) + self.assertEqual(support.fallback_operators, {name: 1}) + self.assertEqual(support.fallback_reasons, {name: {reason}}) + if target == torch.ops.aten.rand_like.default: + self.assertTrue(node.is_impure()) + self.assertEqual(support.unsupported_operators, {}) + + @parameterized.expand(SUPPORT_CLASSES) + def test_requested_fallback(self, _, support_class): + x = torch.empty(2, device="cuda") + target = torch.ops.aten.relu.default + node = self._node(target, x, x) + name = ConverterRegistry.qualified_name_or_str(target) + self.assertIn(node, DYNAMO_CONVERTERS) + for excluded in (name, target): + with self.subTest(excluded=excluded): + support = support_class(torch_executed_ops={excluded}) + self.assertFalse(support.is_node_supported({}, node)) + self.assertEqual(support.fallback_operators, {name: 1}) + self.assertEqual( + support.fallback_reasons, + {name: {"excluded by torch_executed_ops"}}, + ) + + @parameterized.expand(SUPPORT_CLASSES) + def test_rejected_converter_is_not_reported_as_missing(self, _, support_class): + x = torch.empty(2, device="cuda") + target = torch.ops.aten.clone.default + node = self._node(target, x, x) + converters = { + target: [ + ConverterSupport( + converter_implementation=lambda *args: None, + capability_validator=lambda node, settings: False, + ) + ] + } + with patch.object(DYNAMO_CONVERTERS, "registries", [converters]): + self.assertIsNotNone(DYNAMO_CONVERTERS.get_unvalidated(target)) + support = support_class() + self.assertFalse(support.is_node_supported({}, node)) + name = ConverterRegistry.qualified_name_or_str(target) + self.assertEqual( + support.fallback_reasons, + {name: {"no validated TensorRT converter"}}, + ) + + @parameterized.expand(SUPPORT_CLASSES) + def test_same_operator_keeps_multiple_reasons(self, _, support_class): + support = support_class() + target = torch.ops.aten.clone.default + for value in ( + torch.empty((1,) * (trt.Dims.MAX_DIMS + 1), device="cuda"), + torch.empty(2, dtype=torch.complex64, device="cuda"), + ): + self.assertFalse( + support.is_node_supported({}, self._node(target, value, value)) + ) + name = ConverterRegistry.qualified_name_or_str(target) + self.assertEqual(support.fallback_operators, {name: 2}) + self.assertEqual( + support.fallback_reasons, + {name: {"tensor rank exceeds TensorRT limit", "complex tensor dtype"}}, + ) + + @parameterized.expand(SUPPORT_CLASSES) + def test_supported_and_structural_nodes_have_no_fallback(self, _, support_class): + x = torch.empty(2, device="cuda") + node = self._node(torch.ops.aten.relu.default, x, x) + support = support_class() + self.assertTrue(support.is_node_supported({}, node)) + for structural in node.graph.nodes: + if structural.op in ("placeholder", "output"): + support.is_node_supported({}, structural) + self.assertEqual(support.fallback_operators, {}) + self.assertEqual(support.fallback_reasons, {}) + + +if __name__ == "__main__": + run_tests() diff --git a/tests/py/dynamo/partitioning/test_000_fallback_reporting.py b/tests/py/dynamo/partitioning/test_000_fallback_reporting.py new file mode 100644 index 0000000000..ab41daa86f --- /dev/null +++ b/tests/py/dynamo/partitioning/test_000_fallback_reporting.py @@ -0,0 +1,228 @@ +import logging + +import torch +import torch_tensorrt +from parameterized import parameterized +from torch.testing._internal.common_utils import TestCase, run_tests + +SUMMARY_MARKER = "operator(s)" + + +class _SummaryCollector(logging.Filter): + + def __init__(self) -> None: + super().__init__() + self.messages: list[str] = [] + self.levels: list[int] = [] + + def filter(self, record: logging.LogRecord) -> bool: + message = record.getMessage() + if SUMMARY_MARKER in message: + self.messages.append(message) + self.levels.append(record.levelno) + return True + + +class TestFallbackIsReported(TestCase): + """Report fallback at INFO without warning about expected behavior.""" + + @staticmethod + def _six_linear_layers() -> torch.nn.ModuleList: + return torch.nn.ModuleList([torch.nn.Linear(64, 64) for _ in range(6)]) + + @classmethod + def _impure_fallback_module(cls) -> torch.nn.Module: + """Refused on the last path in the support test, and impure, so the older counter + never recorded it.""" + + class ImpureFallback(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.layers = cls._six_linear_layers() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + out = x + for layer in self.layers: + out = torch.relu(layer(out)) + return out + torch.rand_like(out) + + return ImpureFallback().eval().cuda() + + @classmethod + def _complex_fallback_module(cls) -> torch.nn.Module: + """Refused by the complex dtype check, which is one of the earlier returns.""" + + class ComplexFallback(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.layers = cls._six_linear_layers() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + out = x + for layer in self.layers: + out = torch.relu(layer(out)) + return torch.real(torch.fft.fft(out)) + out + + return ComplexFallback().eval().cuda() + + @classmethod + def _fully_supported_module(cls) -> torch.nn.Module: + class FullySupported(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.layers = cls._six_linear_layers() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + out = x + for layer in self.layers: + out = torch.relu(layer(out)) + return out + + return FullySupported().eval().cuda() + + def _compile(self, module, inputs, log_level=logging.INFO, **kwargs): + collector = _SummaryCollector() + logger = logging.getLogger("torch_tensorrt.dynamo._compiler") + previous_level = logger.level + logger.setLevel(log_level) + logger.addFilter(collector) + try: + compiled = torch_tensorrt.dynamo.compile( + torch.export.export(module, tuple(inputs)), + inputs=list(inputs), + min_block_size=1, + enabled_precisions={torch.float32}, + truncate_double=True, + **kwargs, + ) + finally: + logger.removeFilter(collector) + logger.setLevel(previous_level) + self.assertTrue(all(level == logging.INFO for level in collector.levels)) + segments = [name for name, _ in compiled.named_children()] + return segments, collector.messages + + @parameterized.expand( + [ + ( + "impure_refusal", + "_impure_fallback_module", + "rand_like", + "no validated TensorRT converter", + ), + ( + "complex_dtype_refusal", + "_complex_fallback_module", + "fft", + "complex tensor dtype", + ), + ] + ) + def test_fallback_is_reported(self, _, factory, expected_operator, expected_reason): + """Both of these split the graph, and a refusal on any path has to be reported, not + only one of them.""" + inputs = [torch.randn(8, 64, device="cuda")] + segments, messages = self._compile(getattr(self, factory)(), inputs) + self.assertTrue( + any("_run_on_gpu" in segment for segment in segments), + f"expected a PyTorch segment, got {segments}", + ) + self.assertEqual( + len(messages), 1, f"expected exactly one report, got {messages}" + ) + self.assertIn(expected_operator, messages[0]) + self.assertIn(expected_reason, messages[0]) + + def test_fully_supported_module_is_silent(self): + """Nothing fell back, so there is nothing to report.""" + inputs = [torch.randn(8, 64, device="cuda")] + segments, messages = self._compile(self._fully_supported_module(), inputs) + self.assertFalse( + any("_run_on_gpu" in segment for segment in segments), + f"expected no PyTorch segment, got {segments}", + ) + self.assertEqual(messages, []) + + @parameterized.expand([("fast", True), ("global", False)]) + def test_requested_fallback_is_reported(self, _, use_fast_partitioner): + inputs = [torch.randn(8, 64, device="cuda")] + segments, messages = self._compile( + self._fully_supported_module(), + inputs, + torch_executed_ops={"torch.ops.aten.relu.default"}, + use_fast_partitioner=use_fast_partitioner, + ) + self.assertEqual(len(messages), 1) + self.assertIn("torch.ops.aten.relu.default + Operator Count: 6", messages[0]) + self.assertIn("excluded by torch_executed_ops", messages[0]) + + def test_global_partitioner_reports_too(self): + """The global partitioner is the automatic fallback when the fast one raises, so a + user reaches it exactly when they most need to be told something happened.""" + inputs = [torch.randn(8, 64, device="cuda")] + segments, messages = self._compile( + self._impure_fallback_module(), inputs, use_fast_partitioner=False + ) + self.assertEqual( + len(messages), 1, f"expected exactly one report, got {messages}" + ) + + def test_global_partitioner_silent_on_fully_supported(self): + """The global partitioner asks about placeholder and output nodes as well as + operators. Recording those made a fully supported graph report its own inputs and + outputs as fallbacks, so this warns falsely without the callable-node guard.""" + inputs = [torch.randn(8, 64, device="cuda")] + segments, messages = self._compile( + self._fully_supported_module(), inputs, use_fast_partitioner=False + ) + self.assertFalse( + any("_run_on_gpu" in segment for segment in segments), + f"expected no PyTorch segment, got {segments}", + ) + self.assertEqual(messages, []) + + @parameterized.expand([("fast", True), ("global", False)]) + def test_requested_fallback_by_target_is_reported(self, _, use_fast_partitioner): + inputs = [torch.randn(8, 64, device="cuda")] + segments, messages = self._compile( + self._fully_supported_module(), + inputs, + torch_executed_ops={torch.ops.aten.relu.default}, + use_fast_partitioner=use_fast_partitioner, + ) + self.assertEqual(len(messages), 1) + self.assertIn("torch.ops.aten.relu.default + Operator Count: 6", messages[0]) + self.assertIn("excluded by torch_executed_ops", messages[0]) + + @parameterized.expand([("fast", True), ("global", False)]) + def test_mixed_fallback_reasons(self, _, use_fast_partitioner): + inputs = [torch.randn(8, 64, device="cuda")] + _, messages = self._compile( + self._impure_fallback_module(), + inputs, + torch_executed_ops={torch.ops.aten.relu.default}, + use_fast_partitioner=use_fast_partitioner, + ) + self.assertEqual(len(messages), 1) + self.assertIn( + "torch.ops.aten.rand_like.default + Operator Count: 1 " + "(Reasons: no validated TensorRT converter)", + messages[0], + ) + self.assertIn( + "torch.ops.aten.relu.default + Operator Count: 6 " + "(Reasons: excluded by torch_executed_ops)", + messages[0], + ) + self.assertLess(messages[0].index("rand_like"), messages[0].index("relu")) + + def test_warning_level_suppresses_summary(self): + inputs = [torch.randn(8, 64, device="cuda")] + _, messages = self._compile( + self._impure_fallback_module(), inputs, log_level=logging.WARNING + ) + self.assertEqual(messages, []) + + +if __name__ == "__main__": + run_tests()