已合并
fix(triton_experimental): propagate NaN in min/max reductions #44773
伦创建于 19 天前
fix(triton_experimental): propagate NaN in min/max reductions #44773
已合并
共 3 个文件变更+44-4
| @@ -4,9 +4,15 @@ | |||
| 4 | from unittest import mock | 4 | from unittest import mock |
| 5 | 5 | ||
| 6 | import sympy | 6 | import sympy |
| 7 | +import torch | ||
| 7 | from torch._inductor.codegen.triton import IndexingOptions, TritonKernel | 8 | from torch._inductor.codegen.triton import IndexingOptions, TritonKernel |
| 8 | from torch._inductor.fx_passes.control_dependencies import control_deps | 9 | from torch._inductor.fx_passes.control_dependencies import control_deps |
| 9 | -from torch.testing._internal.common_utils import TestCase, run_tests | 10 | +from torch.testing._internal.common_utils import ( |
| 11 | + TestCase, | ||
| 12 | + instantiate_parametrized_tests, | ||
| 13 | + parametrize, | ||
| 14 | + run_tests, | ||
| 15 | +) | ||
| 10 | from torch.utils._ordered_set import OrderedSet | 16 | from torch.utils._ordered_set import OrderedSet |
| 11 | 17 | ||
| 12 | from torch_npu._inductor.triton_experimental import lowering as experimental_lowering | 18 | from torch_npu._inductor.triton_experimental import lowering as experimental_lowering |
| @@ -14,6 +20,7 @@ from torch_npu._inductor.triton_experimental import lowering_override_list | |||
| 14 | from torch_npu._inductor.triton_experimental.codegen import triton as npu_triton_codegen | 20 | from torch_npu._inductor.triton_experimental.codegen import triton as npu_triton_codegen |
| 15 | 21 | ||
| 16 | 22 | ||
| 23 | + | ||
| 17 | class TestTritonExperimentalRegressions(TestCase): | 24 | class TestTritonExperimentalRegressions(TestCase): |
| 18 | def test_control_deps_is_not_replaced_with_fallback(self): | 25 | def test_control_deps_is_not_replaced_with_fallback(self): |
| 19 | self.assertIn(control_deps, lowering_override_list.KEEP_UPSTREAM_LOWERING) | 26 | self.assertIn(control_deps, lowering_override_list.KEEP_UPSTREAM_LOWERING) |
| @@ -86,5 +93,32 @@ class TestTritonExperimentalRegressions(TestCase): | |||
| 86 | self.assertEqual(result.expand_str, "[1, 1]") | 93 | self.assertEqual(result.expand_str, "[1, 1]") |
| 87 | self.assertEqual(result.expand_shape, (1, 1)) | 94 | self.assertEqual(result.expand_shape, (1, 1)) |
| 88 | 95 | ||
| 96 | + | ||
| 97 | + | ||
| 98 | + def test_min_max_reductions_propagate_nan(self, op_name, dim): | ||
| 99 | + # Regression: split reduction's accumulate loop used bare tl.minimum/ | ||
| 100 | + # tl.maximum, which drops NaN on Ascend, so compiled min/max returned | ||
| 101 | + # non-NaN where eager does. The fix adds tl.PropagateNan.ALL. | ||
| 102 | + x = torch.randn(8, 64, 1024, device="npu") * 2000 | ||
| 103 | + x[3, 5, :] = float("nan") | ||
| 104 | + | ||
| 105 | + def reduce_values(t): | ||
| 106 | + # torch.min/max(dim=) return (values, indices); amin/amax return a | ||
| 107 | + # tensor directly. Normalize to values-only so [0] never slices a | ||
| 108 | + # data row. | ||
| 109 | + out = getattr(torch, op_name)(t, dim=dim) | ||
| 110 | + return out[0] if isinstance(out, tuple) else out | ||
| 111 | + | ||
| 112 | + eager_out = reduce_values(x) | ||
| 113 | + | ||
| 114 | + compiled = torch.compile( | ||
| 115 | + reduce_values, | ||
| 116 | + options={"npu_backend": "triton_experimental"}, | ||
| 117 | + ) | ||
| 118 | + compiled_out = compiled(x) | ||
| 119 | + | ||
| 120 | + self.assertEqual(eager_out, compiled_out) | ||
| 121 | + self.assertTrue(bool(compiled_out.isnan().any())) | ||
| 122 | + | ||
| 89 | if __name__ == "__main__": | 123 | if __name__ == "__main__": |
| 90 | run_tests() | 124 | run_tests() |
| @@ -1371,11 +1371,11 @@ class NPUTritonKernelOverrides(TritonKernelOverrides): | |||
| 1371 | 1371 | ||
H | |||
| 1372 | 1372 | ||
| 1373 | def minimum(a, b): | 1373 | def minimum(a, b): |
| 1374 | - return f"tl.minimum({a}, {b})" | 1374 | + return f"tl.minimum({a}, {b}, tl.PropagateNan.ALL)" |
| 1375 | 1375 | ||
| 1376 | 1376 | ||
| 1377 | def maximum(a, b): | 1377 | def maximum(a, b): |
| 1378 | - return f"tl.maximum({a}, {b})" | 1378 | + return f"tl.maximum({a}, {b}, tl.PropagateNan.ALL)" |
| 1379 | 1379 | ||
| 1380 | 1380 | ||
| 1381 | def to_dtype(x, dtype: torch.dtype, src_dtype=None, use_compute_types=True): | 1381 | def to_dtype(x, dtype: torch.dtype, src_dtype=None, use_compute_types=True): |
| @@ -16,6 +16,7 @@ | |||
| 16 | 16 | ||
| 17 | import triton | 17 | import triton |
| 18 | import triton.language as tl | 18 | import triton.language as tl |
| 19 | +from triton.language import core | ||
| 19 | from torch._inductor.runtime import triton_helpers | 20 | from torch._inductor.runtime import triton_helpers |
| 20 | 21 | ||
| 21 | try: | 22 | try: |
| @@ -30,9 +31,14 @@ def max2(a, dim): | |||
| 30 | return tl.max(a, dim, propagate_nan=True) | 31 | return tl.max(a, dim, propagate_nan=True) |
| 31 | 32 | ||
| 32 | 33 | ||
| 34 | + | ||
| 35 | +def _min_prop_nan(a, b): | ||
| 36 | + return core.minimum(a, b, propagate_nan=core.PropagateNan.ALL) | ||
| 37 | + | ||
| 38 | + | ||
| 33 | 39 | ||
| 34 | def min2(a, dim): | 40 | def min2(a, dim): |
| 35 | - return tl.min(a, dim, propagate_nan=True) | 41 | + return tl.reduce(a, dim, _min_prop_nan) |
| 36 | 42 | ||
| 37 | 43 | ||
| 38 | triton_helpers.max2 = max2 | 44 | triton_helpers.max2 = max2 |
bugfix建议同步新增ut用例看护，pr描述中需要贴用例执行结果