[BugFix][Layout] Avoid thread-indexed wide reducer finalize readback

[LeiWang1999]

Summary

• Publish unconstrained wide-plan finalize destinations with participant-wide replicated layouts.
• Prevent fully replicated reducer results from being lowered to local arrays with thread-dependent readback such as reducer[threadIdx.x & 15].
• Keep the inferred layout unchanged when the destination copy chain has another consumer and cannot be overridden safely.

Implementation

• Reuse the finalize-destination copy-chain census to prove that every overridden layout is confined to the reducer publication chain.
• Register the selected layouts before reducer materialization so downstream copies retain static register indices.
• Cover nonzero warp-specialized thread ranges and dtype-converting staging chains.

Validation

• ./format.sh
• cmake --build build -j64 with both the default and CUDA-enabled configurations
• Focused reducer-v2 regression tests:

python -m pytest -q \
  testing/python/language/test_tilelang_language_reducer_v2.py::test_wide_plan_finalize_publication_static_indices_2d \
  testing/python/language/test_tilelang_language_reducer_v2.py::test_wide_plan_finalize_publication_static_indices_1d \
  testing/python/language/test_tilelang_language_reducer_v2.py::test_wide_plan_finalize_publication_legacy_entry \
  testing/python/language/test_tilelang_language_reducer_v2.py::test_wide_plan_finalize_staged_chain_override \
  testing/python/language/test_tilelang_language_reducer_v2.py::test_wide_plan_dst_extra_consumer_not_overridden

Result: 5 passed.

• TileKernels MHC backward benchmark (n=512, h=4096): 15.68 us via CUPTI, with no thread-dependent reducer-result loads in the generated CUDA.

The full reducer-v2 file additionally reported 51 passed and one unrelated target-sensitive packed-narrow codegen assertion (SumOp, 8 expected versus SumOp, 4 on the auto-selected target). That test does not enter the modified wide-plan branch.

Summary

• Fix wide reducer finalize publication layouts.
• Use participant-wide replicated layouts for unconstrained wide-plan finalize destinations.
• Prevent fully replicated reducer results from lowering to thread-indexed local-array readbacks.
• Preserve inferred layouts when copy chains have additional consumers.
• Register selected layouts before reducer materialization.
• Add regression tests for 1D and 2D reductions, legacy syntax, dtype-converting staging chains, additional consumers, warp-specialized ranges, and dual outputs.
• Validate formatting, default and CUDA-enabled builds, five reducer-v2 tests, and the TileKernels MHC backward benchmark.

C++ style / lint notes

• The PR changes C++ implementation code but does not change documented rules in docs/developer_guide/cpp_style.md.
• The C++ API Style Audit (warning only) remains applicable.
• No correctness or build issue is reported from style warnings. Advisory TLCPP003/TLCPP004 findings should not block merge unless they indicate a new API, FFI, or maintainability risk.

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

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: 6c5ec979-8049-4d33-9863-37b1750481ea

📥 Commits

Reviewing files that changed from the base of the PR and between 6b19ebd and bd1f88e.

📒 Files selected for processing (2)

• src/transform/reducer_plan_materialize.cc
• testing/python/language/test_tilelang_language_reducer_v2.py

Included review availability: Your plan includes up to 8 reviews per rolling hour; 7 remain after this review.

---

📝 Walkthrough

Walkthrough

The reducer materialization path now applies validated participant-wide layouts to wide-plan destinations and staging buffers. New tests verify static publication and numerical results across reduction, legacy syntax, staged-copy, and multi-output cases.

Changes

Wide-plan reducer publication

Layer / File(s)	Summary
Destination layout override src/transform/reducer_plan_materialize.cc	Wide-plan fallback validates finalize destinations and copy-only downstream chains. It registers participant-wide layouts for destinations and staging buffers before materialization.
Publication regression coverage testing/python/language/test_tilelang_language_reducer_v2.py	Tests verify static indexing, thread-zero publication guards, legacy reducer syntax, dtype-converting staged copies, multi-output destinations, and numerical results.

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

Merge Risk: ⚪ Minimal · up to bd1f8

This localized reducer layout fix includes focused regression coverage and build validation, and no actionable merge-blocking risk remains after normal checks and review.

Sequence Diagram(s)

sequenceDiagram
  participant WidePlanFallback
  participant TryWideFinalizeDstOverride
  participant LayoutRegistry
  participant ReducerMaterializer
  WidePlanFallback->>TryWideFinalizeDstOverride: validate destination and copy-use chain
  TryWideFinalizeDstOverride->>LayoutRegistry: register participant-wide layouts
  WidePlanFallback->>ReducerMaterializer: materialize with layout overrides
  ReducerMaterializer-->>WidePlanFallback: generated reducer publication

Loading

🚥 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 and concisely describes the main fix: preventing thread-indexed readback for wide reducer finalize results.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

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