TheRock/rocm-libraries/projects/composablekernel/dispatcher/tests/test_multi_abd_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 multi_abd GEMM bridge.

The Tile Engine -> Dispatcher multi_abd bridge relies on one hard invariant:
``GemmKernelConfig.name`` must reproduce, byte-for-byte, the kernel stem that
``unified_gemm_codegen.py`` bakes into the generated kernel (and that the .so
reports at runtime). For multi_abd that stem carries the 4-char (A,B,E,D)
layout plus a ``_multiabd_a{na}_b{nb}_d{nd}_{aop}_{bop}_{cdeop}`` suffix, so
distinct tensor counts / element-wise ops can never collapse onto one kernel.

These tests exercise only pure host-side logic (name generation, the codegen
JSON projection, and the shipped configs/*.json). No GPU, hipcc, or build is
required, so the suite runs green in CPU-only CI.

Run: python3 -m pytest tests/test_multi_abd_bridge.py -v
"""

import ctypes
import json
import sys
import unittest
from pathlib import Path
from unittest import mock

import numpy as np

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

from gemm_utils import (  # noqa: E402
    GemmKernelConfig,
    MultiABDDispatcherLib,
    _output_dtype,
    _dtype_from_kernel_name,
)

# The shipped multi_abd sweep configs the bridge codegens from.
_CONFIG_DIR = (
    REPO_ROOT
    / "tile_engine"
    / "ops"
    / "gemm"
    / "gemm_multi_abd"
    / "configs"
)


def _make_config(**overrides):
    """A canonical multi_abd config; overrides tweak individual fields."""
    kw = dict(
        dtype_a="fp16",
        dtype_b="fp16",
        dtype_c="fp16",
        dtype_acc="fp32",
        layout_a="row",
        layout_b="col",
        layout_c="row",
        layout_d="row",
        variant="multi_abd",
        num_a_tensors=2,
        num_b_tensors=2,
        num_d_tensors=2,
    )
    kw.update(overrides)
    return GemmKernelConfig(**kw)


class TestMultiAbdName(unittest.TestCase):
    """The multi_abd kernel-name contract (the byte-parity invariant)."""

    def test_name_carries_multiabd_suffix(self):
        cfg = _make_config()
        name = cfg.name
        # 4-char (A,B,E,D) layout, not the 3-char C layout.
        self.assertIn("_rcrr_", name)
        # multiabd tensor-count + op suffix, exactly as codegen emits it.
        self.assertTrue(
            name.endswith(
                "_multiabd_a2_b2_d2_PassThrough_PassThrough_PassThrough"
            ),
            name,
        )

    def test_full_stem_is_stable(self):
        # Pin the entire stem so any drift in the naming scheme is caught.
        cfg = _make_config(
            tile_m=128, tile_n=128, tile_k=32,
            wave_m=2, wave_n=2, wave_k=1,
            warp_tile_m=32, warp_tile_n=32, warp_tile_k=16,
            pipeline="compv4", scheduler="intrawave", epilogue="cshuffle",
            pad_m=True, pad_n=True, pad_k=True, persistent=False,
        )
        expected = (
            "gemm_fp16_rcrr_compv4_cshuffle_intrawave"
            "_True_True_True_False"
            "_128x128x32_2x2x1_32x32x16"
            "_multiabd_a2_b2_d2_PassThrough_PassThrough_PassThrough"
        )
        self.assertEqual(cfg.name, expected)

    def test_tensor_counts_change_the_name(self):
        # Distinct tensor counts must not collapse onto one kernel name.
        base = _make_config().name
        more = _make_config(num_a_tensors=3, num_d_tensors=1).name
        self.assertNotEqual(base, more)
        self.assertIn("_a3_", more)
        self.assertIn("_d1_", more)

    def test_elementwise_ops_change_the_name(self):
        base = _make_config().name
        scaled = _make_config(cde_elementwise_op="AddScale").name
        self.assertNotEqual(base, scaled)
        self.assertTrue(scaled.endswith("_PassThrough_PassThrough_AddScale"))

    def test_dtype_recovers_from_name(self):
        # The runner reads the input dtype straight out of the compiled .so
        # name, so every dtype the bridge builds must round-trip.
        for dtype in ("fp16", "bf16"):
            cfg = _make_config(
                dtype_a=dtype, dtype_b=dtype, dtype_c=_output_dtype(dtype),
            )
            self.assertEqual(_dtype_from_kernel_name(cfg.name), dtype)

    def test_name_carries_four_char_layout(self):
        # multi_abd uses the 4-char (A,B,E,D) layout code in the stem; the D
        # char must reflect layout_d independently of the C layout.
        for la, lb, lc, ld in (
            ("row", "col", "row", "row"),
            ("row", "row", "row", "col"),
            ("col", "col", "row", "row"),
        ):
            cfg = _make_config(
                layout_a=la, layout_b=lb, layout_c=lc, layout_d=ld,
            )
            self.assertIn(f"_{cfg.layout4}_", cfg.name)


class TestMultiAbdCodegenJson(unittest.TestCase):
    """The codegen JSON projection must carry the multi_abd block."""

    def test_codegen_json_has_multi_abd_block(self):
        cfg = _make_config(
            num_a_tensors=2, num_b_tensors=2, num_d_tensors=2,
            cde_elementwise_op="AddScale",
        )
        j = cfg.to_codegen_json()
        self.assertIn("multi_abd_config", j)
        mabd = j["multi_abd_config"]
        self.assertEqual(mabd["num_a_tensors"], 2)
        self.assertEqual(mabd["num_b_tensors"], 2)
        self.assertEqual(mabd["num_d_tensors"], 2)
        self.assertEqual(mabd["cde_elementwise_op"], "AddScale")

    def test_codegen_json_core_blocks(self):
        j = _make_config().to_codegen_json()
        self.assertIn("tile_config", j)
        self.assertIn("trait_config", j)
        # dispatcher wave_* is projected onto codegen warp_* (warps per block).
        self.assertEqual(j["tile_config"]["warp_m"], [2])


class TestMultiAbdShippedConfigs(unittest.TestCase):
    """The configs/*.json the bridge codegens from must be well-formed sweeps."""

    def test_config_dir_exists(self):
        self.assertTrue(_CONFIG_DIR.is_dir(), _CONFIG_DIR)

    def test_configs_are_valid_sweeps(self):
        configs = sorted(_CONFIG_DIR.glob("*.json"))
        self.assertGreater(len(configs), 0, "no multi_abd configs shipped")
        for path in configs:
            with self.subTest(config=path.name):
                with open(path) as f:
                    data = json.load(f)
                self.assertIn("tile_config", data)
                self.assertIn("trait_config", data)
                tc = data["tile_config"]
                for key in (
                    "tile_m", "tile_n", "tile_k",
                    "warp_m", "warp_n", "warp_k",
                    "warp_tile_m", "warp_tile_n", "warp_tile_k",
                ):
                    self.assertIn(key, tc, f"{path.name} missing {key}")
                tr = data["trait_config"]
                for key in ("pipeline", "scheduler", "epilogue"):
                    self.assertIn(key, tr, f"{path.name} missing {key}")


class TestMultiAbdAbiMarshalling(unittest.TestCase):
    """GPU-free tests of the C-ABI marshalling in ``MultiABDDispatcherLib``.

    The ``.so`` is mocked (``ctypes.CDLL`` patched), so no generated kernel or
    GPU is needed. These cover the highest-risk-but-previously-untested surface:
    the declared ABI, the pointer/stride array packing, the positional argument
    order, and status forwarding.
    """

    def _make_lib(self, status=0):
        fake = mock.MagicMock()
        fake.dispatcher_run_multi_abd.return_value = status
        with mock.patch("ctypes.CDLL", return_value=fake):
            lib = MultiABDDispatcherLib(Path("/nonexistent/multi_abd.so"))
        return lib, fake

    def test_run_abi_argtypes(self):
        # The declared argtypes must match, in count and key types, the
        # positional arguments run() actually passes (19).
        _, fake = self._make_lib()
        argt = fake.dispatcher_run_multi_abd.argtypes
        self.assertEqual(len(argt), 19)
        for i in range(3):  # A/B/D host-pointer arrays
            self.assertEqual(argt[i], ctypes.POINTER(ctypes.c_void_p))
        self.assertEqual(argt[3], ctypes.c_void_p)  # E host pointer
        for i in (4, 5, 6):  # A/B/D leading-stride arrays
            self.assertEqual(argt[i], ctypes.POINTER(ctypes.c_int64))
        self.assertEqual(argt[-1], ctypes.POINTER(ctypes.c_float))  # time_ms

    def test_run_marshals_pointer_arrays_counts_and_order(self):
        lib, fake = self._make_lib(status=0)
        na, nb, nd = 2, 2, 2
        as_arrays = [np.zeros(4, np.float16) for _ in range(na)]
        bs_arrays = [np.zeros(4, np.float16) for _ in range(nb)]
        ds_arrays = [np.zeros(4, np.float16) for _ in range(nd)]
        e = np.zeros(4, np.float16)
        status, _ = lib.run(
            as_arrays, bs_arrays, ds_arrays, e,
            M=2, N=2, K=2, elem_a=2, elem_b=2, elem_d=2, elem_e=2,
            stride_as=[2, 2], stride_bs=[2, 2], stride_ds=[2, 2], stride_e=2,
        )
        self.assertEqual(status, 0)
        fake.dispatcher_run_multi_abd.assert_called_once()
        args = fake.dispatcher_run_multi_abd.call_args[0]
        self.assertEqual(len(args), 19)
        # A/B/D host-pointer arrays carry exactly one entry per tensor.
        self.assertEqual(len(args[0]), na)
        self.assertEqual(len(args[1]), nb)
        self.assertEqual(len(args[2]), nd)
        # Tensor counts occupy positions 12/13/14; M/N/K positions 15/16/17.
        self.assertEqual((args[12], args[13], args[14]), (na, nb, nd))
        self.assertEqual((args[15], args[16], args[17]), (2, 2, 2))

    def test_run_forwards_nonzero_status(self):
        # A thin shim must surface the C error code (e.g. -3 tensor-count
        # mismatch) verbatim rather than swallowing it.
        lib, fake = self._make_lib(status=-3)
        a = [np.zeros(4, np.float16)]
        status, _ = lib.run(
            a, a, a, np.zeros(4, np.float16),
            M=1, N=1, K=1, elem_a=2, elem_b=2, elem_d=2, elem_e=2,
            stride_as=[1], stride_bs=[1], stride_ds=[1], stride_e=1,
        )
        self.assertEqual(status, -3)


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