diff --git a/docs/autotune_guide.md b/docs/autotune_guide.md new file mode 100644 index 000000000..b962629c5 --- /dev/null +++ b/docs/autotune_guide.md @@ -0,0 +1,64 @@ +# Offline autotune configs + +FlyDSL can serve a previously tuned config without benchmarking. This extends +the direct-JIT autotuner with one opt-in argument: + +```python +@flyc.jit +def launch(x, out, N: fx.Constexpr[int], BLOCK: fx.Constexpr[int]): + ... + + +tuned = autotune( + configs=[Config(BLOCK=128), Config(BLOCK=256)], + key=["N"], + default=lambda x, out, N: Config(BLOCK=256), + artifact_name="my_kernel", +)(launch) +``` + +`artifact_name` enables lookup when `FLYDSL_AUTOTUNE_CONFIG_DIR` is set. Normal +calls use the first available source: + +```text +searched winner cache -> matching artifact -> default -> search +``` + +`FLYDSL_AUTOTUNE=1` bypasses those serving decisions, searches the existing +configs, updates the scratch winner cache, and atomically writes an artifact. +A normal fallback search updates only the scratch cache. +While artifact lookup is active, scratch winners use the same device descriptor +so a same-architecture product cannot shadow the matching artifact. + +## Identity and compatibility + +Artifact identity is the stable `artifact_name`, the declared `key` values, +and the call device's product name, target architecture, and compute-unit +count. Use a globally unique name for each kernel/config schema. The JSON is +self-describing; its filename is an identity digest. + +The declared `key` is the single owner of portable tuning axes. Include every +shape, dtype, layout, or mode that can change the winner. Keep structural knobs +as JIT `Constexpr` parameters on the existing entry point; offline tuning does +not need a build factory or a second key callback. + +Artifacts intentionally do not include a compiler or kernel-source fingerprint. +Treat them as reviewed deployment inputs and retune after a compiler, kernel, +compile-hint, or search-space change that can affect the winner. + +## Failure behavior + +Missing, unreadable, mismatched, or structurally invalid artifacts are ignored, +and normal lookup continues to the default or search path. Artifact config +values cannot overwrite arguments supplied by the caller or declared key axes. +Values must preserve their types when encoded as JSON. `Config.pre_hook` is +process-local code, so it blocks forced artifact generation. + +Once a matching artifact has been accepted, its compile, launch, and runtime +errors propagate normally; FlyDSL does not hide them by retrying the default. +If forced generation cannot establish a device identity or write its artifact, +it fails clearly without caching the winner. + +Generate deployment artifacts on the intended GPU under controlled benchmark +conditions. CI should verify deterministic emit-and-load behavior, not commit a +winner selected from noisy shared-runner timing. diff --git a/docs/index.rst b/docs/index.rst index 08cdd06d0..7cc082474 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -31,6 +31,7 @@ to GPU/ROCDL. layout_system_guide kernel_authoring_guide kernel_tuning_guide + autotune_guide prebuilt_kernels_guide testing_benchmarking_guide cute_layout_algebra_guide diff --git a/kernels/norm/rmsnorm_autotune.py b/kernels/norm/rmsnorm_autotune.py index 7408b9dbc..7901721ca 100644 --- a/kernels/norm/rmsnorm_autotune.py +++ b/kernels/norm/rmsnorm_autotune.py @@ -32,6 +32,7 @@ def _search_configs(input_t, gamma, output, m_in, N, dtype_str="bf16", stream=No configs=_search_configs, key=["m_in", "N", "dtype_str"], default=_default_config, + artifact_name="rmsnorm", )(rmsnorm_direct) diff --git a/python/flydsl/autotune.py b/python/flydsl/autotune.py index 3b4cad77b..bd5853f88 100644 --- a/python/flydsl/autotune.py +++ b/python/flydsl/autotune.py @@ -3,6 +3,7 @@ """FlyDSL autotuner - benchmark multiple kernel configs, pick the fastest.""" +import hashlib import inspect import json import os @@ -10,12 +11,18 @@ from pathlib import Path from typing import Callable, Dict, List +from .utils import env, log +from .utils.file import atomic_write + try: import torch except ImportError: torch = None +_ARTIFACT_VERSION = 1 + + 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") @@ -57,6 +64,38 @@ def _device_fingerprint() -> str: return "" +def _device_descriptor(device=None): + """Portable identity for the device that will run the call.""" + if torch is None: + return None + try: + if not torch.cuda.is_available(): + return None + device = torch.cuda.current_device() if device is None else device + properties = torch.cuda.get_device_properties(device) + name = str(properties.name) + compute_units = getattr(properties, "multi_processor_count", None) + with torch.cuda.device(device): + arch = _device_fingerprint() + except Exception: + return None + if not name or not arch or type(compute_units) is not int or compute_units <= 0: + return None + return {"name": name, "arch": arch, "compute_units": compute_units} + + +def _artifacts_enabled() -> bool: + return bool(env.autotune.config_dir) + + +def _artifact_config_dir() -> Path: + return Path(env.autotune.config_dir).expanduser().resolve() + + +def _canonical_json(value) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + + def _normalize_strides(t) -> tuple: """Bucket strides to {0, 1, other}: the layout *pattern* (broadcast / contiguous / strided) affects the best config, the exact numbers don't.""" @@ -172,6 +211,7 @@ def __init__( post_hook=None, do_bench_fn=None, default=None, + artifact_name=None, ): self.fn = fn # JitFunction instance self.configs = configs @@ -187,11 +227,22 @@ def __init__( self.cache: Dict[tuple, Config] = {} self.default = default - # Infer arg names from the underlying function - if hasattr(fn, "func"): - self.arg_names = list(inspect.signature(fn.func).parameters.keys()) - else: - self.arg_names = list(inspect.signature(fn).parameters.keys()) + # Infer arg names from the underlying function. + source = fn.func if hasattr(fn, "func") else fn + self._signature = inspect.signature(source) + self.arg_names = list(self._signature.parameters.keys()) + + if artifact_name is not None: + if not isinstance(artifact_name, str): + raise TypeError("artifact_name must be a string") + artifact_name = artifact_name.strip() + safe_name = artifact_name.replace("-", "").replace("_", "").replace(".", "") + if not artifact_name.isascii() or not safe_name.isalnum() or len(artifact_name) > 64: + raise ValueError("artifact_name must use 1-64 letters, digits, '.', '_' or '-'") + if len(set(self.key)) != len(self.key) or any(name not in self._signature.parameters for name in self.key): + raise ValueError("artifact keys must be unique kernel parameter names") + self.artifact_name = artifact_name + self._artifact_cache = {} # Disk cache fn_name = getattr(fn, "__name__", None) or getattr(fn, "func", None) @@ -239,6 +290,10 @@ def _make_key(self, args, kwargs): key_vals.append(("_env_", _env_fingerprint())) key_vals.append(("_toolchain_", _toolchain_fingerprint())) key_vals.append(("_device_", _device_fingerprint())) + if self.artifact_name is not None and _artifacts_enabled(): + device = _device_descriptor(self._call_device(args, kwargs)) + descriptor = tuple(sorted(device.items())) if device is not None else None + key_vals.append(("_artifact_", _ARTIFACT_VERSION, self.artifact_name, descriptor)) effective_hints = getattr(self.fn, "_effective_compile_hints", None) if callable(effective_hints): hint_key = tuple( @@ -364,6 +419,149 @@ def _run_config(self, config, args, kwargs): self._reset_tensors(args, merged) return self._run_with_hints(config.compiler_opts(), args, merged) + def _call_device(self, args, kwargs): + if torch is None: + return None + sig_args = dict(zip(self.arg_names, args)) + sig_args.update(kwargs) + return next( + (value.device for value in sig_args.values() if isinstance(value, torch.Tensor) and value.is_cuda), + None, + ) + + def _artifact_ref(self, args, kwargs, *, required): + if self.artifact_name is None or not _artifacts_enabled(): + return None + try: + bound = self._signature.bind_partial(*args, **kwargs) + bound.apply_defaults() + key = {} + for name in self.key: + if name not in bound.arguments: + raise ValueError(f"artifact key {name!r} is missing from the call") + value = bound.arguments[name] + if hasattr(value, "shape"): + value = tuple(value.shape) + elif hasattr(value, "dtype"): + value = str(value.dtype) + key[name] = value + device = _device_descriptor(self._call_device(args, kwargs)) + if device is None: + raise ValueError("call device identity is unavailable") + identity = {"name": self.artifact_name, "key": key, "device": device} + digest = hashlib.sha256(_canonical_json(identity).encode()).hexdigest() + return _artifact_config_dir() / f"{self.artifact_name}-{digest}.json", identity + except (OSError, RuntimeError, TypeError, ValueError, OverflowError, RecursionError) as error: + if required: + raise ValueError(f"cannot generate offline config identity: {error}") from error + log().warning(f"Offline config identity is unavailable: {error}") + return None + + def _validate_artifact_config_types(self, injected): + for name, value in injected.items(): + parameter = self._signature.parameters.get(name) + if parameter is None or parameter.annotation is inspect.Parameter.empty: + continue + annotation = parameter.annotation + coerce = getattr(annotation, "__coerce__", None) + try: + if callable(coerce): + coerce(value) + elif annotation in (bool, int, float, str) and type(value) is not annotation: + raise TypeError(f"expects {annotation.__name__}, got {type(value).__name__}") + except (TypeError, ValueError) as error: + raise ValueError(f"config value for {name!r} does not match its annotation: {error}") from error + + def _decode_artifact_config(self, body, args, kwargs): + if not isinstance(body, dict): + raise ValueError("config must be an object") + if "pre_hook" in body: + raise ValueError("offline configs cannot contain Config.pre_hook") + config = Config.from_dict(body) + try: + call_bound = self._signature.bind_partial(*args, **kwargs) + except TypeError as error: + raise ValueError(f"call does not match the kernel signature: {error}") from error + + call_names = set(call_bound.arguments) + for name, parameter in self._signature.parameters.items(): + if parameter.kind is inspect.Parameter.VAR_KEYWORD and name in call_bound.arguments: + call_names.update(call_bound.arguments[name]) + + injected = config.all_kwargs() + overlap = set(injected).intersection(call_names) | set(body).intersection(self.key) + if overlap: + raise ValueError(f"config would override call arguments: {sorted(overlap)}") + if any(type(value) is not int for value in config.compiler_opts().values()): + raise ValueError("compiler options must be integers") + self._validate_artifact_config_types(injected) + bound_kwargs = dict(kwargs) + bound_kwargs.update(injected) + try: + self._signature.bind(*args, **bound_kwargs) + except TypeError as error: + raise ValueError(f"config does not complete the kernel signature: {error}") from error + return config + + def _load_artifact(self, ref, args, kwargs): + if ref is None: + return None + path, identity = ref + cache_key = str(path) + if cache_key not in self._artifact_cache: + try: + if not path.is_file(): + self._artifact_cache[cache_key] = None + return None + data = json.loads(path.read_text(encoding="utf-8")) + _canonical_json(data) + if not isinstance(data, dict): + raise ValueError("artifact must be an object") + if type(data.get("version")) is not int or data["version"] != _ARTIFACT_VERSION: + raise ValueError("unsupported artifact version") + if _canonical_json(data.get("identity")) != _canonical_json(identity): + raise ValueError("artifact identity does not match") + body = data["config"] + if not isinstance(body, dict): + raise ValueError("config must be an object") + self._artifact_cache[cache_key] = body + except ( + KeyError, + OSError, + json.JSONDecodeError, + TypeError, + ValueError, + OverflowError, + RecursionError, + ) as error: + log().warning(f"Ignoring offline config {path.name}: {error}") + self._artifact_cache[cache_key] = None + return None + + body = self._artifact_cache[cache_key] + if body is None: + return None + try: + return self._decode_artifact_config(body, args, kwargs) + except (TypeError, ValueError, OverflowError, RecursionError) as error: + log().warning(f"Ignoring offline config {path.name}: {error}") + return None + + def _emit_artifact(self, config, ref, args, kwargs): + if config.pre_hook is not None: + raise ValueError("offline configs cannot serialize Config.pre_hook") + path, identity = ref + body = config.to_dict() + if json.loads(_canonical_json(body)) != body: + raise ValueError("offline config values must preserve their types in JSON") + self._decode_artifact_config(body, args, kwargs) + payload = {"version": _ARTIFACT_VERSION, "identity": identity, "config": body} + with atomic_write(path, mode="w", encoding="utf-8") as output: + json.dump(payload, output, indent=2, sort_keys=True, allow_nan=False) + output.write("\n") + self._artifact_cache[str(path)] = body + log().info(f"Wrote offline config {path}") + def __call__(self, *args, **kwargs): key = self._make_key(args, kwargs) force = _tuning_enabled() @@ -371,6 +569,12 @@ def __call__(self, *args, **kwargs): if not force and key in self.cache: return self._run_config(self.cache[key], args, kwargs) + artifact = self._artifact_ref(args, kwargs, required=force) + if not force: + artifact_config = self._load_artifact(artifact, args, kwargs) + if artifact_config is not None: + return self._run_config(artifact_config, args, kwargs) + if not force and self.default is not None: return self._run_config(self.default(*args, **kwargs), args, kwargs) @@ -392,6 +596,8 @@ def __call__(self, *args, **kwargs): best_config, best_time = min(results, key=lambda x: x[1]) print(f"[autotune] best: {best_config} ({best_time:.3f} ms)") + if force and artifact is not None: + self._emit_artifact(best_config, artifact, args, kwargs) self.cache[key] = best_config self._save_disk_cache() @@ -428,6 +634,7 @@ def autotune( post_hook: Callable = None, do_bench: Callable = None, default: Callable = None, + artifact_name: str = None, ): """Autotune decorator for @jit functions. @@ -442,6 +649,9 @@ def myKernel(..., BLOCK: fx.Constexpr[int], ...): the current arguments. default: optional heuristic ``default(*args, **kwargs) -> Config`` used without benchmarking unless ``FLYDSL_AUTOTUNE`` forces a search. + artifact_name: stable name for opt-in config lookup and emission through + ``FLYDSL_AUTOTUNE_CONFIG_DIR``. The existing ``key`` defines the + portable call axes; forced tuning emits a matching artifact. 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 @@ -464,6 +674,7 @@ def decorator(fn): post_hook=post_hook, do_bench_fn=do_bench, default=default, + artifact_name=artifact_name, ) return decorator diff --git a/python/flydsl/compiler/jit_function.py b/python/flydsl/compiler/jit_function.py index a213d9780..9274e4657 100644 --- a/python/flydsl/compiler/jit_function.py +++ b/python/flydsl/compiler/jit_function.py @@ -8,7 +8,6 @@ import os import pickle import pkgutil -import tempfile import threading import time import types @@ -26,6 +25,7 @@ from ..expr.typing import Constexpr, Stream from ..expr.utils.arith import fastmath as fastmath_ctx from ..utils import env, log +from ..utils.file import atomic_write from .ast_rewriter import ASTRewriter from .backends import compile_backend_name, get_backend from .diagnostics import ( @@ -950,7 +950,7 @@ class JitCacheManager: {cache_key}.lock - per-key advisory lock file All disk reads use shared (reader) locks; writes use exclusive locks - with atomic ``tempfile`` + ``os.rename`` to prevent partial reads. + and same-directory atomic replacement to prevent partial reads. """ def __init__(self, cache_dir: Path): @@ -970,20 +970,9 @@ def _lock_file(self, cache_key: str) -> Path: return self.cache_dir / f"{self._safe_key(cache_key)}.lock" @staticmethod - def _atomic_write(cache_file: Path, value: Any) -> None: - """Write *value* atomically via tempfile + rename.""" - cache_file.parent.mkdir(parents=True, exist_ok=True) - fd, tmp = tempfile.mkstemp(dir=str(cache_file.parent), suffix=".tmp") - try: - with os.fdopen(fd, "wb") as f: - pickle.dump(value, f) - os.rename(tmp, str(cache_file)) - except BaseException: - try: - os.unlink(tmp) - except OSError: - pass - raise + def _write_cache_file(cache_file: Path, value: Any) -> None: + with atomic_write(cache_file, mode="wb") as output: + pickle.dump(value, output) def get(self, cache_key: str) -> Optional[Any]: if cache_key in self.memory_cache: @@ -1019,7 +1008,7 @@ def set(self, cache_key: str, value: Any) -> None: if cache_file.exists(): log().debug(f"Cache already exists, skipping write: {cache_file.name}") return - self._atomic_write(cache_file, value) + self._write_cache_file(cache_file, value) log().debug(f"Cache saved: {cache_file.name}") except Exception as e: log().warning(f"Failed to save cache {cache_file}: {e}") @@ -1055,7 +1044,7 @@ def compile_lock(self, cache_key: str): # Cache miss — provide a writer that writes under the already-held lock. def _writer(value): - self._atomic_write(cache_file, value) + self._write_cache_file(cache_file, value) self.memory_cache[cache_key] = value yield (None, _writer) diff --git a/python/flydsl/utils/env.py b/python/flydsl/utils/env.py index 3bb95781f..1839dd7df 100644 --- a/python/flydsl/utils/env.py +++ b/python/flydsl/utils/env.py @@ -216,6 +216,14 @@ def __repr__(self) -> str: return self.help() +class AutotuneEnvManager(EnvManager): + """Autotuning options (``FLYDSL_AUTOTUNE_*`` environment variables).""" + + env_prefix = "AUTOTUNE" + + config_dir = OptStr("", description="Directory for offline config artifacts; empty disables artifacts") + + class CompileEnvManager(EnvManager): """Compile-time options (``FLYDSL_COMPILE_*`` environment variables).""" @@ -292,11 +300,13 @@ class RuntimeEnvManager(EnvManager): ) +autotune = AutotuneEnvManager() compile = CompileEnvManager() debug = DebugEnvManager() runtime = RuntimeEnvManager() __all__ = [ + "autotune", "compile", "debug", "runtime", diff --git a/python/flydsl/utils/file.py b/python/flydsl/utils/file.py new file mode 100644 index 000000000..29dfebaca --- /dev/null +++ b/python/flydsl/utils/file.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +import os +import tempfile +from contextlib import contextmanager +from pathlib import Path +from typing import IO, Iterator, Optional + + +@contextmanager +def atomic_write( + path: Path, + *, + mode: str, + encoding: Optional[str] = None, +) -> Iterator[IO]: + """Write a file through a same-directory temporary and atomically replace it.""" + path.parent.mkdir(parents=True, exist_ok=True) + tmp_path = None + try: + with tempfile.NamedTemporaryFile( + mode=mode, + encoding=encoding, + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as output: + tmp_path = Path(output.name) + yield output + os.replace(tmp_path, path) + finally: + if tmp_path is not None: + tmp_path.unlink(missing_ok=True) diff --git a/tests/kernels/test_rmsnorm_autotune.py b/tests/kernels/test_rmsnorm_autotune.py index 6ca0a3414..65ff1ef02 100644 --- a/tests/kernels/test_rmsnorm_autotune.py +++ b/tests/kernels/test_rmsnorm_autotune.py @@ -3,7 +3,10 @@ """GPU contracts for the direct RMSNorm autotune adopter.""" +import json +import os import re +from pathlib import Path import pytest @@ -26,10 +29,13 @@ @pytest.fixture(autouse=True) def _isolated_tuner(tmp_path, monkeypatch): _rmsnorm_tuner.cache.clear() + _rmsnorm_tuner._artifact_cache.clear() monkeypatch.setattr(_rmsnorm_tuner, "_cache_file", tmp_path / "rmsnorm.json") + monkeypatch.setenv("FLYDSL_AUTOTUNE_CONFIG_DIR", str(tmp_path / "artifacts")) monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) yield _rmsnorm_tuner.cache.clear() + _rmsnorm_tuner._artifact_cache.clear() def _reference(x, g): @@ -88,6 +94,11 @@ def checked_run_config(config, args, kwargs): def test_rmsnorm_autotuned_search_then_cache_hit(monkeypatch): completed = 0 + target = next( + index + for index, config in enumerate(_SEARCH_CONFIGS, start=1) + if config.kwargs["BLOCK_THREADS"] == 128 and config.waves_per_eu == 2 + ) x, g, ref = _inputs(M=8) out = torch.empty_like(x) stream = torch.cuda.Stream() @@ -100,7 +111,7 @@ def bench_once(call, warmup, rep): call() stream.synchronize() completed += 1 - return float(completed) + return 0.0 if completed == target else float(completed) monkeypatch.setattr(_rmsnorm_tuner, "_do_bench", bench_once) monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") @@ -110,6 +121,14 @@ def bench_once(call, warmup, rep): assert completed == len(_SEARCH_CONFIGS) _assert_close(out, ref) + artifacts = list(Path(os.environ["FLYDSL_AUTOTUNE_CONFIG_DIR"]).glob("*.json")) + assert len(artifacts) == 1 + payload = json.loads(artifacts[0].read_text()) + assert payload["identity"]["key"] == {"m_in": x.shape[0], "N": x.shape[1], "dtype_str": "bf16"} + assert payload["identity"]["device"]["name"] == torch.cuda.get_device_name(x.device) + assert payload["config"]["BLOCK_THREADS"] == 128 + assert payload["config"]["waves_per_eu"] == 2 + 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 @@ -128,3 +147,24 @@ def bench_once(call, warmup, rep): assert completed == len(_SEARCH_CONFIGS) _assert_close(cached, ref) + + # A fresh serving decision with no winner cache must load the emitted + # artifact without evaluating the default or search space. + _rmsnorm_tuner.cache.clear() + _rmsnorm_tuner._artifact_cache.clear() + monkeypatch.setattr( + _rmsnorm_tuner, + "default", + lambda *args, **kwargs: pytest.fail("offline artifact should take precedence over default"), + ) + monkeypatch.setattr( + _rmsnorm_tuner, + "configs", + lambda *args, **kwargs: pytest.fail("offline hit should not construct search configs"), + ) + offline = torch.empty_like(x) + rmsnorm_autotuned(x, g, offline, x.shape[0], stream=raw_stream) + stream.synchronize() + + assert completed == len(_SEARCH_CONFIGS) + _assert_close(offline, ref) diff --git a/tests/unit/test_atomic_write.py b/tests/unit/test_atomic_write.py new file mode 100644 index 000000000..f105c557b --- /dev/null +++ b/tests/unit/test_atomic_write.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +import pytest + +from flydsl.utils.file import atomic_write + + +def test_atomic_write_replaces_destination_only_after_success(tmp_path): + destination = tmp_path / "artifact.json" + destination.write_text("old", encoding="utf-8") + + with atomic_write(destination, mode="w", encoding="utf-8") as output: + output.write("new") + assert destination.read_text(encoding="utf-8") == "old" + + assert destination.read_text(encoding="utf-8") == "new" + assert list(tmp_path.iterdir()) == [destination] + + +def test_atomic_write_preserves_destination_and_removes_temp_file_on_failure(tmp_path): + destination = tmp_path / "artifact.json" + destination.write_text("old", encoding="utf-8") + + with pytest.raises(RuntimeError, match="serialize failed"): + with atomic_write(destination, mode="w", encoding="utf-8") as output: + output.write("partial") + raise RuntimeError("serialize failed") + + assert destination.read_text(encoding="utf-8") == "old" + assert list(tmp_path.iterdir()) == [destination] + + +def test_atomic_write_supports_binary_output(tmp_path): + destination = tmp_path / "cache.pkl" + + with atomic_write(destination, mode="wb") as output: + output.write(b"\x00cache") + + assert destination.read_bytes() == b"\x00cache" diff --git a/tests/unit/test_autotune.py b/tests/unit/test_autotune.py index 8e97bf605..8c5fe079b 100644 --- a/tests/unit/test_autotune.py +++ b/tests/unit/test_autotune.py @@ -27,6 +27,8 @@ def _isolate_disk_cache(tmp_path, monkeypatch): """Every test gets a private autotune disk cache so results don't leak across tests (a cached best-config would skip the benchmark loop).""" monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "autotune_cache")) + monkeypatch.delenv("FLYDSL_AUTOTUNE_CONFIG_DIR", raising=False) + monkeypatch.delenv("FLYDSL_AUTOTUNE", raising=False) # ── Fakes ──────────────────────────────────────────────────────────────── @@ -373,6 +375,7 @@ def default(a, out): configs=configs, key=["a"], default=default, + artifact_name="fake-kernel", restore_value=["a"], reset_to_zero=["out"], )(fake_jit) @@ -382,6 +385,7 @@ def default(a, out): assert tuned.reset_to_zero == ["out"] assert tuned.configs is configs assert tuned.default is default + assert tuned.artifact_name == "fake-kernel" # ── two-track default/search ───────────────────────────────────────────── @@ -453,5 +457,258 @@ def bench(call, warmup, rep): assert args[1]._data[0] == 64.0 +# ── offline config artifacts ──────────────────────────────────────────── +class FakeConstexprInt: + @classmethod + def __coerce__(cls, value): + if type(value) is not int: + raise TypeError(f"expects int, got {type(value).__name__}") + return value + + +def _artifact_kernel(a, out, m_in, N, dtype_str, BLOCK_THREADS: FakeConstexprInt, stream: int = 0): + out._data[0] = float(BLOCK_THREADS) + + +def _artifact_default(a, out, m_in, N, dtype_str): + return Config(BLOCK_THREADS=7) + + +def _make_artifact_tuner(**overrides): + options = { + "fn": _artifact_kernel, + "configs": [Config(BLOCK_THREADS=64)], + "key": ["m_in", "N", "dtype_str"], + "default": _artifact_default, + "artifact_name": "rmsnorm", + "do_bench_fn": lambda call, warmup, rep: (call(), 1.0)[1], + } + options.update(overrides) + return _make_tuner(**options) + + +@pytest.fixture +def artifact_dir(tmp_path, monkeypatch): + import importlib + + at = importlib.import_module("flydsl.autotune") + + path = tmp_path / "artifacts" + monkeypatch.setenv("FLYDSL_AUTOTUNE_CONFIG_DIR", str(path)) + monkeypatch.setattr( + at, + "_device_descriptor", + lambda device=None: { + "name": "Test GPU", + "arch": "gfx-test", + "compute_units": 1, + }, + raising=False, + ) + return path + + +def _emit_artifact(monkeypatch, artifact_dir, *, args=None, config=None): + args = args or (FakeTensor((16, 512)), FakeTensor((1,)), 16, 512, "bf16") + monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") + tuner = _make_artifact_tuner(configs=[config or Config(BLOCK_THREADS=64)]) + tuner(*args) + files = list(artifact_dir.glob("*.json")) + assert len(files) == 1 + return files[0], args + + +def test_artifact_lifecycle(monkeypatch, tmp_path, artifact_dir): + path, args = _emit_artifact(monkeypatch, artifact_dir) + payload = json.loads(path.read_text()) + + assert payload["identity"]["key"] == {"m_in": 16, "N": 512, "dtype_str": "bf16"} + assert payload["config"] == {"BLOCK_THREADS": 64} + + monkeypatch.delenv("FLYDSL_AUTOTUNE") + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "fresh-cache")) + + def fail_configs(*_args, **_kwargs): + pytest.fail("artifact lookup evaluated the search space") + + def fail_default(*_args, **_kwargs): + pytest.fail("artifact lookup evaluated the heuristic default") + + loaded = _make_artifact_tuner( + configs=fail_configs, + default=fail_default, + do_bench_fn=lambda *args, **kwargs: pytest.fail("artifact lookup benchmarked"), + ) + out = FakeTensor((1,)) + loaded(args[0], out, *args[2:]) + + assert out._data[0] == 64.0 + assert loaded.cache == {}, "artifact decisions must not enter the searched-winner cache" + + monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") + replacement = _make_artifact_tuner(configs=[Config(BLOCK_THREADS=128)]) + replacement(*args) + assert json.loads(path.read_text())["config"] == {"BLOCK_THREADS": 128} + + +def test_artifact_schema_partitions_searched_winner_cache(monkeypatch, artifact_dir): + args = (FakeTensor((16, 512)), FakeTensor((1,)), 16, 512, "bf16") + monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") + old_schema = _make_artifact_tuner( + artifact_name="rmsnorm-v1", + configs=[Config(BLOCK_THREADS=64)], + ) + old_schema(*args) + + new_schema = _make_artifact_tuner(artifact_name="rmsnorm-v2") + ref = new_schema._artifact_ref(args, {}, required=True) + new_schema._emit_artifact(Config(BLOCK_THREADS=128), ref, args, {}) + + monkeypatch.delenv("FLYDSL_AUTOTUNE") + loaded = _make_artifact_tuner(artifact_name="rmsnorm-v2") + out = FakeTensor((1,)) + loaded(args[0], out, *args[2:]) + + assert out._data[0] == 128.0 + + +def test_artifact_identity_uses_declared_key_and_device(artifact_dir, monkeypatch): + import importlib + + at = importlib.import_module("flydsl.autotune") + + def kernel(a, out, N, BLOCK_THREADS: int, dtype_str: str = "bf16"): + pass + + tuner = _make_tuner(fn=kernel, key=["N", "dtype_str"], artifact_name="defaults") + args = (FakeTensor((8, 512)), FakeTensor((1,)), 512) + first = tuner._artifact_ref(args, {}, required=True) + scratch_key = tuner._make_key(args, {}) + + assert first == tuner._artifact_ref(args, {"dtype_str": "bf16"}, required=True) + assert first != tuner._artifact_ref((args[0], args[1], 1024), {}, required=True) + + monkeypatch.setattr( + at, + "_device_descriptor", + lambda device=None: {"name": "Other GPU", "arch": "gfx-test", "compute_units": 2}, + ) + assert first != tuner._artifact_ref(args, {}, required=True) + assert scratch_key != tuner._make_key(args, {}) + + +def test_cached_artifact_is_revalidated_for_each_call(monkeypatch, tmp_path, artifact_dir): + _, args = _emit_artifact(monkeypatch, artifact_dir) + monkeypatch.delenv("FLYDSL_AUTOTUNE") + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "fresh-cache")) + tuner = _make_artifact_tuner() + ref = tuner._artifact_ref(args, {}, required=True) + + assert tuner._load_artifact(ref, args, {}).to_dict() == {"BLOCK_THREADS": 64} + assert tuner._load_artifact(ref, args, {"BLOCK_THREADS": 32}) is None + assert tuner._load_artifact(ref, args, {}).to_dict() == {"BLOCK_THREADS": 64} + + +@pytest.mark.parametrize("case", ["corrupt", "version", "identity", "config", "override", "type"]) +def test_invalid_artifact_falls_back(monkeypatch, tmp_path, artifact_dir, case): + path, args = _emit_artifact(monkeypatch, artifact_dir) + payload = json.loads(path.read_text()) + if case == "corrupt": + path.write_text("{") + else: + if case == "version": + payload["version"] = 2 + elif case == "identity": + payload["identity"]["key"]["N"] = 1024 + elif case == "config": + payload["config"] = {"UNKNOWN": 64} + elif case == "type": + payload["config"] = {"BLOCK_THREADS": "64"} + else: + payload["config"] = {"BLOCK_THREADS": 64, "N": 1024} + path.write_text(json.dumps(payload)) + monkeypatch.delenv("FLYDSL_AUTOTUNE") + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "fresh-cache")) + + tuner = _make_artifact_tuner( + configs=lambda *_args, **_kwargs: pytest.fail("invalid artifact triggered search"), + do_bench_fn=lambda *args, **kwargs: pytest.fail("invalid artifact benchmarked"), + ) + out = FakeTensor((1,)) + tuner(args[0], out, *args[2:]) + + assert out._data[0] == 7.0 + + +def test_artifact_runtime_failure_is_not_masked_by_default(monkeypatch, tmp_path, artifact_dir): + _, args = _emit_artifact(monkeypatch, artifact_dir) + monkeypatch.delenv("FLYDSL_AUTOTUNE") + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "fresh-cache")) + seen = [] + + def failing_kernel(a, out, m_in, N, dtype_str, BLOCK_THREADS: int): + seen.append(BLOCK_THREADS) + if BLOCK_THREADS == 64: + raise RuntimeError("kernel failed") + out._data[0] = float(BLOCK_THREADS) + + tuner = _make_artifact_tuner(fn=failing_kernel) + + with pytest.raises(RuntimeError, match="kernel failed"): + tuner(*args) + + assert seen == [64] + + +@pytest.mark.parametrize( + "config, message", + [ + pytest.param(Config(BLOCK_THREADS=64, pre_hook=lambda kwargs: None), "pre_hook", id="pre-hook"), + pytest.param(Config(BLOCK_THREADS=(64,)), "preserve their types", id="json-type-change"), + ], +) +def test_unpersistable_artifact_config_blocks_generation(monkeypatch, artifact_dir, config, message): + def kernel(a, out, m_in, N, dtype_str, BLOCK_THREADS): + pass + + monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") + tuner = _make_artifact_tuner(fn=kernel, configs=[config]) + args = (FakeTensor((16, 512)), FakeTensor((1,)), 16, 512, "bf16") + + with pytest.raises(ValueError, match=message): + tuner(*args) + + assert tuner.cache == {} + assert not tuner._cache_file.exists() + assert not list(artifact_dir.glob("*.json")) + + +def test_unavailable_device_identity_falls_back_but_blocks_generation(monkeypatch, artifact_dir): + import importlib + + at = importlib.import_module("flydsl.autotune") + monkeypatch.setattr(at, "_device_descriptor", lambda device=None: None) + args = (FakeTensor((16, 512)), FakeTensor((1,)), 16, 512, "bf16") + tuner = _make_artifact_tuner() + + tuner(*args) + assert args[1]._data[0] == 7.0 + + monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") + with pytest.raises(ValueError, match="cannot generate offline config identity"): + tuner(*args) + + +@pytest.mark.parametrize("name", ["../rmsnorm", 123]) +def test_artifact_name_must_be_safe(name): + with pytest.raises((TypeError, ValueError), match="artifact_name"): + _make_artifact_tuner(artifact_name=name) + + +def test_artifact_key_must_name_a_kernel_parameter(): + with pytest.raises(ValueError, match="artifact keys"): + _make_artifact_tuner(key=["mni"]) + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__, "-v"]))