[FlyOp] Add stages args for MmaMakeFragment to simplify multi-stages allocation

[sjfeng1999]

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Copilot]

Pull request overview

This PR introduces an optional stages attribute/argument for fly.mma.make_fragment to simplify allocating multi-stage MMA fragments, wiring it through the Fly dialect op definition, C++ type inference/lowering, and the Python DSL API. It also updates an example to use the new API and adds an MLIR regression test for the staged layout behavior.

Changes:

• Add stages (optional i32 attr) to fly.mma.make_fragment with updated assembly format and type inference behavior.
• Update layout lowering to allocate the fragment directly from the inferred static layout.
• Expose stages in Python make_fragment_{A,B,C} helpers, update an example, and add a new MLIR FileCheck test.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/mlir/LayoutAlgebra/mma_make_fragment_stages.mlir	New regression test covering staged vs non-staged fragment layout/allocation.
python/flydsl/expr/typing.py	Add keyword-only stages to TiledMma.make_fragment_{A,B,C} and export Numeric.
python/flydsl/expr/primitive.py	Thread stages through mma_make_fragment Python wrapper.
python/flydsl/expr/derived.py	Add make_tiled_copy_tv helper and export it.
lib/Dialect/Fly/Transforms/LayoutLowering.cpp	Change MmaMakeFragmentOp lowering to allocate directly from inferred static layout.
lib/Dialect/Fly/IR/FlyOps.cpp	Implement stages handling in MmaMakeFragmentOp return type inference.
include/flydsl/Dialect/Fly/IR/FlyOps.td	Add optional stages attribute and assembly format to the op definition.
examples/04-preshuffle_gemm.py	Use stages=2 for B fragment allocation and minor runtime/validation tweaks.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.