+ {"".join(sections_html)}
+
+No differences.
' + assert _make_diff_html("x", "x", context=3) == 'No differences.
' + + +def test_no_skipped_phantom_records(monkeypatch, tmp_path): + """Pre-registration is gone: no SKIPPED records, indices global-monotonic.""" + from tilelang.tools.lower_trace import enable, disable + from tilelang.tools.lower_trace import core as _core + from tilelang.tools.lower_trace.core import STATUS_SKIPPED + from tilelang.backend.pass_pipeline import resolve_pipeline + import tilelang.language as T + + monkeypatch.setenv("TL_LOWER_TRACE", "both") + monkeypatch.setenv("TL_LOWER_TRACE_DIR", str(tmp_path)) + + disable() + enable() + + @T.prim_func + def tiny(A: T.Tensor((32,), "float32"), B: T.Tensor((32,), "float32")): + with T.Kernel(32): + tid = T.get_thread_binding() + B[tid] = A[tid] + 1.0 + + mod = tvm.IRModule({"main": tiny}) + target = tvm.target.Target("c") + pipeline = resolve_pipeline(target) + pipeline.lower(mod, target) + + # No phantom/skipped records remain — every record is COMPLETED or FAILED + skipped = [r for r in _core._records if r.status == STATUS_SKIPPED] + assert not skipped, f"Found {len(skipped)} SKIPPED records (pre-registration not removed)" + + # Indices are strictly increasing across all records (global-monotonic) + indices = [r.index for r in _core._records] + assert indices == sorted(indices), f"Indices not ascending: {indices}" + assert len(indices) == len(set(indices)), f"Duplicate indices: {indices}" + + # No phantom LetInline slot when should_force_let_inline() is False + letinline = [r for r in _core._records if "LetInline" in r.name] + assert not letinline, f"Phantom LetInline records found: {letinline}" + + disable() + monkeypatch.delenv("TL_LOWER_TRACE", raising=False) + monkeypatch.delenv("TL_LOWER_TRACE_DIR", raising=False) + + +def test_terminal_mode_no_html(monkeypatch, tmp_path): + """TL_LOWER_TRACE=terminal must not produce an HTML report at process exit. + + Regression guard for the ``_final_report()`` hook: ``_save_raw_files()`` + populates ``_run_dir`` even in terminal-only mode, so without an explicit + ``_should_gen_html()`` check the atexit hook would still emit ``report.html``. + """ + from tilelang.tools.lower_trace import enable, disable + from tilelang.backend.pass_pipeline import resolve_pipeline + import tilelang.language as T + + monkeypatch.setenv("TL_LOWER_TRACE", "terminal") + monkeypatch.setenv("TL_LOWER_TRACE_DIR", str(tmp_path)) + + disable() + enable() + + try: + + @T.prim_func + def tiny(A: T.Tensor((32,), "float32"), B: T.Tensor((32,), "float32")): + with T.Kernel(32): + tid = T.get_thread_binding() + B[tid] = A[tid] + 1.0 + + mod = tvm.IRModule({"main": tiny}) + target = tvm.target.Target("c") + pipeline = resolve_pipeline(target) + pipeline.lower(mod, target) + + # Simulate the atexit hook that fires at process exit. + _core._final_report() + + script_dir = _core._script_dir + assert script_dir is not None, "script_dir should be set after a run" + symlink_report = os.path.join(script_dir, "report.html") + assert not os.path.exists(symlink_report), f"terminal mode must not write report.html symlink at {symlink_report}" + if _core._run_dir is not None: + run_report = os.path.join(_core._run_dir, "report.html") + assert not os.path.exists(run_report), f"terminal mode must not write report.html at {run_report}" + finally: + disable() + + +def test_lower_trace_html_on_failure(tmp_path): + """lower_trace() flushes a partial HTML report even when a pass raises. + + Regression guard: previously an exception from ``p(mod)`` aborted the + function before ``generate_html()`` ran, so ``mode='html'``/``'both'`` + lost the partial trace of completed passes. + """ + from tilelang.tools.lower_trace import lower_trace + + program = _simple_program() + html_path = str(tmp_path / "partial.html") + + class _BoomPass: + def __call__(self, mod): + raise RuntimeError("intentional boom") + + passes = [ + ("Simplify", _noop_pass()), + ("Boom", _BoomPass()), + ] + + with pytest.raises(RuntimeError, match="intentional boom"): + lower_trace(program, passes, mode="html", html_path=html_path) + + assert os.path.exists(html_path), "partial HTML report must be flushed on failure" + with open(html_path) as f: + content = f.read() + assert "Simplify" in content, "completed pass must still appear in partial report" + assert "Boom" in content, "failing pass name must appear in partial report" + assert "FAILED" in content, "failing step must be marked FAILED" + + +# --------------------------------------------------------------------------- +# Codegen edit-and-recompile (Phase 1: _make_patched_source_module for _without_compile) +# --------------------------------------------------------------------------- + + +class _MockCodegenModule: + """Minimal stand-in for a TVM runtime.Module returned by codegen FFIs.""" + + def __init__(self, source: str): + self._source = source + + def inspect_source(self) -> str: + return self._source + + def get_source(self) -> str: + return self._source + + +def _make_mock_build(source: str): + """Return a mock codegen FFI that always produces *source*.""" + + def mock_build(*args, **kwargs): + return _MockCodegenModule(source) + + return mock_build + + +class _MockPatchedModule: + """Stand-in for the CSourceModule returned by _make_patched_source_module.""" + + def __init__(self, source: str): + self._source = source + + def get_source(self) -> str: + return self._source + + def inspect_source(self) -> str: + return self._source + + +def _patched_module_factory(original_module, patched_source): + """Test replacement for _make_patched_source_module (avoids real TVM C++ FFI).""" + return _MockPatchedModule(patched_source) + + +def _setup_trace_overrides(tmp_path, mode="terminal"): + """Set lower_trace overrides for unit testing. + + The autouse ``_isolate_env`` fixture calls ``disable()`` before each test, + so overrides can be set safely inside the test body. + """ + _core._mode_override = mode + _core._trace_dir_override = str(tmp_path) + _core._codegen_output_path_override = str(tmp_path / "codegen.cpp") + _core.reset() + + +def _clear_trace_overrides(): + """Reset overrides (also done by the autouse fixture's ``disable()``).""" + _core._mode_override = _core._UNSET + _core._trace_dir_override = _core._UNSET + _core._codegen_output_path_override = _core._UNSET + _core.reset() + + +@contextlib.contextmanager +def _patch_make_patched_source_module(): + """Patch _make_patched_source_module so tests avoid the real TVM C++ FFI.""" + from unittest.mock import patch + + with patch("tilelang.tools.lower_trace.core._make_patched_source_module", side_effect=_patched_module_factory) as mock: + yield mock + + +def test_codegen_proxy_for_without_compile(tmp_path): + """*_without_compile FFIs return a patched module when user edits codegen.cpp.""" + from tilelang.tools.lower_trace.core import _wrap_codegen_ffi + + source_v1 = "// generated kernel v1\n" + mock_build = _make_mock_build(source_v1) + wrapper = _wrap_codegen_ffi(mock_build, "target.build.tilelang_cuda_without_compile") + + _setup_trace_overrides(tmp_path) + codegen_path = _core._codegen_output_path_override + + with _patch_make_patched_source_module() as mock_factory: + try: + # Run 1: initializes codegen.cpp + .original from codegen output + result1 = wrapper("fake_mod") + assert result1.inspect_source() == source_v1 + + # Edit codegen.cpp (user edit) + edited = "// edited by user\n" + with open(codegen_path, "w") as f: + f.write(edited) + + # Run 2: user edited, codegen unchanged → PATCHED → patched module returned + result2 = wrapper("fake_mod") + assert mock_factory.called, "_make_patched_source_module should be called for _without_compile FFI" + assert result2.get_source() == edited, "Patched module should return the user-edited source" + finally: + _clear_trace_overrides() + + +def test_codegen_proxy_for_source_only_ffi(tmp_path): + """Source-only FFIs without a _without_compile suffix (tilelang_c, webgpu) also return patched module.""" + from tilelang.tools.lower_trace.core import _wrap_codegen_ffi + + source_v1 = "// generated C kernel v1\n" + mock_build = _make_mock_build(source_v1) + wrapper = _wrap_codegen_ffi(mock_build, "target.build.tilelang_c") + + _setup_trace_overrides(tmp_path) + codegen_path = _core._codegen_output_path_override + + with _patch_make_patched_source_module() as mock_factory: + try: + # Run 1: initializes codegen.cpp + .original + wrapper("fake_mod") + + # Edit codegen.cpp + edited = "// edited C kernel\n" + with open(codegen_path, "w") as f: + f.write(edited) + + # Run 2: PATCHED → patched module returned (tilelang_c is in _SOURCE_ONLY_CODEGEN_FFIS) + result2 = wrapper("fake_mod") + assert mock_factory.called, "Expected patched module for source-only FFI tilelang_c" + assert result2.get_source() == edited + finally: + _clear_trace_overrides() + + +def test_codegen_no_proxy_for_full_compile(tmp_path, capsys): + """Full-compile FFIs return the real module (not patched) + NOTE when user edits codegen.cpp.""" + from tilelang.tools.lower_trace.core import _wrap_codegen_ffi + + source_v1 = "// generated kernel v1\n" + mock_build = _make_mock_build(source_v1) + wrapper = _wrap_codegen_ffi(mock_build, "target.build.tilelang_cuda") + + _setup_trace_overrides(tmp_path) + codegen_path = _core._codegen_output_path_override + target = tvm.target.Target("cuda") + + with _patch_make_patched_source_module() as mock_factory: + try: + # Run 1: initializes codegen.cpp + .original + result1 = wrapper("fake_mod", target) + assert result1.inspect_source() == source_v1 + + # Edit codegen.cpp + edited = "// edited by user\n" + with open(codegen_path, "w") as f: + f.write(edited) + + # Run 2: user edited, codegen unchanged → PATCHED + # But full-compile FFI → return real module, NOT patched + capsys.readouterr() # clear prior output + result2 = wrapper("fake_mod", target) + assert not mock_factory.called, "Full-compile FFI must NOT call _make_patched_source_module (would crash tvm_ffi backend)" + assert result2.inspect_source() == source_v1, "Full-compile FFI should return original (unpatched) module" + + captured = capsys.readouterr() + assert "NOT recompiled" in captured.out + assert "nvrtc" in captured.out # backend hint + finally: + _clear_trace_overrides() + + +def test_codegen_conflict_backup(tmp_path): + """CONFLICT: both user edited and codegen changed → backup + regenerate.""" + from tilelang.tools.lower_trace.core import _wrap_codegen_ffi + + source_v1 = "// generated kernel v1\n" + mock_build = _make_mock_build(source_v1) + wrapper = _wrap_codegen_ffi(mock_build, "target.build.tilelang_cuda_without_compile") + + _setup_trace_overrides(tmp_path) + codegen_path = _core._codegen_output_path_override + original_path = codegen_path + ".original" + + with _patch_make_patched_source_module() as mock_factory: + try: + # Run 1: init + wrapper("fake_mod") + + # Edit codegen.cpp (user edit) + with open(codegen_path, "w") as f: + f.write("// user edit\n") + + # Change codegen output (new wrapper with different source) + source_v2 = "// new codegen output v2\n" + wrapper = _wrap_codegen_ffi(_make_mock_build(source_v2), "target.build.tilelang_cuda_without_compile") + + # Run 2: CONFLICT — working != current + result2 = wrapper("fake_mod") + + # .bak files created + assert os.path.exists(codegen_path + ".bak"), "User working copy not backed up" + assert os.path.exists(original_path + ".bak"), "Old baseline not backed up" + + # codegen.cpp regenerated from new codegen + with open(codegen_path) as f: + assert f.read() == source_v2 + + # .original advanced to new codegen + with open(original_path) as f: + assert f.read() == source_v2 + + # No patched module returned (regenerated from new codegen, patched_text=None) + assert not mock_factory.called, "CONFLICT must not call _make_patched_source_module" + assert result2.inspect_source() == source_v2 + finally: + _clear_trace_overrides() + + +def test_codegen_synced(tmp_path): + """SYNCED: user edits match new codegen output → baseline advances, patched module returned.""" + from tilelang.tools.lower_trace.core import _wrap_codegen_ffi + + source_v1 = "// generated kernel v1\n" + mock_build = _make_mock_build(source_v1) + wrapper = _wrap_codegen_ffi(mock_build, "target.build.tilelang_cuda_without_compile") + + _setup_trace_overrides(tmp_path) + codegen_path = _core._codegen_output_path_override + original_path = codegen_path + ".original" + + with _patch_make_patched_source_module() as mock_factory: + try: + # Run 1: init + wrapper("fake_mod") + + # Edit codegen.cpp to match what the new codegen will produce + source_v2 = "// new codegen output v2\n" + with open(codegen_path, "w") as f: + f.write(source_v2) + + # Change codegen output to the same value + wrapper = _wrap_codegen_ffi(_make_mock_build(source_v2), "target.build.tilelang_cuda_without_compile") + + # Run 2: SYNCED + result2 = wrapper("fake_mod") + + # .original advanced to new codegen + with open(original_path) as f: + assert f.read() == source_v2 + + # Patched module returned (patched_text = working_text = source_v2) + assert mock_factory.called, "SYNCED should call _make_patched_source_module for _without_compile FFI" + assert result2.get_source() == source_v2 + finally: + _clear_trace_overrides() + + +def test_codegen_phase_reset_on_inspect_source_failure(tmp_path): + """_current_phase must be reset even if post-codegen tracing raises. + + Regression guard: previously an exception in the codegen post-processing + (inspect_source / file I/O / diff) left _current_phase stuck at "codegen", + misattributing later records. The exception is now caught and warned + (does not propagate), but _current_phase must still be restored. + """ + from tilelang.tools.lower_trace.core import _wrap_codegen_ffi + + class _ExplodingModule: + def inspect_source(self): + raise RuntimeError("inspect_source blew up") + + def mock_build(*args, **kwargs): + return _ExplodingModule() + + wrapper = _wrap_codegen_ffi(mock_build, "target.build.tilelang_cuda_without_compile") + + _setup_trace_overrides(tmp_path) + try: + # The post-codegen tracing exception is caught + warned (not re-raised). + wrapper("fake_mod") + + assert _core._current_phase is None, "_current_phase must be reset after inspect_source failure" + finally: + _clear_trace_overrides() + + +def test_codegen_restores_outer_phase(tmp_path): + """codegen nested in an active pipeline phase must restore it, not clear to None.""" + from tilelang.tools.lower_trace.core import _wrap_codegen_ffi + + source_v1 = "// generated kernel v1\n" + wrapper = _wrap_codegen_ffi(_make_mock_build(source_v1), "target.build.tilelang_cuda_without_compile") + + _setup_trace_overrides(tmp_path) + try: + _core._current_phase = "pipeline_test" + wrapper("fake_mod") + assert _core._current_phase == "pipeline_test", "outer phase must be restored after codegen" + finally: + _core._current_phase = None + _clear_trace_overrides() + + +def test_codegen_record_index_after_nested_pass(tmp_path): + """codegen record index must come after any pass invoked inside original_build. + + Regression guard: pre-allocating the codegen idx before ``original_build`` ran + let an internal traced pass (e.g. ``tir.transform.Simplify``) grab a later + index, so records could appear as N+1 before N. The idx is now allocated + immediately before appending the codegen record. + """ + from tilelang.tools.lower_trace.core import _wrap_codegen_ffi, LowerRecord, STATUS_COMPLETED + + nested_idx = [] + + def mock_build(*args, **kwargs): + with _core._lock: + nested_idx.append(_core._pass_index) + _core._pass_index += 1 + _core._records.append( + LowerRecord( + phase="codegen", + name="internal_simplify", + index=nested_idx[0], + before_text="", + after_text="", + changed=False, + add_lines=0, + del_lines=0, + status=STATUS_COMPLETED, + ) + ) + return _MockCodegenModule("// generated\n") + + wrapper = _wrap_codegen_ffi(mock_build, "target.build.tilelang_cuda_without_compile") + + _setup_trace_overrides(tmp_path) + try: + wrapper("fake_mod") + + codegen_records = [r for r in _core._records if r.name == "codegen"] + assert codegen_records, "no codegen record found" + assert codegen_records[-1].index > nested_idx[0], "codegen index not after nested pass" + + indices = [r.index for r in _core._records] + assert indices == sorted(indices), f"records not in ascending index order: {indices}" + finally: + _clear_trace_overrides() + + +def test_import_time_activation(tmp_path): + """The env hook in tilelang/__init__.py must activate tracing on first import. + + This module clears TL_LOWER_TRACE before importing tilelang, so the rest of + the suite never exercises the import-time activation path. Use a subprocess + to set the env var *before* the first import and verify tracing is on. + """ + import subprocess + import sys + + env = dict(os.environ) + env["TL_LOWER_TRACE"] = "1" + env["TL_LOWER_TRACE_DIR"] = str(tmp_path) + + code = ( + "import tilelang\n" + "from tilelang.tools.lower_trace import core as _core\n" + "assert _core._is_trace_enabled(), 'tracing not enabled at import time'\n" + "assert _core._get_mode() == 'html', 'expected html mode for TL_LOWER_TRACE=1'\n" + ) + result = subprocess.run( + [sys.executable, "-c", code], + env=env, + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, f"subprocess failed (rc={result.returncode}):\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}" + + +if __name__ == "__main__": + tilelang.testing.main() diff --git a/tilelang/__init__.py b/tilelang/__init__.py index a0d1a8eb16..a515e99679 100644 --- a/tilelang/__init__.py +++ b/tilelang/__init__.py @@ -214,6 +214,12 @@ def _load_tile_lang_lib(): from . import rocm as rocm # noqa: F401 from . import metal as metal # noqa: F401 + _lt_value = os.environ.get("TL_LOWER_TRACE", "") + if _lt_value and _lt_value.lower().strip() not in ("0", "false", "no", "off", ""): + from .tools.lower_trace import enable as _lower_trace_enable + + _lower_trace_enable() + del _lt_value del _lazy_load_lib # Install pass diff hook if TILELANG_PASS_DIFF is enabled (zero overhead when off) diff --git a/tilelang/tools/__init__.py b/tilelang/tools/__init__.py index 7a8bde5140..6fd4e98f41 100644 --- a/tilelang/tools/__init__.py +++ b/tilelang/tools/__init__.py @@ -1,2 +1,12 @@ from .plot_layout import plot_layout # noqa: F401 from .Analyzer import * + + +def __getattr__(name): + if name == "lower_trace": + import importlib + + mod = importlib.import_module(".lower_trace", __name__) + globals()["lower_trace"] = mod + return mod + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/tilelang/tools/lower_trace/__init__.py b/tilelang/tools/lower_trace/__init__.py new file mode 100644 index 0000000000..696978024d --- /dev/null +++ b/tilelang/tools/lower_trace/__init__.py @@ -0,0 +1,192 @@ +"""IR Lower Trace — zero-intrusion debug tool for visualizing compilation passes. + +Generates a self-contained HTML page and/or terminal diff showing all passes in the +compilation pipeline with side-by-side IR diff for each pass. + +Usage:: + + TL_LOWER_TRACE=1 python my_kernel.py # HTML report + TL_LOWER_TRACE=terminal python my_kernel.py # terminal diff only + TL_LOWER_TRACE=both python my_kernel.py # both terminal and HTML + +Programmatic API:: + + from tilelang.tools.lower_trace import lower_trace + + lower_trace(func, my_pass, mode="terminal") + lower_trace(func, [pass_a, pass_b], mode="both", html_path="diff.html") +""" + +from __future__ import annotations + +from .core import ( + enable, + disable, + reset, + LowerRecord, + STATUS_COMPLETED, + STATUS_FAILED, + STATUS_SKIPPED, + STATUS_CODEGEN, + _get_pass_display_name, +) + +__all__ = [ + "STATUS_CODEGEN", + "STATUS_COMPLETED", + "STATUS_FAILED", + "STATUS_SKIPPED", + "LowerRecord", + "disable", + "enable", + "lower_trace", + "reset", +] + + +def lower_trace( + func_or_mod, + passes, + *, + mode: str = "terminal", + context: int = 3, + html_path: str = "lower_trace_report.html", +) -> list[dict]: + """Compare IR before and after each pass in a chain. + + Parameters + ---------- + func_or_mod : PrimFunc or IRModule + The starting IR. + passes : Pass or list[Pass] or list[tuple[str, Pass]] + A single pass, a list of passes, or a list of (name, pass) pairs. + mode : {"terminal", "html", "both"} + Output mode. + context : int + Number of context lines in the unified diff (default 3). + html_path : str + Output path for HTML report (default ``lower_trace_report.html``). + + Returns + ------- + list[dict] + One entry per pass step, each containing: + ``name``, ``before_script``, ``after_script``, ``diff_lines``, + ``insertions``, ``deletions``, ``changed``. + """ + from tilelang import tvm + from .diff import unified_diff + + mode = str(mode).strip().lower() + if mode not in ("terminal", "html", "both"): + raise ValueError(f"mode must be one of 'terminal', 'html', 'both', got {mode!r}") + + if isinstance(func_or_mod, tvm.IRModule): + mod = func_or_mod + else: + mod = tvm.IRModule({"main": func_or_mod}) + + # A bare ``(name, pass)`` tuple is a single named pass, not an iterable of + # passes; wrap it so the loop below treats it as one entry. A non-list/tuple + # value (a bare Pass) is likewise wrapped. A list/tuple of passes (or of + # named pairs) is iterated as-is. + if (isinstance(passes, tuple) and len(passes) == 2 and isinstance(passes[0], str)) or not isinstance(passes, (list, tuple)): + passes = [passes] + + named_passes: list[tuple[str, object]] = [] + for p in passes: + if isinstance(p, (list, tuple)) and len(p) == 2: + named_passes.append((str(p[0]), p[1])) + else: + named_passes.append((_get_pass_display_name(p), p)) + + results: list[dict] = [] + + try: + for step_idx, (name, p) in enumerate(named_passes, 1): + before_script = mod.script() + try: + mod = p(mod) + except Exception as e: + results.append( + { + "name": name, + "before_script": before_script, + "after_script": "", + "diff_lines": [], + "insertions": 0, + "deletions": 0, + "changed": False, + "error": str(e), + } + ) + raise + + after_script = mod.script() + + diff_text = unified_diff( + before_script, + after_script, + before_label=f"step {step_idx} before", + after_label=f"step {step_idx} after", + context=context, + color=False, + ) + diff_lines = diff_text.splitlines() if diff_text else [] + + insertions = sum(1 for d in diff_lines if d.startswith("+") and not d.startswith("+++")) + deletions = sum(1 for d in diff_lines if d.startswith("-") and not d.startswith("---")) + changed = insertions > 0 or deletions > 0 + + step_result = { + "name": name, + "before_script": before_script, + "after_script": after_script, + "diff_lines": diff_lines, + "insertions": insertions, + "deletions": deletions, + "changed": changed, + } + results.append(step_result) + + if mode in ("terminal", "both"): + header = f"\n{'=' * 60}\n Pass {step_idx}: {name}\n{'=' * 60}\n" + print(header) + if changed: + colored = unified_diff( + before_script, + after_script, + before_label=f"step {step_idx} before", + after_label=f"step {step_idx} after", + context=context, + color=True, + ) + print(colored, end="") + print(f"\n >>> +{insertions} insertion(s), -{deletions} deletion(s)") + else: + print(" (no changes)") + finally: + if mode in ("html", "both"): + from .html import generate_html + + records = [] + for i, r in enumerate(results): + failed = "error" in r + records.append( + LowerRecord( + phase="lower_trace", + name=r["name"], + index=i, + before_text=r["before_script"], + after_text=r["after_script"], + changed=r["changed"], + add_lines=r["insertions"], + del_lines=r["deletions"], + status=STATUS_FAILED if failed else STATUS_COMPLETED, + error_msg=r.get("error", ""), + ) + ) + generate_html(records, html_path) + print(f"\nHTML report written to: {html_path}") + + return results diff --git a/tilelang/tools/lower_trace/core.py b/tilelang/tools/lower_trace/core.py new file mode 100644 index 0000000000..7777d18b68 --- /dev/null +++ b/tilelang/tools/lower_trace/core.py @@ -0,0 +1,1287 @@ +"""IR lower trace — zero-intrusion debug tool for visualizing compilation passes. + +Monkey-patches ``tvm.ir.transform.Pass.__call__`` and ``PassPipeline.lower`` +to automatically capture IR before/after every pass and generate diff reports. + +This module has **no dependency on ``tilelang.env``**; configuration is read +from ``os.environ`` directly, or passed programmatically via ``enable()``. + +Supports two architectures: +- New: ``PassPipeline.lower`` (each backend registers a pipeline object) +- Old: phase-based functions called from ``tilelang.engine.lower`` + +Usage:: + + TL_LOWER_TRACE=1 python my_kernel.py # HTML report + TL_LOWER_TRACE=terminal python my_kernel.py # terminal diff only + TL_LOWER_TRACE=both python my_kernel.py # both terminal and HTML +""" + +from __future__ import annotations + +import ast +import contextlib +import difflib +import dis +import functools +import inspect +import os +import re +import shutil +import sys +import threading +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from .diff import ( + _ANSI_BOLD, + _ANSI_BLUE, + _ANSI_CYAN, + _ANSI_DIM, + _ANSI_GREEN, + _ANSI_RED, + _ANSI_RESET, + _ANSI_YELLOW, +) + +if TYPE_CHECKING: + from collections.abc import Callable + +STATUS_COMPLETED = "completed" +STATUS_FAILED = "failed" +STATUS_SKIPPED = "skipped" +STATUS_CODEGEN = "codegen" + + +def _get_tvm_ffi(): + """FFI facade exposing ``get_global_func``/``register_global_func``. + + Prefers the new-style ``tvm.ffi`` (unified FFI); falls back to the legacy + ``tvm._ffi`` shipped in ``3rdparty/tvm``, where the register entry is named + ``register_func`` instead of ``register_global_func``. + """ + try: + import tvm.ffi as _ffi + + if hasattr(_ffi, "register_global_func") and hasattr(_ffi, "get_global_func"): + return _ffi + except ImportError: + pass + import tvm._ffi as _ffi + + class _LegacyFFI: + """Adapter exposing the ``register_global_func``/``get_global_func`` API over the legacy ``tvm._ffi`` (where registration is named ``register_func``).""" + + get_global_func = staticmethod(_ffi.get_global_func) + + @staticmethod + def register_global_func(name, func=None, override=False): + """Register a global FFI function via the legacy ``tvm._ffi.register_func`` entry point.""" + return _ffi.register_func(name, func, override=override) + + return _LegacyFFI() + + +def _inspect_module_source(mod): + """Return the source text of a ``tvm.runtime.Module``. + + Prefers the newer ``inspect_source``; falls back to the legacy ``get_source`` + shipped with ``3rdparty/tvm``. Returns ``""`` only when neither hook is + available, so downstream string handling stays safe. + + Exceptions raised by an *available* hook propagate to the caller — this is + required so codegen tracing can observe failures (e.g. a module whose + ``inspect_source`` raises) and reset the current phase correctly. The + fallback to ``get_source`` only happens when the preferred attribute is + missing or not callable. + """ + for _attr in ("inspect_source", "get_source"): + _fn = getattr(mod, _attr, None) + if callable(_fn): + return _fn() or "" + return "" + + +@dataclass +class LowerRecord: + """Result of running a single pass.""" + + phase: str + name: str + index: int + before_text: str + after_text: str + changed: bool + add_lines: int = 0 + del_lines: int = 0 + status: str = STATUS_COMPLETED + error_msg: str = "" + + +_records: list[LowerRecord] = [] +_section_cache: dict[tuple[str, int], str] = {} +_original_pass_call: Callable | None = None +_original_pipeline_lower: object | None = None +_original_codegen_ffis: dict[str, Callable] = {} +_legacy_patched: bool = False +# (target, attr_name, original_or_MISSING, is_dict) — restored by disable() +_legacy_phase_originals: list[tuple[object, str, object, bool]] = [] +_MISSING: object = object() + +_CODEGEN_FFI_NAMES: list[str] = [ + "target.build.tilelang_cuda", + "target.build.tilelang_cuda_without_compile", + "target.build.tilelang_cutedsl", + "target.build.tilelang_cutedsl_without_compile", + "target.build.tilelang_hip", + "target.build.tilelang_hip_without_compile", + "target.build.tilelang_metal", + "target.build.tilelang_c", + "target.build.tilelang_c_host", + "target.build.tilelang_ascend", + "target.build.tilelang_ascend_pto", + "target.build.llvm", + "target.build.webgpu", + "target.build.tilelang_cpp", + "target.build.tilelang_webgpu", +] + +# FFIs whose returned module is consumed *only* via ``get_source()`` — i.e. +# the module holds source text (C / WGSL / etc.), not a compiled binary, and is +# never passed to ``host_mod.import_module()``. For these FFIs a real +# ``CSourceModule`` rebuilt from the user-edited source can be returned in +# place of the original module so that the downstream JIT adapter (NVRTC / +# Cython / CuTeDSL / bisheng) recompiles the user-edited source. +# +# Membership is determined by tracing the call sites in +# ``tilelang.engine.lower``: +# - ``*_without_compile`` FFIs are only called from ``device_codegen_without_compile`` +# (which sets ``enable_device_compile=False`` → ``enable_host_codegen=False``), +# so the result is never passed to ``import_module``. +# - ``tilelang_cpp`` and ``tilelang_webgpu`` produce ``CSourceModule`` / ``WebGPUModule`` +# (source-only, no binary compilation step); they are likewise only called from +# ``device_codegen_without_compile``. +# - ``tilelang_metal`` is dual-use: called from both ``device_codegen`` +# (full-compile → ``import_module``) and ``device_codegen_without_compile`` +# (source-only). Since the wrapper cannot distinguish the two call sites +# at runtime, it is conservatively excluded. +# - ``llvm``, ``tilelang_c_host`` are excluded because they may be reached +# via ``host_codegen`` → ``import_module`` (the module's binary or source +# is consumed by the host runtime module tree). +# - ``tilelang_cuda``, ``tilelang_hip``, ``tilelang_cutedsl`` (full-compile +# variants) produce binary modules consumed via ``import_module``. +# - ``tilelang_ascend`` and ``tilelang_ascend_pto`` are called from +# ``device_codegen`` (not ``device_codegen_without_compile``), but +# ``BuildTileLangAscend`` returns a ``CSourceModuleCreate(code, "c", ...)`` +# — a *source-only* module, not a binary. ``tilelang.engine.lower`` only +# calls ``.get_source()`` on it and never ``import_module()``; the +# ``CythonKernelAdapter`` recompiles that source string via the ``bisheng`` +# compiler. Hence a ``CSourceModule`` rebuilt from the patched source is +# safe and user edits get recompiled. +# +# New FFIs default to *not* being in this set (conservative: return real module +# + NOTE), and must be explicitly added here once their call chain is verified +# to be source-only. +_SOURCE_ONLY_CODEGEN_FFIS: frozenset[str] = frozenset( + { + "target.build.tilelang_cuda_without_compile", + "target.build.tilelang_cutedsl_without_compile", + "target.build.tilelang_hip_without_compile", + "target.build.tilelang_c", + "target.build.webgpu", + "target.build.tilelang_cpp", + "target.build.tilelang_webgpu", + "target.build.tilelang_ascend", + "target.build.tilelang_ascend_pto", + } +) +_current_phase: str | None = None +_pass_index: int = 0 +_auto_flush: bool = False +_script_dir: str | None = None +_run_dir: str | None = None +_lock = threading.RLock() +_run_counter: int = 0 +_atexit_registered: bool = False + +_UNSET: object = object() +_mode_override: str | None | object = _UNSET +_trace_dir_override: str | None | object = _UNSET +_codegen_output_path_override: str | None | object = _UNSET + +# Phase label used for passes that run outside any PassPipeline.lower window +# (e.g. pre-pipeline module passes and tvm.build postproc), so they are still +# captured by the global Pass.__call__ hook. +_UNSCOPED_PHASE = "unscoped" + + +def _parse_lower_trace_mode(value: str | None) -> str | None: + """Parse a TL_LOWER_TRACE-style value into a mode string.""" + if value is None: + return None + v = value.lower().strip() + if v in ("", "0", "false", "no", "off"): + return None + if v in ("1", "true", "yes", "on"): + return "html" + if v in ("terminal", "html", "both"): + return v + return "html" + + +def _get_mode() -> str | None: + """Return the effective trace mode, preferring the programmatic override then the env var.""" + if _mode_override is not _UNSET: + return _mode_override # type: ignore[return-value] + return _parse_lower_trace_mode(os.environ.get("TL_LOWER_TRACE")) + + +def _is_trace_enabled() -> bool: + """Return True when tracing is currently active (mode is not None).""" + return _get_mode() is not None + + +def _should_print_terminal() -> bool: + """Return True when the terminal diff output should be produced.""" + mode = _get_mode() + return mode in ("terminal", "both") + + +def _should_gen_html() -> bool: + """Return True when the HTML report should be produced.""" + mode = _get_mode() + return mode in ("html", "both") + + +def _get_base_trace_dir() -> str: + """Return the configured base trace directory (first level).""" + if _trace_dir_override is not _UNSET and _trace_dir_override: + return _trace_dir_override # type: ignore[return-value] + return os.environ.get("TL_LOWER_TRACE_DIR") or os.path.join(".", "tmp", "lower_trace_dir") + + +def _ensure_script_dir() -> str: + """Return ``No differences.
' + if content_equal and not nl_equal: + before_nl = "present" if before_text.endswith("\n") else "absent" + after_nl = "present" if after_text.endswith("\n") else "absent" + return f'Only trailing newline differs (before: {before_nl}, after: {after_nl}).
' + + before_collapse = [True] * len(before_lines) + after_collapse = [True] * len(after_lines) + for tag, i1, i2, j1, j2 in opcodes: + if tag != "equal": + for i in range(i1, i2): + before_collapse[i] = False + for j in range(j1, j2): + after_collapse[j] = False + for tag, i1, i2, j1, j2 in opcodes: + if tag == "equal": + for k in range(min(context, i2 - i1)): + before_collapse[i1 + k] = False + after_collapse[j1 + k] = False + for k in range(min(context, i2 - i1)): + before_collapse[i2 - 1 - k] = False + after_collapse[j2 - 1 - k] = False + + rows = [] + + for tag, i1, i2, j1, j2 in opcodes: + if tag == "equal": + for i, j in zip(range(i1, i2), range(j1, j2), strict=True): + collapsed = before_collapse[i] + hidden_attr = ' class="row-hidden" data-collapse="1"' if collapsed else "" + ln_l = f'IR before this pass (execution failed):
' + '' + f'{_ir_with_ln}{_esc(rec.before_text) if rec.before_text else ""}'
+ f"This pass did not run (a previous pass failed).
' + f"{_esc(rec.before_text)}\n'
+ f'{_esc(rec.after_text)}\n'
+ f"{_esc(rec.before_text)}'
+ f'{_esc(rec.after_text)}'
+ f"This pass did not modify the IR.
' + '' + '' + f'{_ir_with_ln}{_esc(rec.after_text)}'
+ f"