Skip to content

[TEST][CODEGEN] Fix the test scripts tries to tell numpy a dtype name that it cannot recognise #18430

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tlopex merged 2 commits into from
Nov 11, 2025
Merged

[TEST][CODEGEN] Fix the test scripts tries to tell numpy a dtype name that it cannot recognise #18430

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ea19457
[#18394] The test scripts tries to tell numpy a dtype name that it ca…
cchung100m Nov 10, 2025
285b64d
[#18394] Fix the lint error
cchung100m Nov 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 30 additions & 8 deletions tests/python/codegen/test_target_codegen_cuda_fp4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,11 @@
from tvm.script import tir as T

try:
import ml_dtypes
from ml_dtypes import float4_e2m1fn

ML_DTYPES_AVAILABLE = True
except ImportError:
ml_dtypes = None
ML_DTYPES_AVAILABLE = False


@pytest.mark.parametrize("promoted_dtype", ["float32x2", "float16x2"])
Expand Down Expand Up @@ -63,7 +65,6 @@ def add(
fadd = tvm.compile(sch.mod, target=target)
dev = tvm.device(target, 0)

numpytype = "float4_e2m1fn"
if "x" in native_dtype:
lanes = int(native_dtype.split("x")[-1])
else:
Expand All @@ -75,18 +76,39 @@ def add(
promoted_base_dtype = promoted_dtype

np_shape = (vector_length, lanes) if lanes > 1 else (vector_length,)
a_np = np.random.uniform(low=0, high=5, size=np_shape).astype(numpytype)

# Create test data - either using ml_dtypes if available, or using int8 with valid FP4 values
if ML_DTYPES_AVAILABLE:
a_np = np.random.uniform(low=0, high=5, size=np_shape).astype(float4_e2m1fn)
b_np = np.random.uniform(low=0, high=5, size=np_shape).astype(float4_e2m1fn)
else:
# float4_e2m1fn possible values: [0, 0.5, 1, 1.5, 2, 3, 4, 6]
# We will create int8 arrays with valid FP4 bit patterns
valid_fp4_values = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] # 4-bit values
a_np = np.random.choice(valid_fp4_values, size=np_shape).astype(np.int8)
b_np = np.random.choice(valid_fp4_values, size=np_shape).astype(np.int8)

a = tvm.runtime.empty(shape=(vector_length,), dtype=native_dtype, device=dev)
a.copyfrom(a_np)
b_np = np.random.uniform(low=0, high=5, size=np_shape).astype(numpytype)
b = tvm.runtime.empty(shape=(vector_length,), dtype=native_dtype, device=dev)
b.copyfrom(b_np)
c = tvm.runtime.empty(shape=(vector_length,), dtype=native_dtype, device=dev)
fadd(a, b, c)

tvm.testing.assert_allclose(
c.numpy().astype(promoted_base_dtype), (a_np + b_np).astype(promoted_base_dtype)
)
# For the comparison, we will convert result to the promoted dtype and compare
# Note: When ml_dtypes is not available, we skip the numpy-level computation comparison
# and just verify that the CUDA kernel compiles and executes without error
c_result = c.numpy().astype(promoted_base_dtype)

if ML_DTYPES_AVAILABLE:
# Full comparison when ml_dtypes is available
expected = (a_np + b_np).astype(promoted_base_dtype)
tvm.testing.assert_allclose(c_result, expected)
else:
# When ml_dtypes is not available, we just verify the comparison ran successfully
# by checking that we got a result with the expected shape and dtype
assert c_result.shape == np_shape
assert c_result.dtype == promoted_base_dtype


@tvm.testing.requires_cuda_compute_version(10)
Expand Down