[Tool] Provide layout visualization tool

docs/tutorials/debug_tools_for_tilelang.md

@@ -171,6 +171,30 @@ The output messages will include something like:
 msg='hello world' BlockIdx=(0, 0, 0), ThreadIdx=(0, 0, 0): 0
 ```
 
+### Visual Layout Inference For TileLang
+ The **Visual Layout Inference** tool automatically generates visual diagrams that illustrate the mapping between logical indices, thread IDs, and register file locations.
+
+When TileLang performs layout inference, it determines how fragment buffers are distributed across threads. The visual layout tool captures this information and generates:
+1. **Textual output**: A human-readable description of the layout mapping
+2. **Visual diagrams**: Color-coded plots showing the thread-to-data mapping
+
+The visual layout inference tool is controlled through the `TL_ENABLE_LAYOUT_VISUALIZATION` pass configuration. By default, visualization is **disabled** to avoid performance overhead during compilation.
+
+`TL_ENABLE_LAYOUT_VISUALIZATION` accepts string values to control output formats:
+- "True" or "all": Enabled, generates all formats (PDF, PNG, SVG)
+- "png": Generate PNG format only
+- "pdf": Generate PDF format only
+- "svg": Generate SVG format only
+
+The output messages will include something like:
+```
+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]
+```
+
+
 ## Conclusion
 
 By carefully examining intermediate representations (IR) before final code generation—and by leveraging runtime printing through `T.print`—one can quickly diagnose where index calculations, copy logic, or other kernel operations deviate from the intended behavior. This two-pronged approach (inspecting IR transformations and using runtime prints) is often sufficient for resolving generation and correctness issues in TileLang programs.

examples/visual_layout_inference/visual_layout_inference.py

@@ -0,0 +1,59 @@
+import tilelang
+import tilelang.language as T
+
+tilelang.disable_cache()
+
+
+# use pass_configs to enable layout visualization
+@tilelang.jit(
+    out_idx=[-1], pass_configs={tilelang.PassConfigKey.TL_ENABLE_LAYOUT_VISUALIZATION: "False"})
+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()

pyproject.toml

@@ -42,6 +42,8 @@ dependencies = [
 # mldtypes should be greater than 0.5.1
 # if you want to enable fp4
 fp4 = ["ml-dtypes>=0.5.1"]
+# if you want to enable layout inference visualization
+vis = ["matplotlib"]
 
 [build-system]
 requires = ["cython>=3.0.0", "scikit-build-core"]

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(kEnableLayoutVisualization, ffi::String);
 
 DataType cuTensorMapType() { return DataType::UInt(8, 128); }
 

src/op/builtin.h

@@ -51,6 +51,8 @@ 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 *kEnableLayoutVisualization =
+    "tl.enable_layout_visualization";
 /*!
  * \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,26 @@ 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()
+
+    config_value = pass_ctx.config.get(tilelang.PassConfigKey.TL_ENABLE_LAYOUT_VISUALIZATION.value,
+                                       "")
+
+    if config_value is None:
+        return False
+
+    config_str = str(config_value).strip().lower()
+    return bool(config_str and config_str != "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,64 @@
+import tilelang
+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, formats: str = "png"):
+        super().__init__()
+        self.layout_found = []
+        self.processed_layouts = set()
+        self.formats = formats
+
+    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", formats=self.formats)
+                        self.processed_layouts.add(layout_id)
+
+        self.visit_stmt(op.body)
+
+
+def LayoutVisual():
+
+    def pass_fn(func: tir.PrimFunc, mod, ctx):
+        pass_ctx = tilelang.transform.get_pass_context()
+        config_value = pass_ctx.config.get(
+            tilelang.PassConfigKey.TL_ENABLE_LAYOUT_VISUALIZATION.value)
+
+        config_str = str(config_value).strip().lower()
+        if not config_str or config_str == "false":
+            return func
+        elif config_str == "true":
+            formats = "all"
+        else:
+            formats = config_str
+
+        _LayoutVisualVisitor(formats=formats).visit_stmt(func.body)
+        return func
+
+    return prim_func_pass(pass_fn, opt_level=0)

tilelang/tools/plot_layout.py

@@ -1,11 +1,13 @@
+from __future__ import annotations
 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",
-                verbose: bool = False) -> None:
+                verbose: bool = False,
+                formats: str | list[str] = "png") -> None:
     """
     Plot the layout of a buffer.
 
@@ -21,9 +23,10 @@ def plot_layout(layout: T.Layout,
         The colormap to use for visualization (default is "RdPu").
     verbose : bool, optional
         If True, prints additional information about the mapping (default is False).
-
+    formats : str | list[str], optional
+        The formats to save the image in (default is "png").
     Returns
-    -------
+    -------s
     None
     """
     import os
@@ -82,6 +85,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
@@ -191,17 +200,30 @@ def plot_layout(layout: T.Layout,
     # Save the figure in multiple formats
     plt.tight_layout()
 
-    # Save as PDF
-    pdf_path = tmp_directory / f"{name}.pdf"
-    plt.savefig(pdf_path, bbox_inches="tight")
-    print(f"Saved pdf format into {pdf_path}")
-
-    # Save as PNG
-    png_path = tmp_directory / f"{name}.png"
-    plt.savefig(png_path, bbox_inches="tight", transparent=False, dpi=255)
-    print(f"Saved png format into {png_path}")
-
-    # Save as SVG
-    svg_path = tmp_directory / f"{name}.svg"
-    plt.savefig(svg_path, bbox_inches="tight", format="svg")
-    print(f"Saved svg format into {svg_path}")
+    if isinstance(formats, str):
+        formats_str = formats.strip().lower()
+        if formats_str == 'all':
+            formats_list = ['pdf', 'png', 'svg']
+        elif "," in formats_str:
+            formats_list = [f.strip() for f in formats_str.split(',')]
+        else:
+            formats_list = [formats_str]
+    else:
+        raise TypeError(f"Expected str, but got {type(formats).__name__}. "
+                        f"Please pass a string like 'png', 'pdf', 'svg', 'all', or 'png,pdf'.")
+
+    # Save the figure
+    if 'pdf' in formats_list:
+        pdf_path = tmp_directory / f"{name}.pdf"
+        plt.savefig(pdf_path, bbox_inches="tight")
+        print(f"Saved pdf format into {pdf_path}")
+
+    if 'png' in formats_list:
+        png_path = tmp_directory / f"{name}.png"
+        plt.savefig(png_path, bbox_inches="tight", transparent=False, dpi=255)
+        print(f"Saved png format into {png_path}")
+
+    if 'svg' in formats_list:
+        svg_path = tmp_directory / f"{name}.svg"
+        plt.savefig(svg_path, bbox_inches="tight", format="svg")
+        print(f"Saved svg format into {svg_path}")

tilelang/transform/pass_config.py

@@ -69,6 +69,14 @@ 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_VISUALIZATION = "tl.enable_layout_visualization"
+    """Enable layout inference visualization. Accepts string values:
+    - "" or "false": disabled (default)
+    - "true" or "all": enabled, generate all formats (pdf, png, svg)
+    - "png", "pdf", "svg": enabled, generate specified format
+    - "png,svg": enabled, generate multiple formats (comma-separated)
+    """ ""
+
     TL_STORAGE_REWRITE_DETECT_INPLACE = "tl.storage_rewrite_detect_inplace"
     """Control StorageRewrite inplace detection.