[BugFix] Fix boundary check in loop layout

[tzj-fxz]

As title

Summary by CodeRabbit

• Improvements

  • Runtime guards now only trigger when the inverse mapping is non‑bijective or loop bounds mismatch expected shapes.
  • Added safety checks to ensure fragment/layout indices stay within loop extents, emitting predicates when needed.
  • Inverse/detection logic now iteratively relaxes checks for symbolic shapes and adjusts decision levels.

• Tests

  • Added boundary-check tests validating loop-partition transformation equivalence.
  • Updated transform tests to precompute modules and verify lowering/partitioning behavior.

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

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

📝 Walkthrough

Walkthrough

Add a bounds_match parity check to loop-partition guard logic and make need_guard depend on inverse-map bijectivity or bounds mismatch; emit IfThenElse guards when required. Add loop-extent safety predicates in ParallelOp::InferLayout. Make DetectIterMap iterative with level escalation. Add/adjust TileLang transform tests.

Changes

Cohort / File(s)	Summary
Loop partition update src/transform/loop_partition.cc	Add bounds_match parity check comparing inner-loop extents/minima to inverse-layout inputs; make need_guard depend on inverse bijectivity OR bounds_match; insert IfThenElse guard when necessary; minor formatting tweaks.
ParallelOp layout safety src/op/parallel.cc	When loop_layout_ exists, synthesize a thread-index vector, compute inverse logical indices, and emit LT predicates ensuring each logical index < original loop extent (added before block-size checks).
Layout inverse checks src/layout/layout.cc	Replace single-call DetectIterMap with iterative calls that escalate/adjust the level until errors clear or NoCheck; return final result/level.
TileLang transform tests testing/python/transform/test_tilelang_transform_loop_partition_boundary.py, testing/python/transform/test_tilelang_transform_lower_tile_op.py	Add a boundary-check test validating before/after TileLang loop-partition transforms; precompute before_mod/after_mod in LowerTileOp test and use explicit PassContext for reference construction; compare transformed modules structurally.

Sequence Diagram(s)

sequenceDiagram
  rect rgba(240,248,255,0.5)
  participant TL as TileLang Transform
  participant LP as LoopPartition
  participant LD as LayoutDetector
  participant RG as Guard/PredicateGen
  participant PL as ParallelOp::InferLayout
  end
  TL->>LP: run loop-partition transform
  LP->>LD: DetectIterMap / compute inverse layout
  LD-->>LP: inverse map + level/result
  LP->>LP: compute bounds_match (inner extents vs inverse inputs)
  alt need_guard (not bijective OR bounds mismatch)
    LP->>RG: emit IfThenElse guard (runtime bound checks)
    RG-->>TL: guarded transformed loop
  else bijective & bounds_match
    LP-->>TL: emit unguarded transformed loop
  end
  TL->>PL: InferLayout (with loop_layout_)
  PL->>LD: compute inverse logical indices from synthetic thread index
  PL->>RG: emit LT predicates (logical_i < original_extent) as needed
  RG-->>PL: predicates attached to layout inference

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• [Bugfix] Recover code for flexible parallel #1032: Modifications to src/op/parallel.cc adding predicates tying inverse layout/logical indices to loop extents.
• [Refactor] Relocate layout transformation of ptx_stmatrix  #1689: Related changes in src/op/parallel.cc for loop-extent vs layout/input extent safety checks.
• [Parallel] Support T.Parallel with dynamic extents #990: Overlaps with guard generation and inverse-mapping-based bounds checks in src/transform/loop_partition.cc.

Suggested reviewers

• LeiWang1999
• kurisu6912

Poem

🐇 I hopped through loops to match each bound,

I checked inverse maps where answers are found,
I placed a guard when shapes did not align,
I nudged predicates so indices stay fine,
🥕 A tiny thump — transforms now safe and sound

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 14.29% 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 '[BugFix] Fix boundary check in loop layout' directly relates to the main changes which introduce boundary checks and validation in loop layout handling across multiple files.

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

✨ Finishing touches

• 📝 Generate docstrings

---

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.}