diff --git a/kernels/norm/rmsnorm_autotune.py b/kernels/norm/rmsnorm_autotune.py new file mode 100644 index 000000000..7408b9dbc --- /dev/null +++ b/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, + ) diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 424f06468..115b475cc 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/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) diff --git a/python/flydsl/autotune.py b/python/flydsl/autotune.py index b243d528b..3b4cad77b 100644 --- a/python/flydsl/autotune.py +++ b/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 diff --git a/python/flydsl/compiler/backends/base.py b/python/flydsl/compiler/backends/base.py index d27458119..0f6da1181 100644 --- a/python/flydsl/compiler/backends/base.py +++ b/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. diff --git a/python/flydsl/compiler/backends/rocm.py b/python/flydsl/compiler/backends/rocm.py index c32a328bf..9d6113af2 100644 --- a/python/flydsl/compiler/backends/rocm.py +++ b/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'] @@ -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]) diff --git a/python/flydsl/compiler/jit_function.py b/python/flydsl/compiler/jit_function.py index c2a411dbe..a213d9780 100644 --- a/python/flydsl/compiler/jit_function.py +++ b/python/flydsl/compiler/jit_function.py @@ -45,6 +45,7 @@ effective_fastmath_hint, func_def_location, get_gpu_module_body, + merge_compile_hints, ) from .link_utils import _append_link_lib_options_to_attach_targets, _format_link_lib_options from .protocol import ( @@ -747,14 +748,11 @@ class PipelineConfig: external: bool -def _pipeline_fragments_for_mode(backend) -> PipelineConfig: +def _pipeline_fragments_for_mode(backend, *, compile_hints: dict) -> PipelineConfig: """Return pipeline configuration including optional external split.""" - from .kernel_function import CompilationContext - - hints = CompilationContext.get_compile_hints() - llvm_opts = hints.get("llvm_options") + llvm_opts = compile_hints.get("llvm_options") if _use_external_binary_codegen(): - pre_binary_fragments, binary_fragment = backend.external_binary_pipeline_fragments(compile_hints=hints) + pre_binary_fragments, binary_fragment = backend.external_binary_pipeline_fragments(compile_hints=compile_hints) return PipelineConfig( fragments=[*pre_binary_fragments, binary_fragment], pre_binary=pre_binary_fragments, @@ -763,7 +761,7 @@ def _pipeline_fragments_for_mode(backend) -> PipelineConfig: external=True, ) - fragments = backend.pipeline_fragments(compile_hints=hints) + fragments = backend.pipeline_fragments(compile_hints=compile_hints) return PipelineConfig( fragments=fragments, pre_binary=None, @@ -798,8 +796,10 @@ def compile( backend = get_backend(arch=arch) + compile_hints = CompilationContext.get_compile_hints() module = ir.Module.parse(module.operation.get_asm(enable_debug_info=env.debug.enable_debug_info)) - cfg = _pipeline_fragments_for_mode(backend) + backend.lower_compile_hints(module, compile_hints=compile_hints) + cfg = _pipeline_fragments_for_mode(backend, compile_hints=compile_hints) fragments = cfg.fragments pre_binary_fragments = cfg.pre_binary binary_fragment = cfg.binary_fragment @@ -1203,6 +1203,10 @@ def __get__(self, obj, objtype=None): return self return partial(self.__call__, obj) + def _effective_compile_hints(self): + """Resolve persistent defaults and the current thread-local overlay.""" + return merge_compile_hints(self.compile_hints, CompilationContext.get_compile_hints()) + def _get_global_refs(self, owner_cls=None) -> List[Tuple[str, str, dict]]: """Memoized global-ref discovery (see :func:`_discover_global_refs`).""" cache = self._global_refs_cache @@ -1280,7 +1284,7 @@ def _ensure_cache_manager(self, owner_cls=None): self.cache_manager = JitCacheManager(cache_dir) self.cache_manager.load_all() - def _resolve_and_make_cache_key(self, bound_args): + def _resolve_and_make_cache_key(self, bound_args, *, effective_hints=None): """Resolve raw call values into JitArgument instances *in place* and build the tuple cache key from them. @@ -1298,8 +1302,16 @@ def _resolve_and_make_cache_key(self, bound_args): sig = self._sig # Re-read env vars on every call. key_parts = [("_env_", _cache_invalidating_env_values()), ("_target_", self._backend_target)] - if self.compile_hints: - key_parts.append(("_hints_", tuple(sorted((k, str(v)) for k, v in self.compile_hints.items())))) + if effective_hints is None: + effective_hints = self._effective_compile_hints() + if effective_hints: + hint_key = tuple( + sorted( + (key, type(value).__module__, type(value).__qualname__, repr(value)) + for key, value in effective_hints.items() + ) + ) + key_parts.append(("_hints_", hint_key)) for name, arg in bound_args.items(): param = sig.parameters.get(name) @@ -1349,9 +1361,11 @@ def _globals_key_prefix(self, owner_cls=None) -> tuple: cache[owner_cls] = (("_globals_", tuple(sorted(stable.items()))),) if stable else () return cache[owner_cls] - def _build_full_cache_key(self, bound_arguments, *, owner_cls=None, bound_self=None): + def _build_full_cache_key(self, bound_arguments, *, owner_cls=None, bound_self=None, effective_hints=None): """Build the complete cache key: arg signatures + stable globals snapshot + self type.""" - cache_key = self._globals_key_prefix(owner_cls) + self._resolve_and_make_cache_key(bound_arguments) + cache_key = self._globals_key_prefix(owner_cls) + self._resolve_and_make_cache_key( + bound_arguments, effective_hints=effective_hints + ) if bound_self is not None: cache_key = (("_self_type_", type(bound_self)),) + cache_key return cache_key @@ -1386,7 +1400,14 @@ def __call__(self, *args, **kwargs): bound = sig.bind(*args, **kwargs) bound.apply_defaults() - cache_key = self._build_full_cache_key(bound.arguments, owner_cls=owner_cls, bound_self=bound_self) + # Resolve once so cache identity and compilation use the same options. + effective_hints = self._effective_compile_hints() + cache_key = self._build_full_cache_key( + bound.arguments, + owner_cls=owner_cls, + bound_self=bound_self, + effective_hints=effective_hints, + ) args_tuple = tuple(bound.arguments.values()) @@ -1449,7 +1470,7 @@ def __call__(self, *args, **kwargs): ) raise RuntimeError(msg) - _hints_ctx = CompilationContext.compile_hints(self.compile_hints) if self.compile_hints else nullcontext() + _hints_ctx = CompilationContext.compile_hints(effective_hints) if effective_hints else nullcontext() compiled_func = None # will be set inside lock or compile path diff --git a/python/flydsl/compiler/kernel_function.py b/python/flydsl/compiler/kernel_function.py index ce96d77c2..acbaf4147 100644 --- a/python/flydsl/compiler/kernel_function.py +++ b/python/flydsl/compiler/kernel_function.py @@ -171,6 +171,15 @@ def _normalize_dim(dim: DimType) -> Tuple[DimValueType, DimValueType, DimValueTy # ============================================================================= +def merge_compile_hints(*layers) -> dict: + """Shallow-merge hint layers; later non-None values win.""" + merged = {} + for layer in layers: + if layer: + merged.update((key, value) for key, value in layer.items() if value is not None) + return merged + + class CompilationContext: """Context for tracking compilation state within a @jit function. @@ -188,14 +197,14 @@ class CompilationContext: @classmethod @contextmanager def compile_hints(cls, hints: dict): - """Context manager for setting compiler hints (thread-safe). + """Set thread-local hints, shallow-merging nested contexts. Usage: with CompilationContext.compile_hints({"waves_per_eu": 2}): fn(*args, **kwargs) """ prev = getattr(cls._compile_hints, "data", None) - cls._compile_hints.data = hints + cls._compile_hints.data = merge_compile_hints(prev, hints) try: yield finally: diff --git a/tests/README.md b/tests/README.md index 34ab18bb7..4db6a88a9 100644 --- a/tests/README.md +++ b/tests/README.md @@ -47,6 +47,8 @@ Use the same names as [`python/flydsl/utils/env.py`](../python/flydsl/utils/env. | Compile without execution | `COMPILE_ONLY` | | JIT cache directory | `FLYDSL_RUNTIME_CACHE_DIR` | | Enable/disable JIT disk cache | `FLYDSL_RUNTIME_ENABLE_CACHE` (`0` / `false` to disable; in-memory cache remains active) | +| Force autotuning search | `FLYDSL_AUTOTUNE` (`1` / `true` to ignore the heuristic/cached best) | +| Autotune result cache directory | `FLYDSL_AUTOTUNE_CACHE_DIR` | | IR dump | `FLYDSL_DUMP_IR`, `FLYDSL_DUMP_DIR` | | Device runtime kind | `FLYDSL_RUNTIME_KIND` | | ROCm arch hints (detection helpers) | `FLYDSL_GPU_ARCH`, `HSA_OVERRIDE_GFX_VERSION` | @@ -90,4 +92,3 @@ export FLYDSL_RUNTIME_ENABLE_CACHE=0 # or: rm -rf ~/.flydsl/cache ## MLIR FileCheck tests `tests/mlir/**/*.mlir` checks are driven by **`scripts/run_tests.sh`** (FileCheck + `fly-opt`), not by pytest. Tiering for those may be documented in parallel in this README as the RFC rollout continues; see RFC open questions. - diff --git a/tests/kernels/test_rmsnorm_autotune.py b/tests/kernels/test_rmsnorm_autotune.py new file mode 100644 index 000000000..6ca0a3414 --- /dev/null +++ b/tests/kernels/test_rmsnorm_autotune.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""GPU contracts for the direct RMSNorm autotune adopter.""" + +import re + +import pytest + +pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] + +try: + import torch +except ImportError: + torch = None +if torch is None or not torch.cuda.is_available(): + pytest.skip("CUDA/ROCm not available. Skipping GPU tests.", allow_module_level=True) + +import flydsl.compiler as flyc # noqa: E402 +from kernels.norm.rmsnorm_autotune import _SEARCH_CONFIGS, _rmsnorm_tuner, rmsnorm_autotuned # noqa: E402 +from kernels.norm.rmsnorm_kernel import rmsnorm_direct # noqa: E402 + +EPS = 1e-5 + + +@pytest.fixture(autouse=True) +def _isolated_tuner(tmp_path, monkeypatch): + _rmsnorm_tuner.cache.clear() + monkeypatch.setattr(_rmsnorm_tuner, "_cache_file", tmp_path / "rmsnorm.json") + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + yield + _rmsnorm_tuner.cache.clear() + + +def _reference(x, g): + xf = x.float() + return xf * torch.rsqrt((xf * xf).mean(-1, keepdim=True) + EPS) * g.float() + + +def _inputs(M=32, N=8192): + generator = torch.Generator(device="cuda").manual_seed(0) + x = torch.randn(M, N, device="cuda", dtype=torch.bfloat16, generator=generator) + g = torch.rand(N, device="cuda", dtype=torch.bfloat16, generator=generator) + return x, g, _reference(x, g) + + +def _assert_close(out, ref): + torch.testing.assert_close(out.float(), ref, rtol=0, atol=2e-2) + + +def test_rmsnorm_direct_specializes_known_block_size(): + x, g, ref = _inputs(M=1) + out = torch.empty_like(x) + stream = torch.cuda.current_stream() + + compiled = flyc.compile(rmsnorm_direct, x, g, out, x.shape[0], x.shape[1], "bf16", 512, stream) + stream.synchronize() + artifact = compiled._keepalive + + assert "known_block_size = array" in artifact.source_ir + match = re.search(r"max_flat_workgroup_size\s*=\s*(\d+)", artifact.ir) + assert match is not None and int(match.group(1)) == 512 + _assert_close(out, ref) + + +def test_rmsnorm_autotuned_default_uses_current_stream_and_skips_search(monkeypatch): + monkeypatch.setattr(_rmsnorm_tuner, "_bench_one", lambda *args, **kwargs: pytest.fail("unexpected search")) + x, g, ref = _inputs() + out = torch.empty_like(x) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + observed_streams = [] + original_run_config = _rmsnorm_tuner._run_config + + def checked_run_config(config, args, kwargs): + observed_streams.append(kwargs["stream"].cuda_stream) + return original_run_config(config, args, kwargs) + + monkeypatch.setattr(_rmsnorm_tuner, "_run_config", checked_run_config) + + with torch.cuda.stream(stream): + rmsnorm_autotuned(x, g, out, x.shape[0]) + stream.synchronize() + + assert observed_streams == [stream.cuda_stream] + _assert_close(out, ref) + + +def test_rmsnorm_autotuned_search_then_cache_hit(monkeypatch): + completed = 0 + x, g, ref = _inputs(M=8) + out = torch.empty_like(x) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + raw_stream = stream.cuda_stream + + def bench_once(call, warmup, rep): + nonlocal completed + assert torch.cuda.current_stream().cuda_stream == stream.cuda_stream + call() + stream.synchronize() + completed += 1 + return float(completed) + + monkeypatch.setattr(_rmsnorm_tuner, "_do_bench", bench_once) + monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") + rmsnorm_autotuned(x, g, out, x.shape[0], stream=raw_stream) + stream.synchronize() + + assert completed == len(_SEARCH_CONFIGS) + _assert_close(out, ref) + + call_kwargs = {"N": x.shape[1], "dtype_str": "bf16", "stream": raw_stream} + winner_key = _rmsnorm_tuner._make_key((x, g, out, x.shape[0]), call_kwargs) + assert winner_key in _rmsnorm_tuner.cache + other_m_key = _rmsnorm_tuner._make_key((x, g, out, x.shape[0] + 1), call_kwargs) + assert other_m_key != winner_key + + monkeypatch.delenv("FLYDSL_AUTOTUNE") + monkeypatch.setattr( + _rmsnorm_tuner, + "default", + lambda *args, **kwargs: pytest.fail("cached winner should take precedence over default"), + ) + cached = torch.empty_like(x) + rmsnorm_autotuned(x, g, cached, x.shape[0], stream=raw_stream) + stream.synchronize() + + assert completed == len(_SEARCH_CONFIGS) + _assert_close(cached, ref) diff --git a/tests/unit/test_autotune.py b/tests/unit/test_autotune.py index a4ae8459d..8e97bf605 100644 --- a/tests/unit/test_autotune.py +++ b/tests/unit/test_autotune.py @@ -206,6 +206,21 @@ def test_key_insensitive_to_kwarg_order(): assert k1 == k2 +def test_key_varies_with_effective_compile_hints(): + hints = {"waves_per_eu": 1} + + def fn(a, out, **kw): + pass + + fn._effective_compile_hints = lambda: dict(hints) + tuner = _make_tuner(fn=fn) + args = (FakeTensor((8, 8)), FakeTensor((8, 8))) + first = tuner._make_key(args, {}) + hints["waves_per_eu"] = 2 + + assert tuner._make_key(args, {}) != first + + # ── restore_value (in-place correctness) ──────────────────────────────── def test_restore_value_restores_between_reps(): """A kernel that mutates its input in place must see pristine inputs on @@ -343,14 +358,21 @@ def bench(call, warmup, rep): # ── decorator ──────────────────────────────────────────────────────────── def test_autotune_decorator_wraps_into_autotuner(): - """@autotune returns an Autotuner that forwards restore_value/reset_to_zero.""" + """@autotune returns an Autotuner and forwards its options.""" def fake_jit(a, out, **kw): pass + def configs(a, out): + return [Config(BLOCK=128)] + + def default(a, out): + return Config(BLOCK=64) + tuned = autotune( - configs=[Config(BLOCK=128)], + configs=configs, key=["a"], + default=default, restore_value=["a"], reset_to_zero=["out"], )(fake_jit) @@ -358,7 +380,77 @@ def fake_jit(a, out, **kw): assert isinstance(tuned, Autotuner) assert tuned.restore_value == ["a"] assert tuned.reset_to_zero == ["out"] - assert [c.kwargs["BLOCK"] for c in tuned.configs] == [128] + assert tuned.configs is configs + assert tuned.default is default + + +# ── two-track default/search ───────────────────────────────────────────── +def test_cache_hit_precedes_default_and_search(monkeypatch): + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) + default_calls = 0 + + def fn(a, out, BLOCK): + out._data[0] = float(BLOCK) + + def default(a, out): + nonlocal default_calls + default_calls += 1 + return Config(BLOCK=999) + + def fail_bench(call, warmup, rep): + pytest.fail("normal path benchmarked configs") + + tuner = _make_tuner( + fn=fn, + configs=[Config(BLOCK=64), Config(BLOCK=128)], + default=default, + do_bench_fn=fail_bench, + ) + a = FakeTensor((8,)) + out = FakeTensor((1,)) + args = (a, out) + tuner.cache[tuner._make_key(args, {})] = Config(BLOCK=128) + + tuner(*args) + + assert out._data[0] == 128.0 + assert default_calls == 0 + + tuner.cache.clear() + tuner(*args) + + assert out._data[0] == 999.0 + assert default_calls == 1 + + +def test_force_search_bypasses_cache_and_default(monkeypatch): + monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") + calls = {"configs": 0, "default": 0, "bench": 0} + + def fn(a, out, BLOCK): + out._data[0] = float(BLOCK) + + def configs(a, out): + calls["configs"] += 1 + return [Config(BLOCK=64), Config(BLOCK=128)] + + def default(a, out): + calls["default"] += 1 + return Config(BLOCK=7) + + def bench(call, warmup, rep): + calls["bench"] += 1 + call() + return float(calls["bench"]) + + tuner = _make_tuner(fn=fn, configs=configs, default=default, do_bench_fn=bench) + args = (FakeTensor((8,)), FakeTensor((1,))) + tuner.cache[tuner._make_key(args, {})] = Config(BLOCK=999) + + tuner(*args) + + assert calls == {"configs": 1, "default": 0, "bench": 2} + assert args[1]._data[0] == 64.0 if __name__ == "__main__": diff --git a/tests/unit/test_external_llvm_codegen.py b/tests/unit/test_external_llvm_codegen.py index b430854c0..f1460885f 100644 --- a/tests/unit/test_external_llvm_codegen.py +++ b/tests/unit/test_external_llvm_codegen.py @@ -4,13 +4,20 @@ import json from pathlib import Path +import pytest + from flydsl._mlir import ir +from flydsl._mlir._mlir_libs._mlirDialectsLLVM import translate_module_to_llvmir +from flydsl._mlir.passmanager import PassManager from flydsl.compiler.backends.rocm import RocmBackend from flydsl.compiler.external_llvm import ( _format_llvm_cli_options, external_llvm_fingerprint, run_external_binary_codegen, ) +from flydsl.compiler.jit_function import _create_mlir_context + +pytestmark = [pytest.mark.l1b_target_dialect, pytest.mark.rocm_lower] def _write_executable(path: Path, text: str) -> None: @@ -68,6 +75,81 @@ def test_rocm_external_pipeline_split_matches_full_pipeline(): assert "--amdgpu-num-vgpr=128" in binary +def test_rocm_lower_wpe_preserves_source_default_and_overrides_kernel_entries(): + backend = RocmBackend(RocmBackend.make_target("gfx942")) + src = r"""module { + gpu.module @m { + gpu.func @a() kernel attributes { + rocdl.waves_per_eu = 3 : i32, + passthrough = [["keep", "yes"]] + } { gpu.return } + gpu.func @b() kernel attributes { + passthrough = [["amdgpu-waves-per-eu", "1,1"]] + } { gpu.return } + gpu.func @helper() { gpu.return } + } + }""" + + with ir.Context() as ctx, ir.Location.unknown(ctx): + ctx.load_all_available_dialects() + baseline = ir.Module.parse(src) + baseline_asm = str(baseline) + backend.lower_compile_hints(baseline, compile_hints={"waves_per_eu": 0}) + assert str(baseline) == baseline_asm + + module = ir.Module.parse(src) + backend.lower_compile_hints(module, compile_hints={"waves_per_eu": 2}) + funcs = { + ir.StringAttr(op.attributes["sym_name"]).value: str(op) + for op in module.body.operations[0].regions[0].blocks[0].operations + if op.operation.name == "gpu.func" + } + + for name in ("a", "b"): + assert funcs[name].count("amdgpu-waves-per-eu") == 1 + assert '"amdgpu-waves-per-eu", "2,2"' in funcs[name] + assert "rocdl.waves_per_eu" not in funcs[name] + assert '"keep", "yes"' in funcs["a"] + assert "amdgpu-waves-per-eu" not in funcs["helper"] + + +@pytest.mark.parametrize( + ("value", "error"), + [ + (True, TypeError), + ("2", TypeError), + (-1, ValueError), + ], +) +def test_rocm_lower_wpe_rejects_invalid_values(value, error): + backend = RocmBackend(RocmBackend.make_target("gfx942")) + + with ir.Context() as ctx, ir.Location.unknown(ctx): + module = ir.Module.create() + with pytest.raises(error, match="waves_per_eu"): + backend.lower_compile_hints(module, compile_hints={"waves_per_eu": value}) + + +def test_rocm_wpe_reaches_native_llvm_as_exact_constraint(): + backend = RocmBackend(RocmBackend.make_target("gfx942")) + with _create_mlir_context() as ctx: + module = ir.Module.parse( + """module attributes {gpu.container_module} { + gpu.module @m { + gpu.func @k() kernel attributes {rocdl.waves_per_eu = 1 : i32} { gpu.return } + } + }""", + context=ctx, + ) + backend.lower_compile_hints(module, compile_hints={"waves_per_eu": 2}) + pre_binary, _ = backend.external_binary_pipeline_fragments(compile_hints={}) + PassManager.parse(f"builtin.module({','.join(pre_binary)})", ctx).run(module.operation) + llvm_ir = translate_module_to_llvmir(module.body.operations[0].operation) + + assert '"amdgpu-waves-per-eu"="2,2"' in llvm_ir + assert '"amdgpu-waves-per-eu"="1"' not in llvm_ir + + def test_external_llvm_fingerprint_uses_configured_tools(tmp_path, monkeypatch): llvm_dir = _make_fake_llvm(tmp_path) monkeypatch.setenv("FLYDSL_COMPILE_LLVM_DIR", str(llvm_dir)) diff --git a/tests/unit/test_jit_cache_key.py b/tests/unit/test_jit_cache_key.py index 1b29349f6..377dafaf6 100644 --- a/tests/unit/test_jit_cache_key.py +++ b/tests/unit/test_jit_cache_key.py @@ -6,6 +6,7 @@ import flydsl.compiler as flyc import flydsl.expr as fx from flydsl.compiler.jit_argument import JitArgumentRegistry +from flydsl.compiler.kernel_function import CompilationContext class _FakeCudaStream: @@ -62,3 +63,25 @@ def test_future_annotations_runtime_int32_ignores_value_in_cache_key(): assert key1 == key2 assert ("n", (fx.Int32,)) in key1 assert ("n", (int, 1)) not in key1 + + +def test_thread_local_wpe_overrides_persistent_hint_and_enters_cache_key(): + @flyc.jit + def launch(stream: fx.Stream = fx.Stream(None)): + pass + + launch.compile_hints = {"waves_per_eu": 4} + persistent = _cache_key(launch) + with CompilationContext.compile_hints({"waves_per_eu": 1}): + outer = _cache_key(launch) + assert launch._effective_compile_hints()["waves_per_eu"] == 1 + with CompilationContext.compile_hints({"waves_per_eu": 2}): + inner = _cache_key(launch) + assert launch._effective_compile_hints()["waves_per_eu"] == 2 + assert launch._effective_compile_hints()["waves_per_eu"] == 1 + assert launch._effective_compile_hints()["waves_per_eu"] == 4 + with CompilationContext.compile_hints({"waves_per_eu": "2"}): + invalid_type = _cache_key(launch) + + assert len({persistent, outer, inner}) == 3 + assert invalid_type != inner