mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
cjandcaojian5
21e8c064ac
* Reject unsupported fast-math dtypes * test: cover all unsupported fastmath integer dtypes --------- Co-authored-by: caojian5 <caojian5@huawei.com>
57 lines
1.2 KiB
Python
57 lines
1.2 KiB
Python
import re
|
|
|
|
import pytest
|
|
|
|
import tilelang.language as T
|
|
from tvm import tirx
|
|
|
|
|
|
FASTMATH_INTRINSICS = [
|
|
T.__exp,
|
|
T.__exp10,
|
|
T.__log,
|
|
T.__log2,
|
|
T.__log10,
|
|
T.__tan,
|
|
T.__cos,
|
|
T.__sin,
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"dtype",
|
|
[
|
|
"bool",
|
|
"int8",
|
|
"int16",
|
|
"int32",
|
|
"int64",
|
|
"uint8",
|
|
"uint16",
|
|
"uint32",
|
|
"uint64",
|
|
"float8_e4m3fn",
|
|
"float8_e5m2",
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("intrinsic", FASTMATH_INTRINSICS)
|
|
def test_fastmath_rejects_unsupported_dtype(intrinsic, dtype):
|
|
value = tirx.Var("value", dtype)
|
|
|
|
with pytest.raises(
|
|
TypeError,
|
|
match=rf"T\.{intrinsic.__name__} only supports floating-point inputs, "
|
|
rf"but got {re.escape(dtype)}",
|
|
):
|
|
intrinsic(value)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", ["float16", "bfloat16", "float32", "float64"])
|
|
@pytest.mark.parametrize("intrinsic", FASTMATH_INTRINSICS)
|
|
def test_fastmath_accepts_float_dtype_at_frontend(intrinsic, dtype):
|
|
value = tirx.Var("value", dtype)
|
|
result = intrinsic(value)
|
|
|
|
assert isinstance(result, tirx.Call)
|
|
assert result.dtype == dtype
|