From 8c9ce49e2ab10131282f0d90b6f6a57c8ab912ed Mon Sep 17 00:00:00 2001 From: Patrick Ribbsaeter Date: Sat, 15 Aug 2026 12:50:03 +0200 Subject: [PATCH] [Inductor] Align floating-point sum reduction signed-zero semantics with eager (#187721) In ATen eager and C++ Inductor, summation reductions initialize the accumulator to positive zero (+0.0), evaluating (+0.0) + (-0.0) = +0.0 under IEEE-754. In Triton Inductor reductions and unrolled pointwise reductions, tree reductions and functools.reduce over tensors consisting of negative zeros evaluated (-0.0) + (-0.0) = -0.0, causing F.kl_div(..., log_target=True) and downstream sign-sensitive operations to produce disparate signs between eager and compiled mode. This commit: 1. Adds +0.0 absorption to floating-point sum reductions in Triton codegen (final_reduction) and unrolled pointwise reductions (_unroll_reduction_fn), restoring exact signed-zero parity with eager ATen. 2. Expands regression tests across float16, bfloat16, float32, and float64 covering small unrolled reductions (< unroll threshold), large Triton reductions (>= unroll threshold), and 2D dim-reductions with direct signbit assertions. Fixes #187721 --- test/inductor/test_cpu_repro.py | 39 +++++++++++++++ test/inductor/test_cuda_repro.py | 80 +++++++++++++++++++++++++++++++ torch/_inductor/codegen/triton.py | 3 ++ torch/_inductor/ir.py | 6 ++- 4 files changed, 127 insertions(+), 1 deletion(-) diff --git a/test/inductor/test_cpu_repro.py b/test/inductor/test_cpu_repro.py index 906d6dc..07219ca 100644 --- a/test/inductor/test_cpu_repro.py +++ b/test/inductor/test_cpu_repro.py @@ -7564,6 +7564,45 @@ class CPUReproTests(TestCase): ) self.assertTrue(cuda_storage.has_exceeded_max_reads()) + @parametrize( + "dtype", + (torch.float32, torch.float64), + ) + def test_kl_div_reduction_signbit_parity_187721(self, dtype): + # https://github.com/pytorch/pytorch/issues/187721 + def fn(input_t, target_t): + y = F.kl_div(input_t, target_t, reduction="mean", log_target=True) + sign_val = torch.copysign(torch.tensor(1.0, dtype=torch.float32), y) + return sign_val, y + + input_t = torch.tensor([1000.0, 1000.0], dtype=dtype) + target_t = torch.tensor([-1000.0, -1000.0], dtype=dtype) + + eager_sign, eager_y = fn(input_t, target_t) + compiled_fn = torch.compile(fn, backend="inductor", fullgraph=True) + compiled_sign, compiled_y = compiled_fn(input_t, target_t) + + self.assertEqual(eager_sign, compiled_sign) + self.assertEqual(eager_y.signbit(), compiled_y.signbit()) + self.assertFalse(compiled_y.signbit().item()) + self.assertEqual(compiled_y.dtype, dtype) + + def fn_sum(x): + y = torch.sum(x) + sign_val = torch.copysign(torch.tensor(1.0, dtype=torch.float32), y) + return sign_val, y + + compiled_sum = torch.compile(fn_sum, backend="inductor", fullgraph=True) + + x = torch.tensor([-0.0, -0.0], dtype=dtype) + eager_sign, eager_y = fn_sum(x) + compiled_sign, compiled_y = compiled_sum(x) + + self.assertEqual(eager_sign, compiled_sign) + self.assertEqual(eager_y.signbit(), compiled_y.signbit()) + self.assertFalse(compiled_y.signbit().item()) + self.assertEqual(compiled_y.dtype, dtype) + if __name__ == "__main__": from torch._inductor.test_case import run_tests diff --git a/test/inductor/test_cuda_repro.py b/test/inductor/test_cuda_repro.py index 0941b80..24103cb 100644 --- a/test/inductor/test_cuda_repro.py +++ b/test/inductor/test_cuda_repro.py @@ -3511,6 +3511,86 @@ def triton_poi_fused_add_reflection_pad2d_0(in_ptr0, in_ptr1, out_ptr0, xnumel, self.common(fn4, [y]) + @parametrize( + "dtype", + (torch.float16, torch.bfloat16, torch.float32, torch.float64), + ) + def test_kl_div_reduction_signbit_parity_187721(self, dtype): + # https://github.com/pytorch/pytorch/issues/187721 + def fn(input_t, target_t): + y = F.kl_div(input_t, target_t, reduction="mean", log_target=True) + sign_val = torch.copysign( + torch.tensor(1.0, device=input_t.device, dtype=torch.float32), y + ) + return sign_val, y + + # 1. Small/unrolled reduction path (< unroll threshold) + input_small = torch.tensor([1000.0, 1000.0], dtype=dtype, device=device_type) + target_small = torch.tensor([-1000.0, -1000.0], dtype=dtype, device=device_type) + + eager_sign_s, eager_y_s = fn(input_small, target_small) + compiled_fn = torch.compile(fn, backend="inductor", fullgraph=True) + compiled_sign_s, compiled_y_s = compiled_fn(input_small, target_small) + + self.assertEqual(eager_sign_s, compiled_sign_s) + self.assertEqual(eager_y_s.signbit(), compiled_y_s.signbit()) + self.assertFalse(compiled_y_s.signbit().item()) + self.assertEqual(compiled_y_s.dtype, dtype) + + # 2. Large reduction reaching Triton final_reduction (>= unroll threshold) + input_large = torch.full((1024,), 1000.0, dtype=dtype, device=device_type) + target_large = torch.full((1024,), -1000.0, dtype=dtype, device=device_type) + + eager_sign_l, eager_y_l = fn(input_large, target_large) + compiled_sign_l, compiled_y_l = compiled_fn(input_large, target_large) + + self.assertEqual(eager_sign_l, compiled_sign_l) + self.assertEqual(eager_y_l.signbit(), compiled_y_l.signbit()) + self.assertFalse(compiled_y_l.signbit().item()) + self.assertEqual(compiled_y_l.dtype, dtype) + + def fn_sum(x): + y = torch.sum(x) + sign_val = torch.copysign( + torch.tensor(1.0, device=x.device, dtype=torch.float32), y + ) + return sign_val, y + + compiled_sum = torch.compile(fn_sum, backend="inductor", fullgraph=True) + + # 3. Direct sum: small unrolled reduction (2 elements) + x_small = torch.tensor([-0.0, -0.0], dtype=dtype, device=device_type) + eager_sign, eager_y = fn_sum(x_small) + compiled_sign, compiled_y = compiled_sum(x_small) + + self.assertEqual(eager_sign, compiled_sign) + self.assertEqual(eager_y.signbit(), compiled_y.signbit()) + self.assertFalse(compiled_y.signbit().item()) + self.assertEqual(compiled_y.dtype, dtype) + + # 4. Direct sum: large Triton reduction (1024 elements) + x_large = torch.full((1024,), -0.0, dtype=dtype, device=device_type) + eager_sign_lg, eager_y_lg = fn_sum(x_large) + compiled_sign_lg, compiled_y_lg = compiled_sum(x_large) + + self.assertEqual(eager_sign_lg, compiled_sign_lg) + self.assertEqual(eager_y_lg.signbit(), compiled_y_lg.signbit()) + self.assertFalse(compiled_y_lg.signbit().item()) + self.assertEqual(compiled_y_lg.dtype, dtype) + + # 5. Direct sum: dim=1 reduction retaining outer dimension + def fn_dim(x): + return torch.sum(x, dim=1) + + compiled_dim = torch.compile(fn_dim, backend="inductor", fullgraph=True) + x_dim = torch.full((16, 64), -0.0, dtype=dtype, device=device_type) + eager_dim = fn_dim(x_dim) + compiled_dim_out = compiled_dim(x_dim) + + self.assertEqual(eager_dim.signbit(), compiled_dim_out.signbit()) + self.assertFalse(compiled_dim_out.signbit().any().item()) + self.assertEqual(compiled_dim_out.dtype, dtype) + if __name__ == "__main__": from torch._inductor.test_case import run_tests diff --git a/torch/_inductor/codegen/triton.py b/torch/_inductor/codegen/triton.py index 2f8469f..9807d3f 100644 --- a/torch/_inductor/codegen/triton.py +++ b/torch/_inductor/codegen/triton.py @@ -5491,6 +5491,9 @@ class TritonKernel(SIMDKernel[TritonCSEVariable]): f"{triton_reduction_fn}({value}, {dim}{reduction_ordering})", value.shape, ) + if reduction_type == "sum" and dtype.is_floating_point: + # IEEE-754 (-0.0) + (-0.0) == -0.0, but eager sum initializes to +0.0 (+0.0 + -0.0 == +0.0); absorb negative zero. + result = f"({result} + 0.0)" if result_type is not None: result = f"{result}.to({self.dtype_to_str(result_type)})" diff --git a/torch/_inductor/ir.py b/torch/_inductor/ir.py index 2ca2223..1827930 100644 --- a/torch/_inductor/ir.py +++ b/torch/_inductor/ir.py @@ -1701,7 +1701,7 @@ class Reduction(Loops): combine_fn = get_reduction_combine_fn(reduction_type, src_dtype) def fn(index: Sequence[_IntLike]) -> Any: - return functools.reduce( + res = functools.reduce( combine_fn, ( value_fn(index, rindex) @@ -1710,6 +1710,10 @@ class Reduction(Loops): ) ), ) + if reduction_type == "sum" and src_dtype.is_floating_point: + # IEEE-754 (-0.0) + (-0.0) == -0.0, but eager sum initializes to +0.0 (+0.0 + -0.0 == +0.0); absorb negative zero. + res = ops.add(res, ops.constant(0.0, src_dtype)) + return res value_fn: Callable[[Sequence[_IntLike], Sequence[_IntLike]], Any] if reduction_type in ( -- 2.55.0.windows.3