diff --git a/tests/end_to_end/test_core.py b/tests/end_to_end/test_core.py index b98b7771..d9ab989e 100644 --- a/tests/end_to_end/test_core.py +++ b/tests/end_to_end/test_core.py @@ -3,6 +3,7 @@ import triton import triton.language as tl +from triton import knobs import tilelens from tilelens.clients import Profiler, Sanitizer @@ -79,8 +80,10 @@ def my_gluon_kernel(out, BLOCK: ttgl.constexpr): assert isinstance(my_gluon_kernel, GluonTrace) assert isinstance(my_gluon_kernel.runner, Autotuner) - assert isinstance(my_gluon_kernel.runner.fn, GluonInterpretedFunction) - assert my_gluon_kernel.runner.fn.fn is my_gluon_kernel.base_fn + leaf = my_gluon_kernel.runner.fn + assert isinstance(leaf.interpreted_fn, GluonInterpretedFunction) + assert leaf.interpreted_fn.fn is my_gluon_kernel.base_fn + assert leaf.fn is my_gluon_kernel.jit_fn def test_gluon_trace_handles_heuristics_wrapper(): @@ -99,8 +102,62 @@ def my_gluon_kernel(out, BLOCK: ttgl.constexpr): assert isinstance(my_gluon_kernel, GluonTrace) assert isinstance(my_gluon_kernel.runner, Heuristics) - assert isinstance(my_gluon_kernel.runner.fn, GluonInterpretedFunction) - assert my_gluon_kernel.runner.fn.fn is my_gluon_kernel.base_fn + leaf = my_gluon_kernel.runner.fn + assert isinstance(leaf.interpreted_fn, GluonInterpretedFunction) + assert leaf.interpreted_fn.fn is my_gluon_kernel.base_fn + assert leaf.fn is my_gluon_kernel.jit_fn + + +def test_gluon_autotune_listener_sees_jit_function(monkeypatch): + from triton.experimental import gluon + from triton.experimental.gluon import language as ttgl + + calls = [] + monkeypatch.setattr(knobs.autotuning, "listener", lambda **kw: calls.append(kw)) + + @tilelens.trace("tracer", frontend="gluon") + @triton.autotune( + configs=[ + triton.Config({"BLOCK": 1}, num_warps=1), + triton.Config({"BLOCK": 2}, num_warps=1), + ], + key=[], + ) + @gluon.jit + def my_gluon_kernel(in_ptr, out_ptr, BLOCK: ttgl.constexpr): + ttgl.store(out_ptr, ttgl.load(in_ptr)) + + inp = torch.tensor([42.0]) + out = torch.empty_like(inp) + my_gluon_kernel[(1,)](inp, out) + + assert [call["fn"] for call in calls] == [my_gluon_kernel.jit_fn] + torch.testing.assert_close(out, inp, atol=0, rtol=0) + + +def test_gluon_autotune_retrace_keeps_jit_function(): + from triton.experimental import gluon + from triton.experimental.gluon import language as ttgl + + @triton.autotune( + configs=[ + triton.Config({"BLOCK": 1}, num_warps=1), + triton.Config({"BLOCK": 2}, num_warps=1), + ], + key=[], + ) + @gluon.jit + def my_gluon_kernel(in_ptr, out_ptr, BLOCK: ttgl.constexpr): + ttgl.store(out_ptr, ttgl.load(in_ptr)) + + first = tilelens.trace("tracer", frontend="gluon")(my_gluon_kernel) + second = tilelens.trace("tracer", frontend="gluon")(my_gluon_kernel) + + assert second.jit_fn is first.jit_fn + inp = torch.tensor([42.0]) + out = torch.empty_like(inp) + second[(1,)](inp, out) + torch.testing.assert_close(out, inp, atol=0, rtol=0) def test_gluon_sanitizer_run_preserves_instrumentation_mode(monkeypatch): @@ -368,3 +425,64 @@ def noop_kernel(n, BLOCK: tl.constexpr): traced = tilelens.trace(client=Sanitizer())(noop_kernel) traced[(1,)](n=32) + + +def _autotuned_fill_kernel(**autotune_kwargs): + # A JITFunction, not the InterpretedFunction triton.jit returns in + # interpreter mode, which tests/unit/test_multithreading.py turns on for + # the whole session. + with knobs.runtime.scope(): + knobs.runtime.interpret = False + + # Two configs, so a launch with a new key benchmarks them. + @triton.autotune( + configs=[triton.Config({"BLOCK": 32}), triton.Config({"BLOCK": 64})], + key=["n"], + **autotune_kwargs, + ) + @triton.jit + def fill_kernel(x_ptr, n, BLOCK: tl.constexpr): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + tl.store(x_ptr + offs, 1.0, mask=offs < n) + + return fill_kernel + + +def _fill_grid(meta): + return (triton.cdiv(100, meta["BLOCK"]),) + + +def test_autotune_listener_sees_jit_function(monkeypatch): + calls = [] + monkeypatch.setattr(knobs.autotuning, "listener", lambda **kw: calls.append(kw)) + traced = tilelens.trace("tracer")(_autotuned_fill_kernel()) + + x = torch.zeros(100) + traced[_fill_grid](x, 100) + + assert [call["fn"] for call in calls] == [traced.jit_fn] + assert torch.all(x == 1) + + +def test_autotune_retrace_keeps_jit_function(): + kernel = _autotuned_fill_kernel() + first = tilelens.trace("tracer")(kernel) + second = tilelens.trace("tracer")(kernel) + + assert first.jit_fn is not None + assert second.jit_fn is first.jit_fn + x = torch.zeros(100) + second[_fill_grid](x, 100) + assert torch.all(x == 1) + + +def test_autotune_interpreter_skips_disk_cache(monkeypatch, tmp_path): + monkeypatch.setenv("TRITON_CACHE_DIR", str(tmp_path)) + kernel = _autotuned_fill_kernel(cache_results=True) + traced = tilelens.trace("tracer")(kernel) + + x = torch.zeros(100) + traced[_fill_grid](x, 100) + + assert torch.all(x == 1) + assert not list(tmp_path.rglob("*.autotune.json")) diff --git a/tilelens/core/trace.py b/tilelens/core/trace.py index 28f38301..ca1136e9 100644 --- a/tilelens/core/trace.py +++ b/tilelens/core/trace.py @@ -1,4 +1,4 @@ -from copy import deepcopy +from copy import copy from collections.abc import Callable from typing import Any from ..utils.traceback_utils import CODE_KEYS, get_code_key @@ -65,6 +65,33 @@ def __getitem__(self, grid): ) +class _InterpretedLeaf: + """The interpreted kernel at the bottom of a traced Autotuner/Heuristics chain. + + Everything is the interpreted function's, except ``fn``: that is the + kernel's JITFunction, not the Python function an InterpretedFunction's + ``fn`` is, since Triton's Autotuner follows ``.fn`` from its own down to + the JITFunction it tunes (its disk cache key; ``knobs.autotuning.listener``, + which Triton 3.8 calls after autotuning a launch). + """ + + def __init__(self, jit_fn: Any, interpreted_fn: Any) -> None: + self.jit_fn = jit_fn + self.interpreted_fn = interpreted_fn + + @property + def fn(self) -> Any: + return self.jit_fn + + def run(self, *args, **kwargs): + return self.interpreted_fn.run(*args, **kwargs) + + def __getattr__(self, name: str) -> Any: + if name in ("jit_fn", "interpreted_fn"): # not set yet (e.g. mid-copy) + raise AttributeError(name) + return getattr(self.interpreted_fn, name) + + class KernelTraceSupport: @staticmethod def _is_autotuner(runner: Any) -> bool: @@ -83,14 +110,24 @@ def dummy_benchmarker(fn, quantiles): fn() return (1.0, 1.0, 1.0) - def _interpreter_runner(self, runner: Any, interpreted_fn: Any) -> Any: + def _interpreter_runner( + self, runner: Any, interpreted_fn: Any, jit_fn: Any | None + ) -> Any: + leaf = ( + interpreted_fn + if jit_fn is None + else _InterpretedLeaf(jit_fn, interpreted_fn) + ) if self._is_autotuner(runner): - runner.fn = interpreted_fn + runner.fn = leaf # Kernel Cache: replace the benchmark with a dummy to skip performance testing. runner._do_bench = self.dummy_benchmarker + # Triton only turns the disk cache off for an autotuner decorated in + # interpreter mode; the dummy timings must never be persisted. + runner.cache_results = False return runner if self._is_heuristics(runner): - runner.fn = interpreted_fn + runner.fn = leaf return runner return interpreted_fn @@ -99,7 +136,8 @@ def _warmup_runner(self, runner: Any, jit_fn: Any | None) -> Any | None: return jit_fn if jit_fn is None: return None - warmup_runner = deepcopy(runner) + # No deepcopy: the JITFunction under runner holds an RLock. + warmup_runner = copy(runner) warmup_runner.fn = jit_fn return warmup_runner @@ -144,6 +182,9 @@ def unpack_kernel( ) -> tuple[Any | None, Callable | None, Any | None]: if isinstance(source, TritonTrace): return source.jit_fn, source.base_fn, source.interpreted_fn + if isinstance(source, _InterpretedLeaf): + # The leaf an earlier trace put under this autotuner or heuristics + return source.jit_fn, source.jit_fn.fn, source.interpreted_fn if isinstance(source, JITFunction): base_fn = source.fn return source, base_fn, InterpretedFunction(base_fn) @@ -160,7 +201,7 @@ def unpack_kernel( self.jit_fn, self.base_fn, self.interpreted_fn = unpack_kernel(runner.fn) else: self.jit_fn, self.base_fn, self.interpreted_fn = unpack_kernel(runner) - self.runner = self._interpreter_runner(runner, self.interpreted_fn) + self.runner = self._interpreter_runner(runner, self.interpreted_fn, self.jit_fn) self.warmup_runner = self._warmup_runner(runner, self.jit_fn) self.arg_names = runner.arg_names @@ -304,22 +345,29 @@ def __init__(self, runner: Any, client: str | Client) -> None: if not all(hasattr(runner, attr) for attr in ("fn", "run", "arg_names")): runner = gluon.jit(runner) - def unpack_kernel(source: Any) -> tuple[Callable, Any]: + def unpack_kernel(source: Any) -> tuple[Any, Callable, Any]: if isinstance(source, GluonTrace): - return source.base_fn, source.interpreted_fn + return source.jit_fn, source.base_fn, source.interpreted_fn + if isinstance(source, _InterpretedLeaf): + # The leaf an earlier trace put under this autotuner or heuristics + return source.jit_fn, source.jit_fn.fn, source.interpreted_fn if isinstance(source, Heuristics): return unpack_kernel(source.fn) if all(hasattr(source, attr) for attr in ("fn", "arg_names")): base_fn = source.fn - return base_fn, GluonInterpretedFunction(base_fn, source.arg_names) + return ( + source, + base_fn, + GluonInterpretedFunction(base_fn, source.arg_names), + ) raise TypeError(f"Unsupported runner type: {type(source)}") - self.base_fn, self.interpreted_fn = unpack_kernel( + self.jit_fn, self.base_fn, self.interpreted_fn = unpack_kernel( runner.fn if self._is_autotuner(runner) else runner ) self.fn = runner self.arg_names = runner.arg_names - self.runner = self._interpreter_runner(runner, self.interpreted_fn) + self.runner = self._interpreter_runner(runner, self.interpreted_fn, self.jit_fn) TraceInterface.__init__(self, client) self._copy_callable_attrs(runner, self.base_fn)