Auto schedule

[Denverjin]

Summary by CodeRabbit

• New Features

  • Added shared memory liveness boundary attribute handling in auto-scheduling.
  • Enhanced warpgroup partitioning with segmented control flow management.

• Improvements

  • Improved shared memory allocation analysis with new boundary markers.
  • Extended variable substitution capabilities for IR transformations.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces shared-memory liveness boundary tracking by adding a new attribute constant, enhancing IR structure nodes with variable substitution capabilities and task-node integration, refactoring warpgroup partitioning logic to use segmented control flow with liveness boundaries, and updating shared-memory allocation analysis to recognize the new boundary attribute.

Changes

Cohort / File(s)	Summary
Attribute Definition src/op/builtin.h	Added string constant kSharedMemoryLivenessBoundary in tvm::tl::attr namespace to define attribute key for shared-memory liveness boundary handling.
IR Structure Foundation src/transform/auto_schedule/ir_structure.h	Introduced pure-virtual SubstituteVar interface and implemented it across TaskNode, ControlNode, WrapperNode, ScheduleUnit, and SequenceNode. Added task member to ControlNode and extended variable/region aggregation to include task nodes.
Auto-schedule Implementation src/transform/auto_schedule.cc	Simplified LetDecl conflict resolution, extended IRStructureBuilder::VisitStmt_ to analyze loop bounds for resource usage and populate control_node->task, updated CloneIRStructureWithWarpgroupFilter to remap variables in loop bounds and wrapper values, and refactored ApplyWarpgroupPartitionToIRStructure to use segmented warpgroup construction with liveness boundary markers between segments.
Shared Memory Allocation Analysis src/transform/merge_shared_memory_allocations.cc	Updated SharedMemLinearAccessPatternFinder to recognize kSharedMemoryLivenessBoundary attribute as a new scope boundary for liveness analysis.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Possibly related PRs

• silentCoder-dev/timelang#6: Modifies auto-schedule IR handling around node task ownership and propagation, overlapping with this PR's task-node integration and variable substitution enhancements.

Poem

🐰 Boundaries now bloom with liveness bright,
Variables substitute left and right,
Tasks nest in control nodes clean,
Warpgroups segment what's in between,
Shared memory dances, finally tight! ✨

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 34.38% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check	❓ Inconclusive	The title 'Auto schedule' is vague and does not clearly convey the specific changes made in the pull request, such as adding shared memory liveness boundary handling, refactoring conflict detection, or enhancing IR structure variable substitution.	Consider using a more descriptive title that highlights the main changes, such as 'Add shared memory liveness boundary support for auto-scheduling' or 'Refactor auto-schedule IR structure with variable substitution improvements'.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/transform/auto_schedule/ir_structure.h`:
- Around line 512-519: SubstituteVar currently updates child and task but does
not rewrite wrapper, leaving wrapper->wrapper with stale references; update the
SubstituteVar override to also visit/replace occurrences inside wrapper (e.g.,
call wrapper->SubstituteVar or directly update wrapper->wrapper) so that
WrapperNode's inner pointer is rewritten consistently; ensure the replacement
uses the same old_var/new_var pair and preserves wrapper identity expected by
ConvertIRStructureToStmt (so subsequent LetStmt/AttrStmt reconstruction sees the
updated variables).
In `@src/transform/merge_shared_memory_allocations.cc`:
- Around line 283-286: The visitor currently unwraps AttrStmtNodes inside
VisitWarpSpecializationBody which causes kSharedMemoryLivenessBoundary markers
to be skipped; update VisitWarpSpecializationBody so that when it encounters an
AttrStmtNode whose attr_key == "kSharedMemoryLivenessBoundary" it calls
VisitNewScope(op) (or otherwise dispatches to the same handling as the outer
visitor) instead of recursing directly into attr->body, ensuring the liveness
boundary attr is processed; reference VisitWarpSpecializationBody,
VisitNewScope, and AttrStmtNode to locate and change the unwrap logic so attrs
with key "kSharedMemoryLivenessBoundary" invoke VisitNewScope.

🪄 Autofix (Beta)

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: defaults

Review profile: CHILL

Plan: Pro

Run ID: b0d3e013-a1b9-4dc8-b3ee-27c7ff6e6fd6

📥 Commits

Reviewing files that changed from the base of the PR and between d10b9dc and 58237d2.

📒 Files selected for processing (4)

• src/op/builtin.h
• src/transform/auto_schedule.cc
• src/transform/auto_schedule/ir_structure.h
• src/transform/merge_shared_memory_allocations.cc