[Tool] Provide layout visualization tool

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.