From 72a2610b3f3195e5187b4820050ac7bd2054104a Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Wed, 9 Sep 2026 21:03:52 -0700 Subject: [PATCH 1/2] feat: report operators that fall back to PyTorch ## Problem A model can split into TensorRT engines and PyTorch sections without naming the unsupported operators at the default logging level. Users may only notice the split when performance changes. The existing support record also omits refused operations with side effects, such as random number generation. ## Change Record every refused executable operation in both the fast and global partitioners. Use that record for the dry-run report and one warning from the shared compilation path. Name the operators that caused fallback, but leave out operations the caller explicitly chose to run in PyTorch through `torch_executed_ops`. Fully supported models remain silent. The warning covers graphs that reach partitioning. Small graphs that return earlier keep their existing skipped-compilation message. ## Tests Passed 7/7 focused tests. Three fail without this change because the warning is missing: an unsupported operation with side effects, an unsupported complex data type, and fallback through the global partitioner. The other four tests check that supported models stay silent under both partitioners and that caller-requested fallback stays silent when specified by an operator or its name. Those controls pass before and after. The tests build engines and inspect the graph sections and logs. They do not execute the returned models to compare outputs. Tested on Linux x86_64 with Python 3.12 and TensorRT 11.2. Windows, aarch64, TensorRT-RTX, concurrent logging, nested hierarchical graphs, and downstream log parsers were not tested. --- py/torch_tensorrt/dynamo/_compiler.py | 29 ++- .../partitioning/_adjacency_partitioner.py | 17 ++ .../partitioning/_global_partitioner.py | 18 ++ .../test_000_fallback_reporting.py | 203 ++++++++++++++++++ 4 files changed, 266 insertions(+), 1 deletion(-) create mode 100644 tests/py/dynamo/partitioning/test_000_fallback_reporting.py diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index a5122a8524..9ce1d7d438 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -40,6 +40,9 @@ from torch_tensorrt.dynamo.conversion._ConverterRegistry import ( DYNAMO_CONVERTERS as CONVERTERS, ) +from torch_tensorrt.dynamo.conversion._ConverterRegistry import ( + ConverterRegistry, +) from torch_tensorrt.dynamo.debug._DebuggerConfig import DebuggerConfig from torch_tensorrt.dynamo.debug._supports_debugger import fn_supports_debugger from torch_tensorrt.dynamo.lowering import ( @@ -1440,7 +1443,31 @@ 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 + + # Operators the caller named in torch_executed_ops are left out: that fallback was + # asked for. The set can hold either a qualified name string or an operator target + # object, and fallback_operators is keyed by name, so normalize to names first or a + # target object never matches and the caller is warned about their own choice. + excluded_names = { + ConverterRegistry.qualified_name_or_str(op) + for op in settings.torch_executed_ops + } + reported_fallbacks = { + node_name: count + for node_name, count in supported_ops.fallback_operators.items() + if node_name not in excluded_names + } + if reported_fallbacks: + named = ", ".join( + f"{node_name} + Operator Count: {count}" + for node_name, count in sorted(reported_fallbacks.items()) + ) + logger.warning( + f"{len(reported_fallbacks)} operator(s) have no TensorRT converter and will " + f"run in PyTorch, so this model was split around them: {named}. " + f"Compile with dryrun=True for the full report." + ) # 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..a1ed70f4f8 100644 --- a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py +++ b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py @@ -42,9 +42,21 @@ 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] = {} + # unsupported_operators skips impure nodes, so it cannot answer "did anything fall + # back". This one records every refusal. + self.fallback_operators: Dict[str, int] = {} 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) -> None: + # Only executable operators count as fallbacks. Placeholder and output nodes are + # graph structure, not operators, and recording them makes a fully supported graph + # report its own inputs and outputs as unconverted. + if node.op in CALLABLE_NODE_OPS: + self.fallback_operators[node_name] = ( + self.fallback_operators.get(node_name, 0) + 1 + ) + def is_node_supported( self, submodules: Dict[str, torch.nn.Module], node: torch.fx.Node ) -> bool: @@ -68,6 +80,7 @@ def is_node_supported( "non-target device region", node_name, ) + self._record_fallback(node, node_name) return False if TorchTensorRTOperatorSupport._exceeds_max_tensor_rank(node): @@ -76,6 +89,7 @@ def is_node_supported( self.unsupported_operators[node_name] = ( self.unsupported_operators.get(node_name, 0) + 1 ) + self._record_fallback(node, node_name) return False if TorchTensorRTOperatorSupport._has_complex_dtype(node): @@ -84,6 +98,7 @@ def is_node_supported( self.unsupported_operators[node_name] = ( self.unsupported_operators.get(node_name, 0) + 1 ) + self._record_fallback(node, node_name) return False if TorchTensorRTOperatorSupport._has_bf16_on_turing(node, settings): @@ -106,6 +121,7 @@ def is_node_supported( self.unsupported_operators[node_name] = ( self.unsupported_operators.get(node_name, 0) + 1 ) + self._record_fallback(node, node_name) return False if ( @@ -128,6 +144,7 @@ def is_node_supported( else: self.unsupported_operators[node_name] += 1 + self._record_fallback(node, node_name) 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..f0557837eb 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 ( @@ -149,6 +150,9 @@ def __init__( # Initialize sets of supported/unsupported operators self.supported_operators: Dict[str, int] = {} self.unsupported_operators: Dict[str, int] = {} + # unsupported_operators skips impure nodes, so it cannot answer "did anything + # fall back". This one records every refusal. + self.fallback_operators: Dict[str, int] = {} self.torch_executed_ops: Collection[Target] = torch_executed_ops self._non_target_device_cache: Dict[torch.fx.Node, bool] = {} @@ -279,6 +283,15 @@ def _requires_output_allocator(node: torch.fx.Node) -> bool: "requires_output_allocator", False ) + def _record_fallback(self, node: torch.fx.Node, node_name: str) -> None: + # Only executable operators count as fallbacks. Placeholder and output nodes are + # graph structure, not operators, and recording them makes a fully supported graph + # report its own inputs and outputs as unconverted. + if node.op in CALLABLE_NODE_OPS: + self.fallback_operators[node_name] = ( + self.fallback_operators.get(node_name, 0) + 1 + ) + def is_node_supported( self, submodules: Mapping[str, torch.nn.Module], node: torch.fx.Node ) -> bool: @@ -299,6 +312,7 @@ def is_node_supported( "non-target device region", node_name, ) + self._record_fallback(node, node_name) return False if self._exceeds_max_tensor_rank(node): @@ -307,6 +321,7 @@ def is_node_supported( self.unsupported_operators[node_name] = ( self.unsupported_operators.get(node_name, 0) + 1 ) + self._record_fallback(node, node_name) return False if self._has_complex_dtype(node): @@ -316,6 +331,7 @@ def is_node_supported( self.unsupported_operators[node_name] = ( self.unsupported_operators.get(node_name, 0) + 1 ) + self._record_fallback(node, node_name) return False if self._has_bf16_on_turing(node, settings): @@ -338,6 +354,7 @@ def is_node_supported( self.unsupported_operators[node_name] = ( self.unsupported_operators.get(node_name, 0) + 1 ) + self._record_fallback(node, node_name) return False if ( @@ -360,6 +377,7 @@ def is_node_supported( else: self.unsupported_operators[node_name] += 1 + self._record_fallback(node, node_name) return False def print_support_overview( 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..449631d56b --- /dev/null +++ b/tests/py/dynamo/partitioning/test_000_fallback_reporting.py @@ -0,0 +1,203 @@ +import logging + +import torch +import torch_tensorrt +from parameterized import parameterized +from torch.testing._internal.common_utils import TestCase, run_tests + +WARNING_MARKER = "have no TensorRT converter and will run in PyTorch" + + +class _WarningCollector(logging.Filter): + """A filter rather than a handler, so the record is seen wherever it is emitted from. + + Attaching a handler to the torch_tensorrt logger is not enough: the message comes from a + child logger and torch_tensorrt installs a root handler of its own during compilation. + """ + + def __init__(self) -> None: + super().__init__() + self.messages: list[str] = [] + + def filter(self, record: logging.LogRecord) -> bool: + message = record.getMessage() + if WARNING_MARKER in message: + self.messages.append(message) + return True + + +class TestFallbackIsReported(TestCase): + """A model that falls back partly to PyTorch should say so at default verbosity. + + The partition report exists already but is DEBUG, which is off by default, so a model + that quietly became several engines plus a PyTorch segment looked exactly like one that + compiled whole. + """ + + @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, **kwargs): + collector = _WarningCollector() + loggers = [logging.getLogger()] + [ + logging.getLogger(name) + for name in list(logging.root.manager.loggerDict) + if name.startswith("torch_tensorrt") + ] + for each in loggers: + each.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: + for each in loggers: + each.removeFilter(collector) + segments = [name for name, _ in compiled.named_children()] + return segments, collector.messages + + @parameterized.expand( + [ + ("impure_refusal", "_impure_fallback_module", "rand_like"), + ("complex_dtype_refusal", "_complex_fallback_module", "fft"), + ] + ) + def test_fallback_is_reported(self, _, factory, expected_operator): + """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]) + + 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, []) + + def test_requested_fallback_is_silent(self): + """The caller asked for this operator to stay in PyTorch, so warning about it would + be telling them about their own choice.""" + inputs = [torch.randn(8, 64, device="cuda")] + segments, messages = self._compile( + self._fully_supported_module(), + inputs, + torch_executed_ops={"torch.ops.aten.relu.default"}, + ) + self.assertTrue( + any("_run_on_gpu" in segment for segment in segments), + f"expected a PyTorch segment, got {segments}", + ) + self.assertEqual(messages, []) + + 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, []) + + def test_requested_fallback_by_target_is_silent(self): + """torch_executed_ops accepts an operator target object, not only a qualified name + string. The report filter keys on names, so a target object has to be normalized or + the caller is warned about a fallback they asked for.""" + inputs = [torch.randn(8, 64, device="cuda")] + segments, messages = self._compile( + self._fully_supported_module(), + inputs, + torch_executed_ops={torch.ops.aten.relu.default}, + ) + self.assertTrue( + any("_run_on_gpu" in segment for segment in segments), + f"expected a PyTorch segment, got {segments}", + ) + self.assertEqual(messages, []) + + +if __name__ == "__main__": + run_tests() From ee06164e8a650b6145a68c7ce638374d2111463e Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Thu, 10 Sep 2026 16:42:18 -0700 Subject: [PATCH 2/2] fix: explain PyTorch fallback at info level Fallback is expected behavior, but a warning without a reason makes it hard to distinguish caller choices from unsupported operations. Include observed reasons and requested fallback in the summary, and report it at INFO without changing support decisions. Test Plan: Passed 10 support tests and 12 compilation/logging tests on Linux x86_64 with CUDA. Both partitioners report reasons and requested fallback; fully supported graphs stay silent and WARNING suppresses the summary. Four rank-boundary execution tests also passed with PyTorch output comparisons. The support tests fail without the reason metadata. --- py/torch_tensorrt/dynamo/_compiler.py | 34 ++-- .../partitioning/_adjacency_partitioner.py | 38 ++-- .../partitioning/_global_partitioner.py | 38 ++-- .../partitioning/test_000_fallback_reasons.py | 166 ++++++++++++++++++ .../test_000_fallback_reporting.py | 111 +++++++----- 5 files changed, 296 insertions(+), 91 deletions(-) create mode 100644 tests/py/dynamo/partitioning/test_000_fallback_reasons.py diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index 9ce1d7d438..943ec68c28 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -40,9 +40,6 @@ from torch_tensorrt.dynamo.conversion._ConverterRegistry import ( DYNAMO_CONVERTERS as CONVERTERS, ) -from torch_tensorrt.dynamo.conversion._ConverterRegistry import ( - ConverterRegistry, -) from torch_tensorrt.dynamo.debug._DebuggerConfig import DebuggerConfig from torch_tensorrt.dynamo.debug._supports_debugger import fn_supports_debugger from torch_tensorrt.dynamo.lowering import ( @@ -1445,28 +1442,17 @@ def preserve_module_specs( dryrun_tracker.unsupported_ops = supported_ops.fallback_operators - # Operators the caller named in torch_executed_ops are left out: that fallback was - # asked for. The set can hold either a qualified name string or an operator target - # object, and fallback_operators is keyed by name, so normalize to names first or a - # target object never matches and the caller is warned about their own choice. - excluded_names = { - ConverterRegistry.qualified_name_or_str(op) - for op in settings.torch_executed_ops - } - reported_fallbacks = { - node_name: count - for node_name, count in supported_ops.fallback_operators.items() - if node_name not in excluded_names - } - if reported_fallbacks: - named = ", ".join( - f"{node_name} + Operator Count: {count}" - for node_name, count in sorted(reported_fallbacks.items()) + 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.warning( - f"{len(reported_fallbacks)} operator(s) have no TensorRT converter and will " - f"run in PyTorch, so this model was split around them: {named}. " - f"Compile with dryrun=True for the full report." + 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 diff --git a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py index a1ed70f4f8..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,20 +42,21 @@ 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] = {} - # unsupported_operators skips impure nodes, so it cannot answer "did anything fall - # back". This one records every refusal. + # 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) -> None: - # Only executable operators count as fallbacks. Placeholder and output nodes are - # graph structure, not operators, and recording them makes a fully supported graph - # report its own inputs and outputs as unconverted. + 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 @@ -80,7 +81,7 @@ def is_node_supported( "non-target device region", node_name, ) - self._record_fallback(node, node_name) + self._record_fallback(node, node_name, "explicit non-target device region") return False if TorchTensorRTOperatorSupport._exceeds_max_tensor_rank(node): @@ -89,7 +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) + self._record_fallback(node, node_name, "tensor rank exceeds TensorRT limit") return False if TorchTensorRTOperatorSupport._has_complex_dtype(node): @@ -98,7 +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) + self._record_fallback(node, node_name, "complex tensor dtype") return False if TorchTensorRTOperatorSupport._has_bf16_on_turing(node, settings): @@ -121,7 +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) + self._record_fallback( + node, + node_name, + "data-dependent output shape (fallback_data_dependent_ops=True)", + ) return False if ( @@ -144,7 +149,16 @@ def is_node_supported( else: self.unsupported_operators[node_name] += 1 - self._record_fallback(node, node_name) + 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 f0557837eb..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 @@ -150,9 +150,9 @@ def __init__( # Initialize sets of supported/unsupported operators self.supported_operators: Dict[str, int] = {} self.unsupported_operators: Dict[str, int] = {} - # unsupported_operators skips impure nodes, so it cannot answer "did anything - # fall back". This one records every refusal. + # 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] = {} @@ -283,14 +283,15 @@ def _requires_output_allocator(node: torch.fx.Node) -> bool: "requires_output_allocator", False ) - def _record_fallback(self, node: torch.fx.Node, node_name: str) -> None: - # Only executable operators count as fallbacks. Placeholder and output nodes are - # graph structure, not operators, and recording them makes a fully supported graph - # report its own inputs and outputs as unconverted. + 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 @@ -312,7 +313,7 @@ def is_node_supported( "non-target device region", node_name, ) - self._record_fallback(node, node_name) + self._record_fallback(node, node_name, "explicit non-target device region") return False if self._exceeds_max_tensor_rank(node): @@ -321,7 +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) + self._record_fallback(node, node_name, "tensor rank exceeds TensorRT limit") return False if self._has_complex_dtype(node): @@ -331,7 +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) + self._record_fallback(node, node_name, "complex tensor dtype") return False if self._has_bf16_on_turing(node, settings): @@ -354,7 +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) + self._record_fallback( + node, + node_name, + "data-dependent output shape (fallback_data_dependent_ops=True)", + ) return False if ( @@ -377,7 +382,16 @@ def is_node_supported( else: self.unsupported_operators[node_name] += 1 - self._record_fallback(node, node_name) + 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 index 449631d56b..ab41daa86f 100644 --- a/tests/py/dynamo/partitioning/test_000_fallback_reporting.py +++ b/tests/py/dynamo/partitioning/test_000_fallback_reporting.py @@ -5,34 +5,26 @@ from parameterized import parameterized from torch.testing._internal.common_utils import TestCase, run_tests -WARNING_MARKER = "have no TensorRT converter and will run in PyTorch" +SUMMARY_MARKER = "operator(s)" -class _WarningCollector(logging.Filter): - """A filter rather than a handler, so the record is seen wherever it is emitted from. - - Attaching a handler to the torch_tensorrt logger is not enough: the message comes from a - child logger and torch_tensorrt installs a root handler of its own during compilation. - """ +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 WARNING_MARKER in message: + if SUMMARY_MARKER in message: self.messages.append(message) + self.levels.append(record.levelno) return True class TestFallbackIsReported(TestCase): - """A model that falls back partly to PyTorch should say so at default verbosity. - - The partition report exists already but is DEBUG, which is off by default, so a model - that quietly became several engines plus a PyTorch segment looked exactly like one that - compiled whole. - """ + """Report fallback at INFO without warning about expected behavior.""" @staticmethod def _six_linear_layers() -> torch.nn.ModuleList: @@ -88,15 +80,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return FullySupported().eval().cuda() - def _compile(self, module, inputs, **kwargs): - collector = _WarningCollector() - loggers = [logging.getLogger()] + [ - logging.getLogger(name) - for name in list(logging.root.manager.loggerDict) - if name.startswith("torch_tensorrt") - ] - for each in loggers: - each.addFilter(collector) + 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)), @@ -107,18 +96,29 @@ def _compile(self, module, inputs, **kwargs): **kwargs, ) finally: - for each in loggers: - each.removeFilter(collector) + 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"), - ("complex_dtype_refusal", "_complex_fallback_module", "fft"), + ( + "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): + 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")] @@ -131,6 +131,7 @@ def test_fallback_is_reported(self, _, factory, expected_operator): 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.""" @@ -142,20 +143,18 @@ def test_fully_supported_module_is_silent(self): ) self.assertEqual(messages, []) - def test_requested_fallback_is_silent(self): - """The caller asked for this operator to stay in PyTorch, so warning about it would - be telling them about their own choice.""" + @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.assertTrue( - any("_run_on_gpu" in segment for segment in segments), - f"expected a PyTorch segment, got {segments}", - ) - self.assertEqual(messages, []) + 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 @@ -182,19 +181,45 @@ def test_global_partitioner_silent_on_fully_supported(self): ) self.assertEqual(messages, []) - def test_requested_fallback_by_target_is_silent(self): - """torch_executed_ops accepts an operator target object, not only a qualified name - string. The report filter keys on names, so a target object has to be normalized or - the caller is warned about a fallback they asked for.""" + @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.assertTrue( - any("_run_on_gpu" in segment for segment in segments), - f"expected a PyTorch segment, got {segments}", + 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, [])