mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-25 07:14:16 +08:00
BoHao Chen
940b106108
* [BugFix] Honor nan_propagate in reduce max/min/absmax clear=False write-back MakeUpdate (the clear=False write-back merge) used plain Max/Min for max/min/absmax reducers, silently dropping NaN when nan_propagate=True. Now it emits max_nan/min_nan on fp16/bf16 when nan_propagate is set, mirroring the existing MakeReduce path. GHI: #2697 * [BugFix] Honor nan_propagate in reduce max/min/absmax clear=False write-back MakeUpdate (the clear=False write-back merge) used plain Max/Min for max/min/absmax reducers, silently dropping NaN when nan_propagate=True. Now it emits max_nan/min_nan on fp16/bf16 when nan_propagate is set, mirroring the existing MakeReduce path. GHI: #2697
181 lines
6.9 KiB
Python
181 lines
6.9 KiB
Python
"""Tests for the per-call ``nan_propagate`` kwarg on T.reduce_max / reduce_min /
|
|
reduce_absmax for float16 and bfloat16 buffers (CUDA only)."""
|
|
|
|
import math
|
|
|
|
import torch
|
|
|
|
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
_DTYPES = [("float16", T.float16, torch.float16), ("bfloat16", T.bfloat16, torch.bfloat16)]
|
|
|
|
|
|
def _compile(prim_func):
|
|
return tilelang.compile(prim_func, out_idx=-1, target="cuda")
|
|
|
|
|
|
def _make_reduce_kernel(reduce_fn, length, dtype, *, nan_propagate):
|
|
|
|
@T.prim_func
|
|
def kernel(a: T.Tensor((length,), dtype), out: T.Tensor((1,), dtype)):
|
|
with T.Kernel(1, threads=32):
|
|
frag = T.alloc_fragment((length,), dtype)
|
|
out_frag = T.alloc_fragment((1,), dtype)
|
|
T.copy(a, frag)
|
|
reduce_fn(frag, out_frag, nan_propagate=nan_propagate)
|
|
T.copy(out_frag, out)
|
|
|
|
return kernel
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Source-level checks: confirm the right reducer / intrinsic is emitted.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_reduce_max_default_uses_plain_op():
|
|
k = _compile(_make_reduce_kernel(T.reduce_max, 64, T.float16, nan_propagate=False))
|
|
src = k.get_kernel_source()
|
|
assert "tl::MaxOp" in src and "MaxOpNan" not in src
|
|
assert "__hmax(" in src and "__hmax_nan" not in src
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_reduce_max_nan_propagate_uses_nan_op():
|
|
k = _compile(_make_reduce_kernel(T.reduce_max, 64, T.float16, nan_propagate=True))
|
|
src = k.get_kernel_source()
|
|
assert "tl::MaxOpNan" in src
|
|
assert "__hmax_nan" in src
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_reduce_min_nan_propagate_uses_nan_op():
|
|
k = _compile(_make_reduce_kernel(T.reduce_min, 64, T.bfloat16, nan_propagate=True))
|
|
src = k.get_kernel_source()
|
|
assert "tl::MinOpNan" in src
|
|
assert "__hmin_nan" in src
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_reduce_absmax_nan_propagate_uses_nan_op():
|
|
k = _compile(_make_reduce_kernel(T.reduce_absmax, 64, T.float16, nan_propagate=True))
|
|
src = k.get_kernel_source()
|
|
assert "tl::MaxOpNan" in src
|
|
assert "__hmax_nan" in src
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Runtime behavioral checks: NaN actually propagates only when requested.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_reduce_max_runtime_nan_behavior():
|
|
for _, tl_dtype, torch_dtype in _DTYPES:
|
|
length = 64
|
|
a = torch.arange(length, dtype=torch.float32).to(torch_dtype).cuda()
|
|
a[7] = float("nan")
|
|
|
|
k_default = _compile(_make_reduce_kernel(T.reduce_max, length, tl_dtype, nan_propagate=False))
|
|
k_nan = _compile(_make_reduce_kernel(T.reduce_max, length, tl_dtype, nan_propagate=True))
|
|
|
|
out_default = k_default(a)
|
|
out_nan = k_nan(a)
|
|
|
|
assert not math.isnan(out_default.float().item()), f"{tl_dtype}: default reduce_max should ignore NaN, got {out_default}"
|
|
assert math.isnan(out_nan.float().item()), f"{tl_dtype}: nan_propagate reduce_max should return NaN, got {out_nan}"
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_reduce_min_runtime_nan_behavior():
|
|
for _, tl_dtype, torch_dtype in _DTYPES:
|
|
length = 64
|
|
a = torch.arange(length, dtype=torch.float32).to(torch_dtype).cuda()
|
|
a[13] = float("nan")
|
|
|
|
k_default = _compile(_make_reduce_kernel(T.reduce_min, length, tl_dtype, nan_propagate=False))
|
|
k_nan = _compile(_make_reduce_kernel(T.reduce_min, length, tl_dtype, nan_propagate=True))
|
|
|
|
assert not math.isnan(k_default(a).float().item())
|
|
assert math.isnan(k_nan(a).float().item())
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# clear=False path: the write-back merge (MakeUpdate) must also honor
|
|
# nan_propagate. See GH-2697.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_reduce_clear_false_kernel(reduce_fn, length, dtype, *, nan_propagate):
|
|
|
|
@T.prim_func
|
|
def kernel(a: T.Tensor((length,), dtype), out: T.Tensor((1,), dtype)):
|
|
with T.Kernel(1, threads=32):
|
|
frag = T.alloc_fragment((length,), dtype)
|
|
out_frag = T.alloc_fragment((1,), dtype)
|
|
T.copy(a, frag)
|
|
# Seed the output fragment with -1 so the NaN test is meaningful.
|
|
out_frag[0] = -1.0
|
|
reduce_fn(frag, out_frag, clear=False, nan_propagate=nan_propagate)
|
|
T.copy(out_frag, out)
|
|
|
|
return kernel
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_reduce_max_clear_false_nan_propagate():
|
|
"""clear=False + nan_propagate=True must yield NaN (GH-2697)."""
|
|
for _, tl_dtype, torch_dtype in _DTYPES:
|
|
length = 64
|
|
a = torch.arange(length, dtype=torch.float32).to(torch_dtype).cuda()
|
|
a[7] = float("nan")
|
|
|
|
k_nan = _compile(_make_reduce_clear_false_kernel(T.reduce_max, length, tl_dtype, nan_propagate=True))
|
|
out = k_nan(a)
|
|
assert math.isnan(out.float().item()), f"{tl_dtype}: reduce_max clear=False nan_propagate=True should return NaN, got {out.item()}"
|
|
|
|
k_default = _compile(_make_reduce_clear_false_kernel(T.reduce_max, length, tl_dtype, nan_propagate=False))
|
|
out = k_default(a)
|
|
assert not math.isnan(out.float().item()), f"{tl_dtype}: reduce_max clear=False nan_propagate=False should not return NaN"
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_reduce_min_clear_false_nan_propagate():
|
|
for _, tl_dtype, torch_dtype in _DTYPES:
|
|
length = 64
|
|
a = torch.arange(length, dtype=torch.float32).to(torch_dtype).cuda()
|
|
a[13] = float("nan")
|
|
|
|
k_nan = _compile(_make_reduce_clear_false_kernel(T.reduce_min, length, tl_dtype, nan_propagate=True))
|
|
out = k_nan(a)
|
|
assert math.isnan(out.float().item()), f"{tl_dtype}: reduce_min clear=False nan_propagate=True should return NaN, got {out.item()}"
|
|
|
|
k_default = _compile(_make_reduce_clear_false_kernel(T.reduce_min, length, tl_dtype, nan_propagate=False))
|
|
out = k_default(a)
|
|
assert not math.isnan(out.float().item())
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_reduce_absmax_clear_false_nan_propagate():
|
|
for _, tl_dtype, torch_dtype in _DTYPES:
|
|
length = 64
|
|
a = torch.arange(length, dtype=torch.float32).to(torch_dtype).cuda()
|
|
a[13] = float("nan")
|
|
|
|
k_nan = _compile(_make_reduce_clear_false_kernel(T.reduce_absmax, length, tl_dtype, nan_propagate=True))
|
|
out = k_nan(a)
|
|
assert math.isnan(out.float().item()), (
|
|
f"{tl_dtype}: reduce_absmax clear=False nan_propagate=True should return NaN, got {out.item()}"
|
|
)
|
|
|
|
k_default = _compile(_make_reduce_clear_false_kernel(T.reduce_absmax, length, tl_dtype, nan_propagate=False))
|
|
out = k_default(a)
|
|
assert not math.isnan(out.float().item())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|