TheRock/rocm-libraries/projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py

#!/usr/bin/env python3

# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT

"""CPU-only unit tests for the microscaling (mx) GEMM TileEngine -> Dispatcher bridge.

Locks the config name format, the codegen-JSON projection, the dtype/layout/warp-tile
validity gate, the e8m0 scale codec, the fp8/fp4 quantization round-trips, and the numpy
microscaled reference. No GPU, no hipcc, no Old-TE builder import required.
"""

import sys
import unittest
from pathlib import Path

import numpy as np

_DISP = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(_DISP / "python"))
sys.path.insert(0, str(_DISP / "codegen"))

from mx_gemm_utils import (  # noqa: E402
    SCALE_BLOCK,
    E8M0_ONE,
    MxGemmKernelConfig,
    MxGemmProblem,
    default_fp8_config,
    default_fp4_config,
    e8m0_to_float,
    float_to_e8m0,
    quantize_fp8,
    dequantize_fp8,
    quantize_fp4_packed,
    dequantize_fp4_packed,
    mx_gemm_reference,
)


class TestConfigName(unittest.TestCase):
    def test_fallback_name_prefix(self):
        cfg = default_fp8_config()
        self.assertTrue(cfg._fallback_name().startswith("mx_gemm_fp8_rcr_"))

    def test_fallback_name_encodes_tiles(self):
        cfg = MxGemmKernelConfig(
            datatype="fp4", tile_m=64, tile_n=128, tile_k=256,
            warp_m=1, warp_n=2, warp_k=1,
            warp_tile_m=16, warp_tile_n=16, warp_tile_k=128,
        )
        name = cfg._fallback_name()
        self.assertIn("_fp4_rcr_", name)
        self.assertIn("64x128x256", name)
        self.assertIn("1x2x1", name)
        self.assertIn("16x16x128", name)
        self.assertNotIn(" ", name)

    def test_persistent_suffix_only_when_set(self):
        self.assertNotIn("True", default_fp8_config()._fallback_name())
        cfg = default_fp8_config()
        cfg.persistent = True
        self.assertIn("True", cfg._fallback_name())


class TestCodegenJson(unittest.TestCase):
    def test_projection_roundtrip(self):
        cfg = MxGemmKernelConfig(
            datatype="fp8", tile_m=128, tile_n=128, tile_k=128,
            warp_m=2, warp_n=2, warp_k=1,
            warp_tile_m=16, warp_tile_n=16, warp_tile_k=128,
            k_block_per_cu=3,
            # Pin the arch so to_codegen_config() does not shell out to rocminfo;
            # keeps this a CPU-only test on non-ROCm runners.
            gpu_target="gfx950",
        )
        j = cfg.to_codegen_config()
        self.assertEqual(j["datatype"], "fp8")
        self.assertEqual(j["layout"], "rcr")
        self.assertEqual(j["tile_config"]["tile_k"], 128)
        self.assertEqual(j["tile_config"]["warp_tile_k"], 128)
        self.assertEqual(j["k_block_per_cu"], 3)


class TestValidity(unittest.TestCase):
    def test_default_configs_valid(self):
        self.assertTrue(default_fp8_config().is_valid())
        self.assertTrue(default_fp4_config().is_valid())

    def test_non_rcr_rejected(self):
        cfg = default_fp8_config()
        cfg.layout = "rrr"
        self.assertFalse(cfg.is_valid())

    def test_bad_dtype_rejected(self):
        cfg = default_fp8_config()
        cfg.datatype = "bf16"
        self.assertFalse(cfg.is_valid())

    def test_bf8_rejected(self):
        # Old-TE argparse (choices=["fp4","fp8"]) + validate_gemm_mx never
        # compile bf8/e5m2 for mx_gemm, so the bridge must reject it too. Assert
        # both the config-level gate (is_valid) and, when importable, the
        # codegen-level gate (_validate) refuse dtype="bf8".
        cfg = default_fp8_config()
        cfg.datatype = "bf8"
        # Pin the arch so to_codegen_config() below stays CPU-only (no rocminfo).
        cfg.gpu_target = "gfx950"
        self.assertFalse(cfg.is_valid())

        try:
            from unified_mx_gemm_codegen import _validate  # noqa: E402
        except Exception as exc:  # noqa: BLE001
            self.skipTest(f"codegen import unavailable: {exc}")
        with self.assertRaises(Exception):
            _validate(cfg.to_codegen_config())

    def test_wrong_warp_tile_rejected(self):
        cfg = default_fp8_config()
        cfg.warp_tile_k = 64
        self.assertFalse(cfg.is_valid())

    def test_indivisible_tile_rejected(self):
        cfg = default_fp8_config()
        cfg.tile_m = 100  # not a multiple of warp_m * warp_tile_m (2*16=32)
        self.assertFalse(cfg.is_valid())


class TestProblem(unittest.TestCase):
    def test_scale_k_and_flops(self):
        p = MxGemmProblem(M=64, N=128, K=256)
        self.assertEqual(p.scale_k, 256 // SCALE_BLOCK)
        self.assertEqual(p.flops, 2 * 64 * 128 * 256)

    def test_k_not_multiple_of_32_raises(self):
        with self.assertRaises(ValueError):
            MxGemmProblem(M=32, N=32, K=48)


class TestE8m0Codec(unittest.TestCase):
    def test_one_is_byte_127(self):
        self.assertEqual(int(float_to_e8m0(np.float32(1.0))), E8M0_ONE)
        self.assertEqual(float(e8m0_to_float(E8M0_ONE)), 1.0)

    def test_power_of_two_roundtrip(self):
        for s in (0.25, 0.5, 1.0, 2.0, 4.0, 8.0):
            b = float_to_e8m0(np.float32(s))
            self.assertAlmostEqual(float(e8m0_to_float(b)), s, places=6)

    def test_255_is_nan(self):
        with np.errstate(over="ignore"):
            self.assertTrue(np.isnan(float(e8m0_to_float(255))))

    def test_non_positive_raises(self):
        # Contract is a strictly-positive power-of-two; a 0.0/negative scale is a
        # caller bug and must fail loudly rather than silently encode 1.0.
        for bad in (0.0, -1.0, -0.0):
            with self.assertRaises(ValueError):
                float_to_e8m0(np.float32(bad))
        with self.assertRaises(ValueError):
            float_to_e8m0(np.array([1.0, 0.0, 2.0], np.float32))

    def test_non_finite_raises(self):
        with np.errstate(invalid="ignore"):
            for bad in (np.inf, np.nan):
                with self.assertRaises(ValueError):
                    float_to_e8m0(np.float32(bad))


class TestFp8Codec(unittest.TestCase):
    def test_known_bytes(self):
        self.assertEqual(int(quantize_fp8(np.array([1.0], np.float32))[0]), 0x38)
        self.assertEqual(int(quantize_fp8(np.array([-2.0], np.float32))[0]), 0xC0)

    def test_roundtrip_grid(self):
        grid = np.array([-2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0], np.float32)
        vals = np.tile(grid, (4, 1))
        self.assertTrue(np.array_equal(dequantize_fp8(quantize_fp8(vals)), vals))

    def test_off_grid_raises(self):
        # The vectorized codec must reject values not on the exact e4m3 grid
        # instead of snapping them to a neighbour byte.
        with self.assertRaises(KeyError):
            quantize_fp8(np.array([[0.3]], np.float32))

    def test_neg_zero_collapses_to_zero_byte(self):
        self.assertEqual(int(quantize_fp8(np.array([-0.0], np.float32))[0]), 0x00)

    def test_shape_preserved_2d(self):
        vals = np.full((3, 5), 1.0, np.float32)
        self.assertEqual(quantize_fp8(vals).shape, (3, 5))


class TestFp4Codec(unittest.TestCase):
    def test_pack_two_per_byte(self):
        # one row, K=2 -> a single packed byte; low nibble even-K, high nibble odd-K.
        vals = np.array([[1.0, 2.0]], np.float32)  # codes: 1.0->2, 2.0->4
        packed = quantize_fp4_packed(vals)
        self.assertEqual(packed.shape, (1, 1))
        self.assertEqual(int(packed[0, 0]), (4 << 4) | 2)

    def test_roundtrip_grid(self):
        grid = np.array([-2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0], np.float32)
        rng = np.random.default_rng(0)
        vals = rng.choice(grid, size=(4, 8)).astype(np.float32)
        packed = quantize_fp4_packed(vals)
        self.assertEqual(packed.shape, (4, 4))
        self.assertTrue(np.array_equal(dequantize_fp4_packed(packed, 8), vals))

    def test_off_grid_raises(self):
        # 2.5 exists in fp8 e4m3 but NOT in the fp4 e2m1 grid -> must reject.
        with self.assertRaises(KeyError):
            quantize_fp4_packed(np.array([[2.5, 1.0]], np.float32))


class TestReference(unittest.TestCase):
    def _inputs(self, M, N, K, seed=0):
        rng = np.random.default_rng(seed)
        grid = np.array([-2.0, -1.0, -0.5, 0.5, 1.0, 2.0], np.float32)
        A = rng.choice(grid, size=(M, K)).astype(np.float32)
        B = rng.choice(grid, size=(K, N)).astype(np.float32)
        return A, B

    def test_unit_scales_equals_plain_matmul(self):
        M, N, K = 4, 3, 32
        A, B = self._inputs(M, N, K)
        prob = MxGemmProblem(M=M, N=N, K=K)
        one = int(float_to_e8m0(np.float32(1.0)))
        sa = np.full((M, prob.scale_k), one, np.uint8)
        sb = np.full((N, prob.scale_k), one, np.uint8)
        ref = mx_gemm_reference(A, B, sa, sb, prob).astype(np.float32)
        plain = (A @ B).astype(np.float16).astype(np.float32)
        self.assertTrue(np.allclose(ref, plain, atol=1e-2))

    def test_power_of_two_scale_multiplies(self):
        M, N, K = 4, 3, 32
        A, B = self._inputs(M, N, K, seed=1)
        prob = MxGemmProblem(M=M, N=N, K=K)
        two = int(float_to_e8m0(np.float32(2.0)))
        sa = np.full((M, prob.scale_k), two, np.uint8)
        sb = np.full((N, prob.scale_k), two, np.uint8)
        ref = mx_gemm_reference(A, B, sa, sb, prob).astype(np.float32)
        # scale_a=2 and scale_b=2 -> product scaled by 4.
        expected = (4.0 * (A @ B)).astype(np.float16).astype(np.float32)
        self.assertTrue(np.allclose(ref, expected, atol=1e-1))


class TestCodegenNameContract(unittest.TestCase):
    """Optional: byte-exact name parity when the Old-TE builder is importable."""

    def test_codegen_name_matches_config(self):
        try:
            from unified_mx_gemm_codegen import kernel_name
        except Exception as exc:  # noqa: BLE001
            self.skipTest(f"codegen import unavailable: {exc}")
        cfg = default_fp8_config()
        # Pin the arch so to_codegen_config() stays CPU-only (no rocminfo).
        cfg.gpu_target = "gfx950"
        try:
            name = kernel_name(cfg.to_codegen_config())
        except Exception as exc:  # noqa: BLE001
            self.skipTest(f"Old-TE builder unavailable: {exc}")
        self.assertTrue(name.startswith("mx_gemm_fp8_rcr_"))


if __name__ == "__main__":
    unittest.main()