[Enhancement][Tool] Tree-style pretty ASTPrinter

[SiriusNEO]

Previously we add a trivial ASTPrinter for debugging purpose. This PR improves it to a formal tool in TileLang, which features:

• Tree-style AST printing like rich.
• Use the pass option to control whether it's enabled/disabled.

Usage: annotate the kernel you want to print with:

@tilelang.jit(out_idx=[-1], pass_configs={
    tilelang.PassConfigKey.TL_AST_PRINT_ENABLE: True,
})

The gemm example will be printed like this:

PrimFunc(params=[A_handle, B_handle, C_handle], ret_type=None, buffer_map={A_handle: A, B_handle: B, C_handle: C}, attrs={"global_symbol": "gemm"})
└── body=BlockRealize: with T.block("root"):
├── iter_values: []
├── predicate: T.bool8(True)
└── block: Block: with T.block("root", no_realize=True):
    ├── iter_vars: []
    ├── reads: []
    ├── writes: []
    ├── name_hint: 'root'
    ├── alloc_buffers: []
    ├── match_buffers: []
    ├── annotations: {}
    └── body: AttrStmt: with T.launch_thread("blockIdx.x", 8) as bx:
        ├── node: T.iter_var(bx, T.Range(0, 8), "ThreadIndex", "blockIdx.x")
        ├── attr_key: 'thread_extent'
        ├── value: 8
        └── body: AttrStmt: with T.launch_thread("blockIdx.y", 8) as by:
            ├── node: T.iter_var(by, T.Range(0, 8), "ThreadIndex", "blockIdx.y")
            ├── attr_key: 'thread_extent'
            ├── value: 8
            └── body: AttrStmt: with T.launch_thread("threadIdx.x", 128) as tx:
                ├── node: T.iter_var(tx, T.Range(0, 128), "ThreadIndex", "threadIdx.x")
                ├── attr_key: 'thread_extent'
                ├── value: 128
                └── body: AttrStmt: with T.launch_thread("threadIdx.y", 1) as ty:
                    ├── node: T.iter_var(ty, T.Range(0, 1), "ThreadIndex", "threadIdx.y")
                    ├── attr_key: 'thread_extent'
                    ├── value: 1
                    └── body: AttrStmt: with T.launch_thread("threadIdx.z", 1) as tz:
                        ├── node: T.iter_var(tz, T.Range(0, 1), "ThreadIndex", "threadIdx.z")
                        ├── attr_key: 'thread_extent'
                        ├── value: 1
                        └── body: BlockRealize: with T.block("tilelang_root"):
                            ├── iter_values: []
                            ├── predicate: T.bool8(True)
                            └── block: Block: with T.block("tilelang_root", no_realize=True):
                                ├── iter_vars: []
                                ├── reads: [A[by * 128, 0:993], B[0:993, bx * 128], C[by * 128, bx * 128]]
                                ├── writes: []
                                ├── name_hint: 'tilelang_root'
                                ├── alloc_buffers: [A_shared, B_shared, C_local]
                                ├── match_buffers: []
                                ├── annotations: {}
                                └── body: SeqStmt: C_local = T.Buffer((128, 128), scope="local.fragment")
                                    ├── seq0: Evaluate: C_local = T.Buffer((128, 128), scope="local.fragment")
                                    |   └── value: T.fill(T.region(C_local[0, 0], 2, 128, 128), 0)
                                    ├── seq1: For: for k in T.serial(32, annotations={"num_stages": 3}):
                                    |   ├── loop_var: k
                                    |   ├── min: 0
                                    |   ├── extent: 32
                                    |   ├── kind: 0
                                    |   ├── annotations: {'num_stages': 3}
                                    |   └── body: SeqStmt: A = T.Buffer((1024, 1024), "float16", strides=(1024, 1))
                                    |       ├── seq0: Evaluate: A = T.Buffer((1024, 1024), "float16", strides=(1024, 1))
                                    |       |   └── value: T.copy(T.region(A[by * 128, k * 32], 1, 128, 32), T.region(A_shared[0, 0], 2, 128, 32), -1, T.bool8(False), 0)
                                    |       ├── seq1: Evaluate: B = T.Buffer((1024, 1024), "float16", strides=(1024, 1))
                                    |       |   └── value: T.copy(T.region(B[k * 32, bx * 128], 1, 32, 128), T.region(B_shared[0, 0], 2, 32, 128), -1, T.bool8(False), 0)
                                    |       └── seq2: Evaluate: A_shared = T.Buffer((128, 32), "float16", scope="shared.dyn")
                                    |           └── value: T.gemm_py(T.region(A_shared[0, 0], 1, 128, 32), T.region(B_shared[0, 0], 1, 32, 128), T.region(C_local[0, 0], 3, 128, 128), T.boo...
                                    └── seq2: Evaluate: C_local = T.Buffer((128, 128), scope="local.fragment")
                                        └── value: T.copy(T.region(C_local[0, 0], 1, 128, 128), T.region(C[by * 128, bx * 128], 2, 128, 128), -1, T.bool8(False), 0)

Summary by CodeRabbit

• New Features
  • Introduced enhanced AST pretty-printing with tree-like formatting and improved indentation for debugging and code analysis.
  • Added configuration option to enable automatic AST printing during semantic validation, providing detailed syntax tree visualization for developers.

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

[coderabbitai]

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

Walkthrough

Adds a TileLang AST printing feature: new pass-config key registration in C++, TileLang-side config enum, a conditional hook in the PreLowerSemanticCheck phase, and a visitor-based AST pretty-printer implementation in Python.

Changes

Cohort / File(s)	Summary
C++ Pass Config Registration src/op/builtin.h, src/op/builtin.cc	Declares kASTPrintEnable ("tl.ast_print_enable") and registers it as a Bool pass configuration option in the builtins pass config registry.
Python Pass Config Key tilelang/transform/pass_config.py	Adds TL_AST_PRINT_ENABLE = "tl.ast_print_enable" to PassConfigKey enum for TileLang pass configuration.
AST Printer Enhancement tilelang/analysis/ast_printer.py	Introduces _ASTPrintVisitor (subclass of PyStmtExprVisitor) and refactors ASTPrinter to use visitor-driven, tree-like pretty-printing with field-aware traversal and indentation/connector rendering.
Phase-level Control tilelang/engine/phase.py	Adds should_enable_ast_print(pass_ctx=None) -> bool and conditionally invokes ASTPrinter()(mod) inside PreLowerSemanticCheck when the config flag is enabled.

Sequence Diagram(s)

sequenceDiagram
    autonumber
    participant Phase as PreLowerSemanticCheck
    participant Config as PassContext
    participant Printer as ASTPrinter
    participant Visitor as _ASTPrintVisitor
    participant AST as Module/Stmt

    Phase->>Config: should_enable_ast_print(pass_ctx)?
    Config-->>Phase: TL_AST_PRINT_ENABLE true/false

    rect rgb(200,230,201)
        Note over Phase,Printer: when enabled
        Phase->>Printer: ASTPrinter()(mod)
        Printer->>Printer: print metadata (params, ret, buffers, attrs)
        Printer->>Visitor: instantiate visitor
        Visitor->>AST: traverse nodes
        Visitor-->>Visitor: format tree nodes (connectors, indentation)
        Visitor-->>Printer: return formatted lines
        Printer-->>Phase: completed printing
    end

    Phase->>Phase: continue semantic checks

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Suggested reviewers

• chengyupku

Poem

🐰 I hop through nodes both wide and slight,

I stitch the tree with branches, line and light.
A flag, a visitor, and prints that sing—
Now TileLang's forest shows everything! 🌿✨

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.00% 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 '[Enhancement][Tool] Tree-style pretty ASTPrinter' directly and accurately describes the main change: adding a tree-style AST pretty-printer as a tool enhancement, which aligns with all modified files introducing AST printing functionality controlled via a pass configuration option.

✨ 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]

👋 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]

Actionable comments posted: 0

🧹 Nitpick comments (3)

tilelang/analysis/ast_printer.py (3)

88-91: Unnecessary indent.append after visiting.

Line 90 appends to the indent list after visiting is complete, but the visitor is never used again and the function returns immediately. This appears to be leftover code that has no effect.

🔎 Proposed fix

         visitor.print_stmt_brief(func.body, func_body_prefix)
+        visitor.indent.append(_normal_indent)
         visitor.visit_stmt(func.body)
-        visitor.indent.append(_normal_indent)
         return func

Note: If the intent was to add initial indentation for the body traversal, the append should happen before visit_stmt, not after.

---

32-56: Consider using TIR node's known fields instead of class dict introspection.

The current approach uses stmt.__class__.__dict__.keys() to discover fields, which may miss inherited fields or include unexpected entries. For robustness, consider using TIR's known field accessors or a predefined field list per node type.

However, for a debugging tool, this approach is acceptable as-is.

---

84-84: Optional: Prefix unused parameters with underscore to silence linter warnings.

The mod and ctx parameters are required by the prim_func_pass signature but are unused. You can prefix them with underscores to indicate intentional non-use.

🔎 Proposed fix

-    def pass_fn(func: PrimFunc, mod, ctx) -> PrimFunc:
+    def pass_fn(func: PrimFunc, _mod, _ctx) -> PrimFunc:

📜 Review details

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 2b79a76 and 08fb5f2.

📒 Files selected for processing (5)

• src/op/builtin.cc
• src/op/builtin.h
• tilelang/analysis/ast_printer.py
• tilelang/engine/phase.py
• tilelang/transform/pass_config.py

🧰 Additional context used

🧬 Code graph analysis (1)

tilelang/engine/phase.py (2)

tilelang/transform/__init__.py (1)

• get_pass_context (12-14)

tilelang/analysis/ast_printer.py (1)

• ASTPrinter (79-93)

🪛 Ruff (0.14.10)

tilelang/analysis/ast_printer.py

84-84: Unused function argument: mod

(ARG001)

---

84-84: Unused function argument: ctx

(ARG001)

🔇 Additional comments (6)

src/op/builtin.h (1)

58-58: LGTM!

The new kASTPrintEnable constant is correctly defined and follows the established naming convention and placement of other pass configuration keys in this file.

src/op/builtin.cc (1)

37-37: LGTM!

The pass configuration option is correctly registered with the appropriate Bool type, consistent with other similar enable/disable flags in this file.

tilelang/analysis/ast_printer.py (1)

7-15: LGTM!

The module-level constants for tree formatting are well-organized and clearly named. The line limit of 140 characters and the use of box-drawing characters (├──, └──) provide a clean visual structure.

tilelang/transform/pass_config.py (1)

87-88: LGTM!

The new TL_AST_PRINT_ENABLE configuration key is well-documented and correctly aligned with the C++ counterpart in builtin.h.

tilelang/engine/phase.py (2)

66-69: LGTM!

The should_enable_ast_print function correctly follows the established pattern of other configuration helper functions in this file.

---

121-123: LGTM!

The conditional AST printing is appropriately integrated into PreLowerSemanticCheck, placed before the semantic validation checks. This allows developers to inspect the AST before any validation errors occur.

[Copilot]

Pull request overview

This PR enhances the AST printer in TileLang from a simple debugging tool to a formal tree-style visualization tool similar to the rich library. The feature can be enabled/disabled via pass configuration options and provides hierarchical, indented AST printing with tree connectors for better readability.

Key changes:

• Added pass configuration key TL_AST_PRINT_ENABLE to control AST printing
• Implemented a visitor-based AST printer with tree-style formatting using Unicode box-drawing characters
• Integrated the printer into the pre-lower semantic check phase when enabled

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
tilelang/transform/pass_config.py	Adds the TL_AST_PRINT_ENABLE configuration key for enabling AST printing
tilelang/engine/phase.py	Adds should_enable_ast_print() helper function and conditionally invokes AST printer during pre-lower checks
tilelang/analysis/ast_printer.py	Complete rewrite of AST printer using visitor pattern with tree-style formatting, indentation management, and field traversal logic
src/op/builtin.h	Declares C++ constant for AST print enable configuration
src/op/builtin.cc	Registers the AST print enable pass configuration option in C++

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[coderabbitai]

Actionable comments posted: 0

♻️ Duplicate comments (4)

tilelang/analysis/ast_printer.py (4)

7-14: Module constants are well-defined.

The tree connector characters and indentation settings provide a clear, readable tree structure. The naming convention concern was already noted in a previous review.

---

32-43: Child field tracking may lose earlier child fields if multiple exist.

If a statement class has multiple child fields (e.g., both body and block), only the last encountered one is retained in child_field_name. However, this appears acceptable for TIR statements which typically have a single child field per node type.

---

79-79: Consider adding return type annotation.

For better type safety, add the return type annotation as previously suggested.

---

99-99: Dead code: indent append after traversal.

This line appends to the indent list after visit_stmt has completed. The visitor has finished printing at this point, so this operation has no effect.

🔎 Proposed fix

         visitor.print_stmt_brief(func.body, func_body_prefix)
         visitor.visit_stmt(func.body)
-        visitor.indent.append(_normal_indent)
         return func

🧹 Nitpick comments (1)

tilelang/analysis/ast_printer.py (1)

28-30: String parsing relies on TVM's repr format.

The split(" ") (double space) parsing is tied to TVM's current repr formatting. If the format changes, this may break or produce unexpected output.

Consider adding a fallback or documenting the expected format:

def print_stmt_brief(self, stmt: Stmt, prefix: str) -> None:
    stmt_repr = repr(stmt).splitlines()[0]
    # Extract first segment before double-space (TVM script format)
    stmt_script = stmt_repr.split("  ")[0].strip() if "  " in stmt_repr else stmt_repr.strip()
    self.print_with_clip(prefix + f"{stmt.__class__.__name__}: " + stmt_script)

📜 Review details

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 08fb5f2 and 2d2af0b.

📒 Files selected for processing (1)

• tilelang/analysis/ast_printer.py

🧰 Additional context used

🧬 Code graph analysis (1)

tilelang/analysis/ast_printer.py (2)

src/op/atomic_add.cc (2)

• s (423-423)
• s (423-423)

tilelang/analysis/nested_loop_checker.py (1)

• pass_fn (115-117)

🪛 Ruff (0.14.10)

tilelang/analysis/ast_printer.py

93-93: Unused function argument: mod

(ARG001)

---

93-93: Unused function argument: ctx

(ARG001)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)

• GitHub Check: Test for Python 3.12 with CUDA-12.8 (on self-hosted-nvidia)
• GitHub Check: Test for Python 3.12 with Metal (on macos-latest)
• GitHub Check: Test for Python 3.12 with Nightly-ROCm-7.1 (on self-hosted-amd)

🔇 Additional comments (3)

tilelang/analysis/ast_printer.py (3)

1-5: LGTM!

Imports are appropriate for the visitor-based AST printer implementation.

---

68-76: SeqStmt handling is well-implemented.

The special handling for sequence statements correctly iterates children with proper tree connectors and indentation management.

---

93-100: Pass function signature follows TVM convention.

The mod and ctx parameters are required by the prim_func_pass interface even when unused. The static analysis warning is a false positive for this standard pattern.

[LeiWang1999]

LGTM, Pretty awesome.

[LeiWang1999]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Traverse missing statement children in AST printer

The new tree printer only recurses into fields listed in _child_fields (currently just body, block, seq). TIR statement nodes like IfThenElse use then_case/else_case, and Block can have init, which are statements but won’t be visited or expanded, so whole subtrees disappear from the tree output. This is a regression from the previous ir_transform traversal, and it can mislead debugging whenever kernels include conditionals or other statement children not named body/block/seq. Consider expanding the child field list or delegating traversal to PyStmtExprVisitor for all statement children.

Useful? React with 👍 / 👎.