[Tool] Provide layout visualization tool

examples/visual_layout_inference/visual_layout_inference.py

@@ -0,0 +1,56 @@
+import tilelang
+import tilelang.language as T
+
+
+# use pass_configs to enable layout visualization
+@tilelang.jit(out_idx=[-1], pass_configs={tilelang.PassConfigKey.TL_ENABLE_LAYOUT_VISUAL: True})
+def matmul(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float"):
+
+    @T.prim_func
+    def gemm(
+            A: T.Tensor((M, K), dtype),
+            B: T.Tensor((K, 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)
+
+            T.copy(C_local, C[by * block_M, bx * block_N])
+
+    return gemm
+
+
+def main():
+    kernel = matmul(128, 128, 128, 32, 32, 32)
+
+    import torch
+
+    a = torch.randn(128, 128).cuda().half()
+    b = torch.randn(128, 128).cuda().half()
+
+    c = kernel(a, b)
+
+    ref_c = a @ b
+
+    torch.testing.assert_close(c, ref_c, rtol=1e-2, atol=1e-2)
+    print("All check passed.")
+
+    # print the layout visualization result and save figures to ./tmp.
+    '''
+    C_local layout inference:
+    Shape: [32, 32] -> [8]
+    Thread: _j // 16 * 64 + _i // 16 * 32 + _i % 8 * 4 + _j % 8 // 2
+    Index:  [_j % 16 // 8 * 4 + _i % 16 // 8 * 2 + _j % 2]
+    '''
+
+
+if __name__ == "__main__":
+    main()

src/op/builtin.cc

@@ -34,6 +34,7 @@ TVM_REGISTER_PASS_CONFIG_OPTION(kDisableVectorize256, Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION(kDisableWGMMA, Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION(kDisableShuffleElect, Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION(kStorageRewriteDetectInplace, Bool);
+TVM_REGISTER_PASS_CONFIG_OPTION(kEnableLayoutVisual, Bool);
 
 DataType cuTensorMapType() { return DataType::UInt(8, 128); }
 

src/op/builtin.h

@@ -51,6 +51,7 @@ static constexpr const char *kDisableWGMMA = "tl.disable_wgmma";
 static constexpr const char *kDisableShuffleElect = "tl.disable_shuffle_elect";
 static constexpr const char *kStorageRewriteDetectInplace =
     "tl.storage_rewrite_detect_inplace";
+static constexpr const char *kEnableLayoutVisual = "tl.enable_layout_visual";
 /*!
  * \brief Whether to disable dynamic tail split
  *

tilelang/__init__.py

@@ -137,6 +137,7 @@ def _load_tile_lang_lib():
     transform,  # noqa: F401
     language,  # noqa: F401
     engine,  # noqa: F401
+    tools,  # noqa: F401
 )
 from .autotuner import autotune  # noqa: F401
 from .transform import PassConfigKey  # noqa: F401

tilelang/engine/lower.py

@@ -16,6 +16,7 @@
 from tilelang.engine.param import KernelParam, CompiledArtifact
 from tilelang.utils.target import determine_target
 from tilelang.engine.phase import (
+    LayoutVisual,
     PreLowerSemanticCheck,
     LowerAndLegalize,
     OptimizeForTarget,
@@ -249,6 +250,9 @@ def lower(
     # Phase 1: Lower and legalize the IR
     mod = LowerAndLegalize(mod, target)
 
+    # Visualize the layout
+    LayoutVisual(mod)
+
     # Phase 2: Optimize the IR for the target
     mod = OptimizeForTarget(mod, target)
 

tilelang/engine/phase.py

@@ -67,6 +67,19 @@ def should_force_let_inline(pass_ctx: PassContext | None = None) -> bool:
     return bool(pass_ctx and pass_ctx.config.get(tilelang.PassConfigKey.TL_FORCE_LET_INLINE, False))
 
 
+def should_enable_layout_visual(pass_ctx: PassContext | None = None) -> bool:
+    if pass_ctx is None:
+        pass_ctx = tilelang.transform.get_pass_context()
+    return bool(
+        pass_ctx and pass_ctx.config.get(tilelang.PassConfigKey.TL_ENABLE_LAYOUT_VISUAL, False))
+
+
+def LayoutVisual(mod: IRModule) -> None:
+    """Apply layout visualization pass if enabled."""
+    if should_enable_layout_visual():
+        tilelang.tools.LayoutVisual()(mod)
+
+
 def PreLowerSemanticCheck(mod: IRModule) -> None:
     """
     Check whether the module is valid before lowering. If not, raise a user-friendly error

tilelang/tools/__init__.py

@@ -1,2 +1,3 @@
 from .plot_layout import plot_layout  # noqa: F401
 from .Analyzer import *
+from .layout_visual import LayoutVisual  # noqa: F401

tilelang/tools/layout_visual.py

@@ -0,0 +1,50 @@
+import tilelang.language as T
+from tvm import tir
+from tvm.tir import PyStmtExprVisitor
+
+from tvm.tir.transform import prim_func_pass
+from tilelang.tools.plot_layout import plot_layout
+
+
+def print_layout_format(layout: T.Fragment) -> str:
+    input_shape = layout.get_input_shape()
+    output_shape = layout.get_output_shape()
+    lines = [
+        f"  Shape: {input_shape} -> {output_shape}", f"  Thread: {layout.forward_thread}",
+        f"  Index:  {layout.forward_index}"
+    ]
+
+    return "\n".join(lines)
+
+
+@tir.functor.visitor
+class _LayoutVisualVisitor(PyStmtExprVisitor):
+
+    def __init__(self):
+        super().__init__()
+        self.layout_found = []
+        self.processed_layouts = set()
+
+    def visit_block_(self, op: tir.Block) -> None:
+        if "layout_map" in op.annotations:
+            layout_map = op.annotations["layout_map"]
+
+            for key, layout in layout_map.items():
+                if isinstance(layout, T.Fragment):
+                    layout_id = str(layout)
+                    if layout_id not in self.processed_layouts:
+                        print(f"{key} layout inference:")
+                        print(print_layout_format(layout))
+                        plot_layout(layout, name=f"{key}_layout")
+                        self.processed_layouts.add(layout_id)
+
+        self.visit_stmt(op.body)
+
+
+def LayoutVisual():
+
+    def pass_fn(func: tir.PrimFunc, mod, ctx):
+        _LayoutVisualVisitor().visit_stmt(func.body)
+        return func
+
+    return prim_func_pass(pass_fn, opt_level=0)

tilelang/tools/plot_layout.py

@@ -1,7 +1,7 @@
 import tilelang.language as T
 
 
-def plot_layout(layout: T.Layout,
+def plot_layout(layout: T.Fragment,
                 save_directory="./tmp",
                 name: str = "layout",
                 colormap: str = "RdPu",
@@ -82,6 +82,12 @@ def plot_layout(layout: T.Layout,
     raw_colors = [cmap(i) for i in range(num_threads)]
     colors = raw_colors.copy()
 
+    # Show the distribution of registers in each thread of a warp.
+    warp_size = 32
+    spectral_camp = plt.get_cmap("hsv", warp_size * 6)
+    for i in range(warp_size):
+        colors[i] = spectral_camp(i * 6)
+
     # Determine the number of rows and columns in the input shape
     nrows, ncols = input_shape
     # Adjust figure size to maintain square cells

tilelang/transform/pass_config.py

@@ -69,6 +69,9 @@ class PassConfigKey(str, Enum):
     TL_FORCE_LET_INLINE = "tl.force_let_inline"
     """Force TileLang to inline let bindings during simplification. Default: False"""
 
+    TL_ENABLE_LAYOUT_VISUAL = "tl.enable_layout_visual"
+    """Enable layout inference visualization. Default: False"""
+
     TL_STORAGE_REWRITE_DETECT_INPLACE = "tl.storage_rewrite_detect_inplace"
     """Control StorageRewrite inplace detection.