[Feature] Add compiler pass timing profiling by penguin-wwy · Pull Request #2622 · tile-ai/tilelang

penguin-wwy · 2026-07-11T18:36:33Z

Summary

Add compiler pass timing enabled through environment variables or PassConfig.
Report inclusive/self time, threshold filtering, Top-N passes, and compilation phase context.
Support regular JIT and grouped compilation, including failure-path reporting.

Why not reuse TVM's PassTimingInstrument?

Data persistence: TVM clears profiling data in ExitPassContext, so render() must be called
inside the with PassContext block.
Custom output: TileLang needs custom sorting, threshold filtering, Top-N reporting, and phase
grouping.
Extensibility: This can later include IR size statistics, skipped-pass counts, and other metrics.

Example

2026-07-12 02:31:39  [TileLang:tilelang.pass_timing:INFO] (pass_timing.py:219): 
================================================================================================
TileLang Pass Timing Report
Context: stage=jit-lower, kernel=matmul, backend=cython
================================================================================================
Total: 0.1970s (55 passes)
------------------------------------------------------------------------------------------------
  #  Pass Name                                    Inclusive       Self   Incl %  Self %
------------------------------------------------------------------------------------------------
  1  pass_fn                                        0.0002s    0.0002s    0.1%    0.1%
  2  pass_fn                                        0.0004s    0.0004s    0.2%    0.2%
  3  tirx.BindTarget                                0.0011s    0.0011s    0.6%    0.6%
  4  tl.MaterializeKernelLaunch                     0.0001s    0.0001s    0.0%    0.0%
  5  pass_fn                                        0.0008s    0.0008s    0.4%    0.4%
  6  tl.LegalizeNegativeIndex                       0.0030s    0.0030s    1.5%    1.5%
  7  tl.VerifyParallelLoop                          0.0001s    0.0001s    0.0%    0.0%
  8  tl.InjectAssumes                               0.0001s    0.0001s    0.1%    0.1%
  9  tl.Simplify                                    0.0048s    0.0048s    2.4%    2.4%
 10  tl.LayoutReducer                               0.0016s    0.0016s    0.8%    0.8%
 11  tl.IfStmtBinding                               0.0001s    0.0001s    0.0%    0.0%
 12  tl.Simplify                                    0.0042s    0.0042s    2.1%    2.1%
 13  tl.LayoutInference                             0.0114s    0.0114s    5.8%    5.8%
 14  tl.LowerTileOp                                 0.0448s    0.0448s   22.7%   22.7%
 15  pass_fn                                        0.0011s    0.0011s    0.6%    0.6%
 16  tl.LegalizeVectorizedLoop                      0.0207s    0.0207s   10.5%   10.5%
 17  tl.LegalizeSafeMemoryAccess                    0.0056s    0.0056s    2.9%    2.9%
 18  tl.LowerAccessPtr                              0.0000s    0.0000s    0.0%    0.0%
 19  tl.Simplify                                    0.0080s    0.0080s    4.1%    4.1%
 20  tl.HoistNonRestrictParams                      0.0000s    0.0000s    0.0%    0.0%
 21  tl.PlanAndUpdateBufferAllocationLocation       0.0014s    0.0014s    0.7%    0.7%
 22  tl.HoistGlobalBufferAllocations                0.0000s    0.0000s    0.0%    0.0%
 23  tl.LowerOpaqueBlock                            0.0002s    0.0002s    0.1%    0.1%
 24  tl.Simplify                                    0.0063s    0.0063s    3.2%    3.2%
 25  tirx.NarrowDataType                            0.0019s    0.0019s    1.0%    1.0%
 26  tl.FlattenBuffer                               0.0053s    0.0053s    2.7%    2.7%
 27  tl.ConfigIndexBitwidth                         0.0025s    0.0025s    1.3%    1.3%
 28  tirx.Simplify                                  0.0060s    0.0060s    3.1%    3.1%
 29  tl.VectorizeLoop                               0.0001s    0.0001s    0.1%    0.1%
 30  tir.StorageRewrite                             0.0004s    0.0004s    0.2%    0.2%
 31  tl.LoopUnswitching                             0.0001s    0.0001s    0.0%    0.0%
 32  tl.UnrollLoop                                  0.0001s    0.0001s    0.1%    0.1%
 33  s_tir.RenormalizeSplitPattern                  0.0042s    0.0042s    2.1%    2.1%
 34  tirx.Simplify                                  0.0048s    0.0048s    2.5%    2.5%
 35  tirx.RemoveNoOp                                0.0027s    0.0027s    1.4%    1.4%
 36  s_tir.HoistIfThenElse                          0.0093s    0.0000s    4.7%    0.0%
 37    s_tir.InsertHoistIfThenElse                  0.0020s    0.0020s    1.0%    1.0%
 38    tirx.Simplify                                0.0047s    0.0047s    2.4%    2.4%
 39    tirx.RemoveNoOp                              0.0026s    0.0026s    1.3%    1.3%
 40  tirx.VerifyMemory                              0.0000s    0.0000s    0.0%    0.0%
 41  tirx.AnnotateEntryFunc                         0.0000s    0.0000s    0.0%    0.0%
 42  s_tir.InferFragment                            0.0000s    0.0000s    0.0%    0.0%
 43  tl.AnnotateDeviceRegions                       0.0000s    0.0000s    0.0%    0.0%
 44  tl.SplitHostDevice                             0.0008s    0.0006s    0.4%    0.3%
 45    tirx.ConvertSSA                              0.0001s    0.0001s    0.1%    0.1%
 46  tl.AnnotateReadOnlyParams                      0.0001s    0.0001s    0.0%    0.0%
 47  tl.MergeIfStmt                                 0.0001s    0.0001s    0.0%    0.0%
 48  tl.MakePackedAPI                               0.0105s    0.0105s    5.3%    5.3%
 49  tl.Simplify                                    0.0242s    0.0242s   12.3%   12.3%
 50  tl.LowerDeviceKernelLaunch                     0.0003s    0.0003s    0.1%    0.1%
 51  tirx.Filter                                    0.0001s    0.0001s    0.0%    0.0%
 52  tirx.Filter                                    0.0000s    0.0000s    0.0%    0.0%
 53  tl.LowerIntrin                                 0.0020s    0.0020s    1.0%    1.0%
 54  tirx.Simplify                                  0.0045s    0.0045s    2.3%    2.3%
 55  pass_fn                                        0.0006s    0.0006s    0.3%    0.3%
------------------------------------------------------------------------------------------------
Top 10 Slowest Passes by Inclusive Time:
   1. tl.LowerTileOp                                   0.0448s  22.7%
   2. tl.Simplify                                      0.0242s  12.3%
   3. tl.LegalizeVectorizedLoop                        0.0207s  10.5%
   4. tl.LayoutInference                               0.0114s   5.8%
   5. tl.MakePackedAPI                                 0.0105s   5.3%
   6. s_tir.HoistIfThenElse                            0.0093s   4.7%
   7. tl.Simplify                                      0.0080s   4.1%
   8. tl.Simplify                                      0.0063s   3.2%
   9. tirx.Simplify                                    0.0060s   3.1%
  10. tl.LegalizeSafeMemoryAccess                      0.0056s   2.9%
================================================================================================

Summary

Adds configurable compiler pass-timing profiling for regular JIT and grouped compilation.

Highlights

Records inclusive and self pass durations with nested-pass tracking.
Supports threshold filtering, contextual reports, and failure-path reporting.
Enables profiling through environment variables or PassConfig.
Adds profiling to lowering, host, and device compilation phases.
Exposes timing records and instrumentation through tilelang.utils.
Adds environment parsing, configuration keys, and comprehensive timing tests.

C++ style / lint notes

Touches C++ registration and declarations in src/op/builtin.cc and src/op/builtin.h.
No changes are indicated for rules documented in docs/developer_guide/cpp_style.md.
The C++ API Style Audit (warning only) may report advisory findings; these should remain separate from correctness, build, and test issues.

github-actions · 2026-07-11T18:36:41Z

👋 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 · 2026-07-11T18:39:17Z

📝 Walkthrough

Walkthrough

TileLang adds pass profiling configuration, threshold validation, nested pass timing records, formatted reports, and timing instrumentation to JIT and grouped compilation phases. Profiling can be enabled through pass configuration or environment variables.

Changes

Pass profiling

Layer / File(s)	Summary
Profiling configuration contracts `src/op/builtin.h`, `src/op/builtin.cc`, `tilelang/transform/pass_config.py`, `tilelang/env.py`, `testing/python/components/test_tilelang_env.py`	Adds pass profiling keys, environment-backed enablement and thresholds, threshold validation, configuration precedence, and parsing tests.
Pass timing engine and reporting `tilelang/utils/pass_timing.py`, `tilelang/utils/__init__.py`, `testing/python/debug/test_pass_timing.py`	Adds nested pass timing records, self-duration calculation, timing instrument construction, threshold-filtered reports, exit reporting, public exports, and instrumentation tests.
JIT lowering integration `tilelang/jit/kernel.py`	Builds profiling instruments for JIT lowering and reports timing for the lower phase.
Grouped compilation integration `tilelang/autotuner/grouped_compile.py`	Creates phase-specific instruments and reports timing for grouped-lower, grouped-device, and grouped-host compilation.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Configuration
  participant JITKernel
  participant PassContext
  participant TileLangPassTimingInstrument
  participant TimingReport
  Configuration->>JITKernel: provide profiling flag and threshold
  JITKernel->>PassContext: install pass instruments
  PassContext->>TileLangPassTimingInstrument: record pass callbacks
  TileLangPassTimingInstrument->>TimingReport: produce timing report on exit

Suggested reviewers: LeiWang1999, lucifer1004

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 36.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: adding compiler pass timing profiling.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

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

coderabbitai

🧹 Nitpick comments (1)

tilelang/jit/kernel.py (1)

239-255: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Extract shared instrument-setup helper to eliminate duplication.

The base-instrument construction (lines 239–242) and threshold-resolution logic (lines 244–251) are duplicated in tilelang/autotuner/grouped_compile.py (lines 44–56). A shared factory would keep both call sites in sync if the setup logic evolves.

♻️ Proposed shared helper in tilelang/utils/pass_timing.py

def make_pass_instrument_factory(pass_configs: Mapping[str, Any]) -> Callable[[], tuple[list, object | None]]:
    """Return a zero-arg factory that produces fresh (instruments, timing_instrument) tuples."""
    from tilelang import env
    from tilelang.transform import PassConfigKey

    base_pass_instruments: list = []
    if pass_configs.get(PassConfigKey.TL_ENABLE_DUMP_IR):
        dump_ir_path = pass_configs.get(PassConfigKey.TL_DUMP_IR_DIR, "./dump_ir")
        base_pass_instruments.append(tvm.ir.instrument.DumpIR(dump_dir=dump_ir_path))

    profile_threshold_ms: float | None = None
    if pass_configs.get(PassConfigKey.TL_PASS_PROFILE) or env.is_pass_profile_enabled():
        profile_threshold_ms = resolve_pass_profile_threshold_ms(
            pass_configs,
            PassConfigKey.TL_PASS_PROFILE_THRESHOLD_MS,
            env.get_pass_profile_threshold_ms,
        )

    def create_pass_instruments():
        return build_pass_instruments(base_pass_instruments, profile_threshold_ms)

    return create_pass_instruments

Usage in kernel.py:

-        base_pass_instruments = []
-        if pass_configs.get(PassConfigKey.TL_ENABLE_DUMP_IR):
-            dump_ir_path = pass_configs.get(PassConfigKey.TL_DUMP_IR_DIR, "./dump_ir")
-            base_pass_instruments.append(tvm.ir.instrument.DumpIR(dump_dir=dump_ir_path))
-
-        # Pass timing instrument
-        profile_threshold_ms = None
-        if pass_configs.get(PassConfigKey.TL_PASS_PROFILE) or env.is_pass_profile_enabled():
-            profile_threshold_ms = resolve_pass_profile_threshold_ms(
-                pass_configs,
-                PassConfigKey.TL_PASS_PROFILE_THRESHOLD_MS,
-                env.get_pass_profile_threshold_ms,
-            )
-        pass_instruments, timing_instrument = build_pass_instruments(
-            base_pass_instruments,
-            profile_threshold_ms,
-        )
+        create_pass_instruments = make_pass_instrument_factory(pass_configs)
+        pass_instruments, timing_instrument = create_pass_instruments()

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tilelang/jit/kernel.py` around lines 239 - 255, Extract the duplicated
dump-IR and pass-profile setup from the kernel pass-instrument flow and grouped
compilation into a shared make_pass_instrument_factory helper in pass_timing.py.
Have the factory resolve configuration once and return a zero-argument creator
that produces fresh instrument/timing tuples via build_pass_instruments, then
update the affected call sites to use it while preserving existing configuration
behavior.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tilelang/jit/kernel.py`:
- Around line 239-255: Extract the duplicated dump-IR and pass-profile setup
from the kernel pass-instrument flow and grouped compilation into a shared
make_pass_instrument_factory helper in pass_timing.py. Have the factory resolve
configuration once and return a zero-argument creator that produces fresh
instrument/timing tuples via build_pass_instruments, then update the affected
call sites to use it while preserving existing configuration behavior.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: e490ab81-6db4-4cdc-8b95-30facf493cf5

📥 Commits

Reviewing files that changed from the base of the PR and between bd76473 and dda55f0.

📒 Files selected for processing (10)

src/op/builtin.cc
src/op/builtin.h
testing/python/components/test_tilelang_env.py
testing/python/debug/test_pass_timing.py
tilelang/autotuner/grouped_compile.py
tilelang/env.py
tilelang/jit/kernel.py
tilelang/transform/pass_config.py
tilelang/utils/__init__.py
tilelang/utils/pass_timing.py

[Feature] Add compiler pass timing profiling

dda55f0

coderabbitai Bot reviewed Jul 11, 2026

View reviewed changes

LeiWang1999 requested a review from SiriusNEO July 13, 2026 09:05

SiriusNEO approved these changes Jul 16, 2026

View reviewed changes

SiriusNEO merged commit 4a5cf09 into tile-ai:main Jul 16, 2026
6 of 7 checks passed

This was referenced Aug 6, 2026

[Enhancement] Show kernel name in pass timing report #2905

Merged

[Debug][TIR][JIT] Unify pass instrumentation per compilation #2923

Merged

Navigation Menu

[Feature] Add compiler pass timing profiling - #2622

[Feature] Add compiler pass timing profiling#2622
SiriusNEO merged 1 commit into
tile-ai:mainfrom
penguin-wwy:pass_dev

penguin-wwy commented Jul 11, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

github-actions Bot commented Jul 11, 2026

Uh oh!

coderabbitai Bot commented Jul 11, 2026 •

edited

Loading

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

penguin-wwy commented Jul 11, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Highlights

C++ style / lint notes

Uh oh!

github-actions Bot commented Jul 11, 2026

Uh oh!

coderabbitai Bot commented Jul 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

penguin-wwy commented Jul 11, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 11, 2026 •

edited

Loading