[Bugfix] Improve robustness of ProveFragmentContains with fully replicated layout

[LeiWang1999]

as title.

Summary by CodeRabbit

• Refactor
  • Optimized fragment replication handling with early exit logic for improved performance when fragments are fully replicated.
  • Refined parameter handling in fragment validation operations for cleaner code organization.

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

[coderabbitai]

📝 Walkthrough

Walkthrough

An optimization in ProveFragmentContains adds an early exit when a large fragment is completely replicated, and adjusts how the check_forward_index flag is passed at caller sites—one explicitly set to true, another relying on default behavior.

Changes

Cohort / File(s)	Summary
ProveFragmentContains Optimization src/op/parallel.cc	Added early exit condition: returns true immediately if the large fragment is completely replicated. Updated caller sites to explicitly pass check_forward_index flag (true in one call) and delegate without explicit flag passing in the ChooseBestCandidate contains helper.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• [Enhancement] Improve equality checks in layout nodes and fragment validation #1573: Directly modifies ProveFragmentContains and its usage of the check_forward_index flag at multiple caller sites.
• [BugFix] Corrected when proving loop layout contains a fragment buffer layout #1708: Changes how ProveFragmentContains is called in src/op/parallel.cc, including flag parameter passing adjustments.
• [Bugfix] Bind thread range for fragment inference in Parallel strict layout inference stage. #1359: Adjusts forward index and forward variable handling in src/op/parallel.cc, overlapping with the fragment replication logic.

Poem

🐰 A rabbit hops with early cheer,
When fragments fully replicate here!
No need to check the forward way,
Early exit saves the day! ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% 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 clearly and specifically identifies the main change: a bug fix that improves robustness of ProveFragmentContains when handling fully replicated layouts, which directly aligns with the code changes.

_{✏️ 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.}

[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! 🚀