[Debug][TIR][JIT] Unify pass instrumentation per compilation - #2923
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe PR replaces process-wide tracing hooks with shared TVM ChangesPass-event instrumentation migration
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant CompilationEntry
participant CompilePassInstrumentationSession
participant PassPipeline
participant PassInstrument
participant LowerTraceSession
participant BackendCodegen
CompilationEntry->>CompilePassInstrumentationSession: create compilation session
CompilationEntry->>PassPipeline: execute lowering
PassPipeline->>PassInstrument: emit nested pass callbacks
PassInstrument->>LowerTraceSession: record pass events
PassPipeline->>BackendCodegen: request generated source
BackendCodegen->>LowerTraceSession: record codegen event
CompilePassInstrumentationSession-->>CompilationEntry: finalize session
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
testing/python/debug/test_lower_trace.py (1)
140-143: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMake the instrument assertion independent of other registered providers.
build_pass_instrumentsreturns every registered provider's instrument. The assertionlen(instruments) == 1fails if another tool (for example the pass visualizer) registers a provider during the session. Assert on the lower-trace instrument instead of the total count.♻️ Proposed change
instruments, timing = build_pass_instruments([], threshold_ms=None) assert timing is None - assert len(instruments) == 1 - assert isinstance(instruments[0].observer, _core._LowerTraceObserver) + trace_instruments = [ + item for item in instruments if isinstance(getattr(item, "observer", None), _core._LowerTraceObserver) + ] + assert len(trace_instruments) == 1 assert Pass.__call__ is original_pass_call🤖 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 `@testing/python/debug/test_lower_trace.py` around lines 140 - 143, Update the test around build_pass_instruments to stop asserting the total instruments length; locate the instrument whose observer is an instance of _core._LowerTraceObserver and assert that it exists, while preserving the timing assertion and lower-trace behavior.testing/python/debug/test_pass_events.py (1)
196-205: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd a timeout to the barrier wait.
barrier.wait()has no timeout. If the pool cannot start both worker threads, the test blocks forever and the CI job hangs until the global timeout. A timeout converts the hang into aBrokenBarrierErrorfailure.♻️ Proposed change
def read_phase(name): with pass_phase(name): - barrier.wait() + barrier.wait(timeout=30) return current_pass_phase()🤖 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 `@testing/python/debug/test_pass_events.py` around lines 196 - 205, Update the barrier synchronization in test_pass_phase_is_isolated_between_threads by supplying a finite timeout to barrier.wait(), so failure to start both worker threads raises BrokenBarrierError instead of hanging indefinitely.
🤖 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.
Inline comments:
In `@tilelang/tools/lower_trace/core.py`:
- Around line 213-216: Reorder the union members in _mode_override,
_trace_dir_override, and _codegen_output_path_override so None appears last,
satisfying Ruff RUF036 while preserving their existing types and sentinel
behavior.
In `@tilelang/utils/pass_events.py`:
- Around line 234-272: The shared StackedPassInstrument keeps _stack and
_next_sequence as non-context-local state, allowing concurrent lowering to
corrupt pass tracking; move both into ContextVar-backed state (or otherwise
isolate them per thread) while preserving enter_pass_ctx, exit_pass_ctx, and
run_before_pass behavior. In tilelang/utils/pass_events.py lines 234-272, update
the StackedPassInstrument state access; in tilelang/tools/lower_trace/core.py
lines 863-878, verify whether PassContext.current() is thread-local in the
installed TVM version and, if not, create a separate instrument per lowering
thread.
---
Nitpick comments:
In `@testing/python/debug/test_lower_trace.py`:
- Around line 140-143: Update the test around build_pass_instruments to stop
asserting the total instruments length; locate the instrument whose observer is
an instance of _core._LowerTraceObserver and assert that it exists, while
preserving the timing assertion and lower-trace behavior.
In `@testing/python/debug/test_pass_events.py`:
- Around line 196-205: Update the barrier synchronization in
test_pass_phase_is_isolated_between_threads by supplying a finite timeout to
barrier.wait(), so failure to start both worker threads raises
BrokenBarrierError instead of hanging indefinitely.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 9075b9f5-6c1d-41bc-a91e-e4689e8bc60a
📒 Files selected for processing (11)
docs/tools/lower_trace.mdtesting/python/debug/test_lower_trace.pytesting/python/debug/test_pass_events.pytilelang/backend/pass_pipeline/pipeline.pytilelang/jit/adapter/wrapper.pytilelang/tools/lower_trace/__init__.pytilelang/tools/lower_trace/core.pytilelang/tools/lower_trace/html.pytilelang/tools/pass_visualizer/core.pytilelang/utils/pass_events.pytilelang/utils/pass_timing.py
| _UNSET: object = object() | ||
| _mode_override: str | None | object = _UNSET | ||
| _trace_dir_override: str | None | object = _UNSET | ||
| _codegen_output_path_override: str | None | object = _UNSET |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Reorder the sentinel union to satisfy Ruff RUF036.
Ruff reports RUF036 on these three annotations. Move None to the end of each union.
♻️ Proposed change
_UNSET: object = object()
-_mode_override: str | None | object = _UNSET
-_trace_dir_override: str | None | object = _UNSET
-_codegen_output_path_override: str | None | object = _UNSET
+_mode_override: str | object | None = _UNSET
+_trace_dir_override: str | object | None = _UNSET
+_codegen_output_path_override: str | object | None = _UNSET📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| _UNSET: object = object() | |
| _mode_override: str | None | object = _UNSET | |
| _trace_dir_override: str | None | object = _UNSET | |
| _codegen_output_path_override: str | None | object = _UNSET | |
| _UNSET: object = object() | |
| _mode_override: str | object | None = _UNSET | |
| _trace_dir_override: str | object | None = _UNSET | |
| _codegen_output_path_override: str | object | None = _UNSET |
🧰 Tools
🪛 Ruff (0.16.1)
[warning] 214-214: None not at the end of the type union.
Move None to the end of the type union
(RUF036)
[warning] 215-215: None not at the end of the type union.
Move None to the end of the type union
(RUF036)
[warning] 216-216: None not at the end of the type union.
Move None to the end of the type union
(RUF036)
🤖 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/tools/lower_trace/core.py` around lines 213 - 216, Reorder the union
members in _mode_override, _trace_dir_override, and
_codegen_output_path_override so None appears last, satisfying Ruff RUF036 while
preserving their existing types and sentinel behavior.
Source: Linters/SAST tools
LeiWang1999
changed the title
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (7)
tilelang/tools/lower_trace/core.py (2)
713-726: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSort
__all__to satisfy Ruff RUF022.Ruff reports RUF022 on this list. Apply the isort-style ordering, or add a
# noqa: RUF022if the current grouping is intentional.🤖 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/tools/lower_trace/core.py` around lines 713 - 726, Sort the entries in __all__ using isort-style ordering to satisfy Ruff RUF022, preserving all existing exports; only add # noqa: RUF022 if the current grouping is intentionally required.Source: Linters/SAST tools
314-322: 🚀 Performance & Scalability | 🔵 Trivial | ⚖️ Poor tradeoffConsider making nested-pass capture configurable.
create_pass_instrumentsetscapture_nested=True, and_LowerTraceObserver.pass_startedstringifies the full module for every captured pass. Nested C++ passes multiply the record count, so the session retains one full before-text and after-text string per nested pass, andsave_raw_fileswrites two files for each. On a large kernel this increases peak memory and trace I/O substantially compared to top-level-only capture. Expose the flag through_LowerTraceConfigso users can opt out.Also applies to: 536-539
🤖 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/tools/lower_trace/core.py` around lines 314 - 322, Expose a capture_nested option on _LowerTraceConfig and use that configured value in create_pass_instrument instead of hard-coding True. Ensure the setting propagates through the existing configuration path so users can disable nested-pass capture while preserving enabled behavior and the current default.testing/python/debug/test_pass_events.py (1)
204-214: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the rejected reuse configuration.
compile_pass_instrumentationraisesValueErrorwhen a caller passestoolsor setsinclude_default_tools=Falsewhile reusing an active session. No test exercises that branch. Add a case so the guard cannot regress silently.♻️ Proposed addition
def test_reuse_rejects_tool_overrides(): with compile_pass_instrumentation(name="outer"): with pytest.raises(ValueError): with compile_pass_instrumentation(name="nested", tools=[_RecordingTool("x", [])]): pass with pytest.raises(ValueError): with compile_pass_instrumentation(name="nested", include_default_tools=False): pass🤖 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 `@testing/python/debug/test_pass_events.py` around lines 204 - 214, Add a test near test_nested_helpers_reuse_the_owning_compile_session that opens an outer compile_pass_instrumentation session, then asserts nested reuse raises ValueError when tools is provided and separately when include_default_tools=False.tilelang/utils/pass_events.py (1)
137-148: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winRecord the suppressed finalization error.
If
erroris not None,finish_erroris discarded silently. A tool that fails during finalization then leaves no trace, which makes tool bugs hard to diagnose during an already failing compilation. Attach the finalization failure to the compilation error or log it.♻️ Proposed change
# Never hide the compilation failure that tools were asked to observe. - if error is None and finish_error is not None: - raise finish_error + if finish_error is None: + return + if error is None: + raise finish_error + error.add_note(f"pass-instrumentation tool finalization also failed: {finish_error!r}")🤖 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/utils/pass_events.py` around lines 137 - 148, Update PassEvents.finish so the first captured finish_error is preserved when a compilation error is already present: attach it to error or log it through the existing error-reporting mechanism, while retaining the current behavior of raising finish_error when error is None. Ensure finalization failures are never silently discarded.tilelang/tools/pass_visualizer/core.py (1)
438-446: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueGuard against silently discarding an earlier instrument.
create_pass_instrumentoverwritesself.instrumenton every call. The compile-session contract allows one call perPassContext, andcreate_pass_instruments()can run more than once in a session. Todaybuild_pass_datacreates only one instrumented context, so no capture is lost. If a second instrumented context is added later, the first context's records disappear from the report with no error. Keep the created instruments in a list, or assert single use.♻️ Proposed change
class StructureTreePassTool(PassInstrumentationTool): """Per-viewer tool that creates its PassContext-local capture instrument.""" def __init__(self) -> None: self.instrument: StructureTreePassInstrument | None = None + self.instruments: list[StructureTreePassInstrument] = [] def create_pass_instrument(self) -> StructureTreePassInstrument: self.instrument = StructureTreePassInstrument() + self.instruments.append(self.instrument) return self.instrument🤖 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/tools/pass_visualizer/core.py` around lines 438 - 446, Update StructureTreePassTool.create_pass_instrument so repeated calls cannot silently replace the previously created instrument: either retain every StructureTreePassInstrument in a collection for reporting, or explicitly reject calls after the first with an assertion. Ensure build_pass_data and any consumers use the retained instruments consistently.testing/python/debug/test_lower_trace.py (1)
641-671: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winGive the barrier a timeout so a regression fails instead of hanging.
threading.Barrier(2)has no timeout.run_codegencurrently calls the build function outside any shared lock, so both threads reach the barrier. If a future change moves the codegen-output lock aroundnext_call, one thread blocks on the lock and never reaches the barrier. The test then hangs the CI job instead of failing.♻️ Proposed change
- barrier = threading.Barrier(2) + barrier = threading.Barrier(2, timeout=30)🤖 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 `@testing/python/debug/test_lower_trace.py` around lines 641 - 671, Update the Barrier construction in test_concurrent_sessions_serialize_a_shared_codegen_path to include a finite timeout, so synchronization regressions raise a test failure instead of hanging. Keep the existing two-party barrier behavior unchanged.tilelang/backend/pass_pipeline/pipeline.py (1)
38-42: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThree copies of the session-plus-instrument idiom. Each site reads
current_compile_pass_instrumentation()before openingcompile_pass_instrumentation, then attaches instruments only when the session is new. That ordering rule is correct but implicit, and a future edit that readshas_sessionafter entering the session would silently attach duplicate instruments. Add one helper (for examplemanaged_compile_session(name)) intilelang/utils/pass_events.pyand call it from all three sites.
tilelang/backend/pass_pipeline/pipeline.py#L38-L42: replace thehas_sessionblock withwith managed_compile_session(f"pipeline-{self.name}"), pass_pipeline(self.name):and drop the now-unusednullcontext,current_compile_pass_instrumentation, andinstrument_current_pass_contextimports.tilelang/engine/lower.py#L310-L314: replace thehas_sessionblock inlower_to_host_device_irwithwith managed_compile_session("lower-to-host-device-ir"):.tilelang/engine/lower.py#L374-L377: replace thehas_sessionblock inlowerwithwith managed_compile_session("lower"):, then remove the module-level imports at lines 7 and 23-27 that become unused.🤖 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/backend/pass_pipeline/pipeline.py` around lines 38 - 42, Centralize the session-ordering logic in a new managed_compile_session(name) helper in tilelang/utils/pass_events.py, ensuring the existing instrumentation check occurs before opening the compile session. Update tilelang/backend/pass_pipeline/pipeline.py lines 38-42 to use it and remove the now-unused imports; update tilelang/engine/lower.py lines 310-314 and 374-377 to use managed_compile_session("lower-to-host-device-ir") and managed_compile_session("lower") respectively, then remove imports made unused by both replacements.
🤖 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.
Inline comments:
In `@testing/python/debug/test_pass_events.py`:
- Around line 231-252: Update compile_one in
test_compile_sessions_are_isolated_between_threads to pass a finite timeout to
barrier.wait(), ensuring the test fails promptly if the other worker exits
before reaching the barrier while preserving the existing synchronization
behavior.
In `@tilelang/tools/lower_trace/core.py`:
- Line 574: Update PassEventObserver.passes_incomplete in the affected observer
implementation to accept Sequence[IncompletePass] instead of
list[IncompletePass], matching the base contract and allowing tuple inputs from
StackedPassInstrument.abort; add Sequence to the TYPE_CHECKING imports if
needed.
---
Nitpick comments:
In `@testing/python/debug/test_lower_trace.py`:
- Around line 641-671: Update the Barrier construction in
test_concurrent_sessions_serialize_a_shared_codegen_path to include a finite
timeout, so synchronization regressions raise a test failure instead of hanging.
Keep the existing two-party barrier behavior unchanged.
In `@testing/python/debug/test_pass_events.py`:
- Around line 204-214: Add a test near
test_nested_helpers_reuse_the_owning_compile_session that opens an outer
compile_pass_instrumentation session, then asserts nested reuse raises
ValueError when tools is provided and separately when
include_default_tools=False.
In `@tilelang/backend/pass_pipeline/pipeline.py`:
- Around line 38-42: Centralize the session-ordering logic in a new
managed_compile_session(name) helper in tilelang/utils/pass_events.py, ensuring
the existing instrumentation check occurs before opening the compile session.
Update tilelang/backend/pass_pipeline/pipeline.py lines 38-42 to use it and
remove the now-unused imports; update tilelang/engine/lower.py lines 310-314 and
374-377 to use managed_compile_session("lower-to-host-device-ir") and
managed_compile_session("lower") respectively, then remove imports made unused
by both replacements.
In `@tilelang/tools/lower_trace/core.py`:
- Around line 713-726: Sort the entries in __all__ using isort-style ordering to
satisfy Ruff RUF022, preserving all existing exports; only add # noqa: RUF022 if
the current grouping is intentionally required.
- Around line 314-322: Expose a capture_nested option on _LowerTraceConfig and
use that configured value in create_pass_instrument instead of hard-coding True.
Ensure the setting propagates through the existing configuration path so users
can disable nested-pass capture while preserving enabled behavior and the
current default.
In `@tilelang/tools/pass_visualizer/core.py`:
- Around line 438-446: Update StructureTreePassTool.create_pass_instrument so
repeated calls cannot silently replace the previously created instrument: either
retain every StructureTreePassInstrument in a collection for reporting, or
explicitly reject calls after the first with an assertion. Ensure
build_pass_data and any consumers use the retained instruments consistently.
In `@tilelang/utils/pass_events.py`:
- Around line 137-148: Update PassEvents.finish so the first captured
finish_error is preserved when a compilation error is already present: attach it
to error or log it through the existing error-reporting mechanism, while
retaining the current behavior of raising finish_error when error is None.
Ensure finalization failures are never silently discarded.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: af13b578-88eb-4055-aa32-1516753aace3
📒 Files selected for processing (17)
docs/tools/lower_trace.mdtesting/python/debug/test_lower_trace.pytesting/python/debug/test_pass_events.pytilelang/autotuner/grouped_compile.pytilelang/backend/device_codegen.pytilelang/backend/host_codegen.pytilelang/backend/pass_pipeline/pipeline.pytilelang/engine/lower.pytilelang/jit/adapter/wrapper.pytilelang/jit/kernel.pytilelang/tools/lower_trace/core.pytilelang/tools/pass_visualizer/README.mdtilelang/tools/pass_visualizer/__init__.pytilelang/tools/pass_visualizer/core.pytilelang/tools/pass_visualizer/viewer.pytilelang/utils/pass_events.pytilelang/utils/pass_timing.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tilelang/jit/adapter/wrapper.py
| def test_compile_sessions_are_isolated_between_threads(): | ||
| barrier = threading.Barrier(2) | ||
|
|
||
| def compile_one(label): | ||
| tool = _RecordingTool(label, []) | ||
| with compile_pass_instrumentation( | ||
| name=label, | ||
| tools=[tool], | ||
| include_default_tools=False, | ||
| ) as session: | ||
| marker = create_pass_instruments()[0] | ||
| barrier.wait() | ||
| assert current_compile_pass_instrumentation() is session | ||
| return session, tool, marker | ||
|
|
||
| with ThreadPoolExecutor(max_workers=2) as pool: | ||
| left, right = list(pool.map(compile_one, ("left", "right"))) | ||
|
|
||
| assert left[0] is not right[0] | ||
| assert left[1] is not right[1] | ||
| assert left[2] is not right[2] | ||
| assert current_compile_pass_instrumentation() is None |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Bound the barrier wait to prevent a CI hang.
barrier.wait() has no timeout. If one worker raises before it reaches the barrier, for example when create_pass_instruments() returns an empty list, the other worker blocks forever. pool.map then never completes and ThreadPoolExecutor.__exit__ joins indefinitely, so the test hangs instead of failing.
🛡️ Proposed fix
- marker = create_pass_instruments()[0]
- barrier.wait()
+ marker = create_pass_instruments()[0]
+ barrier.wait(timeout=30)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def test_compile_sessions_are_isolated_between_threads(): | |
| barrier = threading.Barrier(2) | |
| def compile_one(label): | |
| tool = _RecordingTool(label, []) | |
| with compile_pass_instrumentation( | |
| name=label, | |
| tools=[tool], | |
| include_default_tools=False, | |
| ) as session: | |
| marker = create_pass_instruments()[0] | |
| barrier.wait() | |
| assert current_compile_pass_instrumentation() is session | |
| return session, tool, marker | |
| with ThreadPoolExecutor(max_workers=2) as pool: | |
| left, right = list(pool.map(compile_one, ("left", "right"))) | |
| assert left[0] is not right[0] | |
| assert left[1] is not right[1] | |
| assert left[2] is not right[2] | |
| assert current_compile_pass_instrumentation() is None | |
| def test_compile_sessions_are_isolated_between_threads(): | |
| barrier = threading.Barrier(2) | |
| def compile_one(label): | |
| tool = _RecordingTool(label, []) | |
| with compile_pass_instrumentation( | |
| name=label, | |
| tools=[tool], | |
| include_default_tools=False, | |
| ) as session: | |
| marker = create_pass_instruments()[0] | |
| barrier.wait(timeout=30) | |
| assert current_compile_pass_instrumentation() is session | |
| return session, tool, marker | |
| with ThreadPoolExecutor(max_workers=2) as pool: | |
| left, right = list(pool.map(compile_one, ("left", "right"))) | |
| assert left[0] is not right[0] | |
| assert left[1] is not right[1] | |
| assert left[2] is not right[2] | |
| assert current_compile_pass_instrumentation() is None |
🤖 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 `@testing/python/debug/test_pass_events.py` around lines 231 - 252, Update
compile_one in test_compile_sessions_are_isolated_between_threads to pass a
finite timeout to barrier.wait(), ensuring the test fails promptly if the other
worker exits before reaching the barrier while preserving the existing
synchronization behavior.
| label = f"{phase}/{pass_name}" | ||
| print_diff(before_text, after_text, f"{label} (before)", f"{label} (after)") | ||
|
|
||
| def passes_incomplete(self, passes: list[IncompletePass], error: BaseException | None) -> None: |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Widen the passes_incomplete parameter type to match the base contract.
PassEventObserver.passes_incomplete declares passes: Sequence[IncompletePass], and StackedPassInstrument.abort passes a tuple. This override narrows the type to list, which a type checker reports as an incompatible override.
♻️ Proposed change
- def passes_incomplete(self, passes: list[IncompletePass], error: BaseException | None) -> None:
+ def passes_incomplete(self, passes: Sequence[IncompletePass], error: BaseException | None) -> None:Add Sequence to the TYPE_CHECKING import block at line 52:
if TYPE_CHECKING:
- from collections.abc import Callable, Generator
+ from collections.abc import Callable, Generator, Sequence📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def passes_incomplete(self, passes: list[IncompletePass], error: BaseException | None) -> None: | |
| def passes_incomplete(self, passes: Sequence[IncompletePass], error: BaseException | None) -> None: |
🤖 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/tools/lower_trace/core.py` at line 574, Update
PassEventObserver.passes_incomplete in the affected observer implementation to
accept Sequence[IncompletePass] instead of list[IncompletePass], matching the
base contract and allowing tuple inputs from StackedPassInstrument.abort; add
Sequence to the TYPE_CHECKING imports if needed.
8c0157d to
595114a
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tilelang/autotuner/grouped_compile.py (1)
98-98: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winRemove the unused
normalized_target_hostbinding.Line 98 unpacks
normalized_target_host, but this helper does not use it. Ruff reports RUF059 for this binding. Replace it with_or use the value.Proposed fix
- host_mod, device_mod, params, normalized_target, normalized_target_host = lower_to_host_device_ir( + host_mod, device_mod, params, normalized_target, _ = lower_to_host_device_ir(🤖 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/autotuner/grouped_compile.py` at line 98, Update the unpacking assignment from lower_to_host_device_ir in the grouped compilation flow to replace the unused normalized_target_host binding with _, while preserving the other returned values and existing behavior.Source: Linters/SAST tools
🧹 Nitpick comments (1)
tilelang/engine/lower.py (1)
140-151: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider extracting the repeated session preamble into one helper.
The three-line pattern
has_session = .../compile_pass_instrumentation(...)/nullcontext() if has_session else instrument_current_pass_context()is repeated inlower_to_host_device_ir(Lines 147-150),lower_with_context(Lines 171-174), andlower(Lines 243-246).PassPipeline.lowerintilelang/backend/pass_pipeline/pipeline.pyrepeats it as well. A single context manager intilelang/utils/pass_events.pywould keep the "attach instruments only when I own the session" rule in one place.♻️ Sketch of the shared helper
Add to
tilelang/utils/pass_events.py:`@contextlib.contextmanager` def managed_compile_session(name: str) -> Generator[CompilePassInstrumentationSession, None, None]: """Own a compile session and attach its instruments only when created here.""" has_session = current_compile_pass_instrumentation() is not None with compile_pass_instrumentation(name=name) as session: attach = nullcontext() if has_session else instrument_current_pass_context() with attach: yield sessionThen each call site collapses to:
- has_session = current_compile_pass_instrumentation() is not None - with compile_pass_instrumentation(name="lower-to-host-device-ir"): - attach_instruments = nullcontext() if has_session else instrument_current_pass_context() - with attach_instruments: - return _lower_to_host_device_ir_in_session(func_or_mod, context, runtime_only) + with managed_compile_session("lower-to-host-device-ir"): + return _lower_to_host_device_ir_in_session(func_or_mod, context, runtime_only)🤖 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/engine/lower.py` around lines 140 - 151, Extract the repeated compile-session setup into a shared managed_compile_session context manager in pass_events.py, preserving the rule that instruments are attached only when the helper creates the session. Replace the duplicated preamble in lower_to_host_device_ir, lower_with_context, lower, and PassPipeline.lower with this helper while keeping each method’s existing lowering logic and session name unchanged.
🤖 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.
Inline comments:
In `@testing/python/debug/test_lower_trace.py`:
- Around line 638-668: Add a finite timeout to the barrier.wait() call inside
run_one so worker failures cannot leave the other thread blocked indefinitely.
Keep the existing barrier synchronization and concurrent test behavior
unchanged.
---
Outside diff comments:
In `@tilelang/autotuner/grouped_compile.py`:
- Line 98: Update the unpacking assignment from lower_to_host_device_ir in the
grouped compilation flow to replace the unused normalized_target_host binding
with _, while preserving the other returned values and existing behavior.
---
Nitpick comments:
In `@tilelang/engine/lower.py`:
- Around line 140-151: Extract the repeated compile-session setup into a shared
managed_compile_session context manager in pass_events.py, preserving the rule
that instruments are attached only when the helper creates the session. Replace
the duplicated preamble in lower_to_host_device_ir, lower_with_context, lower,
and PassPipeline.lower with this helper while keeping each method’s existing
lowering logic and session name unchanged.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 639f39fc-cd91-4cde-bdd5-d9377282f0ae
📒 Files selected for processing (9)
testing/python/debug/test_lower_trace.pytilelang/autotuner/grouped_compile.pytilelang/backend/device_codegen.pytilelang/backend/host_codegen.pytilelang/backend/pass_pipeline/pipeline.pytilelang/engine/lower.pytilelang/jit/kernel.pytilelang/tools/pass_visualizer/core.pytilelang/utils/pass_timing.py
🚧 Files skipped from review as they are similar to previous changes (5)
- tilelang/backend/host_codegen.py
- tilelang/backend/device_codegen.py
- tilelang/jit/kernel.py
- tilelang/backend/pass_pipeline/pipeline.py
- tilelang/tools/pass_visualizer/core.py
| def test_concurrent_sessions_serialize_a_shared_codegen_path(tmp_path): | ||
| """An explicit shared edit path is transactional across compile sessions.""" | ||
| from concurrent.futures import ThreadPoolExecutor | ||
| import threading | ||
|
|
||
| output_path = tmp_path / "shared.cpp" | ||
| barrier = threading.Barrier(2) | ||
| source = "// identical generated source\n" | ||
|
|
||
| def run_one(label): | ||
| trace = _core.LowerTraceSession( | ||
| mode="terminal", | ||
| trace_dir=str(tmp_path / label), | ||
| codegen_output=str(output_path), | ||
| ) | ||
|
|
||
| def build(_mod): | ||
| barrier.wait() | ||
| return _MockCodegenModule(source) | ||
|
|
||
| _run_codegen(trace, build, "target.build.tilelang_c") | ||
| return trace | ||
|
|
||
| with ThreadPoolExecutor(max_workers=2) as pool: | ||
| left, right = pool.map(run_one, ("left", "right")) | ||
|
|
||
| assert output_path.read_text() == source | ||
| assert (tmp_path / "shared.cpp.original").read_text() == source | ||
| assert (tmp_path / "shared.cpp.latest").read_text() == source | ||
| assert [record.index for record in left.records] == [0] | ||
| assert [record.index for record in right.records] == [0] |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Add a timeout to the barrier so a failure cannot hang the suite.
barrier.wait() has no timeout. If one worker raises before it reaches the barrier (for example during LowerTraceSession construction or inside run_codegen before next_call()), the other worker blocks forever and the test run hangs instead of failing. A timeout converts that case into a BrokenBarrierError.
🛡️ Proposed fix
def build(_mod):
- barrier.wait()
+ barrier.wait(timeout=30)
return _MockCodegenModule(source)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def test_concurrent_sessions_serialize_a_shared_codegen_path(tmp_path): | |
| """An explicit shared edit path is transactional across compile sessions.""" | |
| from concurrent.futures import ThreadPoolExecutor | |
| import threading | |
| output_path = tmp_path / "shared.cpp" | |
| barrier = threading.Barrier(2) | |
| source = "// identical generated source\n" | |
| def run_one(label): | |
| trace = _core.LowerTraceSession( | |
| mode="terminal", | |
| trace_dir=str(tmp_path / label), | |
| codegen_output=str(output_path), | |
| ) | |
| def build(_mod): | |
| barrier.wait() | |
| return _MockCodegenModule(source) | |
| _run_codegen(trace, build, "target.build.tilelang_c") | |
| return trace | |
| with ThreadPoolExecutor(max_workers=2) as pool: | |
| left, right = pool.map(run_one, ("left", "right")) | |
| assert output_path.read_text() == source | |
| assert (tmp_path / "shared.cpp.original").read_text() == source | |
| assert (tmp_path / "shared.cpp.latest").read_text() == source | |
| assert [record.index for record in left.records] == [0] | |
| assert [record.index for record in right.records] == [0] | |
| def test_concurrent_sessions_serialize_a_shared_codegen_path(tmp_path): | |
| """An explicit shared edit path is transactional across compile sessions.""" | |
| from concurrent.futures import ThreadPoolExecutor | |
| import threading | |
| output_path = tmp_path / "shared.cpp" | |
| barrier = threading.Barrier(2) | |
| source = "// identical generated source\n" | |
| def run_one(label): | |
| trace = _core.LowerTraceSession( | |
| mode="terminal", | |
| trace_dir=str(tmp_path / label), | |
| codegen_output=str(output_path), | |
| ) | |
| def build(_mod): | |
| barrier.wait(timeout=30) | |
| return _MockCodegenModule(source) | |
| _run_codegen(trace, build, "target.build.tilelang_c") | |
| return trace | |
| with ThreadPoolExecutor(max_workers=2) as pool: | |
| left, right = pool.map(run_one, ("left", "right")) | |
| assert output_path.read_text() == source | |
| assert (tmp_path / "shared.cpp.original").read_text() == source | |
| assert (tmp_path / "shared.cpp.latest").read_text() == source | |
| assert [record.index for record in left.records] == [0] | |
| assert [record.index for record in right.records] == [0] |
🤖 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 `@testing/python/debug/test_lower_trace.py` around lines 638 - 668, Add a
finite timeout to the barrier.wait() call inside run_one so worker failures
cannot leave the other thread blocked indefinitely. Keep the existing barrier
synchronization and concurrent test behavior unchanged.
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
tilelang/tools/pass_timing.py (1)
314-316: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument that
sequenceis per-run in the flattenedrecordslist.Each
TileLangPassTimingInstrumentrestartsnext_sequenceat 0. This property concatenates records from several runs, sosequencevalues repeat across contexts. A consumer that sorts the flattened list bysequencewill interleave records from different PassContexts. Add a short docstring that states the ordering contract, or pair each record with its context.♻️ Proposed docstring
`@property` def records(self) -> list[PassTimingRecord]: + """Records for every PassContext, grouped by run in creation order. + + ``PassTimingRecord.sequence`` restarts at 0 for each run, so do not + sort the returned list by ``sequence`` across runs. + """ return [record for run in self._runs for record in run.timing.records]🤖 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/tools/pass_timing.py` around lines 314 - 316, Add a concise docstring to the records property in the timing collection class stating that records are flattened in run order and that each record’s sequence is scoped to its individual run, so sequence values may repeat across runs. Preserve the existing flattening 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.
Inline comments:
In `@tilelang/autotuner/grouped_compile.py`:
- Around line 34-36: Update the compile-pass instrumentation logic in
tilelang/autotuner/grouped_compile.py lines 34-36 and tilelang/jit/kernel.py
lines 215-217 to detect an already active session and reuse it without supplying
a new timing tool; only the root session should configure timing. Add regression
coverage for nested compilation under an active outer session with
TL_PASS_PROFILE=True.
In `@tilelang/tools/pass_timing.py`:
- Line 343: Normalize the TL_PASS_PROFILE value in create_pass_timing_tool
before evaluating it, so string values are enabled only for the established
explicit-on forms and values such as "0" or "false" remain disabled. Preserve
the existing env.is_pass_profile_enabled() fallback and profiling guard
behavior.
- Around line 181-191: Update _create_tvm_instrument to construct the pass
timing instrument through TVM’s public PassInstrument API or pass_instrument
decorator instead of _ffi_instrument_api.PassInstrument. Preserve the existing
callbacks from _PassTimingState, including state.enter_pass_ctx,
state.exit_pass_ctx, state.run_before_pass, and state.run_after_pass, while
removing the private FFI dependency.
---
Nitpick comments:
In `@tilelang/tools/pass_timing.py`:
- Around line 314-316: Add a concise docstring to the records property in the
timing collection class stating that records are flattened in run order and that
each record’s sequence is scoped to its individual run, so sequence values may
repeat across runs. Preserve the existing flattening behavior.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 462d5ff7-5b2e-4ae8-bcdf-908b0157a93c
📒 Files selected for processing (12)
testing/python/debug/test_lower_trace.pytesting/python/debug/test_pass_events.pytesting/python/debug/test_pass_timing.pytesting/python/jit/test_tilelang_jit_adapter_wrapper.pytilelang/autotuner/grouped_compile.pytilelang/jit/adapter/utils.pytilelang/jit/adapter/wrapper.pytilelang/jit/kernel.pytilelang/tools/pass_timing.pytilelang/utils/__init__.pytilelang/utils/pass_events.pytilelang/utils/pass_timing.py
🚧 Files skipped from review as they are similar to previous changes (2)
- testing/python/debug/test_lower_trace.py
- tilelang/utils/pass_events.py
1 participant
Summary
Changes
Validation
Notes