mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-12 05:10:18 +08:00
Lei Wang
d85e0c6e96
* Enhance plot_layout function to support both Fragment and Layout types for visualization. Update parameters for colormap and formats, and introduce helper functions for format parsing and saving plots. Improve documentation for clarity on usage and expected input types. * lint fix * Refactor swizzle layout functions to use dedicated layout creators. Replace inline 2D swizzle functions with calls to `make_full_bank_swizzled_layout`, `make_half_bank_swizzled_layout`, and `make_quarter_bank_swizzled_layout` for improved clarity and maintainability in layout generation. * Remove outdated documentation from layout_swizzle.py and ensure plots are closed after saving in plot_layout.py for better resource management.
25 lines
980 B
Python
25 lines
980 B
Python
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")
|