[Tool] Add tool to print fragment in thread value view

[kurisu6912]

This pr add a tool to plot fragment layout in thread and index view.

Summary by CodeRabbit

Release Notes

• New Features
  • Added a fragment mapping visualization tool showing source-to-(thread,local) mappings as color-coded grids with text annotations.
  • Supports replicated fragments and customizable index labels for clearer mapping views.
  • Exportable to multiple image formats and configurable output directories.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

📝 Walkthrough

Walkthrough

Added a new plotting helper plot_fragment_tv to tilelang/tools/plot_layout.py that renders a fragment's source-index → (thread, local) coordinate mapping as annotated, color-coded 2D grids; supports replication, multiple output formats, and optional file output. Also added itertools import.

Changes

Cohort / File(s)	Summary
New Fragment Visualization Function tilelang/tools/plot_layout.py	Added plot_fragment_tv(...): builds source index arrays, calls fragment forward-mapping methods (thread/local indices), generates string labels, renders color-coded matplotlib grids with annotations, supports replication and multiple output formats, and writes outputs when a path is provided. Also added top-level itertools import.

Sequence Diagram

sequenceDiagram
    participant Caller
    participant plot_fragment_tv
    participant Fragment
    participant NumPy
    participant Matplotlib
    participant FileSystem

    Caller->>plot_fragment_tv: invoke plotting helper (frag, options)
    plot_fragment_tv->>Fragment: query fragment dims / replication
    plot_fragment_tv->>NumPy: construct src index grids (cartesian product)
    NumPy-->>plot_fragment_tv: src_flat_idx / src_idx arrays
    plot_fragment_tv->>Fragment: map_forward_thread / map_forward_index for each src idx
    Fragment-->>plot_fragment_tv: (thread, local) coordinates
    plot_fragment_tv->>NumPy: assemble label matrix and color mapping
    plot_fragment_tv->>Matplotlib: render annotated grid (figure/axes)
    Matplotlib-->>plot_fragment_tv: figure
    plot_fragment_tv->>FileSystem: create output dir (if needed) and save file(s)
    FileSystem-->>Caller: written image(s)

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• [Tool] Provide layout visualization tool #1353: Similar changes to tilelang/tools/plot_layout.py introducing Fragment-based plotting utilities.

Suggested reviewers

• LeiWang1999
• oraluben

Poem

🐰
I hop through indices, row by row,
Mapping threads where data flows,
Colors, labels, grids in sight,
Fragments plotted, neat and light,
A little rabbit cheers—plot bright!

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The pull request title clearly and specifically describes the main change: adding a tool to visualize fragments in thread-value view, which directly matches the new plot_fragment_tv function added to the codebase.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 4

🤖 Fix all issues with AI agents

In `@tilelang/tools/plot_layout.py`:
- Around line 234-235: The docstring for the parameter save_directory is
incorrect (says default "./tmp" while the function signature uses None); update
the docstring in plot_layout.py to reflect the actual default of save_directory
(None) or change the function signature to match the documented default—modify
the parameter description for save_directory in the relevant function/method's
docstring so it accurately states "default is None" (or adjust the default in
the signature to "./tmp" if that was intended) and ensure the text mentions the
actual behavior used by the code.
- Around line 280-285: The loop can divide by zero when mx =
np.max(src_flat_idx); modify the normalization before calling cmap: compute a
safe denominator (e.g., denom = mx if mx > 0 else 1) and use
cmap(src_flat_idx[i, j] / denom) so values are well-defined when src_flat_idx is
all zeros; update the code around variables mx, src_flat_idx, cmap, and the loop
over num_local_dim/num_thread_dim that calls plt.text with src_idx_str
accordingly.
- Around line 266-275: The loop-local variable "name" in plot_layout.py shadows
the function parameter "name" (the output filename), causing the filename to be
overwritten; rename the loop variable to something non-conflicting (e.g.,
"idx_name" or "idx_label") in both occurrences inside the two loops where it's
assigned (the list comprehension building the index string and the identical
code in the else branch), and update any subsequent uses (src_idx_str[...]=...)
to use the new variable so the function parameter "name" remains the filename.
- Around line 290-297: The figure created by this function is never closed when
save_directory is provided, which can leak memory; after the loop that saves
formats (the block using save_directory, save_dir, formats and
plt.tight_layout()), call plt.close() (or plt.close(fig) if a Figure object
named fig is available) to release the figure, or alternatively return the
Figure object so the caller can close it; update the save block in
plot_layout.py to close the figure after saving.

🧹 Nitpick comments (4)

tilelang/tools/plot_layout.py (4)

3-3: Redundant import of itertools at line 49.

The module-level import here makes the local import itertools inside plot_layout (line 49) redundant. Consider removing the local import for consistency.

---

283-283: Unused variable a from tuple unpacking.

The alpha channel a is unpacked but never used.

♻️ Proposed fix: prefix with underscore

-            r, g, b, a = cmap(src_flat_idx[i, j] / mx)
+            r, g, b, _a = cmap(src_flat_idx[i, j] / mx)

---

264-265: Prefer spread syntax over list concatenation.

Per static analysis (RUF005), [rep, *list(item)] is cleaner than [rep] + list(item).

♻️ Proposed fix

-                th = frag.map_forward_thread([rep] + list(item))[0].value
-                dst_idx = frag.map_forward_index([rep] + list(item))[0].value
+                th = frag.map_forward_thread([rep, *item])[0].value
+                dst_idx = frag.map_forward_index([rep, *item])[0].value

---

295-296: Inconsistent formats handling compared to plot_layout.

The existing plot_layout function supports "all" and comma-separated format strings (e.g., "png,pdf"), while this function only handles plain strings or lists. Consider aligning the behavior for consistency.