[Tool] Provide layout visualization tool

[Cunxiao2002]

Summary by CodeRabbit

• New Features

  • Added layout visualization tooling with configurable output formats (png, pdf, svg, or all).
  • Added a global configuration key to enable and control layout visualization during compilation.

• Documentation

  • Added a debugging guide describing layout visualization usage and outputs.
  • Added an example demonstrating layout visualization for a GEMM kernel.

• Chores

  • Added optional visualization dependency (matplotlib).

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

[coderabbitai[bot]]

[!NOTE]

Other AI code review bot(s) detected

CodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review.

Walkthrough

A new layout visualization feature is added to TileLang, enabling visual and textual output of inferred layouts during compilation. Pass config options control enablement and output formats, with new Python API functions querying and applying visualization. A core analysis pass scans IR for layout annotations and generates plots via an enhanced plot_layout utility supporting multiple formats (txt, png, pdf, svg).

Changes

Cohort / File(s)	Summary
Pass Configuration src/op/builtin.cc, src/op/builtin.h, tilelang/transform/pass_config.py	Added two new pass config options: TL_LAYOUT_VISUALIZATION_ENABLE (Bool, default False) and TL_LAYOUT_VISUALIZATION_FORMATS (String, accepts "txt", "png", "pdf", "svg", "all", or comma-separated values). Registered in both C++ and Python config key enums.
Core Visualization Analysis tilelang/analysis/layout_visual.py	New analysis module implementing print_fragment_format() helper and _LayoutVisualVisitor pass that scans IR blocks for layout_map annotations, prints layouts, and generates plots. Provides public LayoutVisual(formats) entry point registered as a TVM prim_func_pass.
Engine Phase Integration tilelang/engine/phase.py	Added three new public functions: should_enable_layout_visual() (reads config, defaults to False), get_layout_visual_formats() (reads config, validates against allowed formats, expands "all", defaults to ["txt"], supports comma-separated lists), and LayoutVisual() wrapper. Modified LowerAndLegalize() to invoke visualization post-LayoutInference().
Plot Layout Utility tilelang/tools/plot_layout.py	Updated plot_layout() signature from layout: T.Layout to layout: T.Fragment, added formats parameter (str | list[str], default "png"). Introduced warp-size-based color augmentation with HSV spectrum, conditional multi-format saves (pdf/png/svg), format validation with TypeError for non-string inputs.
Public Exports tilelang/__init__.py, tilelang/analysis/__init__.py	Exposed tools module in tilelang package; added LayoutVisual export from tilelang.analysis.
Documentation & Examples docs/tutorials/debug_tools_for_tilelang.md, examples/visual_layout_inference/visual_layout_inference.py	Added new tutorial section describing layout visualization tool, configuration, supported formats, and example outputs. Introduced example GEMM kernel demonstrating visualization via pass_configs with PyTorch validation.
Dependencies pyproject.toml	Added optional dependency group vis with matplotlib.

Sequence Diagram

sequenceDiagram
    participant Compiler as Compilation Pipeline
    participant Phase as LowerAndLegalize
    participant Detector as should_enable_layout_visual
    participant Formatter as get_layout_visual_formats
    participant Visitor as LayoutVisual Pass
    participant IR as IR Module
    participant Plotter as plot_layout
    
    Compiler->>Phase: invoke LowerAndLegalize(mod)
    Phase->>Phase: LayoutInference(mod)
    Phase->>Detector: check enable?
    Detector-->>Phase: True/False
    
    alt Visualization Enabled
        Phase->>Formatter: get formats
        Formatter-->>Phase: ["png", "pdf", ...] or ["txt"]
        Phase->>Visitor: LayoutVisual(mod)
        Visitor->>IR: scan for layout_map annotations
        loop For each Fragment layout
            Visitor->>Visitor: print_fragment_format(layout)
            Visitor->>Plotter: plot_layout(layout, formats)
            Plotter-->>Plotter: save txt/png/pdf/svg
        end
    else Visualization Disabled
        Phase->>Phase: continue (no-op)
    end

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~22 minutes

Areas requiring close attention:

• Format validation and normalization logic in get_layout_visual_formats() — ensure comma-separated parsing, case-handling, and "all" expansion are correct
• Layout tracking mechanism in _LayoutVisualVisitor to prevent duplicate processing of layouts
• Consistency between C++ pass config registration and Python config key enum definitions
• Integration between plot_layout() signature change and existing callers; HSV color augmentation logic for large thread counts
• Example verification: GEMM kernel correctness and visualization output plausibility

Poem

🐰 A layout blooms in colors bright,
With visualizations shining light,
From fragments deep to formats free—
Plot, txt, pdf, for all to see!
The compiler's eye now sees the way,
Layout visualization saves the day! 🎨✨

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 30.43% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'Provide layout visualization tool' directly corresponds to the PR's primary objective of adding a layout visualization feature across multiple files and modules.

✨ 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.}

[github-actions[bot]]

👋 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! 🚀

[LeiWang1999]

@Cunxiao2002 thanks, would you mind providing an example for us to reproduce the tool.

[Cunxiao2002]

@LeiWang1999 ok, I will.

[oraluben]

I didn't know it's related to tvm's type🤔it seems not to make much sense to have string True/False

[Cunxiao2002]

@oraluben It seems pass_config can also accept a parameter of a specific type. In that case, I think we should stop using True or False for the check and just pass the image type directly as the argument.

[Cunxiao2002]

@oraluben fixed==

[LeiWang1999]

@codex review