#!/usr/bin/env python3 # Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT """CPU-only unit tests for python/gemm_utils.py. Locks in the bit-level helpers that the TE -> Dispatcher GEMM bridge relies on: * bf16 <-> uint16 encoding (round-to-nearest-even), since numpy has no native bf16 and the runner carries bf16 as a uint16 bit pattern. * fp8 (E4M3) / bf8 (E5M2) FNUZ <-> uint8 encoding, used for the gfx942 8-bit float surface. The decode must be exact to the device format; the encode only needs to land on the nearest representable byte. * dtype / layout parsing from the compiled kernel name, which drives how the runner lays out host buffers. No GPU is touched -- all functions under test are pure host-side logic. Run: python3 -m pytest tests/test_gemm_utils.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 / "python")) import numpy as np # noqa: E402 from gemm_utils import ( # noqa: E402 GemmKernelConfig, _fp32_to_bf16_u16, _bf16_u16_to_fp32, _fp32_to_fp8_u8, _fp8_u8_to_fp32, _fp32_to_bf8_u8, _bf8_u8_to_fp32, _fnuz_decode_table, _output_dtype, _dtype_from_kernel_name, _layout_from_kernel_name, ) class TestBf16Encoding(unittest.TestCase): """bf16 = top 16 bits of fp32 with round-to-nearest-even.""" def test_exactly_representable_roundtrip(self): # Values whose low 16 fp32 mantissa bits are zero are exact in bf16. exact = np.array([0.0, 1.0, -1.0, 2.0, 0.5, -0.5, 4.0, 256.0], dtype=np.float32) out = _bf16_u16_to_fp32(_fp32_to_bf16_u16(exact)) np.testing.assert_array_equal(out, exact) def test_roundtrip_within_bf16_tolerance(self): rng = np.random.default_rng(0) x = (rng.standard_normal(10000) * 100.0).astype(np.float32) out = _bf16_u16_to_fp32(_fp32_to_bf16_u16(x)) # bf16 has 8 bits of significand -> relative error <= 2^-8. rel = np.abs(out - x) / (np.abs(x) + 1e-30) self.assertLessEqual(float(rel.max()), 2.0 ** -8) def test_round_to_nearest_even_ties(self): # Tie halfway between bf16 1.0 (0x3F80, even) and 0x3F81 (odd): # fp32 0x3F808000 must round DOWN to the even neighbor 0x3F80. tie_down = np.array([0x3F808000], dtype=np.uint32).view(np.float32) self.assertEqual(int(_fp32_to_bf16_u16(tie_down)[0]), 0x3F80) # Tie halfway between 0x3F81 (odd) and 0x3F82 (even): # fp32 0x3F818000 must round UP to the even neighbor 0x3F82. tie_up = np.array([0x3F818000], dtype=np.uint32).view(np.float32) self.assertEqual(int(_fp32_to_bf16_u16(tie_up)[0]), 0x3F82) def test_special_values(self): inf = np.array([np.inf, -np.inf], dtype=np.float32) out = _bf16_u16_to_fp32(_fp32_to_bf16_u16(inf)) self.assertTrue(np.isinf(out[0]) and out[0] > 0) self.assertTrue(np.isinf(out[1]) and out[1] < 0) nan = np.array([np.nan], dtype=np.float32) out_nan = _bf16_u16_to_fp32(_fp32_to_bf16_u16(nan)) self.assertTrue(np.isnan(out_nan[0])) def test_dtype_and_size(self): u16 = _fp32_to_bf16_u16(np.zeros(4, dtype=np.float32)) self.assertEqual(u16.dtype, np.uint16) self.assertEqual(u16.itemsize, 2) # must match sizeof(bf16_t) on device class TestFp8Bf8Encoding(unittest.TestCase): """fp8 E4M3 / bf8 E5M2 in the FNUZ format used by gfx942. The decode is the load-bearing half (it must equal the device value for a byte); the encode must land on the nearest representable byte and saturate. """ def test_format_ranges(self): # FNUZ maxima: E4M3 -> 2^7 * 1.875 = 240; E5M2 -> 2^15 * 1.75 = 57344. t43 = _fnuz_decode_table(4, 3) t52 = _fnuz_decode_table(5, 2) self.assertEqual(float(np.nanmax(t43)), 240.0) self.assertEqual(float(np.nanmin(t43)), -240.0) self.assertEqual(float(np.nanmax(t52)), 57344.0) self.assertEqual(float(np.nanmin(t52)), -57344.0) def test_zero_and_nan_slots(self): # 0x00 is +0; the negative-zero slot 0x80 is the lone NaN (FNUZ). for tab in (_fnuz_decode_table(4, 3), _fnuz_decode_table(5, 2)): self.assertEqual(float(tab[0x00]), 0.0) self.assertTrue(np.isnan(tab[0x80])) def test_exactly_representable_roundtrip(self): exact = np.array([0.0, 0.5, 1.0, -1.0, 2.0, -2.0, 1.5, -0.25, 4.0, 8.0], dtype=np.float32) np.testing.assert_array_equal( _fp8_u8_to_fp32(_fp32_to_fp8_u8(exact)), exact) np.testing.assert_array_equal( _bf8_u8_to_fp32(_fp32_to_bf8_u8(exact)), exact) def test_decode_is_consistent_with_encode(self): # The parity contract: ref multiplies decode(encode(x)), so the pair must # be self-consistent and every encoded byte must decode finite. rng = np.random.default_rng(1) x = (rng.standard_normal(5000) * 0.1).astype(np.float32) for enc, dec in ((_fp32_to_fp8_u8, _fp8_u8_to_fp32), (_fp32_to_bf8_u8, _bf8_u8_to_fp32)): d = dec(enc(x)) self.assertTrue(np.all(np.isfinite(d))) def test_saturates_no_inf(self): # FNUZ has no infinity: huge magnitudes clamp to the finite max. big = np.array([1e30, -1e30], dtype=np.float32) self.assertEqual(float(_fp8_u8_to_fp32(_fp32_to_fp8_u8(big))[0]), 240.0) self.assertEqual(float(_bf8_u8_to_fp32(_fp32_to_bf8_u8(big))[1]), -57344.0) def test_dtype_and_size(self): for enc in (_fp32_to_fp8_u8, _fp32_to_bf8_u8): u8 = enc(np.zeros(4, dtype=np.float32)) self.assertEqual(u8.dtype, np.uint8) self.assertEqual(u8.itemsize, 1) # must match sizeof(fp8_t/bf8_t) class TestOutputDtype(unittest.TestCase): """Output (C) element dtype must mirror the codegen's get_output_dtype.""" def test_mapping(self): self.assertEqual(_output_dtype("fp16"), "fp16") self.assertEqual(_output_dtype("bf16"), "bf16") self.assertEqual(_output_dtype("fp8"), "fp16") self.assertEqual(_output_dtype("bf8"), "fp16") self.assertEqual(_output_dtype("int8"), "int32") class TestKernelNameParsing(unittest.TestCase): """The runner reads dtype + layout straight from the compiled .so name.""" _NAME = ("gemm_bf16_rcr_compv3_cshuffle_intrawave_" "False_False_False_False_64x64x64_4x1x1_16x16x16") def test_dtype_from_name(self): self.assertEqual(_dtype_from_kernel_name(self._NAME), "bf16") self.assertEqual( _dtype_from_kernel_name("gemm_fp16_rrr_compv4_cshuffle_intrawave"), "fp16", ) def test_dtype_fallback(self): # Malformed / single-token name falls back to fp16. self.assertEqual(_dtype_from_kernel_name("gemm"), "fp16") def test_layout_from_name(self): self.assertEqual(_layout_from_kernel_name(self._NAME), "rcr") for lay in ("rrr", "ccr", "crr", "rcc"): name = f"gemm_fp16_{lay}_compv3_cshuffle_intrawave" self.assertEqual(_layout_from_kernel_name(name), lay) def test_layout_fallback(self): # A token that is not a 3-char r/c string falls back to rcr. self.assertEqual( _layout_from_kernel_name("gemm_fp16_xyz_compv3"), "rcr" ) self.assertEqual(_layout_from_kernel_name("gemm"), "rcr") class TestConfigNameContract(unittest.TestCase): """GemmKernelConfig.name is the single source of truth tying config -> codegen -> runtime; parsing it back must recover dtype and layout.""" def test_name_roundtrips_through_parsers(self): for dtype in ("fp16", "bf16", "fp8", "bf8", "int8"): for la, lb, lc in (("row", "col", "row"), ("row", "row", "row"), ("col", "col", "row"), ("col", "row", "row")): cfg = GemmKernelConfig( dtype_a=dtype, dtype_b=dtype, dtype_c=_output_dtype(dtype), dtype_acc=("int32" if dtype == "int8" else "fp32"), layout_a=la, layout_b=lb, layout_c=lc, ) name = cfg.name self.assertEqual(_dtype_from_kernel_name(name), dtype) self.assertEqual(_layout_from_kernel_name(name), cfg.layout) if __name__ == "__main__": unittest.main()