tilelang/testing/python/math/test_math_ieee_math.py

import tilelang
import tilelang.language as T
import torch
import tilelang.testing
import pytest

ROUNDING_MODES = ["rn", "rz", "ru", "rd"]


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,
):
    """
    Test IEEE-compliant math operations with specified rounding modes.
    """

    # Define the appropriate function based on operation type to avoid TVM parsing conflicts
    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
    else:  # Single argument operations

        @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="cuda",
        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}")


def test_rounding_mode_validation():
    """Test that invalid rounding modes raise ValueError"""

    # Test with invalid rounding mode
    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")


@tilelang.testing.requires_cuda
@pytest.mark.parametrize("mode", ROUNDING_MODES, ids=ROUNDING_MODES)
def test_ieee_add_all_rounding_modes(mode):
    """Test IEEE addition with all rounding modes"""
    run_ieee_math_test("ieee_add", T.ieee_add, rounding_mode=mode, run_execution=mode == "rn")
    print(f"✓ ieee_add with {mode} passed")


@tilelang.testing.requires_cuda
@pytest.mark.parametrize("mode", ROUNDING_MODES, ids=ROUNDING_MODES)
def test_ieee_sub_all_rounding_modes(mode):
    """Test IEEE subtraction with all rounding modes"""
    run_ieee_math_test("ieee_sub", T.ieee_sub, rounding_mode=mode, run_execution=mode == "rn")
    print(f"✓ ieee_sub with {mode} passed")


@tilelang.testing.requires_cuda
@pytest.mark.parametrize("mode", ROUNDING_MODES, ids=ROUNDING_MODES)
def test_ieee_mul_all_rounding_modes(mode):
    """Test IEEE multiplication with all rounding modes"""
    run_ieee_math_test("ieee_mul", T.ieee_mul, rounding_mode=mode, run_execution=mode == "rn")
    print(f"✓ ieee_mul with {mode} passed")


@tilelang.testing.requires_cuda
@pytest.mark.parametrize("mode", ROUNDING_MODES, ids=ROUNDING_MODES)
def test_ieee_fmaf_all_rounding_modes(mode):
    """Test IEEE fused multiply-add with all rounding modes"""
    run_ieee_math_test("ieee_fmaf", T.ieee_fmaf, rounding_mode=mode, run_execution=mode == "rn")
    print(f"✓ ieee_fmaf with {mode} passed")


@tilelang.testing.requires_cuda
@pytest.mark.parametrize("mode", ROUNDING_MODES, ids=ROUNDING_MODES)
def test_ieee_frcp_all_rounding_modes(mode):
    """Test IEEE reciprocal with all rounding modes"""
    run_ieee_math_test("ieee_frcp", T.ieee_frcp, rounding_mode=mode, run_execution=mode == "rn")
    print(f"✓ ieee_frcp with {mode} passed")


@tilelang.testing.requires_cuda
@pytest.mark.parametrize("mode", ROUNDING_MODES, ids=ROUNDING_MODES)
def test_ieee_fsqrt_all_rounding_modes(mode):
    """Test IEEE square root with all rounding modes"""
    run_ieee_math_test("ieee_fsqrt", T.ieee_fsqrt, rounding_mode=mode, run_execution=mode == "rn")
    print(f"✓ ieee_fsqrt with {mode} passed")


@tilelang.testing.requires_cuda
def test_ieee_frsqrt_rn_only():
    """Test IEEE reciprocal square root (round to nearest only)"""

    @T.prim_func
    def main(
        A: T.Tensor((32, 32), T.float32),
        B: T.Tensor((32, 32), T.float32),
    ):
        with T.Kernel(T.ceildiv(32, 16), T.ceildiv(32, 16), threads=128) as (bx, by):
            for i, j in T.Parallel(16, 16):
                B[by * 16 + i, bx * 16 + j] = T.ieee_frsqrt(A[by * 16 + i, bx * 16 + j])

    kernel = tilelang.compile(
        main,
        out_idx=[1],
        target="cuda",
        pass_configs={
            tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: False,
        },
    )

    print("\n=== Testing ieee_frsqrt (rn only) ===")
    print("✓ ieee_frsqrt compilation test passed")

    # Test numerical execution
    a = torch.abs(torch.randn(32, 32, device="cuda", dtype=torch.float32)) + 0.1

    try:
        result = kernel(a)
        assert result is not None
        print("✓ ieee_frsqrt numerical execution test passed")
    except Exception as e:
        print(f"Warning: ieee_frsqrt execution failed: {e}")


@tilelang.testing.requires_cuda
@pytest.mark.parametrize("mode", ROUNDING_MODES, ids=ROUNDING_MODES)
def test_ieee_fdiv_all_rounding_modes(mode):
    """Test IEEE division with all rounding modes"""
    run_ieee_math_test("ieee_fdiv", T.ieee_fdiv, rounding_mode=mode, run_execution=mode == "rn")
    print(f"✓ ieee_fdiv with {mode} passed")


if __name__ == "__main__":
    tilelang.testing.main()