diff --git a/projects/composablekernel/dispatcher/python/gemm_utils.py b/projects/composablekernel/dispatcher/python/gemm_utils.py index c3c13485bce8..95f5444bddd5 100644 --- a/projects/composablekernel/dispatcher/python/gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/gemm_utils.py @@ -377,6 +377,30 @@ def cleanup(self) -> None: # ============================================================================ +def _fp32_to_bf16_u16(x: np.ndarray) -> np.ndarray: + """Encode fp32 -> bfloat16 bit pattern in a uint16 array (round-to-nearest-even). + + numpy has no native bf16, but the C ABI only cares about the 2-byte memory + layout (sizeof(bf16_t) == 2 == sizeof(uint16)). Truncating the low 16 bits of + the fp32 representation with round-to-nearest-even matches ck_tile's bf16. + """ + u32 = np.ascontiguousarray(x, dtype=np.float32).view(np.uint32) + # round-to-nearest-even: add (lsb-of-kept-bits + 0x7FFF) before truncating + rounding = ((u32 >> 16) & 1) + np.uint32(0x7FFF) + return ((u32 + rounding) >> 16).astype(np.uint16) + + +def _bf16_u16_to_fp32(u16: np.ndarray) -> np.ndarray: + """Decode a uint16 bf16 bit pattern back to fp32 (low 16 mantissa bits zero).""" + return (u16.astype(np.uint32) << 16).view(np.float32) + + +def _dtype_from_kernel_name(name: str) -> str: + """Extract the dtype token from a kernel name like ``gemm___...``.""" + parts = name.split("_") + return parts[1] if len(parts) > 1 else "fp16" + + class GpuGemmRunner: """High-level runner: construct from a .so path, call run(A, B, problem). @@ -402,16 +426,30 @@ def run( M, N, K = problem.M, problem.N, problem.K # A is row-major MxK; B is supplied KxN and stored column-major (the - # 'c' in rcr), matching how the kernel expects its operands. - A_h = np.ascontiguousarray(A, dtype=np.float16) - B_h = np.ascontiguousarray(B.T, dtype=np.float16) - C_h = np.zeros((M, N), dtype=np.float16) + # 'c' in rcr), matching how the kernel expects its operands. The element + # dtype is dictated by the compiled kernel (encoded in its name); the C + # ABI sizes its buffers from sizeof(ADataType), so the host buffers must + # match it byte-for-byte. + dtype = _dtype_from_kernel_name(self._kernel_name) + if dtype == "bf16": + A_h = _fp32_to_bf16_u16(A) + B_h = _fp32_to_bf16_u16(np.ascontiguousarray(B.T)) + C_h = np.zeros((M, N), dtype=np.uint16) + else: # fp16 (default) + A_h = np.ascontiguousarray(A, dtype=np.float16) + B_h = np.ascontiguousarray(B.T, dtype=np.float16) + C_h = np.zeros((M, N), dtype=np.float16) status, time_ms = self.lib.run(A_h, B_h, C_h, M, N, K) + if dtype == "bf16": + C_out = _bf16_u16_to_fp32(C_h) + else: + C_out = C_h + tflops = (problem.flops / (time_ms * 1e-3)) / 1e12 if time_ms > 0 else 0.0 return GemmResult( - output=C_h, + output=C_out, time_ms=time_ms, status=status, tflops=tflops, diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py index 4b90a99cc2ea..ed296e35d6f8 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py @@ -47,7 +47,7 @@ # column-major B (rcr). Constrain the CLI to what the runner actually supports so # a mismatched --dtype/--layout fails fast instead of silently benchmarking the # wrong thing. Later phases widen these sets as the runner gains signatures. -SUPPORTED_DTYPES = ("fp16",) +SUPPORTED_DTYPES = ("fp16", "bf16") SUPPORTED_LAYOUTS = ("rcr",) diff --git a/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py b/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py index 23fca859d4a4..e0c71d768b12 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py +++ b/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py @@ -37,8 +37,10 @@ def _run_one(idx, so_path, prob_dict, kernel_name): problem = GemmProblem.from_dict(prob_dict) np.random.seed(42) - A = (np.random.randn(problem.M, problem.K) * 0.1).astype(np.float16) - B = (np.random.randn(problem.K, problem.N) * 0.1).astype(np.float16) + # Generate fp32 source; the runner encodes to the kernel's real dtype + # (fp16 or bf16) based on the compiled kernel name. + A = (np.random.randn(problem.M, problem.K) * 0.1).astype(np.float32) + B = (np.random.randn(problem.K, problem.N) * 0.1).astype(np.float32) # CRITICAL: load the library ONLY inside this subprocess. runner = GpuGemmRunner(lib_path=so_path)