Fix shared memory boundary detection bug

[Denverjin]

Summary by CodeRabbit

• Refactor
  • Improved internal code transformation pipeline for GPU scheduling and memory allocation optimization.
  • Enhanced handling of statement nesting and shared memory boundary detection during automatic code generation.

[coderabbitai]

📝 Walkthrough

Walkthrough

The changes refactor internal IR transformation logic across three files: a new statement flattener preprocessor for sequence statements, refactored warpgroup partitioning with unified segment merging, and recursive boundary marker detection for shared-memory liveness tracking.

Changes

Cohort / File(s)	Summary
Statement Flattening src/transform/auto_schedule.cc	Introduced SeqStmtFlattener mutator that recursively inlines nested SeqStmt nodes. Integrated into ReNestLetStmts to flatten statements before let/attribute nesting operations.
Warpgroup Partitioning src/transform/auto_schedule/warpgroup_partition.cc	Replaced conditional segmentation with unconditional pass and added logic to merge non-Control segments around Control boundaries. Added shared-memory liveness boundary wrapping via AttrStmt.
Shared-Memory Boundary Detection src/transform/merge_shared_memory_allocations.cc	Enhanced VisitBoundedNewScopes with recursive VisitBodyFindBoundary to traverse nested LetStmt and AttrStmt chains when locating boundary markers, replacing direct-child-only logic.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~22 minutes

Possibly related PRs

• Fix shared memory liveness detection and run formatting #8: Modifies shared-memory liveness boundary handling in warpgroup partitioning that now inserts AttrStmt boundary markers.
• Introduce letstmt defined variables in auto schedule pipeline #3: Alters let/attribute statement nesting behavior through the new SeqStmtFlattener preprocessor step.
• Add warp specialization benchmarks and improve efficiency #10: Updates boundary and sequence statement traversal logic in merge_shared_memory_allocations.cc alongside this PR's changes.

Poem

🐰 Seq statements flatten, neat and small,
Nesting rules apply to all,
Control and segments merge with care,
Boundaries hidden—we'll find them there! ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

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

✅ 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 'Fix shared memory boundary detection bug' directly corresponds to the primary change in src/transform/merge_shared_memory_allocations.cc, where boundary detection logic was updated to recursively traverse wrapper statements.

_{✏️ 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.cc`:
- Around line 873-901: Add an explicit include for the functional header so
std::function is not relied on transitively: in the file containing
SeqStmtFlattener, add `#include` <functional> near the other includes; this makes
the use of std::function<void(const Stmt&)> in the Flatten lambda and the
SeqStmtFlattener class declaration explicit and self-contained.
In `@src/transform/auto_schedule/warpgroup_partition.cc`:
- Around line 1184-1246: The ICHECK_EQ on wg0_segments.size() vs
wg1_segments.size() incorrectly aborts valid asymmetric clones created by
CloneIRStructureWithWarpgroupFilter; instead detect size mismatch before the
segmented-splitting loop and fall back to the old unsplit path by constructing
the full-stmt clones and calling MakeWarpgroupIf on them (e.g., use
ConvertIRStructureToStmt(wg0_structure.get(), outer_enable_epi) and
ConvertIRStructureToStmt(wg1_structure.get(), outer_enable_epi) or their
SegmentToStmt equivalents), set if_then_else to that result and skip the
segmented splitting logic; remove or replace the ICHECK_EQ with this conditional
fallback so the code only performs segmented splitting when wg0_segments.size()
== wg1_segments.size().

🪄 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: 0f072465-d7a3-4ae6-8b9b-f701c80ec1dc

📥 Commits

Reviewing files that changed from the base of the PR and between e2cb01b and 2d2acb7.

📒 Files selected for processing (3)

• src/transform/auto_schedule.cc
• src/transform/auto_schedule/warpgroup_partition.cc
• src/transform/merge_shared_memory_allocations.cc