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

#!/usr/bin/env python3

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

"""CPU-only unit tests for the grouped GEMM variant of the TE -> Dispatcher bridge.

Grouped GEMM has a genuinely different ABI from regular GEMM (a list of
sub-problems, a launch(descs, stream) signature, and an internally-owned device
workspace), so it is generated by a dedicated path in the unified GEMM codegen
and launched by grouped_gemm_ctypes_lib.cpp. These tests lock in the host-side
contract of that path without touching a GPU:

  * the GROUPED variant generates the grouped launch/kernel with fully-qualified
    ck_tile:: types (so the header does not silently depend on a using-directive);
  * the grouped ctypes bridge includes the standard headers it uses directly
    rather than relying on transitive includes from the force-included kernel;
  * the FNUZ decode table shared by the fp8/bf8 path is cached and read-only.

No GPU is touched -- codegen is pure string formatting and the helpers are pure.
Run: python3 -m pytest tests/test_grouped_gemm_codegen.py -v
"""

import sys
import unittest
from pathlib import Path

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

from codegen_common import TileConfig  # noqa: E402

from unified_gemm_codegen import (  # noqa: E402
    GemmVariant,
    KernelConfig,
    TraitConfig,
    CKTileKernelGenerator,
)

from gemm_utils import _fnuz_decode_table  # noqa: E402

GROUPED_CTYPES_LIB = (
    DISPATCHER_DIR / "bindings" / "ctypes" / "grouped_gemm_ctypes_lib.cpp"
)


def _grouped_config() -> KernelConfig:
    """A minimal valid grouped-GEMM kernel config (fp16/rcr, cshuffle)."""
    tile = TileConfig(
        tile_m=128, tile_n=128, tile_k=32,
        warp_m=2, warp_n=2, warp_k=1,
        warp_tile_m=32, warp_tile_n=32, warp_tile_k=16,
    )
    trait = TraitConfig(
        pipeline="compv3", epilogue="cshuffle", scheduler="intrawave",
        pad_m=False, pad_n=False, pad_k=False, persistent=False,
    )
    return KernelConfig(tile=tile, trait=trait, variant=GemmVariant.GROUPED)


class TestGroupedVariant(unittest.TestCase):
    def test_enum_value(self):
        self.assertEqual(GemmVariant.GROUPED.value, "grouped")


class TestGroupedCodegen(unittest.TestCase):
    """The generated grouped header must be self-contained and correctly typed."""

    def setUp(self):
        gen = CKTileKernelGenerator("fp16", "rcr")
        self.src = gen.generate(_grouped_config())

    def test_generates_grouped_launch_and_kernel(self):
        # The grouped path emits the (descs, stream) launch and the grouped kernel.
        self.assertIn(
            "launch(const std::vector<ck_tile::GroupedGemmHostArgs<>>&", self.src
        )
        self.assertIn("ck_tile::GroupedGemmKernel<", self.src)

    def test_device_memory_is_ck_tile_qualified(self):
        # DeviceMem lives in the ck_tile namespace; the workspace allocation must
        # be fully qualified so the header does not depend on a using-directive.
        self.assertIn("ck_tile::DeviceMem workspace_dev", self.src)
        # And no bare, unqualified `DeviceMem <name>` declaration should remain.
        import re

        self.assertIsNone(
            re.search(r"(?<!:)\bDeviceMem\s+workspace_dev", self.src),
            "found an unqualified DeviceMem declaration in generated grouped header",
        )

    def test_grouped_includes_are_present(self):
        # The grouped variant pulls in the headers its launch needs directly.
        for inc in (
            "ck_tile/host/device_memory.hpp",
            "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp",
        ):
            self.assertIn(inc, self.src)


class TestGroupedCtypesLibIncludes(unittest.TestCase):
    """The grouped ctypes bridge must include what it uses, not lean on the
    transitively-included kernel header (include order is otherwise fragile)."""

    def test_includes_array_header(self):
        text = GROUPED_CTYPES_LIB.read_text()
        # std::array<..., 0>{} descriptors are built here; <array> must be direct.
        self.assertIn("std::array", text)
        self.assertIn("#include <array>", text)


class TestFnuzTableCaching(unittest.TestCase):
    """The 256-entry FNUZ decode table is pure per (exp_bits, mant_bits): it must
    be cached and handed out read-only so callers cannot mutate the shared copy."""

    def test_table_is_cached_same_object(self):
        self.assertIs(_fnuz_decode_table(4, 3), _fnuz_decode_table(4, 3))
        self.assertIs(_fnuz_decode_table(5, 2), _fnuz_decode_table(5, 2))

    def test_table_is_read_only(self):
        table = _fnuz_decode_table(4, 3)
        self.assertFalse(table.flags.writeable)
        with self.assertRaises(ValueError):
            table[0] = 1.0


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