mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
137 lines
4.6 KiB
Python
137 lines
4.6 KiB
Python
"""CPU codegen float literal regression tests.
|
|
|
|
Covers two paths through CodeGenTileLangC::VisitExpr_(FloatImmNode*) that have
|
|
broken before:
|
|
|
|
1. Finite fp16 constants must lower to the cast form `(half)1.5e+00f`
|
|
(legal C/C++), NOT the `1.5e+00h` form copied from Metal — the `h`
|
|
suffix is MSL-only and g++ rejects it
|
|
("no matching literal operator ... operator\"\"h").
|
|
2. +/-inf and nan must lower to the C99 macros INFINITY / NAN, NOT the
|
|
base class's "-inff" / "nanf" output (IEEE text + 'f' suffix), which
|
|
g++ rejects ("use of undeclared identifier 'inff'"). This path is
|
|
exercised by reduce max/min identity init values.
|
|
|
|
Both were latent because no prior c-target test emitted a finite fp16
|
|
FloatImm or an inf identity. reduce max/min introduced the inf path; the
|
|
fp16 finite-constant regression was introduced by copying Metal's FloatImm
|
|
override verbatim and fixed by delegating finite values back to the base
|
|
class.
|
|
"""
|
|
|
|
import re
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
from tilelang import tvm
|
|
|
|
|
|
def _compile_c(func, out_idx):
|
|
with tvm.target.Target("c"):
|
|
return tilelang.compile(
|
|
func,
|
|
out_idx=out_idx,
|
|
target="c",
|
|
target_host="c",
|
|
execution_backend="cython",
|
|
)
|
|
|
|
|
|
def _source(func) -> str:
|
|
with tvm.target.Target("c"):
|
|
artifact = tilelang.lower(func)
|
|
return artifact.kernel_source
|
|
|
|
|
|
def test_fp16_finite_constant_arithmetic():
|
|
"""Reproducer: T.float16(1.5) + T.float16(2.25) must compile and run.
|
|
|
|
Regression for the `1.5e+00h` literal bug: the Metal-derived override
|
|
appended an 'h' suffix for fp16, which is invalid in C/C++.
|
|
"""
|
|
|
|
@T.prim_func
|
|
def main(A: T.Tensor((4,), "float16"), B: T.Tensor((4,), "float16")):
|
|
with T.Kernel(1):
|
|
for i in T.grid(4):
|
|
B[i] = A[i] + T.float16(1.5) + T.float16(2.25)
|
|
|
|
# Must compile cleanly (the bug failed at g++ compile step).
|
|
kernel = _compile_c(main, out_idx=[1])
|
|
|
|
# The generated source must use the legal cast form, not the 'h' suffix.
|
|
src = _source(main)
|
|
assert "(half)" in src, f"expected (half) cast form in source:\n{src}"
|
|
assert not re.search(r"[0-9]h\b", src), f"source must not contain an 'h' literal suffix (MSL-only):\n{src}"
|
|
|
|
a = torch.tensor([0.0, 1.0, 2.0, 3.0], dtype=torch.float16)
|
|
out = kernel(a)
|
|
expected = a + torch.tensor(3.75, dtype=torch.float16)
|
|
torch.testing.assert_close(out, expected)
|
|
|
|
|
|
@pytest.mark.parametrize("op", ["max", "min"])
|
|
def test_fp16_reduce_inf_identity(op):
|
|
"""fp16 reduce_max/min identity (+/-inf) must lower to INFINITY, not inff.
|
|
|
|
The base class prints "-inf" + 'f' => "-inff", which g++ rejects.
|
|
"""
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((4, 8), "float16"),
|
|
B: T.Tensor((4,), "float16"),
|
|
):
|
|
with T.Kernel(1):
|
|
src = T.alloc_local((4, 8), "float16")
|
|
dst = T.alloc_local((4,), "float16")
|
|
for i, j in T.grid(4, 8):
|
|
src[i, j] = A[i, j]
|
|
if op == "max":
|
|
T.reduce_max(src, dst, dim=1, clear=True)
|
|
else:
|
|
T.reduce_min(src, dst, dim=1, clear=True)
|
|
for i in T.grid(4):
|
|
B[i] = dst[i]
|
|
|
|
# Source must use the C99 macro form, not the broken literal.
|
|
src = _source(main)
|
|
assert "INFINITY" in src, f"expected INFINITY macro for {op} identity in source:\n{src}"
|
|
assert "inff" not in src, f"source must not contain 'inff':\n{src}"
|
|
|
|
kernel = _compile_c(main, out_idx=[1])
|
|
a = torch.randn((4, 8), dtype=torch.float16)
|
|
out = kernel(a)
|
|
expected = a.amax(dim=1) if op == "max" else a.amin(dim=1)
|
|
torch.testing.assert_close(out, expected)
|
|
|
|
|
|
def test_float32_finite_constant_sanity():
|
|
"""float32 finite constants must keep the 'f' suffix after delegation.
|
|
|
|
Guards against over-narrowing the override: delegating finite values to
|
|
the base class must still produce `1.5e+00f` for float32. An input
|
|
tensor is threaded through so the torch wrapper can infer the CPU
|
|
device (a no-arg call defaults to CUDA).
|
|
"""
|
|
|
|
@T.prim_func
|
|
def main(A: T.Tensor((1,), "float32"), B: T.Tensor((1,), "float32")):
|
|
with T.Kernel(1):
|
|
B[0] = A[0] + T.cast(1.5, "float32") + T.cast(2.25, "float32")
|
|
|
|
src = _source(main)
|
|
assert re.search(r"[0-9]f\b", src), f"expected a float32 'f' suffix literal in source:\n{src}"
|
|
|
|
kernel = _compile_c(main, out_idx=[1])
|
|
a = torch.tensor([0.0], dtype=torch.float32)
|
|
out = kernel(a)
|
|
torch.testing.assert_close(out, torch.tensor([3.75], dtype=torch.float32))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|