[Enhancement] Improve plot_layout visualization for Layouts

examples/plot_layout/layout_swizzle.py

@@ -0,0 +1,32 @@
+from tilelang.layout import (
+    make_full_bank_swizzled_layout,
+    make_half_bank_swizzled_layout,
+    make_quarter_bank_swizzled_layout,
+)
+from tilelang.tools import plot_layout
+
+element_size = 16  # float16 = 16 bits
+
+
+# ---- Plot the swizzle patterns ----
+
+# 1. Quarter-bank (32B) — 1-bit XOR — 8x16
+# Rows 0-3: identity; Rows 4-7: two 8-element halves swap
+layout = make_quarter_bank_swizzled_layout(8, 16, element_size)
+print(f"Quarter-bank swizzle (8x16, fp16): {layout}")
+plot_layout(layout, name="swizzle_quarter_8x16")
+
+# 2. Half-bank (64B) — 2-bit XOR — 8x32
+layout = make_half_bank_swizzled_layout(8, 32, element_size)
+print(f"Half-bank swizzle (8x32, fp16): {layout}")
+plot_layout(layout, name="swizzle_half_8x32")
+
+# 3. Full-bank (128B) — 3-bit XOR — 8x64
+layout = make_full_bank_swizzled_layout(8, 64, element_size)
+print(f"Full-bank swizzle (8x64, fp16): {layout}")
+plot_layout(layout, name="swizzle_full_8x64")
+
+# 4. Full-bank (128B) — multi-tile: 32x64
+layout = make_full_bank_swizzled_layout(32, 64, element_size)
+print(f"Full-bank swizzle (32x64, fp16): {layout}")
+plot_layout(layout, name="swizzle_full_32x64")

examples/plot_layout/layout_transform.py

@@ -0,0 +1,24 @@
+import tilelang.language as T
+from tilelang.tools import plot_layout
+
+# --- Example 1: Simple 2D Transpose (4x4) ---
+transpose_layout = T.Layout([4, 4], lambda i, j: (j, i))
+print("Transpose 4x4:", transpose_layout)
+plot_layout(transpose_layout, name="transpose_4x4")
+
+# --- Example 2: Larger Transpose (8x8) ---
+transpose_8x8 = T.Layout([8, 8], lambda i, j: (j, i))
+print("Transpose 8x8:", transpose_8x8)
+plot_layout(transpose_8x8, name="transpose_8x8")
+
+# --- Example 3: 3D → 2D reshape + transpose ---
+# (i, j, k) with shape [2, 4, 8] → (k, i*4+j)
+reshape_layout = T.Layout([2, 4, 8], lambda i, j, k: (k, i * 4 + j))
+print("Reshape 3D [2,4,8] -> [8,8]:", reshape_layout)
+plot_layout(reshape_layout, name="reshape_3d_to_2d")
+
+# --- Example 4: Interleave layout ---
+# Even rows from first half, odd rows from second half
+interleave = T.Layout([8, 4], lambda i, j: (i % 4 * 2 + i // 4, j))
+print("Interleave [8,4]:", interleave)
+plot_layout(interleave, name="interleave_8x4")