[2/5] autotune: add opt-in search and RMSNorm waves-per-EU tuning (#770)

kernels/norm/rmsnorm_autotune.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+"""Two-track RMSNorm autotuning through the normal direct JIT path."""
+
+from flydsl.autotune import Config, autotune
+from kernels.norm.rmsnorm_common import BLOCK_THREADS
+from kernels.norm.rmsnorm_kernel import SMALL_N_THRESHOLD, rmsnorm_direct
+
+_SEARCH_CONFIGS = (
+    Config(BLOCK_THREADS=128),
+    Config(BLOCK_THREADS=128, waves_per_eu=1),
+    Config(BLOCK_THREADS=128, waves_per_eu=2),
+    Config(BLOCK_THREADS=256),
+    Config(BLOCK_THREADS=256, waves_per_eu=1),
+    Config(BLOCK_THREADS=512),
+    Config(BLOCK_THREADS=512, waves_per_eu=2),
+)
+
+
+def _default_config(*_args, **_kwargs):
+    return Config(BLOCK_THREADS=BLOCK_THREADS)
+
+
+def _search_configs(input_t, gamma, output, m_in, N, dtype_str="bf16", stream=None):
+    if N <= SMALL_N_THRESHOLD:
+        return [_default_config()]
+    return list(_SEARCH_CONFIGS)
+
+
+_rmsnorm_tuner = autotune(
+    configs=_search_configs,
+    key=["m_in", "N", "dtype_str"],
+    default=_default_config,
+)(rmsnorm_direct)
+
+
+def rmsnorm_autotuned(input_t, gamma, output, m_in, dtype_str="bf16", stream=None):
+    import torch
+
+    with torch.cuda.device(input_t.device):
+        launch_stream = torch.cuda.current_stream() if stream is None else stream
+        return _rmsnorm_tuner(
+            input_t,
+            gamma,
+            output,
+            m_in,
+            N=int(input_t.shape[-1]),
+            dtype_str=dtype_str,
+            stream=launch_stream,
+        )

kernels/norm/rmsnorm_kernel.py

@@ -47,6 +47,9 @@
 
 KERNEL_NAME = "rmsnorm"
 
+# The small-N path derives its own block geometry and is not tuned.
+SMALL_N_THRESHOLD = 2048
+
 
 def _store_yscale(scale_copy_atom, yscale_div, index, val):
     r = fx.make_rmem_tensor(1, fx.Float32)
@@ -67,20 +70,24 @@ def _quant_dtype_max(dtype_str: str) -> float:
     raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')")
 
 
-def build_rmsnorm_module(N: int, dtype_str: str, store_rstd: bool = False, eps: float = EPS):
-    if N <= 2048:
+def build_rmsnorm_module(
+    N: int, dtype_str: str, store_rstd: bool = False, eps: float = EPS, BLOCK_THREADS: int = BLOCK_THREADS
+):
+    if N <= SMALL_N_THRESHOLD:
         return _build_rmsnorm_large_m_small_n_module(N, dtype_str, store_rstd, eps)
 
     arch = get_rocm_arch()
     USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95")
 
+    # BLOCK_THREADS controls storage, tiling, and launch geometry.
     tile_cols = BLOCK_THREADS * VEC_WIDTH
     RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE)
     elem_bits = 32 if dtype_str == "f32" else 16
+    _kernel_kwargs = {} if BLOCK_THREADS <= 256 else {"known_block_size": [BLOCK_THREADS, 1, 1]}
 
     SharedStorage = _make_reduction_storage(RED_SLOTS)
 
-    @flyc.kernel
+    @flyc.kernel(**_kernel_kwargs)
     def rmsnorm_kernel(
         Input: fx.Tensor,
         Gamma: fx.Tensor,
@@ -300,6 +307,22 @@ def launch_rmsnorm(
     return launch_rmsnorm
 
 
+@flyc.jit
+def rmsnorm_direct(
+    Input: fx.Tensor,
+    Gamma: fx.Tensor,
+    Output: fx.Tensor,
+    m_in: fx.Int32,
+    N: fx.Constexpr[int],
+    dtype_str: fx.Constexpr[str],
+    BLOCK_THREADS: fx.Constexpr[int],
+    stream: fx.Stream = fx.Stream(None),
+):
+    """Specialize the existing RMSNorm factory through JIT Constexpr inputs."""
+    launch = build_rmsnorm_module(N, dtype_str, BLOCK_THREADS=BLOCK_THREADS)
+    launch(Input, Gamma, Output, m_in, stream)
+
+
 def _build_rmsnorm_large_m_small_n_module(N: int, dtype_str: str, store_rstd: bool = False, eps: float = EPS):
     BLOCK_N = 1 << (N - 1).bit_length()
     BLOCK_M = max(min(16384 // BLOCK_N, 32), 8)

python/flydsl/autotune.py

@@ -6,6 +6,7 @@
 import inspect
 import json
 import os
+from contextlib import nullcontext
 from pathlib import Path
 from typing import Callable, Dict, List
 
@@ -15,6 +16,11 @@
     torch = None
 
 
+def _tuning_enabled() -> bool:
+    """Whether to bypass cached/default configs and run a fresh search."""
+    return os.environ.get("FLYDSL_AUTOTUNE", "").strip().lower() in ("1", "true", "yes", "on")
+
+
 def _env_fingerprint() -> tuple:
     """Sorted cache-invalidating env vars (reuses the JIT's canonical list)."""
     try:
@@ -165,6 +171,7 @@ def __init__(
         pre_hook=None,
         post_hook=None,
         do_bench_fn=None,
+        default=None,
     ):
         self.fn = fn  # JitFunction instance
         self.configs = configs
@@ -178,6 +185,7 @@ def __init__(
         self.post_hook = post_hook
         self._do_bench = do_bench_fn or do_bench
         self.cache: Dict[tuple, Config] = {}
+        self.default = default
 
         # Infer arg names from the underlying function
         if hasattr(fn, "func"):
@@ -231,6 +239,15 @@ def _make_key(self, args, kwargs):
         key_vals.append(("_env_", _env_fingerprint()))
         key_vals.append(("_toolchain_", _toolchain_fingerprint()))
         key_vals.append(("_device_", _device_fingerprint()))
+        effective_hints = getattr(self.fn, "_effective_compile_hints", None)
+        if callable(effective_hints):
+            hint_key = tuple(
+                sorted(
+                    (key, type(value).__module__, type(value).__qualname__, repr(value))
+                    for key, value in effective_hints().items()
+                )
+            )
+            key_vals.append(("_compile_hints_", hint_key))
 
         return tuple(str(v) for v in key_vals)
 
@@ -275,35 +292,56 @@ def _prune(self, configs, args, kwargs):
             return self.prune_configs_by(configs, sig_args)
         return configs
 
+    def _stream_context(self, args, kwargs):
+        """Use an explicit torch/raw stream for benchmark events and setup ops."""
+        if torch is None:
+            return nullcontext()
+        sig_args = dict(zip(self.arg_names, args))
+        sig_args.update(kwargs)
+        stream = sig_args.get("stream")
+        if isinstance(stream, torch.cuda.Stream):
+            return torch.cuda.stream(stream)
+        if isinstance(stream, int):
+            device = next(
+                (value.device for value in sig_args.values() if isinstance(value, torch.Tensor) and value.is_cuda),
+                None,
+            )
+            stream = (
+                torch.cuda.default_stream(device) if stream == 0 else torch.cuda.ExternalStream(stream, device=device)
+            )
+            return torch.cuda.stream(stream)
+        return nullcontext()
+
     def _bench_one(self, config, args, kwargs):
         """Compile and benchmark one config. Returns time in ms."""
         merged_kwargs = dict(kwargs)
         merged_kwargs.update(config.all_kwargs())
         compiler_opts = config.compiler_opts()
 
-        # Snapshot once before any rep runs, so restores are from pristine input.
-        snapshot = self._snapshot_tensors(args, merged_kwargs)
-
-        def kernel_call():
-            # Order: restore/reset the inputs first, THEN run the pre_hooks, so a
-            # hook that sets up state (incl. mutating a tensor) isn't clobbered
-            # by the restore. Each benchmark rep starts from clean inputs.
-            self._restore_tensors(snapshot)
-            self._reset_tensors(args, merged_kwargs)
-            if config.pre_hook:
-                config.pre_hook(merged_kwargs)
-            if self.pre_hook:
-                self.pre_hook(merged_kwargs)
-            self._run_with_hints(compiler_opts, args, merged_kwargs)
-            if self.post_hook:
-                self.post_hook(merged_kwargs)
+        with self._stream_context(args, merged_kwargs):
+            # Snapshot once before any rep runs, so restores are from pristine input.
+            snapshot = self._snapshot_tensors(args, merged_kwargs)
 
-        try:
-            return self._do_bench(kernel_call, warmup=self.warmup, rep=self.rep)
-        finally:
-            # Leave the caller's tensors as a single clean run would.
-            if snapshot:
+            def kernel_call():
+                # Order: restore/reset the inputs first, THEN run the pre_hooks, so a
+                # hook that sets up state (incl. mutating a tensor) isn't clobbered
+                # by the restore. Each benchmark rep starts from clean inputs.
                 self._restore_tensors(snapshot)
+                self._reset_tensors(args, merged_kwargs)
+                if config.pre_hook:
+                    config.pre_hook(merged_kwargs)
+                if self.pre_hook:
+                    self.pre_hook(merged_kwargs)
+                self._run_with_hints(compiler_opts, args, merged_kwargs)
+                if self.post_hook:
+                    self.post_hook(merged_kwargs)
+
+            try:
+                return self._do_bench(kernel_call, warmup=self.warmup, rep=self.rep)
+            finally:
+                # Leave the caller's tensors as a single clean run would.
+                if snapshot:
+                    self._restore_tensors(snapshot)
 
     def _run_with_hints(self, compiler_opts, args, kwargs):
         """Run the kernel with optional compiler hints. Import is deferred so
@@ -322,16 +360,22 @@ def _run_config(self, config, args, kwargs):
         clean run (restore_value tensors are already restored by _bench_one)."""
         merged = dict(kwargs)
         merged.update(config.all_kwargs())
-        self._reset_tensors(args, merged)
-        return self._run_with_hints(config.compiler_opts(), args, merged)
+        with self._stream_context(args, merged):
+            self._reset_tensors(args, merged)
+            return self._run_with_hints(config.compiler_opts(), args, merged)
 
     def __call__(self, *args, **kwargs):
         key = self._make_key(args, kwargs)
-        if key in self.cache:
+        force = _tuning_enabled()
+
+        if not force and key in self.cache:
             return self._run_config(self.cache[key], args, kwargs)
 
-        # Benchmark all configs
-        configs = self._prune(self.configs, args, kwargs)
+        if not force and self.default is not None:
+            return self._run_config(self.default(*args, **kwargs), args, kwargs)
+
+        configs = self.configs(*args, **kwargs) if callable(self.configs) else self.configs
+        configs = self._prune(configs, args, kwargs)
         print(f"[autotune] tuning {len(configs)} configs...")
         results = []
         for i, config in enumerate(configs):
@@ -373,7 +417,7 @@ def _save_disk_cache(self):
 
 
 def autotune(
-    configs: List[Config],
+    configs,
     key: List[str] = None,
     warmup: int = 5,
     rep: int = 25,
@@ -383,6 +427,7 @@ def autotune(
     pre_hook: Callable = None,
     post_hook: Callable = None,
     do_bench: Callable = None,
+    default: Callable = None,
 ):
     """Autotune decorator for @jit functions.
 
@@ -393,6 +438,10 @@ def myKernel(..., BLOCK: fx.Constexpr[int], ...):
             ...
 
     Args:
+        configs: sequence of :class:`Config`, or a callable returning one for
+            the current arguments.
+        default: optional heuristic ``default(*args, **kwargs) -> Config`` used
+            without benchmarking unless ``FLYDSL_AUTOTUNE`` forces a search.
         restore_value: tensor args the kernel mutates in place (output overlaps
             input, or accumulation). Snapshotted and restored before each bench
             rep so every config is measured on identical inputs. Required when
@@ -414,6 +463,7 @@ def decorator(fn):
             pre_hook=pre_hook,
             post_hook=post_hook,
             do_bench_fn=do_bench,
+            default=default,
         )
 
     return decorator

python/flydsl/compiler/backends/base.py

@@ -69,6 +69,10 @@ def pipeline_fragments(self, *, compile_hints: dict) -> List[str]:
         """
         ...
 
+    def lower_compile_hints(self, module, *, compile_hints: dict) -> None:
+        """Optionally materialize backend-specific hints before the pipeline."""
+        return None
+
     def external_binary_pipeline_fragments(self, *, compile_hints: dict) -> Tuple[List[str], str]:
         """Split the pipeline for external device binary code generation.
 

python/flydsl/compiler/backends/rocm.py

@@ -100,6 +100,26 @@ def pipeline_fragments(self, *, compile_hints: dict) -> List[str]:
     def external_binary_pipeline_fragments(self, *, compile_hints: dict) -> Tuple[List[str], str]:
         return self._pipeline_parts(compile_hints=compile_hints)
 
+    def lower_compile_hints(self, module, *, compile_hints: dict) -> None:
+        """Materialize a scalar waves-per-EU override on kernel entries."""
+        waves_per_eu = compile_hints.get("waves_per_eu")
+        if waves_per_eu is None:
+            return
+        if isinstance(waves_per_eu, bool) or not isinstance(waves_per_eu, int):
+            raise TypeError(f"waves_per_eu must be a non-negative int, got {waves_per_eu!r}")
+        if waves_per_eu < 0:
+            raise ValueError(f"waves_per_eu must be >= 0, got {waves_per_eu}")
+        if waves_per_eu == 0:
+            return
+
+        with module.context:
+            for func_op in _iter_gpu_kernel_funcs(module):
+                # rocdl.waves_per_eu expresses a minimum. Replace it with the exact
+                # min/max LLVM passthrough for an explicit compile-hint override.
+                if "rocdl.waves_per_eu" in func_op.attributes:
+                    del func_op.attributes["rocdl.waves_per_eu"]
+                _set_passthrough(func_op, "amdgpu-waves-per-eu", f"{waves_per_eu},{waves_per_eu}")
+
     def gpu_module_targets(self) -> List[str]:
         chip = self.target.arch
         return [f'#rocdl.target<chip = "{chip}">']
@@ -120,3 +140,30 @@ def jit_runtime_lib_basenames(self) -> List[str]:
             "libfly_jit_runtime.so",
             "libmlir_c_runner_utils.so",
         ]
+
+
+def _iter_gpu_kernel_funcs(module):
+    """Yield entry ``gpu.func`` ops, excluding device helpers."""
+    for top in module.body.operations:
+        if top.operation.name != "gpu.module":
+            continue
+        for op in top.regions[0].blocks[0].operations:
+            if op.operation.name == "gpu.func" and "gpu.kernel" in op.attributes:
+                yield op
+
+
+def _set_passthrough(func_op, key: str, value: str) -> None:
+    """Replace one LLVM passthrough key while preserving unrelated entries."""
+    from ..._mlir import ir
+
+    def _entry_key(entry):
+        try:
+            pair = ir.ArrayAttr(entry)
+            return ir.StringAttr(pair[0]).value if len(pair) else None
+        except (TypeError, ValueError):
+            return None
+
+    new_entry = ir.ArrayAttr.get([ir.StringAttr.get(key), ir.StringAttr.get(value)])
+    existing = func_op.attributes["passthrough"] if "passthrough" in func_op.attributes else None
+    kept = [entry for entry in existing if _entry_key(entry) != key] if existing is not None else []
+    func_op.attributes["passthrough"] = ir.ArrayAttr.get([*kept, new_entry])