[ROCDL] Add make_tiled_tdm_atom op and tdm_partition

examples/06-cdna5_tensor_copy.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 FlyDSL Project Contributors
+
+"""Point-to-point tensor copy with the CDNA5 TDM engine (gfx1250).
+
+``TENSOR_LOAD_TO_LDS`` and ``TENSOR_STORE_FROM_LDS`` move a whole tile between global
+memory and LDS on their own: no VGPRs, no per-lane addressing, EXEC ignored. A copy is
+therefore just the two of them back to back, with an ``s_wait_tensorcnt`` in between.
+"""
+
+import torch
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+
+TM, TN = 128, 64
+
+
+@flyc.kernel
+def tdm_copy_kernel(
+    A: fx.Tensor,
+    B: fx.Tensor,
+    tm: fx.Constexpr[int] = TM,
+    tn: fx.Constexpr[int] = TN,
+):
+    # The LDS tile is packed (no row padding), which is what the store needs: TDM
+    # drains LDS with the tile stride and has no de-padding of its own.
+    lds = fx.SharedAllocator().allocate(fx.Array[fx.Float16, tm * tn]).peek()
+    smem_layout = fx.make_layout((tm, tn), (tn, 1))
+    smem_tensor = fx.make_view(lds.ptr, smem_layout)
+
+    # One atom per direction -- the direction is a property of the instruction.
+    # Each carries its own tensor's pointer/stride/extent, so a caller cannot
+    # pair an atom with a coordinate from a different tensor.
+    tdm_load_atom, tdmA = fx.rocdl.cdna5.make_tiled_tdm_atom(fx.rocdl.TensorLoad(), A, smem_layout, (tm, tn))
+    tdm_store_atom, tdmB = fx.rocdl.cdna5.make_tiled_tdm_atom(fx.rocdl.TensorStore(), B, smem_layout, (tm, tn))
+
+    # Taking this block's tile is zipped_divide + slice.
+    blk_tdmA = fx.zipped_divide(tdmA, (tm, tn))[None, (fx.block_idx.x, fx.block_idx.y)]
+    blk_tdmB = fx.zipped_divide(tdmB, (tm, tn))[None, (fx.block_idx.x, fx.block_idx.y)]
+
+    # One layout cuts both sides, so they keep describing the same elements.
+    # There is no thread index in it -- TDM is issued by a single wave, so every
+    # lane sees one partition.
+    #
+    #  The warp coordinate says which warp's share this is: one warp does the
+    # whole tile here, so it is the trivial `0` over a size-1 layout. Splitting
+    # the tile between N warps is `num_warps=N` on the atom and this warp's
+    # index over `make_layout(N)` here.
+    warp_crd, warp_layout = 0, fx.make_layout(1, 1)
+
+    tAsA, tAgA = fx.rocdl.cdna5.tdm_partition(tdm_load_atom, warp_crd, warp_layout, smem_tensor, blk_tdmA)
+    tBsB, tBgB = fx.rocdl.cdna5.tdm_partition(tdm_store_atom, warp_crd, warp_layout, smem_tensor, blk_tdmB)
+
+    fx.copy(tdm_load_atom, tAgA, tAsA)
+    fx.rocdl.s_wait_tensorcnt(0)
+    fx.barrier()
+
+    fx.copy(tdm_store_atom, tBsB, tBgB)
+    fx.rocdl.s_wait_tensorcnt(0)
+
+
+@flyc.jit
+def tdm_tensor_copy(
+    A: fx.Tensor,
+    B: fx.Tensor,
+    m: fx.Int32,
+    n: fx.Int32,
+    stream: fx.Stream = fx.Stream(None),
+):
+    grid = ((m + TM - 1) // TM, (n + TN - 1) // TN, 1)
+    tdm_copy_kernel(A, B).launch(grid=grid, block=(32, 1, 1), stream=stream)
+
+
+# Deliberately not a multiple of the tile in either dim: the atom is built with
+# `init_boundary_check=True` by default, so the descriptor clamps the ragged edge tiles and the
+# lowering derives the clamp from the same coordinate that moved the base address.
+M, N = 128 * 2 + 40, 64 * 3 + 24
+
+A = torch.arange(M * N, dtype=torch.int32).reshape(M, N).to(torch.float16).cuda()
+B = torch.zeros(M, N, dtype=torch.float16).cuda()
+
+tdm_tensor_copy(A, B, M, N)
+
+torch.cuda.synchronize()
+ok = torch.equal(A.cpu(), B.cpu())
+print(f"Result correct: {ok}")

include/flydsl/Conversion/FlyToROCDL/Passes.td

@@ -13,18 +13,3 @@ def FlyToROCDLConversionPass : Pass<"convert-fly-to-rocdl"> {
     "ROCDL::ROCDLDialect"
   ];
 }
-
-def FlyROCDLClusterAttrPass : Pass<"fly-rocdl-cluster-attr"> {
-  let summary = "Inject amdgpu-cluster-dims into llvm.func passthrough";
-  let description = [{
-    Converts the `rocdl.cluster_dims` discardable attribute on `llvm.func`
-    into an LLVM IR `passthrough` entry containing `amdgpu-cluster-dims`.
-
-    This is a hack to work around the upstream ROCDL dialect not supporting
-    cluster_dims translation.  Run this pass inside `gpu.module(...)` after
-    `convert-gpu-to-rocdl` so that `llvm.func` ops already exist.
-  }];
-  let dependentDialects = [
-    "LLVM::LLVMDialect"
-  ];
-}

include/flydsl/Dialect/Fly/IR/FlyInterfaces.td

@@ -34,6 +34,20 @@ def Fly_StatefulOpTypeInterface : TypeInterface<"StatefulOpTypeInterface"> {
       "Build the default state value for this op.",
       "::mlir::Value", "getDefaultState",
       (ins "::mlir::OpBuilder &":$builder, "::mlir::Location":$loc)>,
+    InterfaceMethod<
+      "Build the initial state value for an atom materialized with `args`.",
+      "::mlir::Value", "getAtomState",
+      (ins "::mlir::OpBuilder &":$builder, "::mlir::Location":$loc,
+           "::mlir::ValueRange":$args),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        if (!args.empty()) {
+          ::mlir::emitError(loc) << "this atom takes no construction arguments, got "
+                                 << args.size();
+          return nullptr;
+        }
+        return $_type.getDefaultState(builder, loc);
+      }]>,
     InterfaceMethod<
       "Insert the given value to a field of this stateful op. "
       "Returns the updated struct value after setting the field.",

include/flydsl/Dialect/Fly/IR/FlyOps.td

@@ -347,9 +347,16 @@ def Fly_MakeMmaAtomOp : Fly_Op<"make_mma_atom", [Pure]> {
 }
 
 def Fly_MakeCopyAtomOp : Fly_Op<"make_copy_atom", [Pure]> {
-  let arguments = (ins I32Attr:$valBits);
+  let arguments = (ins Variadic<AnyType>:$args, I32Attr:$valBits);
   let results = (outs Fly_CopyAtom:$result);
-  let assemblyFormat = "attr-dict `:` qualified(type($result))";
+  let assemblyFormat = [{
+    (`(` $args^ `:` type($args) `)`)? attr-dict `:` qualified(type($result))
+  }];
+  let builders = [
+    OpBuilder<(ins "::mlir::Type":$result, "int32_t":$valBits), [{
+      build($_builder, $_state, result, ::mlir::ValueRange{}, valBits);
+    }]>
+  ];
 }
 
 def Fly_AtomSetValueOp : Fly_Op<"atom.set_value", [Pure, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
@@ -359,7 +366,8 @@ def Fly_AtomSetValueOp : Fly_Op<"atom.set_value", [Pure, DeclareOpInterfaceMetho
 }
 
 def Fly_CopyAtomCall : Fly_Op<"copy_atom_call"> {
-  let arguments = (ins Fly_CopyAtom:$copyAtom, Fly_MemRef:$src, Fly_MemRef:$dst, Optional<Fly_MemRef>:$pred);
+  let arguments = (ins Fly_CopyAtom:$copyAtom, Fly_TensorLikeType:$src,
+                       Fly_TensorLikeType:$dst, Optional<Fly_MemRef>:$pred);
 }
 def Fly_MmaAtomCall : Fly_Op<"mma_atom_call"> {
   let arguments = (ins Fly_MmaAtom:$mmaAtom, Fly_MemRef:$d, Fly_MemRef:$a, Fly_MemRef:$b, Fly_MemRef:$c);
@@ -416,7 +424,8 @@ def Fly_MmaMakeFragmentOp : Fly_Op<"mma.make_fragment", [Pure, DeclareOpInterfac
 }
 
 def Fly_CopyOp : Fly_Op<"copy"> {
-  let arguments = (ins AnyType:$copyAtom, Fly_MemRef:$src, Fly_MemRef:$dst, Optional<Fly_MemRef>:$pred);
+  let arguments = (ins AnyType:$copyAtom, Fly_TensorLikeType:$src, Fly_TensorLikeType:$dst,
+                       Optional<Fly_MemRef>:$pred);
   let assemblyFormat = "`(` $copyAtom `,` $src `,` $dst (`,` $pred^)? `)` attr-dict `:` functional-type(operands, results)";
 }
 def Fly_GemmOp : Fly_Op<"gemm"> {

include/flydsl/Dialect/FlyROCDL/CMakeLists.txt

@@ -1 +1,2 @@
 add_subdirectory(IR)
+add_subdirectory(Transforms)

include/flydsl/Dialect/FlyROCDL/IR/Atom.td

@@ -11,12 +11,10 @@ def FlyROCDL_AtomStateField : I32EnumAttr<"AtomStateField", "", [
   I32EnumAttrCase<"ImmOffset",          1, "imm_offset">,
   I32EnumAttrCase<"ScaleA",             2, "scale_a">,
   I32EnumAttrCase<"ScaleB",             3, "scale_b">,
-  I32EnumAttrCase<"WorkgroupMask",      4, "workgroup_mask">
-  // NOTE: the gfx1250 TDM N-D descriptor's per-dim tensor extent (OOB) and per-dim
-  // tensor stride are NOT shared enum fields. They are resolved privately by
-  // CopyOpGFX1250TDMType from the field names "extent_0".."extent_4" /
-  // "stride_0".."stride_3" (set via fly.atom.set_value), keeping this cross-atom
-  // vocabulary free of TDM's per-dim sprawl. See GFX1250/CopyAtom.cpp (tdmGeomSlot).
+  // Introduced in CDNA5.TensorLoad/Store.
+  I32EnumAttrCase<"WorkgroupMask",      4, "workgroup_mask">,
+  I32EnumAttrCase<"AtomicBarrierAddr",  5, "atomic_barrier_addr">,
+  I32EnumAttrCase<"EarlyTimeout",       6, "early_timeout">
 ]> {
   let genSpecializedAttr = 0;
   let cppNamespace = FlyROCDL_Dialect.cppNamespace;

include/flydsl/Dialect/FlyROCDL/IR/CopyAtom.td

@@ -88,4 +88,58 @@ def FlyROCDL_CopyOpGFX1250TDM : FlyROCDL_StatefulCopyOp<"CopyOpGFX1250TDM", "gfx
   let genVerifyDecl = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// CopyOp CDNA5 — N-D TDM (Tensor Data Mover) async Global <-> LDS copy (1-5D),
+// addressed by a tile coordinate.
+//===----------------------------------------------------------------------===//
+
+class FlyROCDL_CopyOpCDNA5TDM<string typeName, string typeMnemonic>
+    : FlyROCDL_StatefulCopyOp<typeName, typeMnemonic, [
+        DeclareTypeInterfaceMethods<Fly_StatefulOpTypeInterface, ["getAtomState"]>
+      ]> {
+  dag tdmParams = (ins
+    ArrayRefParameter<"int32_t">:$tileShape,
+    // The unit the *descriptor* counts. Only its width reaches the hardware (`data_size`).
+    "Type":$dataType,
+    // Where each mode of the *global* tensor sends its *tdm* basis axis, congruent
+    // with that tensor's shape.
+    //
+    // Being an attribute is a real restriction, not just an encoding choice: a leaf is
+    // `scale E axis`, so the scale is a compile-time integer. One axis per mode is exact
+    // while the map is 1:1, which it is up to five modes.
+    // TODO(rank>5, dynamic strides)
+    "::mlir::fly::IntTupleAttr":$tensor2tdm,
+    // Whether this atom arrives on the HW auto-barrier (descriptor config bit 18).
+    // Only the enable is a type property: *which* barrier is an LDS address, and that
+    // is atom state. With this off the atom has no `atomic_barrier_addr` field to set.
+    DefaultValuedParameter<"bool", "false">:$atomicBarrier,
+    DefaultValuedParameter<"int32_t", "0">:$cacheModifier,
+    DefaultValuedParameter<"int32_t", "1">:$iterCount
+  );
+  let genVerifyDecl = 1;
+}
+
+def FlyROCDL_CopyOpCDNA5TensorLoad : FlyROCDL_CopyOpCDNA5TDM<"CopyOpCDNA5TensorLoad", "cdna5.tensor_load"> {
+  let parameters = !con(tdmParams, (ins
+    DefaultValuedParameter<"int32_t", "0">:$padInterval,
+    DefaultValuedParameter<"int32_t", "0">:$padAmount
+  ));
+
+  let assemblyFormat = [{
+    `<` `shape` `=` `[` $tileShape `]` `,` `elem` `=` $dataType
+    `,` `tensor2tdm` `=` $tensor2tdm
+    (`,` struct($atomicBarrier, $cacheModifier, $iterCount, $padInterval, $padAmount)^)? `>`
+  }];
+}
+
+def FlyROCDL_CopyOpCDNA5TensorStore : FlyROCDL_CopyOpCDNA5TDM<"CopyOpCDNA5TensorStore", "cdna5.tensor_store"> {
+  let parameters = tdmParams;
+
+  let assemblyFormat = [{
+    `<` `shape` `=` `[` $tileShape `]` `,` `elem` `=` $dataType
+    `,` `tensor2tdm` `=` $tensor2tdm
+    (`,` struct($atomicBarrier, $cacheModifier, $iterCount)^)? `>`
+  }];
+}
+
 #endif // FLYROCDL_COPYATOM

include/flydsl/Dialect/FlyROCDL/IR/Ops.td

@@ -29,4 +29,35 @@ def FlyROCDL_GetBufferRsrcOp : FlyROCDL_Op<"get_buffer_rsrc",
   let assemblyFormat = "`(` $ptr `)` attr-dict `:` functional-type($ptr, $result)";
 }
 
+class FlyROCDL_MakeTiledTdmAtomOp<string mnemonic>
+    : FlyROCDL_Op<mnemonic, [Pure, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+  dag tdmArgs = (ins
+    AnyType:$tensor,
+    AnyType:$smemLayout,
+    AnyType:$tiler,
+    DefaultValuedOptionalAttr<I32Attr, "1">:$numWarps,
+    DefaultValuedOptionalAttr<BoolAttr, "true">:$initBoundaryCheck,
+    DefaultValuedOptionalAttr<I32Attr, "0">:$cacheModifier,
+    DefaultValuedOptionalAttr<BoolAttr, "false">:$atomicBarrier,
+    OptionalAttr<TypeAttr>:$internalType
+  );
+  let results = (outs AnyType:$atom, AnyType:$coordTensor);
+
+  let assemblyFormat = [{
+    $tensor `,` $smemLayout `,` $tiler attr-dict `:` type(operands)
+  }];
+}
+
+def FlyROCDL_MakeTiledTdmLoadAtomOp
+    : FlyROCDL_MakeTiledTdmAtomOp<"make_tiled_tdm_load_atom"> {
+  let summary = "Build a CDNA5 TDM load atom and the coordinate tensor that addresses it.";
+  let arguments = tdmArgs;
+}
+
+def FlyROCDL_MakeTiledTdmStoreAtomOp
+    : FlyROCDL_MakeTiledTdmAtomOp<"make_tiled_tdm_store_atom"> {
+  let summary = "Build a CDNA5 TDM store atom and the coordinate tensor that addresses it.";
+  let arguments = tdmArgs;
+}
+
 #endif // FLYROCDL_OPS

include/flydsl/Dialect/FlyROCDL/Transforms/CMakeLists.txt

@@ -0,0 +1,4 @@
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name FlyROCDL)
+
+add_mlir_generic_tablegen_target(FlyROCDLTransformPassIncGen)

include/flydsl/Dialect/FlyROCDL/Transforms/Passes.h

@@ -0,0 +1,24 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 FlyDSL Project Contributors
+
+#ifndef FLYDSL_DIALECT_FLYROCDL_TRANSFORMS_PASSES_H
+#define FLYDSL_DIALECT_FLYROCDL_TRANSFORMS_PASSES_H
+
+#include "mlir/Pass/Pass.h"
+
+#include "flydsl/Dialect/Fly/IR/FlyDialect.h"
+#include "flydsl/Dialect/FlyROCDL/IR/Dialect.h"
+
+namespace mlir {
+namespace fly_rocdl {
+
+#define GEN_PASS_DECL
+#include "flydsl/Dialect/FlyROCDL/Transforms/Passes.h.inc"
+
+#define GEN_PASS_REGISTRATION
+#include "flydsl/Dialect/FlyROCDL/Transforms/Passes.h.inc"
+
+} // namespace fly_rocdl
+} // namespace mlir
+
+#endif // FLYDSL_DIALECT_FLYROCDL_TRANSFORMS_PASSES_H

include/flydsl/Dialect/FlyROCDL/Transforms/Passes.td

@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 FlyDSL Project Contributors
+
+#ifndef FLYROCDL_PASSES
+#define FLYROCDL_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def FlyROCDLExpandOpsPass : Pass<"fly-rocdl-expand-ops"> {
+  let summary = "Expand fly_rocdl target specific ops into the ops";
+  let dependentDialects = [
+    "arith::ArithDialect",
+    "mlir::fly::FlyDialect"
+  ];
+}
+
+def FlyROCDLClusterAttrPass : Pass<"fly-rocdl-cluster-attr"> {
+  let summary = "Inject amdgpu-cluster-dims into llvm.func passthrough";
+  let description = [{
+    Converts the `rocdl.cluster_dims` discardable attribute on `llvm.func`
+    into an LLVM IR `passthrough` entry containing `amdgpu-cluster-dims`.
+
+    This is a hack to work around the upstream ROCDL dialect not supporting
+    cluster_dims translation.  Run this pass inside `gpu.module(...)` after
+    `convert-gpu-to-rocdl` so that `llvm.func` ops already exist -- it rewrites
+    an attribute on those, and so is not part of the Fly -> ROCDL conversion.
+  }];
+  let dependentDialects = [
+    "LLVM::LLVMDialect"
+  ];
+}
+
+#endif // FLYROCDL_PASSES

include/flydsl/Dialect/FlyROCDL/Utils/TdmAtomBuilder.h

@@ -0,0 +1,38 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 FlyDSL Project Contributors
+
+#ifndef FLYDSL_DIALECT_FLYROCDL_UTILS_TDMATOMBUILDER_H
+#define FLYDSL_DIALECT_FLYROCDL_UTILS_TDMATOMBUILDER_H
+
+#include "flydsl/Dialect/FlyROCDL/IR/Dialect.h"
+#include "flydsl/Dialect/FlyROCDL/Utils/TdmGeometry.h"
+
+namespace mlir::fly_rocdl {
+
+/// Read a builder's operands and attributes into a `tdm::Request`, run the derivation, and
+/// report the descriptor's data type.
+template <typename AdaptorT>
+FailureOr<tdm::Geometry> deriveTdmAtom(AdaptorT adaptor, tdm::Request &request, Type &dataType,
+                                       function_ref<InFlightDiagnostic()> emitError);
+
+extern template FailureOr<tdm::Geometry>
+deriveTdmAtom<MakeTiledTdmLoadAtomOpAdaptor>(MakeTiledTdmLoadAtomOpAdaptor, tdm::Request &, Type &,
+                                             function_ref<InFlightDiagnostic()>);
+extern template FailureOr<tdm::Geometry>
+deriveTdmAtom<MakeTiledTdmStoreAtomOpAdaptor>(MakeTiledTdmStoreAtomOpAdaptor, tdm::Request &,
+                                              Type &, function_ref<InFlightDiagnostic()>);
+
+FailureOr<Type> tdmLoadOpType(MLIRContext *ctx, const tdm::Geometry &geometry, Type dataType,
+                              fly::IntTupleAttr tensor2tdm, bool atomicBarrier,
+                              int32_t cacheModifier, function_ref<InFlightDiagnostic()> emitError);
+FailureOr<Type> tdmStoreOpType(MLIRContext *ctx, const tdm::Geometry &geometry, Type dataType,
+                               fly::IntTupleAttr tensor2tdm, bool atomicBarrier,
+                               int32_t cacheModifier, function_ref<InFlightDiagnostic()> emitError);
+
+FailureOr<fly::LayoutType> tdmPartitionLayout(Type atomType, Type stensorType, Type gtensorType,
+                                              int32_t numWarps,
+                                              function_ref<InFlightDiagnostic()> emitError);
+
+} // namespace mlir::fly_rocdl
+
+#endif // FLYDSL_DIALECT_FLYROCDL_UTILS_TDMATOMBUILDER_H