[ROCDL] Add make_tiled_tdm_atom op and tdm_partition

[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 first-class CDNA5 (gfx1250) Tensor Data Mover (TDM) “tiled atom” construction, including new FlyROCDL ops/types, a dedicated expansion pass to materialize runtime descriptor arguments and normalize boundary-check state, and supporting changes so copy paths can be addressed by coordinate tensors (TMA-style) instead of memrefs.

Changes:

• Add CDNA5 TDM load/store copy atom types plus geometry derivation, partitioning, and lowering support (C++ + Python bindings).
• Add fly_rocdl.make_tiled_tdm_{load,store}_atom builders and a fly-rocdl-expand-ops pass to expand them and normalize boundary_check state.
• Generalize Fly copy/copy_atom_call infrastructure and ROCm pipelines to support coord-tensor operands and construction-argument-bearing copy atoms.

Reviewed changes

Copilot reviewed 30 out of 30 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/mlir/Conversion/tdm_cdna5.mlir	New FileCheck coverage for CDNA5 TDM descriptor construction, boundary-check folding, iteration, recast, and barrier behavior.
python/flydsl/expr/rocdl/cdna5.py	Python API for CDNA5 TDM load/store atoms, tiled-atom builder, and tdm_partition helper.
python/flydsl/expr/primitive.py	Update make_copy_atom call signature to pass explicit args list; avoid unnecessary set_value when no kwargs in copy().
python/flydsl/compiler/backends/rocm.py	Insert fly-rocdl-expand-ops into the ROCm compilation pipeline.
lib/Dialect/FlyROCDL/Transforms/ExpandOps.cpp	New pass to expand tiled-TDM atom builders and normalize boundary_check → boundary_check_axes.
lib/Dialect/FlyROCDL/Transforms/ClusterAttr.cpp	New pass to translate rocdl.cluster_dims into LLVM passthrough amdgpu-cluster-dims.
lib/Dialect/FlyROCDL/Ops.cpp	Implement type inference for make_tiled_tdm_{load,store}_atom ops.
lib/Dialect/FlyROCDL/CMakeLists.txt	Wire up new CDNA5 sources and FlyROCDL transform pass tablegen target.
lib/Dialect/FlyROCDL/CDNA5/TdmGeometry.cpp	Implement CDNA5 TDM geometry derivation and partition-layout construction.
lib/Dialect/FlyROCDL/CDNA5/TdmAtomBuilder.cpp	Builder-side derivation helpers, partition-layout entrypoint, and type constructors for CDNA5 TDM load/store ops.
lib/Dialect/FlyROCDL/CDNA5/CopyAtom.cpp	Implement CDNA5 TDM atom lowering, state layout, boundary-check semantics, padding/iteration encoding, and ROCDL intrinsic emission.
lib/Dialect/Fly/Utils/IntTupleUtils.cpp	Update basis→tuple value materialization to match attribute-side behavior and support dynamic basis values.
lib/Dialect/Fly/Transforms/LayoutLowering.cpp	Allow coord-tensor operands in copy lowering and improve “single atom call” emission heuristics.
lib/Dialect/Fly/Transforms/ConvertAtomCallToSSAForm.cpp	Treat coord-tensor operands as non-promotable and avoid invalid memref casts.
lib/Conversion/FlyToROCDL/FlyToROCDL.cpp	Convert coord tensors to int-tuples, simplify make_view lowering, validate construction args for stateless atoms, and relax copy lowering to allow coord tensors.
lib/CAPI/Dialect/FlyROCDL/FlyROCDLDialect.cpp	Register FlyROCDL passes via the dialect’s pass registry entrypoint.
lib/Bindings/Python/FlyROCDLExtension.cpp	Add Python bindings for CDNA5 TDM types and expose tdm_partition_layout.
include/flydsl/Dialect/FlyROCDL/Utils/TdmGeometry.h	New public API for CDNA5 TDM geometry/partitioning utilities and constants.
include/flydsl/Dialect/FlyROCDL/Utils/TdmAtomBuilder.h	New public API for deriving tiled-TDM atoms and computing partition layouts.
include/flydsl/Dialect/FlyROCDL/Transforms/Passes.td	Define fly-rocdl-expand-ops and fly-rocdl-cluster-attr passes.
include/flydsl/Dialect/FlyROCDL/Transforms/Passes.h	Declare/register FlyROCDL transform passes.
include/flydsl/Dialect/FlyROCDL/Transforms/CMakeLists.txt	Tablegen target for FlyROCDL transform pass declarations/registration.
include/flydsl/Dialect/FlyROCDL/IR/Ops.td	Add make_tiled_tdm_{load,store}_atom op definitions.
include/flydsl/Dialect/FlyROCDL/IR/CopyAtom.td	Add CDNA5 cdna5.tensor_load / cdna5.tensor_store copy-op types and parameters.
include/flydsl/Dialect/FlyROCDL/IR/Atom.td	Extend shared atom state fields to include CDNA5 TDM barrier/timeout state.
include/flydsl/Dialect/FlyROCDL/CMakeLists.txt	Include FlyROCDL Transforms subdirectory in the build.
include/flydsl/Dialect/Fly/IR/FlyOps.td	Add args operands to fly.make_copy_atom; broaden copy/copy_atom_call to accept “tensor-like” operands.
include/flydsl/Dialect/Fly/IR/FlyInterfaces.td	Add getAtomState(builder, loc, args) to the stateful op type interface with a safe default implementation.
include/flydsl/Conversion/FlyToROCDL/Passes.td	Remove cluster-attr pass definition from conversion passes (now owned by FlyROCDL transforms).
examples/06-cdna5_tensor_copy.py	New end-to-end example using CDNA5 tiled-TDM atoms + partitioning helpers.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.