[BugFix][Transform] Preserve nested cast-store guards

[JayceSu98]

Fixes #2616

Summary

DecoupleTypeCast moves cast stores into generated copy loops. When a source store was nested under several if statements, the transform kept only one guard. The generated copy could therefore overwrite lanes that the source program left untouched.

The pass now combines nested no-else guards. If the surrounding control flow cannot be represented by one predicate, including an if/else, it leaves that store undecoupled. The pass may move a conversion, but it must preserve which lanes are written; that check belongs in the TileLang transform rather than in individual kernels.

Changes

• Collect and conjoin nested root-level if predicates around cast stores.
• Preserve the combined predicate on generated copy stores.
• Treat any unrepresentable statement control flow, including if/else, as ineligible for decoupling.
• Add nested-guard, branch-shape, and runtime sentinel regressions.

Review Notes

• Nested no-else guards are combined with logical conjunction.
• Unsupported branch shapes take the original path rather than a partially predicated rewrite.
• The behavior is intentionally conservative; this PR does not attempt general control-flow cloning.

Validation

• testing/python/issue/test_tilelang_issue_decouple_type_cast_nested_guards.py preserves sentinel values in unwritten FP8 lanes on H100.
• A100 runtime case is skipped by the existing compute-capability requirement for FP8.
• Relevant DecoupleTypeCast suite: 22 passed, 1 skipped.
• Python compilation and git diff --check passed.

[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

The change updates DecoupleTypeCast to conjoin nested guards, skip unsupported branch-specific control flow, and validate the resulting condition before generating copy loops. Regression tests cover both transfer directions, if/else handling, and float8 lane preservation.

Changes

Nested guard preservation

Layer / File(s)	Summary
Guard extraction and transformation tilelang/transform/decouple_type_cast.py	Nested condition-only guards are conjoined, unsupported branch control flow is skipped, and validated conditions drive generated copy loops.
Regression coverage testing/python/transform/test_tilelang_transform_decouple_type_cast.py, testing/python/issue/test_tilelang_issue_2616.py	Tests cover nested guard propagation, if/else handling, and preservation of unwritten float8 lanes.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: leiwang1999

Sequence Diagram(s)

sequenceDiagram
  participant VectorizedLoopTransformation
  participant extract_if_condition
  participant GeneratedCopyLoops
  VectorizedLoopTransformation->>extract_if_condition: Validate and combine nested predicates
  extract_if_condition-->>VectorizedLoopTransformation: Return condition or no result
  VectorizedLoopTransformation->>GeneratedCopyLoops: Create guarded copy-from and copy-to loops

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

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

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The implementation conjoins nested guards, skips unsupported branch shapes, and adds regressions directly addressing issue #2616.
Out of Scope Changes check	✅ Passed	All changes are limited to the DecoupleTypeCast fix and focused regression tests for the linked issue.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main fix: preserving guards for nested cast-store transformations.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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