diff --git a/tests/end_to_end/test_wrapper.py b/tests/end_to_end/test_wrapper.py index 32958077..e0d8d181 100644 --- a/tests/end_to_end/test_wrapper.py +++ b/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 diff --git a/tests/unit/test_wrapper.py b/tests/unit/test_wrapper.py index 12c47896..176327b9 100644 --- a/tests/unit/test_wrapper.py +++ b/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"): diff --git a/tilelens/core/frontend/base.py b/tilelens/core/frontend/base.py index b8442f56..c581cc92 100644 --- a/tilelens/core/frontend/base.py +++ b/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: diff --git a/tilelens/core/frontend/gluon.py b/tilelens/core/frontend/gluon.py index 7f7a1460..80deb181 100644 --- a/tilelens/core/frontend/gluon.py +++ b/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) diff --git a/tilelens/core/frontend/nki.py b/tilelens/core/frontend/nki.py index 6b0820f1..ed4c3698 100644 --- a/tilelens/core/frontend/nki.py +++ b/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 diff --git a/tilelens/core/frontend/nki_beta2.py b/tilelens/core/frontend/nki_beta2.py index 3f47d716..100f3ff5 100644 --- a/tilelens/core/frontend/nki_beta2.py +++ b/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 diff --git a/tilelens/core/frontend/triton.py b/tilelens/core/frontend/triton.py index 30b4846a..324bdb10 100644 --- a/tilelens/core/frontend/triton.py +++ b/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) diff --git a/tilelens/core/trace.py b/tilelens/core/trace.py index abdd11a1..7702f754 100644 --- a/tilelens/core/trace.py +++ b/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