From 07433dad11a8e448f32c36a7f258997ab48970d0 Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Wed, 9 Sep 2026 21:03:50 -0700 Subject: [PATCH] fix: support changing input shapes in log1p ## Problem `log1p` runs in PyTorch when input dimensions can vary between calls. Its TensorRT converter also adds one before converting integers to floating point. A large integer can overflow at that step, giving a wrong result. ## Change Enable dynamic shapes and convert integer and boolean inputs to float32 before adding one. Keep int8 and uint8 inputs in PyTorch because the engine cannot use them on this path. Keep float64 inputs there too unless `truncate_double=True` allows float32 arithmetic. With that option, the wrapper converts the engine output back to float64. The input-type check does not require fixed dimensions. Supported types can still run in TensorRT when the shape changes. ## Tests Passed 29/29 tests covering fixed and changing shapes, integer overflow, and unsupported input types. The new dynamic input-type cases check whether an engine was built, then compare shape, data type, and values with PyTorch at batch sizes 1, 3, and 6. Without the new input-type check, exactly the three fallback cases fail. The float64 case with truncation enabled passes with and without that check. Tested on Linux x86_64 with Python 3.12, TensorRT 11.2, and the native runtime. Windows, aarch64, TensorRT-RTX, TensorRT 10.x, and the Python-only runtime were not tested. The converter still uses `log(1 + x)`. Its existing loss of accuracy near zero is not fixed here. --- .../dynamo/conversion/aten_ops_converters.py | 21 +++- .../dynamo/conversion/impl/unary/ops.py | 11 ++ tests/py/dynamo/conversion/test_log1p.py | 113 ++++++++++++++++++ 3 files changed, 144 insertions(+), 1 deletion(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 2a1ce14d2d5..55e3dcfb8d1 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -2126,7 +2126,26 @@ def aten_ops_log10( ) -@dynamo_tensorrt_converter(torch.ops.aten.log1p.default) +def log1p_validator(node: Node, settings: Optional[CompilationSettings] = None) -> bool: + input_node = node.args[0] + input_meta = input_node.meta.get("tensor_meta") + if input_meta is None: + input_meta = input_node.meta.get("val") + if input_meta is None: + return True + # Casting inside the engine cannot repair an unsupported input binding. + if input_meta.dtype in (torch.int8, torch.uint8): + return False + return input_meta.dtype != torch.float64 or ( + settings is not None and settings.truncate_double + ) + + +@dynamo_tensorrt_converter( + torch.ops.aten.log1p.default, + capability_validator=log1p_validator, + supports_dynamic_shapes=True, +) def aten_ops_log1p( ctx: ConversionContext, target: Target, diff --git a/py/torch_tensorrt/dynamo/conversion/impl/unary/ops.py b/py/torch_tensorrt/dynamo/conversion/impl/unary/ops.py index 12f60514574..5e4b661435c 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/unary/ops.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/unary/ops.py @@ -130,6 +130,17 @@ def log1p( """ Computes log(1 + x) for each element of the input tensor. """ + # Cast before the add, not after. TensorRT's log accepts no integer type, and an + # integer add would also wrap: 1 + INT32_MAX goes negative and log returns NaN. + if isinstance(input_val, TRTTensor) and input_val.dtype not in ( + trt.float32, + trt.float16, + trt.bfloat16, + ): + input_val = cast_trt_tensor( + ctx, input_val, trt.float32, f"{name}_input_cast", target, source_ir + ) + one_plus_x = impl.elementwise.add( ctx, target, source_ir, f"{name}_add", input_val, 1 ) diff --git a/tests/py/dynamo/conversion/test_log1p.py b/tests/py/dynamo/conversion/test_log1p.py index 7e59cc16d3d..5c01730d446 100644 --- a/tests/py/dynamo/conversion/test_log1p.py +++ b/tests/py/dynamo/conversion/test_log1p.py @@ -1,5 +1,6 @@ import torch import torch.nn as nn +import torch_tensorrt from parameterized import parameterized from torch.testing._internal.common_utils import run_tests from torch_tensorrt import Input @@ -68,6 +69,118 @@ def forward(self, input): inputs, ) + @parameterized.expand( + [ + ((1,), (3,), (5,)), + ((1, 20), (2, 20), (3, 20)), + ((2, 3, 4), (3, 4, 5), (4, 5, 6)), + ((2, 3, 4, 5), (3, 5, 5, 6), (4, 5, 6, 7)), + ] + ) + def test_log1p_float_dynamic_shape(self, min_shape, opt_shape, max_shape): + class Log1p(nn.Module): + def forward(self, input): + return torch.ops.aten.log1p.default(input) + + input_specs = [ + Input( + dtype=torch.float32, + min_shape=min_shape, + opt_shape=opt_shape, + max_shape=max_shape, + ), + ] + self.run_test_with_dynamic_shape( + Log1p(), + input_specs, + use_dynamo_tracer=True, + ) + + @parameterized.expand( + [ + ("int32", torch.int32), + ("int64", torch.int64), + ("bool", torch.bool), + ] + ) + def test_log1p_integer_input(self, _, dtype): + """TensorRT's log accepts no integer type, so log1p has to cast. It also has to + cast before adding one, or a large integer wraps and the log returns NaN.""" + + class Log1p(nn.Module): + def forward(self, input): + return torch.ops.aten.log1p.default(input) + + if dtype is torch.bool: + inputs = [torch.tensor([True, False, True, True])] + else: + inputs = [torch.tensor([3, 1, 2, 1], dtype=dtype)] + self.run_test( + Log1p(), + inputs, + use_dynamo_tracer=True, + ) + + def test_log1p_int32_near_max(self): + """1 + INT32_MAX wraps negative in the integer domain, and the log of a negative + number is NaN, so this returns a wrong answer rather than failing.""" + + class Log1p(nn.Module): + def forward(self, input): + return torch.ops.aten.log1p.default(input) + + inputs = [torch.tensor([2147483647, 100, 10, 1], dtype=torch.int32)] + self.run_test( + Log1p(), + inputs, + use_dynamo_tracer=True, + ) + + @parameterized.expand( + [ + ("float64_fallback", torch.float64, False, 0), + ("int8_fallback", torch.int8, False, 0), + ("uint8_fallback", torch.uint8, False, 0), + ("float64_truncated", torch.float64, True, 1), + ("float32", torch.float32, False, 1), + ("float16", torch.float16, False, 1), + ("bfloat16", torch.bfloat16, False, 1), + ("int32", torch.int32, False, 1), + ("int64", torch.int64, False, 1), + ("bool", torch.bool, False, 1), + ] + ) + def test_log1p_dynamic_dtype(self, _, dtype, truncate_double, engine_count): + class Log1p(nn.Module): + def forward(self, x): + return torch.ops.aten.log1p.default(x) + + module = Log1p().eval().cuda() + inputs = (torch.randint(0, 5, (3, 5), device="cuda").to(dtype),) + batch = torch.export.Dim("batch", min=1, max=6) + exported = torch.export.export(module, inputs, dynamic_shapes=({0: batch},)) + compiled = torch_tensorrt.dynamo.compile( + exported, + inputs=inputs, + min_block_size=1, + truncate_double=truncate_double, + ) + self.assertEqual( + sum("_run_on_acc" in name for name, _ in compiled.named_children()), + engine_count, + ) + for size in (1, 3, 6): + with self.subTest(size=size): + x = torch.randint(0, 5, (size, 5), device="cuda").to(dtype) + expected = module(x) + if truncate_double: + # The wrapper restores float64 after float32 engine arithmetic. + torch.testing.assert_close( + compiled(x), expected, rtol=1e-5, atol=1e-6 + ) + else: + torch.testing.assert_close(compiled(x), expected) + if __name__ == "__main__": run_tests()