mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
Chennes
c6294f07e3
common.h defines a bfloat16_t __hfma wrapper that forwards to the __nv_bfloat16 __hfma overload. CUDA only declares that overload for __CUDA_ARCH__ >= 800, so on SM75 merely parsing the wrapper fails with "no instance of overloaded function __hfma". Because common.h is included by every generated kernel, this breaks all SM75 builds, not only bf16 ones. Guard the native path on SM80+ and use an fp32 FMA followed by bf16 conversion on earlier architectures, following the existing pre-SM80 bf16 fallback strategy. Add a compile-only regression test that builds a bf16 ieee_fmaf kernel for an explicit SM75 target. Existing IEEE math tests use target="cuda", so an SM80+ runner exercises only the native branch.
252 lines
8.8 KiB
Python
252 lines
8.8 KiB
Python
import tilelang
|
||
import tilelang.language as T
|
||
import torch
|
||
import tilelang.testing
|
||
import pytest
|
||
|
||
ROUNDING_MODES = ["rn", "rz", "ru", "rd"]
|
||
|
||
# (dtype, rounding_mode) pairs that must compile and produce correct results.
|
||
# - fp32/fp64: all four IEEE rounding modes are natively supported by CUDA.
|
||
# - fp16/bf16: only round-to-nearest-even is available.
|
||
IEEE_DTYPE_MODES = [
|
||
(T.float32, "rn"),
|
||
(T.float32, "rz"),
|
||
(T.float32, "ru"),
|
||
(T.float32, "rd"),
|
||
(T.float64, "rn"),
|
||
(T.float64, "rz"),
|
||
(T.float64, "ru"),
|
||
(T.float64, "rd"),
|
||
(T.float16, "rn"),
|
||
(T.bfloat16, "rn"),
|
||
]
|
||
|
||
# Ops that carry a rounding_mode parameter (everything except frsqrt).
|
||
IEEE_OPS_WITH_RM = [
|
||
("ieee_add", T.ieee_add),
|
||
("ieee_sub", T.ieee_sub),
|
||
("ieee_mul", T.ieee_mul),
|
||
("ieee_fmaf", T.ieee_fmaf),
|
||
("ieee_frcp", T.ieee_frcp),
|
||
("ieee_fsqrt", T.ieee_fsqrt),
|
||
("ieee_fdiv", T.ieee_fdiv),
|
||
]
|
||
|
||
# All 8 ops (including frsqrt).
|
||
IEEE_ALL_OPS = IEEE_OPS_WITH_RM + [("ieee_frsqrt", T.ieee_frsqrt)]
|
||
|
||
|
||
def run_ieee_math_test(
|
||
mathop_name,
|
||
mathop_func,
|
||
rounding_mode="rn",
|
||
M=32,
|
||
N=32,
|
||
block_M=16,
|
||
block_N=16,
|
||
dtype=T.float32,
|
||
run_execution=True,
|
||
target="cuda",
|
||
):
|
||
"""
|
||
Test IEEE-compliant math operations with specified rounding modes.
|
||
"""
|
||
|
||
if mathop_name == "ieee_fmaf":
|
||
|
||
@T.prim_func
|
||
def main_func(
|
||
A: T.Tensor((M, N), dtype),
|
||
B: T.Tensor((M, N), dtype),
|
||
C: T.Tensor((M, N), dtype),
|
||
D: T.Tensor((M, N), dtype),
|
||
):
|
||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
|
||
for i, j in T.Parallel(block_M, block_N):
|
||
D[by * block_M + i, bx * block_N + j] = mathop_func(
|
||
A[by * block_M + i, bx * block_N + j],
|
||
B[by * block_M + i, bx * block_N + j],
|
||
C[by * block_M + i, bx * block_N + j],
|
||
rounding_mode,
|
||
)
|
||
|
||
out_idx = [3]
|
||
num_inputs = 3
|
||
elif mathop_name in ["ieee_add", "ieee_sub", "ieee_mul", "ieee_fdiv"]:
|
||
|
||
@T.prim_func
|
||
def main_func(
|
||
A: T.Tensor((M, N), dtype),
|
||
B: T.Tensor((M, N), dtype),
|
||
C: T.Tensor((M, N), dtype),
|
||
):
|
||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
|
||
for i, j in T.Parallel(block_M, block_N):
|
||
C[by * block_M + i, bx * block_N + j] = mathop_func(
|
||
A[by * block_M + i, bx * block_N + j], B[by * block_M + i, bx * block_N + j], rounding_mode
|
||
)
|
||
|
||
out_idx = [2]
|
||
num_inputs = 2
|
||
elif mathop_name == "ieee_frsqrt": # no rounding_mode parameter
|
||
|
||
@T.prim_func
|
||
def main_func(
|
||
A: T.Tensor((M, N), dtype),
|
||
B: T.Tensor((M, N), dtype),
|
||
):
|
||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
|
||
for i, j in T.Parallel(block_M, block_N):
|
||
B[by * block_M + i, bx * block_N + j] = mathop_func(A[by * block_M + i, bx * block_N + j])
|
||
|
||
out_idx = [1]
|
||
num_inputs = 1
|
||
else: # Single argument with rounding_mode (frcp, fsqrt)
|
||
|
||
@T.prim_func
|
||
def main_func(
|
||
A: T.Tensor((M, N), dtype),
|
||
B: T.Tensor((M, N), dtype),
|
||
):
|
||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
|
||
for i, j in T.Parallel(block_M, block_N):
|
||
B[by * block_M + i, bx * block_N + j] = mathop_func(A[by * block_M + i, bx * block_N + j], rounding_mode)
|
||
|
||
out_idx = [1]
|
||
num_inputs = 1
|
||
|
||
# Test compilation
|
||
kernel = tilelang.compile(
|
||
main_func,
|
||
out_idx=out_idx,
|
||
target=target,
|
||
pass_configs={
|
||
tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: False,
|
||
},
|
||
)
|
||
|
||
print(f"\n=== Testing {mathop_name} with rounding mode {rounding_mode} ===")
|
||
print(f"✓ {mathop_name} compilation test passed")
|
||
|
||
if not run_execution:
|
||
return
|
||
|
||
# Test numerical execution
|
||
torch_dtype = dtype.as_torch()
|
||
a = torch.randn(M, N, device="cuda", dtype=torch_dtype)
|
||
|
||
if num_inputs >= 2:
|
||
b = torch.randn(M, N, device="cuda", dtype=torch_dtype)
|
||
if num_inputs == 3:
|
||
c = torch.randn(M, N, device="cuda", dtype=torch_dtype)
|
||
|
||
# Ensure positive values for functions that need them
|
||
if mathop_name in ["ieee_frcp", "ieee_fsqrt"]:
|
||
a = torch.abs(a) + 0.1
|
||
elif mathop_name == "ieee_fdiv":
|
||
b = torch.abs(b) + 0.1 # Avoid division by zero
|
||
|
||
# Execute kernel
|
||
try:
|
||
if num_inputs == 1:
|
||
result = kernel(a)
|
||
elif num_inputs == 2:
|
||
result = kernel(a, b)
|
||
else: # num_inputs == 3
|
||
result = kernel(a, b, c)
|
||
|
||
assert result is not None
|
||
print(f"✓ {mathop_name} numerical execution test passed")
|
||
except Exception as e:
|
||
print(f"Warning: {mathop_name} execution failed: {e}")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Rounding-mode validation (API-level, no GPU required)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_rounding_mode_validation():
|
||
"""Test that invalid rounding modes raise ValueError"""
|
||
with pytest.raises(ValueError, match="Invalid rounding mode"):
|
||
T.ieee_add(1.0, 2.0, "invalid_mode")
|
||
with pytest.raises(ValueError, match="Invalid rounding mode"):
|
||
T.ieee_mul(1.0, 2.0, "xy")
|
||
with pytest.raises(ValueError, match="Invalid rounding mode"):
|
||
T.ieee_fsqrt(4.0, "bad_mode")
|
||
print("✓ Rounding mode validation test passed")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Numerical tests — every (op × dtype × rounding_mode) valid combination
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@tilelang.testing.requires_cuda
|
||
@pytest.mark.parametrize("op_name,op_func", IEEE_OPS_WITH_RM)
|
||
@pytest.mark.parametrize("dtype,mode", IEEE_DTYPE_MODES)
|
||
def test_ieee_with_rounding_mode(op_name, op_func, dtype, mode):
|
||
"""Compile + run every IEEE op across fp32 / fp64 / fp16 / bf16 and all rounding
|
||
modes that the hardware supports."""
|
||
run_ieee_math_test(op_name, op_func, rounding_mode=mode, dtype=dtype, run_execution=(mode == "rn"))
|
||
|
||
|
||
# ieee_frsqrt — only round-to-nearest (no rounding_mode parameter).
|
||
|
||
|
||
@tilelang.testing.requires_cuda
|
||
@pytest.mark.parametrize("dtype", [T.float32, T.float16, T.bfloat16])
|
||
def test_ieee_frsqrt(dtype):
|
||
run_ieee_math_test("ieee_frsqrt", T.ieee_frsqrt, dtype=dtype)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Rejection tests — combinations that must fail at codegen time
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@tilelang.testing.requires_cuda
|
||
def test_ieee_frsqrt_fp64_rejected():
|
||
"""fp64 has no rsqrt intrinsic — codegen must FATAL."""
|
||
with pytest.raises(Exception, match="frsqrt is not supported for float64"):
|
||
run_ieee_math_test("ieee_frsqrt", T.ieee_frsqrt, dtype=T.float64)
|
||
|
||
|
||
@tilelang.testing.requires_cuda
|
||
def test_ieee_half_non_rn_rejected():
|
||
"""fp16/bf16 only supports round-to-nearest-even."""
|
||
with pytest.raises(Exception, match="Only rounding mode 'rn' is available for half precision"):
|
||
run_ieee_math_test("ieee_add", T.ieee_add, rounding_mode="rz", dtype=T.float16)
|
||
with pytest.raises(Exception, match="Only rounding mode 'rn' is available for half precision"):
|
||
run_ieee_math_test("ieee_add", T.ieee_add, rounding_mode="rz", dtype=T.bfloat16)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Architecture regression — the bf16 __hfma wrapper in common.h must keep
|
||
# compiling for pre-SM80 targets. See the bf16 __hfma guard in
|
||
# src/tl_templates/cuda/common.h.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@tilelang.testing.requires_cuda
|
||
def test_ieee_fmaf_bf16_compiles_for_sm75():
|
||
"""bf16 ieee_fmaf lowers to __hfma, whose native __nv_bfloat16 overload
|
||
only exists on SM80+. common.h is pulled into every generated kernel, so an
|
||
unguarded wrapper breaks *all* SM75 builds, not just bf16 ones. Compile (do
|
||
not run) a bf16 fmaf kernel explicitly for sm_75 to pin the fallback path.
|
||
Compile-only, so it runs on any CUDA runner regardless of the host GPU arch
|
||
-- the parametrized tests above use target="cuda" and would only exercise
|
||
the SM80+ branch on a modern runner."""
|
||
run_ieee_math_test(
|
||
"ieee_fmaf",
|
||
T.ieee_fmaf,
|
||
dtype=T.bfloat16,
|
||
target={"kind": "cuda", "arch": "sm_75"},
|
||
run_execution=False,
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
tilelang.testing.main()
|