[BugFix][CUDA] Restore cp.async for predicated zero-fill copies

[hebo1221]

Fixes #2759.

A bounds-checked global-to-shared copy with a zero-fill else branch writes the same shared-memory location on every path, but PipelinePlanning currently rejects it as a conditional async producer.

This recognizes only that total-copy pattern. Loop bounds, the condition, and all source and destination indices must be pure; partial copies, nonzero fallbacks, and state-dependent guards or indices remain synchronous.

Tested:

• 48 passed, 3 skipped across the affected pipeline planning and lowering tests
• GB10 runtime check for cp.async emission and zero-filled tail values
• original sm90 reproducer emits 32 cp.async.cg.shared.global sites and 32 LDGSTS sites, matching v0.1.8

Summary

• Restores cp.async lowering for pure, predicated global-to-shared copies with total zero-fill fallbacks.
• Keeps partial copies, nonzero fallbacks, missing fallbacks, state-dependent guards, and state-dependent indices synchronous.
• Adds total_zero_fill_copy metadata to PipelineStageInfo.
• Adds pipeline-planning tests and a CUDA regression test for zero-filled tail elements.

Validation

• 48 tests passed and 3 were skipped.
• GB10 runtime validation confirmed cp.async generation and correct zero-filled tail values.
• The sm90 reproducer emitted 32 cp.async.cg.shared.global instructions and 32 LDGSTS sites.

C++ style / lint notes

• The PR changes C++ in src/transform/pipeline_planning.cc.
• The PR does not change docs/developer_guide/cpp_style.md.
• That guide documents the C++ API style audit.
• CI runs C++ API Style Audit (warning only).
• No audit findings were provided. TLCPP003/TLCPP004 warnings are advisory and should not block merging without a clear API, FFI, or maintainability risk.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: bbc6def8-0b7f-4212-a585-74acd33b5406

📥 Commits

Reviewing files that changed from the base of the PR and between 5a9b2a5 and e4ef146.

📒 Files selected for processing (3)

• src/transform/pipeline_planning.cc
• testing/python/issue/test_tilelang_issue_2759.py
• testing/python/transform/test_tilelang_transform_pipeline_planning.py

---

📝 Walkthrough

Walkthrough

Pipeline planning now detects pure parallel guarded global-to-shared copies with zero-fill fallbacks, records their stage metadata, and allows them as asynchronous producers. Planner and CUDA tests cover accepted and rejected patterns.

Changes

Predicated zero-fill copy support

Layer / File(s)	Summary
Zero-fill copy pattern analysis src/transform/pipeline_planning.cc	Structural matching identifies pure parallel global-to-shared copies with guarded zero-fill branches. PipelineStageInfo records the classification.
Async producer integration and validation src/transform/pipeline_planning.cc, testing/python/transform/test_tilelang_transform_pipeline_planning.py, testing/python/issue/test_tilelang_issue_2759.py	Classified conditional copies become asynchronous producers. Tests verify stage annotations, rejected variants, cp_async_gs generation, and zero-filled output.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant PipelineStageAnalysis
  participant ZeroFillCopyMatcher
  participant AsyncProducerEligibility
  participant CUDATestKernel
  PipelineStageAnalysis->>ZeroFillCopyMatcher: inspect guarded parallel copy
  ZeroFillCopyMatcher-->>PipelineStageAnalysis: classify total zero-fill copy
  PipelineStageAnalysis->>AsyncProducerEligibility: evaluate conditional stage
  AsyncProducerEligibility-->>PipelineStageAnalysis: allow asynchronous producer
  PipelineStageAnalysis->>CUDATestKernel: generate predicated pipeline copy
  CUDATestKernel-->>PipelineStageAnalysis: emit cp_async_gs and zero-fill behavior

Loading

Possibly related PRs

• tile-ai/tilelang#2721: Related to predicated asynchronous global-to-shared copies with zero-fill fallbacks.
• tile-ai/tilelang#2842: Related to cp.async zero-fill handling in a different legalization path.

Suggested reviewers: leiwang1999, siriusneo

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly describes restoring cp.async for predicated zero-fill copies, which is the main change.
Linked Issues check	✅ Passed	The changes address issue #2759 by restoring async lowering for pure predicated zero-fill copies while keeping unsupported cases synchronous.
Out of Scope Changes check	✅ Passed	The implementation and tests are directly related to restoring cp.async for predicated zero-fill copies.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

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

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