[FIX] Keep the JITFunction reachable from traced autotuners

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"))

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)