[FIX] Keep device functions JITFunctions under the CLI

tests/end_to_end/test_wrapper.py

@@ -5,6 +5,8 @@
 from pathlib import Path
 from string import Template
 import shutil
+import pytest
+import torch
 from tilelens.wrapper import (
     SANITIZER_COMMAND,
     PROFILER_COMMAND,
@@ -328,3 +330,45 @@ def test_cli_rejects_trace_decorator(tmp_path: Path):
     assert (
         "CLI wrapper" in proc.stderr
     ), f"Expected error message about CLI wrapper conflict in stderr, got:\n{proc.stderr}"
+
+
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_profiler_cli_instruments_device_functions(tmp_path: Path, device):
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("the profiler's compile needs a GPU")
+    script = tmp_path / "device_functions.py"
+    script.write_text(
+        """import torch
+import triton
+import triton.language as tl
+
+@triton.jit
+def add_one(value):
+    for _ in range(1):
+        value += 1
+    return value
+
+@triton.jit
+def kernel(ptr, BLOCK: tl.constexpr):
+    offsets = tl.arange(0, BLOCK)
+    tl.store(ptr + offsets, add_one(tl.load(ptr + offsets)))
+
+"""
+        + f"kernel[(1,)](torch.zeros(8, device={device!r}), BLOCK=8)\n",
+        encoding="utf-8",
+    )
+    # On CPU nothing can compile, so the buffer-load check is off there.
+    env = {
+        **os.environ,
+        "PROFILER_DISABLE_BUFFER_LOAD_CHECK": "1" if device == "cpu" else "0",
+    }
+    proc = subprocess.run(
+        [PROFILER_COMMAND, str(script)],
+        capture_output=True,
+        text=True,
+        env=env,
+        timeout=60,
+    )
+    assert proc.returncode == 0, proc.stderr
+    # The loop inside the device function is instrumented.
+    assert "Total for-loops detected: 1" in proc.stdout

tests/unit/test_wrapper.py

@@ -5,6 +5,7 @@
 from unittest.mock import MagicMock, patch
 
 import tilelens
+import triton
 from tilelens.core.config import config as cfg
 from tilelens.core.trace import TraceInterface
 from tilelens.wrapper import (
@@ -163,6 +164,36 @@ def test_create_patched_jit_with_kwargs_and_custom_original_jit():
     assert result == "final_kernel"
 
 
+def test_patched_jit_preserves_device_function_dependencies(_isolate_cli_active):
+    cfg.cli_active = True
+    patched_jit = create_patched_jit(profiler_wrapper, triton.jit, frontend="triton")
+
+    @patched_jit
+    def add_one(value):
+        return value + 1
+
+    @patched_jit
+    def add_two(value):
+        return add_one(add_one(value))
+
+    # Triton's dependency hashing accepts only JITFunction callees.
+    assert isinstance(add_one, triton.JITFunction)
+    assert add_two.cache_key
+
+
+def test_cli_jit_function_warmup_skips_the_trace(_isolate_cli_active):
+    cfg.cli_active = True
+
+    @triton.jit
+    def kernel(x):
+        pass
+
+    original_run = MagicMock()
+    kernel.run = original_run
+    profiler_wrapper(kernel).run(grid=(1,), warmup=True)
+    original_run.assert_called_once_with(grid=(1,), warmup=True)
+
+
 # ======== create_patched_autotune Tests ===========
 
 
@@ -243,6 +274,11 @@ def test_trace_decorator_works_when_cli_inactive(_isolate_cli_active):
     assert isinstance(result, TraceInterface)
 
 
+def test_trace_rejects_an_unknown_frontend():
+    with pytest.raises(ValueError, match="Unknown frontend"):
+        tilelens.trace(frontend="bogus")(lambda: None)
+
+
 def test_apply_wrapper_rejects_non_cli_invocation():
     """apply_sanitizer/apply_profiler should raise when called from Python, not CLI."""
     with pytest.raises(RuntimeError, match="must be used as a CLI tool"):

tilelens/core/frontend/base.py

@@ -2,10 +2,13 @@
 from contextlib import nullcontext
 from dataclasses import dataclass
 from importlib import import_module
-from typing import Any
+from typing import Any, TYPE_CHECKING
 
 from tilelens.core.data import Op
 
+if TYPE_CHECKING:
+    from tilelens.core.client import Client
+
 
 _MISSING = object()
 
@@ -115,6 +118,12 @@ def from_namespaces(
     def maybe_yield_for_multism(self) -> None:
         return None
 
+    def wrap_kernel(self, kernel: Any, client: "str | Client") -> Any:
+        """Create this frontend's traced callable, including CLI adaptation."""
+        raise NotImplementedError(
+            f"Frontend '{self.name}' does not implement kernel wrapping"
+        )
+
     def run_op_overrider(
         self,
         op: Callable,
@@ -206,7 +215,13 @@ def register_frontend(frontend: Frontend) -> Frontend:
 
 def get_frontend(name: str) -> Frontend:
     if name not in FRONTEND_REGISTRY:
-        import_module(f"tilelens.core.frontend.{name}")
+        module_name = f"tilelens.core.frontend.{name}"
+        try:
+            import_module(module_name)
+        except ModuleNotFoundError as exc:
+            if exc.name != module_name:
+                raise
+            raise ValueError(f"Unknown frontend: {name}") from exc
     try:
         return FRONTEND_REGISTRY[name]
     except KeyError as exc:

tilelens/core/frontend/gluon.py

@@ -566,6 +566,11 @@ def __init__(self):
             namespaces=definition.namespaces,
         )
 
+    def wrap_kernel(self, kernel, client):
+        from ..trace import GluonTrace
+
+        return GluonTrace(kernel, client)
+
     @staticmethod
     def symbolic_ops_for_op_type(op_type: type[Op]) -> tuple[str, ...]:
         return TritonFrontend.symbolic_ops_for_op_type(op_type)

tilelens/core/frontend/nki.py

@@ -93,6 +93,11 @@ def __init__(self):
             namespaces=definition.namespaces,
         )
 
+    def wrap_kernel(self, kernel, client):
+        from ..trace import NKITrace
+
+        return NKITrace(kernel, client, beta2=False)
+
     def patch_lang(self, fn, client_manager: Any = None) -> _LangPatchScope:
         from tilelens.core.simulation.nki import nki_patch_lang
 

tilelens/core/frontend/nki_beta2.py

@@ -76,6 +76,11 @@ def __init__(self):
             namespaces=definition.namespaces,
         )
 
+    def wrap_kernel(self, kernel, client):
+        from ..trace import NKITrace
+
+        return NKITrace(kernel, client, beta2=True)
+
     def patch_lang(self, fn, client_manager: Any = None) -> _LangPatchScope:
         from tilelens.core.simulation.nki_beta2 import nki_patch_lang
 

tilelens/core/frontend/triton.py

@@ -1084,7 +1084,34 @@ def _grid_executor_call(self, grid_executor, *args_dev, **kwargs):
         finally:
             self._current_client_manager = previous_client_manager
 
+    def wrap_kernel(self, kernel, client):
+        from ..trace import TritonTrace
+
+        traced = TritonTrace(kernel, client)
+        if not cfg.cli_active or not isinstance(kernel, JITFunction):
+            return traced
+
+        # Device helpers must remain JITFunctions for dependency hashing and
+        # code generation. Instrument their launch path instead of replacing
+        # the callable that another kernel references.
+        original_run = kernel.run
+
+        def run(*args, **kwargs):
+            # JITFunction.warmup re-enters run(warmup=True) for the real
+            # compile; tracing it would recurse.
+            if kwargs.get("warmup", False):
+                return original_run(*args, **kwargs)
+            return traced.run(*args, **kwargs)
+
+        kernel.run = run
+        kernel._tilelens_trace = traced
+        return kernel
+
     def _jit_function_call(self, jit_function, *args, **kwargs):
+        traced = getattr(jit_function, "_tilelens_trace", None)
+        if traced is not None:
+            return traced(*args, **kwargs)
+
         if (
             self._current_client_manager is not None
             and not self._jit_function_uses_core_module(jit_function)

tilelens/core/trace.py

@@ -8,6 +8,7 @@
 from ..clients.race_detector.race_detector import NullRaceDetector
 from .client import ClientManager, Client
 from .data import Launch
+from .frontend.base import get_frontend
 import types
 
 
@@ -409,7 +410,10 @@ def _is_race_detector_client(selected: str | Client) -> bool:
         return isinstance(selected, NullRaceDetector)
 
     def decorator(kernel) -> TritonTrace | NKITrace | GluonTrace | Any:
-        if cfg.cli_active and isinstance(kernel, TraceInterface):
+        if cfg.cli_active and (
+            isinstance(kernel, TraceInterface)
+            or isinstance(getattr(kernel, "_tilelens_trace", None), TraceInterface)
+        ):
             raise RuntimeError(
                 "@tilelens.trace() decorator cannot be used together with "
                 "CLI wrapper (e.g., tile-sanitizer / tile-profiler). "
@@ -435,15 +439,7 @@ def decorator(kernel) -> TritonTrace | NKITrace | GluonTrace | Any:
 
         trace_source(kernel)
 
-        # First-time wrapping
-        if frontend in ("nki", "nki_beta2"):
-            return NKITrace(kernel, client, beta2=("beta2" in frontend))
-        elif frontend == "gluon":
-            return GluonTrace(kernel, client)
-        elif frontend == "triton":
-            return TritonTrace(kernel, client)
-        else:
-            raise ValueError(f"Unknown frontend: {frontend}")
+        return get_frontend(frontend).wrap_kernel(kernel, client)
 
     return decorator