hgemm: add split-k support for layout version

README.md

@@ -13,6 +13,18 @@ The goal of this repository is to show how to write high-performance GEMM kernel
 
 ![HGEMM BF16 benchmark vs Torch hipBLAS](images/hgemm_benchmark.svg)
 
+## FlyDSL install from source
+
+```bash
+pip uninstall -y flydsl
+git clone git@github.com:ROCm/FlyDSL.git
+cd FlyDSL ; git checkout 9a5c08e77355f915ad35965bea8ea88f0af33bf3
+git submodule sync && git submodule update --init --recursive
+bash scripts/build_llvm.sh -j64
+bash scripts/build.sh -j64
+pip install -e .
+```
+
 ## Run Tests And Benchmarks
 
 ```bash

gemm_tune.py

@@ -1,3 +1,7 @@
+import os
+
+os.environ.setdefault("KINETO_LOG_LEVEL", "6")
+
 import json
 import torch
 import itertools
@@ -10,7 +14,7 @@
 from dataclasses import dataclass
 from flydsl.runtime.device import get_rocm_arch
 
-from kernels.hgemm_layout_gfx950 import hgemm, make_hgemm_gfx950_param
+from kernels.hgemm_layout_gfx950 import hgemm, make_hgemm_param_and_validate
 
 gpu_arch = get_rocm_arch()
 base_dir = Path(__file__).resolve().parent
@@ -25,6 +29,7 @@ class Args:
     n: int
     k: int
     layout: str
+    enable_split_k: bool = False
 
 
 @dataclass
@@ -40,6 +45,45 @@ class TunedArgs:
     tflops: float
 
 
+@dataclass(frozen=True)
+class GemmTileIoUPruner:
+    m: int
+    n: int
+    k: int
+    keep_ratio: float
+
+    @staticmethod
+    def _ceil_div(value, divisor):
+        return (value + divisor - 1) // divisor
+
+    def _split_k_padded(self, block_k, split_k):
+        working_k = self._ceil_div(self.k, split_k)
+        padded_k = 0
+        for split_idx in range(split_k):
+            remaining_k = max(self.k - split_idx * working_k, 0)
+            part_k = min(working_k, remaining_k)
+            if part_k > 0:
+                padded_k += self._ceil_div(part_k, block_k) * block_k
+        return padded_k
+
+    def _config_iou(self, config):
+        padded_m = self._ceil_div(self.m, config["block_m"]) * config["block_m"]
+        padded_n = self._ceil_div(self.n, config["block_n"]) * config["block_n"]
+        padded_k = self._split_k_padded(config["block_k"], config["split_k"])
+        return (self.m * self.n * self.k) / (padded_m * padded_n * padded_k)
+
+    def prune(self, configs):
+        if not configs:
+            return configs
+        config_ious = [self._config_iou(config) for config in configs]
+        threshold = max(config_ious) * self.keep_ratio
+        return [
+            config
+            for config, iou in zip(configs, config_ious)
+            if iou >= threshold
+        ]
+
+
 def empty_layout_matrix(rows, cols, dtype, is_t):
     if is_t:
         return torch.empty((cols, rows), dtype=dtype, device="cuda").t()
@@ -78,7 +122,13 @@ def tuning_benchmark(args, kwargs={}, niters=50):
     c_ref = create_outputs(args)[0]
     torch.addmm(bias, a, b, out=c_ref)
     hgemm(a, b, c, bias=bias, user_kwargs=kwargs, layout=args.layout)
-    tol = float(args.k) / 2048 * 6e-1
+    tol = (
+        float(args.k)
+        / 2048
+        * 6e-1
+        * kwargs.get("split_k", 1)
+        * kwargs.get("k_waves", 1)
+    )
     is_allclose = torch.allclose(c, c_ref, atol=tol, rtol=tol)
     assert is_allclose
     # performance bench
@@ -107,37 +157,72 @@ def tuning_benchmark(args, kwargs={}, niters=50):
     return duration
 
 
-def hgemm_get_configs():
+def hgemm_get_configs(args):
+    split_k_candidates = [1]
+    if args.enable_split_k:
+        split_k_candidates.extend(
+            split_k for split_k in range(2, 10) if args.k % split_k == 0
+        )
     selections = {
         "block_m": [16, 32, 48, 64, 80, 96, 128, 256],
         "block_n": [16, 32, 64, 80, 96, 128, 256],
         "block_k": [64, 128, 256],
         "stages": [i for i in range(2, 10)],
+        "split_k": split_k_candidates,
         "m_waves": [1, 2, 4],
         "n_waves": [1, 2, 4],
+        "k_waves": [1, 2],
         "group_m": [0, 4],
         "use_half_tile_interleaved": [False, True],
     }
     keys = selections.keys()
     values = selections.values()
     configs = [dict(zip(keys, combo)) for combo in itertools.product(*values)]
+    keep_ratio = 0.75 if args.m <= 32 else 0.85 if args.m <= 128 else 0.95
+    configs = GemmTileIoUPruner(
+        args.m,
+        args.n,
+        args.k,
+        keep_ratio,
+    ).prune(configs)
     valid_configs = []
+    is_large_gemm = args.m >= 4096 and args.n >= 4096 and args.k >= 4096
     for config in configs:
-        if not config["use_half_tile_interleaved"]:
-            mma_m_iters = config["block_m"] // config["m_waves"] // 16
-            mma_n_iters = config["block_n"] // config["n_waves"] // 16
-            if mma_m_iters > 4 or mma_n_iters > 4:
+        if is_large_gemm:
+            if not (
+                config["use_half_tile_interleaved"]
+                and config["block_m"] == 256
+                and config["block_n"] == 256
+                and config["block_k"] == 64
+                and config["stages"] == 2
+                and config["split_k"] == 1
+                and config["m_waves"] == 2
+                and config["n_waves"] == 4
+                and config["k_waves"] == 1
+            ):
                 continue
+        else:
+            if not config["use_half_tile_interleaved"]:
+                mma_m_iters = config["block_m"] // config["m_waves"] // 16
+                mma_n_iters = config["block_n"] // config["n_waves"] // 16
+                if mma_m_iters > 4 or mma_n_iters > 4:
+                    continue
         try:
-            make_hgemm_gfx950_param(**config)
-            valid_configs.append(config)
+            param = make_hgemm_param_and_validate(
+                args.m,
+                args.n,
+                args.k,
+                config,
+            )
+            if param is not None:
+                valid_configs.append(config)
         except Exception:
             pass
     return valid_configs
 
 
 def tune_single(args):
-    configs = hgemm_get_configs()
+    configs = hgemm_get_configs(args)
     best_duration = float(1e10)
     best_idx = 0
     pbar = tqdm(total=len(configs), desc=f"{args}")
@@ -171,8 +256,18 @@ def tune_all(
     dtype,
     out_prefix,
     layout,
+    enable_split_k=False,
 ):
     mnks = [
+        # splitk
+        # (32, 384, 7168),
+        # (32, 384, 16384),
+        # (800, 384, 7168),
+        # (32, 7168, 2048),
+        # (8, 7168, 2048),
+        # (8, 5120, 2880),
+        # (32, 2880, 2048),
+        # normal
         (8, 4096, 4096),
         (16, 4096, 4096),
         (32, 4096, 4096),
@@ -201,6 +296,7 @@ def tune_all(
                 n=mnk[1],
                 k=mnk[2],
                 layout=layout,
+                enable_split_k=enable_split_k,
             )
             result = tune_single(args)
             result = vars(result)
@@ -219,6 +315,12 @@ def tune_all(
     )
     parser.add_argument("--single", action="store_true")
     parser.add_argument("--tune_all", action="store_true")
+    parser.add_argument(
+        "--enable_split_k",
+        "--enable-split-k",
+        action="store_true",
+        help="include valid split_k values greater than 1 in tuning",
+    )
     parser.add_argument("--m", type=int, default=4096)
     parser.add_argument("--n", type=int, default=4096)
     parser.add_argument("--k", type=int, default=4096)
@@ -229,7 +331,12 @@ def tune_all(
     if args.single:
         tune_single(args)
     elif args.tune_all:
-        tune_all(args.dtype, args.out, args.layout)
+        tune_all(
+            args.dtype,
+            args.out,
+            args.layout,
+            enable_split_k=args.enable_split_k,
+        )
 
     # rm -rf ~/.flydsl/ ; python3 gemm_tune.py --single --dtype bf16 --m 1024 --n 1024 --k 1024
     # rm -rf ~/.flydsl/ ; python3 gemm_tune.py --single --dtype bf16 --m 2048 --n 2048 --k 2048
@@ -243,3 +350,4 @@ def tune_all(
     # rm -rf ~/.flydsl/ ; python3 gemm_tune.py --single --dtype bf16 --m 32 --n 384 --k 7168
 
     # rm -rf ~/.flydsl/ ; python3 gemm_tune.py --tune_all
+    # rm -rf ~/.flydsl/ ; python3 gemm_tune.py --tune_all --enable_split_k

kernels/common.py

@@ -0,0 +1,39 @@
+from threading import Lock
+from typing import Any, Callable
+
+_compiled_cache_lock = Lock()
+
+
+def run_cached(
+    jit_func: Any,
+    *compile_args: Any,
+    constexpr_param: Any,
+    compiler: Callable[..., Any],
+    dispatch_args: tuple[Any, ...],
+) -> Any:
+    """Cache a layout-dynamic FlyDSL dispatcher by constexpr param."""
+    cache_key = constexpr_param.__cache_signature__()
+    compiled_cache = getattr(jit_func, "_compiled_cache", None)
+    if compiled_cache is not None:
+        compiled = compiled_cache.get(cache_key)
+        if compiled is not None:
+            compiled(*dispatch_args)
+            return compiled
+
+    dispatch_after_wait = False
+    with _compiled_cache_lock:
+        compiled_cache = getattr(jit_func, "_compiled_cache", None)
+        if compiled_cache is None:
+            compiled_cache = {}
+            jit_func._compiled_cache = compiled_cache
+
+        compiled = compiled_cache.get(cache_key)
+        if compiled is None:
+            compiled = compiler(jit_func, *compile_args)
+            compiled_cache[cache_key] = compiled
+        else:
+            dispatch_after_wait = True
+
+    if dispatch_after_wait:
+        compiled(*dispatch_args)
+    return compiled