From 3578b8b905d544691880d9ae424bad3e8e7e8302 Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Wed, 9 Sep 2026 21:03:50 -0700 Subject: [PATCH 1/2] fix: reject unsupported operators when full compilation is required ## Problem `require_full_compilation=True` asks the compiler to reject a model with unsupported operators. A small model can return early, before that check happens. The call then succeeds even though some operations cannot run in TensorRT. ## Change Check for unsupported operators before returning early. If full compilation is required, raise an error that reports how many operators are unsupported. `dryrun` still reports the problem without raising. Preserve the existing behavior for empty models and small models whose operators are all supported. Below `min_block_size`, the minimum operation count, they can return successfully without building a TensorRT engine. The tests and description now distinguish successful return from engine construction. ## Tests Passed 7/7 focused tests. They check rejection, the error count, dry-run behavior, and successful return for small supported and empty models. Larger supported models are checked for both an engine and output matching ordinary PyTorch. Without the early check, the two rejection tests fail. The small supported model still returns without an engine. Tested on Linux x86_64 with Python 3.12, TensorRT 11.2, and the native runtime. Windows, aarch64, TensorRT-RTX, the Python-only runtime, and TorchScript entry points were not tested. --- py/torch_tensorrt/dynamo/_compiler.py | 16 +- .../test_000_require_full_compilation.py | 141 ++++++++++++++++++ 2 files changed, 156 insertions(+), 1 deletion(-) create mode 100644 tests/py/dynamo/partitioning/test_000_require_full_compilation.py diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index a5122a8524..2824b381a6 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -1310,10 +1310,24 @@ def compile_module( ) # If the number of supported operations is 0 or less than the block size, skip the subgraph - # TODO: Add condition to second expression below when require_full_compilation is added if num_supported_ops == 0 or ( num_supported_ops < settings.min_block_size and not settings.dryrun ): + # Only refuse when an operator genuinely has no converter. A graph whose every + # operator converts is fully supported however few of them there are, and all + # three partitioners deliberately disregard min_block_size in that case, so + # raising here would contradict them. dryrun is documented as the way to inspect + # what would fall back, so it stays non fatal. + if ( + settings.require_full_compilation + and num_supported_ops < total_ops + and not settings.dryrun + ): + raise AssertionError( + f"require_full_compilation=True was specified, but " + f"{total_ops - num_supported_ops} of {total_ops} operations in this " + f"subgraph have no TensorRT converter" + ) logger.warning( f"{num_supported_ops} supported operations detected in subgraph containing {total_ops} computational nodes. " f"Skipping this subgraph, since min_block_size was detected to be {settings.min_block_size}" diff --git a/tests/py/dynamo/partitioning/test_000_require_full_compilation.py b/tests/py/dynamo/partitioning/test_000_require_full_compilation.py new file mode 100644 index 0000000000..7ae10da8bb --- /dev/null +++ b/tests/py/dynamo/partitioning/test_000_require_full_compilation.py @@ -0,0 +1,141 @@ +import torch +import torch_tensorrt +from parameterized import parameterized +from torch.testing._internal.common_utils import TestCase, run_tests + + +class TestRequireFullCompilation(TestCase): + """The early return must reject unsupported operators with full compilation required. + + Fully supported graphs below min_block_size still return without building an engine. + These controls check that they are not rejected, not that they run in TensorRT. + """ + + @staticmethod + def _six_linear_layers(): + return torch.nn.ModuleList([torch.nn.Linear(64, 64) for _ in range(6)]) + + @classmethod + def _no_converter_module(cls): + """Small, and its only non-trivial operator has no converter.""" + + class NoConverter(torch.nn.Module): + def forward(self, x): + return x + torch.normal( + 0.0, 1.0, size=x.shape, device=x.device, dtype=x.dtype + ) + + return NoConverter().eval().cuda() + + @classmethod + def _small_fully_convertible_module(cls): + """Every operator converts, and there are fewer than min_block_size of them.""" + + class SmallFullyConvertible(torch.nn.Module): + def forward(self, x): + return torch.relu(x + 1) + + return SmallFullyConvertible().eval().cuda() + + @classmethod + def _large_fully_convertible_module(cls): + class LargeFullyConvertible(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 LargeFullyConvertible().eval().cuda() + + @staticmethod + def _compile(module, inputs, **kwargs): + return torch_tensorrt.dynamo.compile( + torch.export.export(module, tuple(inputs)), + inputs=list(inputs), + require_full_compilation=True, + enabled_precisions={torch.float32}, + truncate_double=True, + **kwargs, + ) + + @staticmethod + def _segments(compiled): + return [name for name, _ in compiled.named_children() if "_run_on" in name] + + def test_operator_without_converter_raises(self): + """The case the flag exists for: something in the graph cannot be converted.""" + inputs = [torch.randn(8, 64, device="cuda")] + with self.assertRaisesRegex(AssertionError, "have no TensorRT converter"): + self._compile(self._no_converter_module(), inputs) + + def test_message_counts_the_unconvertible_operators(self): + """The count has to name the operators that cannot be converted. Counting the + convertible ones instead produced "only 2 of 2 operations are convertible", which + contradicts itself.""" + inputs = [torch.randn(8, 64, device="cuda")] + with self.assertRaises(AssertionError) as raised: + self._compile(self._no_converter_module(), inputs) + message = str(raised.exception) + self.assertIn("require_full_compilation=True", message) + self.assertIn("1 of 2 operations", message) + self.assertIn("have no TensorRT converter", message) + + def test_small_fully_convertible_module_is_not_rejected(self): + """The existing early return for a supported graph is still allowed.""" + inputs = [torch.randn(8, 64, device="cuda")] + compiled = self._compile(self._small_fully_convertible_module(), inputs) + torch.testing.assert_close( + compiled(*inputs), + self._small_fully_convertible_module()(*inputs), + rtol=5e-3, + atol=5e-3, + ) + + def test_empty_graph_is_not_rejected(self): + """A module that returns its input has nothing to convert and nothing to refuse.""" + + class Identity(torch.nn.Module): + def forward(self, x): + return x + + inputs = [torch.randn(8, 64, device="cuda")] + compiled = self._compile(Identity().eval().cuda(), inputs) + torch.testing.assert_close(compiled(*inputs), inputs[0]) + + def test_dryrun_reports_instead_of_raising(self): + """dryrun is documented as the way to inspect what would fall back, so it must + not become fatal.""" + inputs = [torch.randn(8, 64, device="cuda")] + compiled = self._compile(self._no_converter_module(), inputs, dryrun=True) + self.assertIsNotNone(compiled) + + @parameterized.expand( + [("default_min_block_size", {}), ("min_block_size_1", {"min_block_size": 1})] + ) + def test_large_fully_convertible_module_builds_one_engine(self, _, kwargs): + """Guards against over correction. Also asserts an engine was really built, since + a check for the absence of a PyTorch segment passes on an empty segment list.""" + inputs = [torch.randn(8, 64, device="cuda")] + module = self._large_fully_convertible_module() + compiled = self._compile(module, inputs, **kwargs) + segments = self._segments(compiled) + self.assertTrue( + any("_run_on_acc" in segment for segment in segments), + f"expected a TensorRT engine, got {segments}", + ) + self.assertFalse( + any("_run_on_gpu" in segment for segment in segments), + f"expected no PyTorch segment, got {segments}", + ) + torch.testing.assert_close( + compiled(*inputs), module(*inputs), rtol=5e-3, atol=5e-3 + ) + + +if __name__ == "__main__": + run_tests() From 726b774749b8a72ddeca7888c0ee666ab0da92ab Mon Sep 17 00:00:00 2001 From: PyTorch Bot Date: Wed, 23 Sep 2026 13:09:01 -0700 Subject: [PATCH 2/2] test: handle fully compiled TensorRT modules The compiler can return a TensorRT module directly when it compiles the whole graph. The full-compilation test only looked for named child segments, so it missed that engine. Count TensorRT modules including the root. Keep checking that exactly one engine exists, no PyTorch segment remains, and the output matches ordinary PyTorch. --- .../partitioning/test_000_require_full_compilation.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/py/dynamo/partitioning/test_000_require_full_compilation.py b/tests/py/dynamo/partitioning/test_000_require_full_compilation.py index 7ae10da8bb..25a9a5ff14 100644 --- a/tests/py/dynamo/partitioning/test_000_require_full_compilation.py +++ b/tests/py/dynamo/partitioning/test_000_require_full_compilation.py @@ -2,6 +2,7 @@ import torch_tensorrt from parameterized import parameterized from torch.testing._internal.common_utils import TestCase, run_tests +from torch_tensorrt.dynamo.runtime import TorchTensorRTModule class TestRequireFullCompilation(TestCase): @@ -124,10 +125,12 @@ def test_large_fully_convertible_module_builds_one_engine(self, _, kwargs): module = self._large_fully_convertible_module() compiled = self._compile(module, inputs, **kwargs) segments = self._segments(compiled) - self.assertTrue( - any("_run_on_acc" in segment for segment in segments), - f"expected a TensorRT engine, got {segments}", - ) + engines = [ + child + for child in compiled.modules() + if isinstance(child, TorchTensorRTModule) + ] + self.assertEqual(len(engines), 1) self.assertFalse( any("_run_on_gpu" in segment for segment in segments), f"expected no PyTorch segment, got {segments}",