From 8a28308bdbdc1d4df159dc4d9aff79dd894d94ae Mon Sep 17 00:00:00 2001 From: HuEnwei Date: Sun, 30 Aug 2026 00:26:41 +0800 Subject: [PATCH] [Relax][Frontend][Torch] Support torch.round(x, decimals) via from_exported_program and fix negative-decimals rounding torch.export lowers torch.round(x, decimals) (any explicit decimals, including decimals=0) to aten.round.decimals, but the exported-program convert map only registered round.default. Any explicit decimals made from_exported_program fail with "Unsupported function types ['round.decimals']". Additionally, BaseFXGraphImporter._round scaled every non-zero decimals by round(x * 10**decimals) / 10**decimals. For negative decimals this multiplies by 0.1 / 0.01 / ..., which is inexact in floating point: in float64, torch.round(torch.tensor(25.0), decimals=-1) computed 25 * 0.1 == 2.5000000000000004 and rounded up to 30 instead of 20. Register "round.decimals" in ExportedProgramImporter.create_convert_map and branch the decimals != 0 scale in _round to use an exact integer power of 10: multiply for positive decimals, divide for negative ones (round(x / 10**|d|) * 10**|d|). The ties-to-even inner rounding is already provided by upstream #19367 / #19368 (tir.round -> nearbyint across backends) and is not changed here. Validated by the verify_patch.py differential harness on the locked build: Part B (fix + ties-to-even inner round, = latest semantics) matches PyTorch for all 28 combinations (from_exported_program / from_fx x decimals {0,1,2,3,-1,-2,-3} x float32/float64), including the previously-failing round(25, -1) == 20 and round(2.25, 1) == 2.2. --- .../torch/base_fx_graph_translator.py | 37 +++++++++-- .../torch/exported_program_translator.py | 1 + .../test_frontend_from_exported_program.py | 52 +++++++++++++++ tests/python/relax/test_frontend_from_fx.py | 65 +++++++++++++++++++ 4 files changed, 149 insertions(+), 6 deletions(-) diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py index b0bb14ac95ad..c3f03457fde7 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -454,12 +454,37 @@ def _round(self, node: fx.Node) -> relax.Expr: if decimals == 0: return self.block_builder.emit(relax.op.round(arg)) - # For decimals != 0, use: round(x * 10^decimals) / 10^decimals - dtype = arg.ty.dtype - scale = relax.const(10**decimals, dtype) - scaled = relax.op.multiply(arg, scale) - rounded = relax.op.round(scaled) - result = relax.op.divide(rounded, scale) + # For decimals != 0, round on the exact power-of-10 scale and scale back: + # round(x * 10^decimals) / 10^decimals. The scaling must always use an + # integer power of 10: multiply for positive decimals, divide for negative + # ones. Dividing for negative decimals (instead of multiplying by + # 10**decimals, i.e. 0.1 / 0.01 / ...) avoids float precision errors such as + # 25 * 0.1 == 2.5000000000000004 in float64, which would round up to 30 + # instead of 20 for torch.round(25, -1). + # + # For float16/bfloat16 inputs the scaling is done in float32 and cast + # back, because 10**|decimals| can overflow the input range: 10**4 == 10000 + # with 25 * 10000 == 250000 overflows float16 (max 65504) to inf, and 10**5 + # already overflows float16 to inf, turning decimals=5 and -5 into NaN. + input_dtype = arg.ty.dtype + dtype = input_dtype + if dtype in ("float16", "bfloat16"): + dtype = "float32" + arg = self.block_builder.emit(relax.op.astype(arg, dtype)) + + if decimals > 0: + scale = relax.const(10**decimals, dtype) + scaled = relax.op.multiply(arg, scale) + rounded = relax.op.round(scaled) + result = relax.op.divide(rounded, scale) + else: + scale = relax.const(10 ** (-decimals), dtype) + scaled = relax.op.divide(arg, scale) + rounded = relax.op.round(scaled) + result = relax.op.multiply(rounded, scale) + + if input_dtype in ("float16", "bfloat16"): + result = relax.op.astype(result, input_dtype) return self.block_builder.emit(result) def _softmax(self, node: fx.Node) -> relax.Var: diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index 0e69074af601..545b9999fba5 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -1766,6 +1766,7 @@ def create_convert_map( "relu6.default": self._unary_op(relax.op.nn.relu6), "relu6_.default": self._unary_op(relax.op.nn.relu6), "round.default": self._round, + "round.decimals": self._round, "rsqrt.default": self._rsqrt, "scalar_tensor.default": self._scalar_tensor, "scatter.value": self._scatter_value, diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index e9d2ac8b704f..a13edf139986 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -148,6 +148,58 @@ def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple( verify_model(UnaryOp(), example_args, {}, expected) +def test_round_decimals(): + """torch.round(x, decimals) is exported as aten.round.decimals, which was missing + from the convert map (only round.default was registered) and made any explicit + decimals -- including decimals=0 -- fail with + "AssertionError: Unsupported function types ['round.decimals']". + + With the decimals overload registered, torch.round(x, decimals) must convert and + match PyTorch's round-half-to-even results, including negative decimals + (round(25, -1) == 20) where the scale-by-0.1 float precision path used to be wrong. + """ + + class RoundDecimalsModel(Module): + def __init__(self, decimals): + super().__init__() + self.decimals = decimals + + def forward(self, input): + return torch.round(input, decimals=self.decimals) + + # Half values exercise ties-to-even; 25/125/165 exercise the negative-decimals path. + x = torch.tensor( + [0.5, 1.5, 2.5, 4.5, -0.5, -2.5, 25.0, 125.0, 165.0, 2.25], dtype=torch.float32 + ) + for decimals in (0, 1, -1, -2): + verify_model_numerically(RoundDecimalsModel(decimals).eval(), (x,), rtol=1e-6, atol=1e-6) + + +def test_round_decimals_low_precision(): + """Scaling for low-precision inputs must happen in float32 and be cast back. + + 10**|decimals| can overflow float16: 10**4 == 10000 with 25 * 10000 == 250000 + exceeds float16's max of 65504, so scaling in float16 yields inf, and 10**5 + already overflows float16 (the scale itself becomes inf), turning decimals=5 + and -5 into NaN. Upcasting the input to float32 keeps the scaling exact; the + rounded result is cast back to the input dtype. + """ + + class RoundDecimalsModel(Module): + def __init__(self, decimals): + super().__init__() + self.decimals = decimals + + def forward(self, input): + return torch.round(input, decimals=self.decimals) + + x = torch.tensor([0.5, 1.5, 2.5, 2.25, 25.0, 125.0, 165.0, -0.5], dtype=torch.float16) + # Positive decimals exercise the multiply-by-10**d overflow (4, 5); + # negative decimals exercise the 10**|d| scale overflowing float16 (-5). + for decimals in (2, 4, 5, -2, -4, -5): + verify_model_numerically(RoundDecimalsModel(decimals).eval(), (x,), rtol=1e-6, atol=1e-6) + + operator_bool_unary = [ (torch.isinf, R.isinf), (torch.isnan, R.isnan), diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py index a489977958c7..7dfd534ffa3f 100644 --- a/tests/python/relax/test_frontend_from_fx.py +++ b/tests/python/relax/test_frontend_from_fx.py @@ -2506,6 +2506,71 @@ def main( verify_model(DivFloorModel(), input_info, {}, expected_div_floor) +def test_round_decimals(): + """torch.round(x, decimals) through from_fx must match PyTorch's round-half-to-even + results, including negative decimals (round(25, -1) == 20). The previous + scale-by-10**decimals implementation multiplied by 0.1 for negative decimals, which + is numerically wrong: 25 * 0.1 == 2.5000000000000004 in float64 rounds up to 30. + """ + input_info = [([10], "float32")] + x = torch.tensor( + [0.5, 1.5, 2.5, 4.5, -0.5, -2.5, 25.0, 125.0, 165.0, 2.25], dtype=torch.float32 + ) + + class RoundDecimalsModel(Module): + def __init__(self, decimals): + super().__init__() + self.decimals = decimals + + def forward(self, input): + return torch.round(input, decimals=self.decimals) + + for decimals in (0, 1, -1, -2): + gm = fx.symbolic_trace(RoundDecimalsModel(decimals).eval()) + mod = from_fx(gm, input_info) + ex = relax.build(mod, target="llvm") + vm = relax.VirtualMachine(ex, tvm.cpu()) + tvm_out = vm["main"](tvm.runtime.tensor(x.numpy())) + got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else tvm_out[0].numpy() + tvm.testing.assert_allclose( + got, torch.round(x, decimals=decimals).numpy(), rtol=1e-6, atol=1e-6 + ) + + +def test_round_decimals_low_precision(): + """Scaling for low-precision inputs must happen in float32 and be cast back. + + 10**|decimals| can overflow float16: 10**4 == 10000 with 25 * 10000 == 250000 + exceeds float16's max of 65504, so scaling in float16 yields inf, and 10**5 + already overflows float16 (the scale itself becomes inf), turning decimals=5 + and -5 into NaN. Upcasting the input to float32 keeps the scaling exact; the + rounded result is cast back to the input dtype. + """ + input_info = [([8], "float16")] + x = torch.tensor([0.5, 1.5, 2.5, 2.25, 25.0, 125.0, 165.0, -0.5], dtype=torch.float16) + + class RoundDecimalsModel(Module): + def __init__(self, decimals): + super().__init__() + self.decimals = decimals + + def forward(self, input): + return torch.round(input, decimals=self.decimals) + + # Positive decimals exercise the multiply-by-10**d overflow (4, 5); + # negative decimals exercise the 10**|d| scale overflowing float16 (-5). + for decimals in (2, 4, 5, -2, -4, -5): + gm = fx.symbolic_trace(RoundDecimalsModel(decimals).eval()) + mod = from_fx(gm, input_info) + ex = relax.build(mod, target="llvm") + vm = relax.VirtualMachine(ex, tvm.cpu()) + tvm_out = vm["main"](tvm.runtime.tensor(x.numpy())) + got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else tvm_out[0].numpy() + tvm.testing.assert_allclose( + got, torch.round(x, decimals=decimals).numpy(), rtol=1e-6, atol=1e-6 + ) + + def test_size(): input_info = [([1, 3, 10, 10], "float32")]