From 7c18c8918a36cb9ec4f0c1ede558b32508e3a103 Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Wed, 9 Sep 2026 21:03:53 -0700 Subject: [PATCH] fix: support dynamic inputs in clone and _to_copy ## Problem Copying a model input can fall back to PyTorch or fail during conversion when its dimensions can vary between calls. The special `clone` and `_to_copy` converters used at model boundaries did not declare support for those dynamic shapes. ## Change Enable dynamic shapes for both converters. Their implementations do not need fixed dimensions. Add checks that leave unsupported input types and memory layouts in PyTorch, including copies whose non-contiguous layout the TensorRT layer cannot preserve. Test float64 and uint8 fallback separately. Each case checks that no engine was built and that the result's data type and values match ordinary PyTorch. ## Tests Passed 7/7 focused tests. The float32 `clone` and `_to_copy` cases build an engine and run at batch sizes 1, 3, and 6. Other cases check channels-last memory layout, float64 and uint8 fallback, and float64 input with conversion to float32 enabled. Tested on Linux x86_64 with Python 3.12, TensorRT 11.2, and the native runtime. Windows, aarch64, TensorRT-RTX, TensorRT 10.x, the Python-only runtime, and export/save/load were not tested. Dynamic dimensions of size zero are outside this change. --- .../dynamo/conversion/aten_ops_converters.py | 79 ++++++++- .../test_clone_placeholder_dynamic.py | 150 ++++++++++++++++++ 2 files changed, 227 insertions(+), 2 deletions(-) create mode 100644 tests/py/dynamo/conversion/test_clone_placeholder_dynamic.py diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 2a1ce14d2d..01881cd4fb 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -1811,13 +1811,88 @@ def aten_ops_clone_copy_dtype( ) +def _dynamic_placeholder_copy_supported( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + """Reject the cases this converter cannot serve once it accepts a dynamic input. + + Each of these ran correctly in PyTorch before the placeholder path declared dynamic + shape support, so they fall back rather than fail: + + A memory_format the copy cannot preserve. The layer produces standard contiguous + strides, so a channels-last copy would return the right values with the wrong layout, + silently. + + An input dtype the engine cannot bind. float64 needs truncate_double, and without it + the binding expects float32 and rejects the caller's tensor at the first call. uint8 + aborts the build outright. + + A non-contiguous input with no memory_format. Default clone preserves the input strides, + so a channels-last or transposed input keeps its layout in eager, but the copy the layer + builds is contiguous. That returns the right values with the wrong strides, silently. + """ + if node.kwargs.get("memory_format") is not None: + _LOGGER.debug( + f"{node.target} with memory_format={node.kwargs['memory_format']} cannot " + "preserve strides, falling back" + ) + return False + + input_node = node.args[0] if node.args else None + input_meta = input_node.meta.get("val") if isinstance(input_node, Node) else None + if not isinstance(input_meta, torch.Tensor): + return True + + if not input_meta.is_contiguous(): + # No memory_format was given (checked above), so this is a preserve-format copy of a + # non-contiguous input, which the contiguous layer cannot reproduce. + _LOGGER.debug( + f"{node.target} preserves the strides of a non-contiguous input, which the " + "copy cannot, falling back" + ) + return False + + if input_meta.dtype == torch.uint8: + _LOGGER.debug( + f"{node.target} with a uint8 input is not supported, falling back" + ) + return False + if input_meta.dtype == torch.float64 and not ( + settings is not None and settings.truncate_double + ): + _LOGGER.debug( + f"{node.target} with a float64 input needs truncate_double=True, falling back" + ) + return False + + return True + + +def _clone_placeholder_validator( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + return is_only_operator_on_placeholder( + node, settings + ) and _dynamic_placeholder_copy_supported(node, settings) + + +def _to_copy_placeholder_validator( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + return to_copy_dtype_validator(placeholder_only=True)( + node, settings + ) and _dynamic_placeholder_copy_supported(node, settings) + + @dynamo_tensorrt_converter( torch.ops.aten.clone.default, - capability_validator=is_only_operator_on_placeholder, + capability_validator=_clone_placeholder_validator, + supports_dynamic_shapes=True, ) @dynamo_tensorrt_converter( torch.ops.aten._to_copy.default, - capability_validator=to_copy_dtype_validator(placeholder_only=True), + capability_validator=_to_copy_placeholder_validator, + supports_dynamic_shapes=True, ) def aten_ops_clone_copy_placeholder( ctx: ConversionContext, diff --git a/tests/py/dynamo/conversion/test_clone_placeholder_dynamic.py b/tests/py/dynamo/conversion/test_clone_placeholder_dynamic.py new file mode 100644 index 0000000000..ce7b7b7b72 --- /dev/null +++ b/tests/py/dynamo/conversion/test_clone_placeholder_dynamic.py @@ -0,0 +1,150 @@ +import torch +import torch_tensorrt +from parameterized import parameterized +from torch.testing._internal.common_utils import TestCase, run_tests + + +class TestClonePlaceholderDynamicShape(TestCase): + """clone and _to_copy each have two registrations, a general one and one for the case + where the node takes a graph input and its result is a graph output. Both call the same + dtype cast, which reads no shapes, so both must accept a dynamic input. + + The cases the cast cannot serve have to keep falling back, since they run correctly in + PyTorch: a memory_format it cannot preserve, and an input dtype the engine cannot bind. + """ + + @staticmethod + def _compile(module, inputs, dim_max=6, **kwargs): + batch = torch.export.Dim("batch", min=1, max=dim_max) + exported = torch.export.export(module, inputs, dynamic_shapes=({0: batch},)) + return torch_tensorrt.dynamo.compile( + exported, arg_inputs=inputs, min_block_size=1, **kwargs + ) + + @staticmethod + def _engines(compiled): + return sum(1 for name, _ in compiled.named_children() if "_run_on_acc" in name) + + @parameterized.expand( + [ + ("clone", lambda x: torch.ops.aten.clone.default(x)), + ( + "to_copy", + lambda x: torch.ops.aten._to_copy.default(x, dtype=torch.float16), + ), + ] + ) + def test_placeholder_copy_with_dynamic_dim(self, _, operation): + """Both registrations are fixed, so both need a case. Reverting either flag alone + left the other one's test green.""" + + class OnlyCopy(torch.nn.Module): + def forward(self, x): + return operation(x) + + module = OnlyCopy().eval().cuda() + inputs = (torch.randn(3, 5, device="cuda"),) + compiled = self._compile(module, inputs) + + self.assertEqual( + self._engines(compiled), + 1, + f"expected one engine, got {[n for n, _ in compiled.named_children()]}", + ) + # The declared range is 1 to 6, so run both ends of it as well as the middle. + for size in (1, 3, 6): + sized = (torch.randn(size, 5, device="cuda"),) + torch.testing.assert_close(compiled(*sized), module(*sized)) + + def test_channels_last_falls_back(self): + """The cast produces standard contiguous strides, so a channels-last copy would + return the right values with the wrong layout. Comparing values alone cannot see + that, which is why this asserts on the strides.""" + + class CloneChannelsLast(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.clone.default( + x, memory_format=torch.channels_last + ) + + module = CloneChannelsLast().eval().cuda() + inputs = ( + torch.randn(3, 4, 5, 3, device="cuda").to( + memory_format=torch.channels_last + ), + ) + compiled = self._compile(module, inputs) + result = compiled(*inputs) + + self.assertTrue( + result.is_contiguous(memory_format=torch.channels_last), + f"expected channels-last strides, got {result.stride()}", + ) + torch.testing.assert_close(result, module(*inputs)) + + def test_default_clone_of_channels_last_falls_back(self): + """A default clone with no memory_format still preserves the input strides, so a + channels-last input keeps its layout in eager. Export represents this with empty + kwargs, so it is not caught by the explicit-memory_format check; the copy the layer + builds is contiguous, which is the same silent layout change one level down.""" + + class CloneDefault(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.clone.default(x) + + module = CloneDefault().eval().cuda() + inputs = ( + torch.randn(3, 4, 5, 3, device="cuda").to( + memory_format=torch.channels_last + ), + ) + compiled = self._compile(module, inputs) + result = compiled(*inputs) + + self.assertTrue( + result.is_contiguous(memory_format=torch.channels_last), + f"expected channels-last strides, got {result.stride()}", + ) + torch.testing.assert_close(result, module(*inputs)) + + @parameterized.expand([("float64", torch.float64), ("uint8", torch.uint8)]) + def test_unbindable_input_dtype_falls_back(self, _, dtype): + """float64 needs truncate_double, and without it the binding expects float32 and + rejects the caller's tensor. uint8 aborts the build. Both run in PyTorch.""" + + class OnlyClone(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.clone.default(x) + + module = OnlyClone().eval().cuda() + if dtype == torch.uint8: + inputs = (torch.randint(0, 255, (3, 5), dtype=dtype, device="cuda"),) + else: + inputs = (torch.randn(3, 5, dtype=dtype, device="cuda"),) + + compiled = self._compile(module, inputs) + result = compiled(*inputs) + + self.assertEqual(self._engines(compiled), 0) + self.assertEqual(result.dtype, dtype) + torch.testing.assert_close(result, module(*inputs)) + + def test_float64_converts_when_truncation_is_allowed(self): + """With truncate_double the engine can bind it, so it should not fall back.""" + + class OnlyClone(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.clone.default(x) + + module = OnlyClone().eval().cuda() + inputs = (torch.randn(3, 5, dtype=torch.float64, device="cuda"),) + compiled = self._compile(module, inputs, truncate_double=True) + + self.assertEqual(self._engines(compiled), 1) + torch.testing.assert_close( + compiled(*inputs).double(), module(*inputs), rtol=1e-3, atol=1e-3 + ) + + +if __name__ == "__main__": + run_tests()