mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-20 06:30:07 +08:00
Lei Wang
5c88121e3e
* Refactor pass visualizer to use PassInstrument * Address pass visualizer review feedback
300 lines
10 KiB
Python
300 lines
10 KiB
Python
# type: ignore
|
|
"""Tests for the pass_visualizer debugging tool.
|
|
|
|
Covers:
|
|
- Core helpers: build_module, inspect_structure, StructureTreePassInstrument
|
|
- Real-pipeline structure capture and per-pass diffing (build_pass_data via viewer)
|
|
- HTML / text emission (emit_html, emit_txt)
|
|
|
|
These run the CUDA lowering prologue on a small kernel. An explicit cuda target
|
|
is passed so the pipeline does not depend on auto target detection.
|
|
"""
|
|
|
|
import pytest
|
|
|
|
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
from tilelang import tvm
|
|
|
|
from tilelang.tools.pass_visualizer import (
|
|
build_module,
|
|
capture_structure,
|
|
inspect_structure,
|
|
PassStructureRecord,
|
|
StructureTreePassInstrument,
|
|
)
|
|
from tilelang.tools.pass_visualizer.viewer import (
|
|
build_pass_data,
|
|
emit_html,
|
|
emit_txt,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _gemm_relu_kernel():
|
|
"""A small fused GEMM + bias + ReLU @tilelang.jit kernel."""
|
|
|
|
@tilelang.jit(out_idx=[-1])
|
|
def gemm_relu(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float32"):
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((M, K), dtype),
|
|
B: T.Tensor((K, N), dtype),
|
|
bias: T.Tensor((N,), dtype),
|
|
C: T.Tensor((M, N), dtype),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
|
|
A_shared = T.alloc_shared((block_M, block_K), dtype)
|
|
B_shared = T.alloc_shared((block_K, block_N), dtype)
|
|
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
|
|
|
T.clear(C_local)
|
|
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
|
|
T.copy(A[by * block_M, k * block_K], A_shared)
|
|
T.copy(B[k * block_K, bx * block_N], B_shared)
|
|
T.gemm(A_shared, B_shared, C_local)
|
|
|
|
for i, j in T.Parallel(block_M, block_N):
|
|
C_local[i, j] = T.max(C_local[i, j] + bias[bx * block_N + j], 0)
|
|
|
|
T.copy(C_local, C[by * block_M, bx * block_N])
|
|
|
|
return main
|
|
|
|
return gemm_relu
|
|
|
|
|
|
def _build_small_module():
|
|
"""Build an IRModule + resolved cuda target for the small kernel."""
|
|
kernel = _gemm_relu_kernel()
|
|
func = kernel.get_tir(M=128, N=128, K=128, block_M=64, block_N=64, block_K=32)
|
|
cuda_target = tvm.target.Target("cuda", host="llvm")
|
|
return build_module(func, target=cuda_target)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Core helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_structure_tree_instrument_records_real_passes():
|
|
"""The instrument records top-level passes using their PassInfo names."""
|
|
mod, target = _build_small_module()
|
|
instrument = StructureTreePassInstrument()
|
|
|
|
with tvm.transform.PassContext(instruments=[instrument]), target:
|
|
mod = tvm.tirx.transform.BindTarget(target)(mod)
|
|
tilelang.transform.MaterializeKernelLaunch()(mod)
|
|
|
|
records = instrument.ordered_records()
|
|
assert [record.name for record in records] == ["tirx.BindTarget", "tl.MaterializeKernelLaunch"]
|
|
assert instrument.input_lines
|
|
assert any("PrimFunc" in line for line in instrument.input_lines)
|
|
|
|
|
|
def test_structure_tree_instrument_ignores_nested_passes():
|
|
"""Nested implementation passes do not duplicate linear browser stages."""
|
|
mod, _target = _build_small_module()
|
|
instrument = StructureTreePassInstrument()
|
|
|
|
@tvm.transform.module_pass(opt_level=0, name="test.Outer")
|
|
def outer_pass(mod, _ctx):
|
|
return tvm.tirx.transform.Simplify()(mod)
|
|
|
|
with tvm.transform.PassContext(instruments=[instrument]):
|
|
outer_pass(mod)
|
|
|
|
records = instrument.ordered_records()
|
|
assert [record.name for record in records] == ["test.Outer"]
|
|
|
|
|
|
def test_pass_structure_record_changed_uses_structure_snapshot():
|
|
"""The changed flag describes the structure-tree diff shown by the viewer."""
|
|
unchanged = PassStructureRecord(
|
|
name="test.Unchanged",
|
|
sequence=0,
|
|
before_lines=["PrimFunc", " body"],
|
|
after_lines=["PrimFunc", " body"],
|
|
)
|
|
changed = PassStructureRecord(
|
|
name="test.Changed",
|
|
sequence=1,
|
|
before_lines=["PrimFunc", " body"],
|
|
after_lines=["PrimFunc", " changed body"],
|
|
)
|
|
|
|
assert not unchanged.changed
|
|
assert changed.changed
|
|
|
|
|
|
def test_structure_tree_instrument_cleans_up_after_failure():
|
|
"""A missing after-pass callback is retained as a diagnostic, not stale state."""
|
|
mod, _target = _build_small_module()
|
|
instrument = StructureTreePassInstrument()
|
|
|
|
@tvm.transform.module_pass(opt_level=0, name="test.Fail")
|
|
def failing_pass(_mod, _ctx):
|
|
raise RuntimeError("expected failure")
|
|
|
|
with pytest.raises(RuntimeError, match="expected failure"), tvm.transform.PassContext(instruments=[instrument]):
|
|
failing_pass(mod)
|
|
|
|
assert instrument.records == []
|
|
assert instrument.incomplete_passes == ["test.Fail"]
|
|
|
|
# Entering a new context resets the incomplete frame and records normally.
|
|
with tvm.transform.PassContext(instruments=[instrument]):
|
|
tvm.tirx.transform.Simplify()(mod)
|
|
assert [record.name for record in instrument.ordered_records()] == ["tirx.Simplify"]
|
|
assert instrument.incomplete_passes == []
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_inspect_structure_renders_tree(capsys):
|
|
"""inspect_structure prints the SBlock tree with tile ops expanded by field."""
|
|
mod, _target = _build_small_module()
|
|
inspect_structure(mod)
|
|
out = capsys.readouterr().out
|
|
|
|
assert "PrimFunc" in out
|
|
assert "SBlock" in out
|
|
# gemm is expanded by field name, not printed as one positional line.
|
|
assert "T.gemm" in out
|
|
assert "a_region" in out
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_capture_structure_returns_lines():
|
|
"""capture_structure turns inspect_structure output into text lines."""
|
|
mod, _target = _build_small_module()
|
|
lines = capture_structure(mod)
|
|
|
|
assert isinstance(lines, list)
|
|
assert len(lines) > 0
|
|
assert any("PrimFunc" in ln for ln in lines)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# build_pass_data + emission (file-driven)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_build_pass_data_and_emit(tmp_path):
|
|
"""End-to-end: observe the real prologue and emit HTML/txt."""
|
|
import os
|
|
|
|
kernel_path = os.path.join(
|
|
os.path.dirname(tilelang.tools.pass_visualizer.__file__),
|
|
"examples",
|
|
"gemm_relu.py",
|
|
)
|
|
with open(kernel_path) as f:
|
|
source = f.read()
|
|
|
|
kwargs = {"M": 128, "N": 128, "K": 128, "block_M": 64, "block_N": 64, "block_K": 32}
|
|
name, stages = build_pass_data(kernel_path, None, "auto", kwargs, source)
|
|
|
|
assert name == "gemm_relu"
|
|
# source + (input) + at least one pass.
|
|
assert len(stages) >= 3
|
|
for st in stages:
|
|
assert "name" in st and "flag" in st and "rows" in st
|
|
flags = {st["flag"] for st in stages}
|
|
assert "source" in flags
|
|
assert "input" in flags
|
|
stage_names = [st["name"].split("] ", 1)[-1] for st in stages[2:]]
|
|
assert "LayoutInference" in stage_names
|
|
assert "LowerTileOp" in stage_names
|
|
assert "AddWrapperForSingleBufStore" in stage_names
|
|
assert "DecoupleTypeCast" in stage_names
|
|
assert "pass_fn" not in stage_names
|
|
|
|
html = emit_html(name, stages)
|
|
assert "Pass browser" in html
|
|
assert "gemm_relu" in html
|
|
|
|
txt = emit_txt(name, stages)
|
|
assert "kernel: gemm_relu" in txt
|
|
assert "T.gemm" in txt
|
|
|
|
|
|
def test_build_pass_data_rejects_non_cuda_target():
|
|
"""The focused viewer rejects backends whose pass pipeline it cannot dispatch."""
|
|
import os
|
|
|
|
kernel_path = os.path.join(
|
|
os.path.dirname(tilelang.tools.pass_visualizer.__file__),
|
|
"examples",
|
|
"gemm_relu.py",
|
|
)
|
|
with open(kernel_path) as f:
|
|
source = f.read()
|
|
|
|
kwargs = {"M": 128, "N": 128, "K": 128, "block_M": 64, "block_N": 64, "block_K": 32}
|
|
with pytest.raises(ValueError, match="currently supports only CUDA targets"):
|
|
build_pass_data(kernel_path, None, "llvm", kwargs, source)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_build_pass_data_observes_canonical_prologue(monkeypatch):
|
|
"""A newly introduced real pipeline pass appears without a viewer pass list."""
|
|
import os
|
|
|
|
kernel_path = os.path.join(
|
|
os.path.dirname(tilelang.tools.pass_visualizer.__file__),
|
|
"examples",
|
|
"gemm_relu.py",
|
|
)
|
|
with open(kernel_path) as f:
|
|
source = f.read()
|
|
|
|
@tvm.transform.module_pass(opt_level=0, name="test.InjectedPipelinePass")
|
|
def injected_pass(mod, _ctx):
|
|
return mod
|
|
|
|
def injected_prologue(mod, _target):
|
|
return injected_pass(mod)
|
|
|
|
monkeypatch.setattr("tilelang.tools.pass_visualizer.viewer.CUDAPassPipelineBodyPrologue", injected_prologue)
|
|
kwargs = {"M": 128, "N": 128, "K": 128, "block_M": 64, "block_N": 64, "block_K": 32}
|
|
_name, stages = build_pass_data(kernel_path, None, "cuda", kwargs, source)
|
|
|
|
assert [stage["name"] for stage in stages] == ["source code", "(input)", "[01] InjectedPipelinePass"]
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_build_pass_data_respects_jit_pass_configs(tmp_path):
|
|
"""Conditional stages use the analyzed JITImpl's real PassContext config."""
|
|
import os
|
|
|
|
kernel_path = os.path.join(
|
|
os.path.dirname(tilelang.tools.pass_visualizer.__file__),
|
|
"examples",
|
|
"gemm_relu.py",
|
|
)
|
|
with open(kernel_path) as f:
|
|
source = f.read()
|
|
source = source.replace(
|
|
"@tilelang.jit(out_idx=[-1])",
|
|
'@tilelang.jit(out_idx=[-1], pass_configs={"tl.disable_warp_specialized": True})',
|
|
)
|
|
configured_path = tmp_path / "gemm_relu_no_ws.py"
|
|
configured_path.write_text(source)
|
|
|
|
kwargs = {"M": 128, "N": 128, "K": 128, "block_M": 64, "block_N": 64, "block_K": 32}
|
|
_name, stages = build_pass_data(str(configured_path), None, "cuda", kwargs, source)
|
|
|
|
stage_names = [st["name"] for st in stages]
|
|
assert not any("ProducerConsumerWarpSpecialized" in name for name in stage_names)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|