Support wgmma_gemm

[Denverjin]

Summary by CodeRabbit

Release Notes

• New Features
  • Extended auto-schedule resource analysis to recognize WGMMA GEMM operations and properly classify them as tensor core compute workloads
  • Added WGMMA GEMM barrier rewriting support to the auto-scheduler
  • Updated example benchmark to use WGMMA GEMM operations for improved performance

[coderabbitai]

Caution

Review failed

Pull request was closed or merged during review

📝 Walkthrough

Walkthrough

Support for WGMMA GEMM operations is added across the TileLang framework. The FlashMLA benchmark kernel switches to using T.wgmma_gemm() for attention-score and positional encoding multiplications. The auto-schedule resource analysis and barrier rewriting logic are extended to recognize and handle these new WGMMA GEMM operation types.

Changes

Cohort / File(s)	Summary
FlashMLA Benchmark examples/auto_schedule/flashmla_benchmark.py	Replaced three T.gemm() calls with T.wgmma_gemm() in the _flashattn_body kernel's _flashmla_body loop for Step 1 and Step 2 attention-score matrix multiplications and positional encoding operations, preserving all call arguments.
Auto-schedule Resource Analysis src/transform/auto_schedule.cc	Extended resource analysis to recognize tl.tileop.wgmma_gemm_py and tl.tileop.wgmma_gemm operation types alongside existing GEMM variants, marking them for tensor core compute detection and contributing to the collected tensor core shape set.
Barrier Rewriting src/transform/auto_schedule/barrier.h	Updated barrier rewriting logic to handle WGMMA GEMM operations by recognizing them and replacing the mbar argument (index 16) when present, consistent with existing GEMM barrier handling.

Possibly related PRs

• Add warp specialization benchmarks and improve efficiency #10 — Main PR and this PR both modify the FlashMLA benchmark to use WGMMA GEMM ops and update auto-schedule/barrier handling to recognize these new GEMM operation types.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Poem

🐰 Hopping through tiles with WGMMA cheer,
New GEMM types appear, crystal clear!
Barriers rewired, schedules blessed,
Tensor cores dance in FlashMLA's nest! ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 16.67% 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 'Support wgmma_gemm' clearly and concisely describes the main change: adding support for the wgmma_gemm operation across the codebase (examples, auto-schedule resource analysis, and barrier rewriting).

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}