Gfx1250 tdm gather scatter

NOTICE

@@ -0,0 +1,41 @@
+FlyDSL
+Copyright (c) 2025 FlyDSL Project Contributors
+
+This product includes software developed as part of the FlyDSL project,
+licensed under the Apache License, Version 2.0 (see LICENSE).
+
+------------------------------------------------------------------------
+Third-party components
+------------------------------------------------------------------------
+
+Triton (https://github.com/triton-lang/triton) — MIT License
+
+The gfx1250 TDM (Tensor Data Mover) copy-atom lowering in
+  lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp
+is a port of the AMD gfx1250 TDM descriptor bitfield layout, the N-D warp
+distribution, and the row-gather descriptor packing from Triton's AMD backend
+(third_party/amd/lib/TritonAMDGPUToLLVM/TDMUtility.cpp and
+.../backend/include/TDMCommon.h: createTDMDescriptor / fillTDMDescriptor /
+tdmGetWarpDistribution). Triton is used under the MIT License:
+
+    Copyright 2018-2020 Philippe Tillet
+    Copyright 2020-2022 OpenAI
+
+    Permission is hereby granted, free of charge, to any person obtaining
+    a copy of this software and associated documentation files (the
+    "Software"), to deal in the Software without restriction, including
+    without limitation the rights to use, copy, modify, merge, publish,
+    distribute, sublicense, and/or sell copies of the Software, and to
+    permit persons to whom the Software is furnished to do so, subject to
+    the following conditions:
+
+    The above copyright notice and this permission notice shall be
+    included in all copies or substantial portions of the Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+    EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+    MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+    IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+    CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+    TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+    SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

README.md

@@ -406,7 +406,7 @@ FlyDSL's design is inspired by ideas from several projects:
 - [NVIDIA CUTLASS](https://github.com/NVIDIA/cutlass) — CuTe layout algebra concepts (BSD-3-Clause parts only; no EULA-licensed code was referenced)
 - [ROCm Composable Kernel](https://github.com/ROCm/composable_kernel) — tile-based kernel design patterns for AMD GPUs
 - [ROCm AIter](https://github.com/ROCm/aiter) — test infrastructure and performance comparison baselines (MIT)
-- [Triton](https://github.com/triton-lang/triton) — Python DSL for GPU kernel authoring
+- [Triton](https://github.com/triton-lang/triton) — Python DSL for GPU kernel authoring; the gfx1250 TDM copy-atom descriptor packing is ported from Triton's AMD backend (MIT, see [`NOTICE`](NOTICE))
 - [HipKittens](https://github.com/HazyResearch/HipKittens) — minimal, opinionated C++ embedded primitives for fast AMD AI kernels (part of the ThunderKittens family)
 
 ## 📄 License

examples/05-gather_scatter.py

@@ -21,6 +21,18 @@
     offset tensor : (TV, Rest...)
     pred tensor   : (TV, Rest...)  optional
 
+The same ``fx.gather`` / ``fx.scatter`` entry points also drive hardware
+whole-tile gather on gfx1250: pass a TDM gather atom
+(``fx.rocdl.make_tdm_gather_atom(src)``, ``copy_rank == 2``) instead of a
+per-element ``UniversalCopy`` atom and the call lowers to a single
+``rocdl.tensor.load.to.lds`` / ``store.from.lds`` gather. In that mode the
+``offset`` argument is the row-index operand (an i16/i32 tensor whose element
+type sets the index width and whose length sets the row count), ``base_iter``
+only supplies the global shape/direction token (the base pointer is atom
+state), and no per-instance loop is emitted. See
+``tests/mlir/Conversion/tdm_gather_gfx1250.mlir`` for the lowering and
+``fx.rocdl.make_tdm_gather_atom`` for the atom builder. The example kernels
+below use the arch-neutral ``UniversalCopy`` path so they run everywhere.
 """
 
 import torch

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

@@ -80,30 +80,38 @@ def Fly_CopyOpTypeInterface : TypeInterface<"CopyOpTypeInterface"> {
     InterfaceMethod<"", "::mlir::Attribute", "getThrBitLayoutDst", (ins)>,
     InterfaceMethod<"", "::mlir::Attribute", "getThrBitLayoutRef", (ins)>,
     InterfaceMethod<
-      "Emit the lowering IR for a copy_atom_call with this CopyOp type.",
+      "Emit the lowering IR for a copy_atom_call with this CopyOp type. "
+      "`indicesMemTy`/`indices` carry the optional gather/scatter row-index "
+      "operand (null when absent); atoms that do not gather ignore them.",
       "::mlir::LogicalResult",
       "emitAtomCall",
       (ins "::mlir::OpBuilder &":$builder,
            "::mlir::Location":$loc,
            "::mlir::Type":$copyAtomTy,
            "::mlir::Type":$srcMemTy,
            "::mlir::Type":$dstMemTy,
+           "::mlir::Type":$indicesMemTy,
            "::mlir::Value":$atomVal,
            "::mlir::Value":$src,
-           "::mlir::Value":$dst)>,
+           "::mlir::Value":$dst,
+           "::mlir::Value":$indices)>,
     InterfaceMethod<
-      "Emit the lowering IR for a predicated copy_atom_call with this CopyOp type.",
+      "Emit the lowering IR for a predicated copy_atom_call with this CopyOp type. "
+      "`indicesMemTy`/`indices` carry the optional gather/scatter row-index "
+      "operand (null when absent); atoms that do not gather ignore them.",
       "::mlir::LogicalResult",
       "emitAtomCall",
       (ins "::mlir::OpBuilder &":$builder,
            "::mlir::Location":$loc,
            "::mlir::Type":$copyAtomTy,
            "::mlir::Type":$srcMemTy,
            "::mlir::Type":$dstMemTy,
+           "::mlir::Type":$indicesMemTy,
            "::mlir::Type":$predMemTy,
            "::mlir::Value":$atomVal,
            "::mlir::Value":$src,
            "::mlir::Value":$dst,
+           "::mlir::Value":$indices,
            "::mlir::Value":$pred)>,
     InterfaceMethod<
       "Emit SSA-form copy: reads src, returns vector result.",

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

@@ -357,8 +357,13 @@ def Fly_AtomSetValueOp : Fly_Op<"atom.set_value", [Pure, DeclareOpInterfaceMetho
   let assemblyFormat = "`(` $atom `,` $field `,` $value `)` attr-dict `:` functional-type(operands, results)";
 }
 
-def Fly_CopyAtomCall : Fly_Op<"copy_atom_call"> {
-  let arguments = (ins Fly_CopyAtom:$copyAtom, Fly_MemRef:$src, Fly_MemRef:$dst, Optional<Fly_MemRef>:$pred);
+def Fly_CopyAtomCall : Fly_Op<"copy_atom_call", [AttrSizedOperandSegments]> {
+  let arguments = (ins Fly_CopyAtom:$copyAtom, Fly_MemRef:$src, Fly_MemRef:$dst,
+                       Optional<Fly_MemRef>:$indices, Optional<Fly_MemRef>:$pred);
+  // `indices` is keyword-led (no leading comma) so it stays distinguishable from
+  // the bare comma-led optional `pred`; both may appear, indices then pred.
+  let assemblyFormat = "`(` $copyAtom `,` $src `,` $dst (`indices` `=` $indices^)? (`,` $pred^)? `)` "
+                       "attr-dict `:` functional-type(operands, results)";
 }
 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);
@@ -414,9 +419,13 @@ def Fly_MmaMakeFragmentOp : Fly_Op<"mma.make_fragment", [Pure, DeclareOpInterfac
   let assemblyFormat = "`(` $operand_id `,` $tiled_mma `,` $input (`,` `stages` `=` $stages^)? `)` attr-dict `:` functional-type(operands, results)";
 }
 
-def Fly_CopyOp : Fly_Op<"copy"> {
-  let arguments = (ins AnyType:$copyAtom, Fly_MemRef:$src, Fly_MemRef:$dst, Optional<Fly_MemRef>:$pred);
-  let assemblyFormat = "`(` $copyAtom `,` $src `,` $dst (`,` $pred^)? `)` attr-dict `:` functional-type(operands, results)";
+def Fly_CopyOp : Fly_Op<"copy", [AttrSizedOperandSegments]> {
+  let arguments = (ins AnyType:$copyAtom, Fly_MemRef:$src, Fly_MemRef:$dst,
+                       Optional<Fly_MemRef>:$indices, Optional<Fly_MemRef>:$pred);
+  // `indices` is keyword-led (no leading comma) so it stays distinguishable from
+  // the bare comma-led optional `pred`; both may appear, indices then pred.
+  let assemblyFormat = "`(` $copyAtom `,` $src `,` $dst (`indices` `=` $indices^)? (`,` $pred^)? `)` "
+                       "attr-dict `:` functional-type(operands, results)";
 }
 def Fly_GemmOp : Fly_Op<"gemm"> {
   let arguments = (ins AnyType:$mmaAtom, Fly_MemRef:$d, Fly_MemRef:$a, Fly_MemRef:$b, Fly_MemRef:$c,

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

@@ -269,10 +269,11 @@ def Fly_CopyAtom : Fly_Type<"CopyAtom", "copy_atom", [
     ::mlir::Attribute getThrValLayoutRef();
 
     ::mlir::LogicalResult emitAtomCall(OpBuilder &builder, Location loc, Type copyAtomTy,
-                                       Type srcMemTy, Type dstMemTy, Value atomVal, Value src, Value dst) const;
+                                       Type srcMemTy, Type dstMemTy, Type indicesMemTy, Value atomVal,
+                                       Value src, Value dst, Value indices) const;
     ::mlir::LogicalResult emitAtomCall(OpBuilder &builder, Location loc, Type copyAtomTy,
-                                       Type srcMemTy, Type dstMemTy, Type predMemTy, Value atomVal, Value src, Value dst,
-                                       Value pred) const;
+                                       Type srcMemTy, Type dstMemTy, Type indicesMemTy, Type predMemTy,
+                                       Value atomVal, Value src, Value dst, Value indices, Value pred) const;
 
     ::mlir::FailureOr<::mlir::Value> emitAtomCallSSA(OpBuilder &builder, Location loc, Type resultTy,
                                                      Type copyAtomTy, Type srcTy, Type dstTy,

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

@@ -31,6 +31,14 @@ def FlyROCDL_AtomStateField : I32EnumAttr<"AtomStateField", "", [
   let cppNamespace = FlyROCDL_Dialect.cppNamespace;
 }
 
+// gfx1250 TDM copy mode: whole-tile contiguous move (tiled) vs row gather/scatter
+// selected by an explicit row-index operand (gather). Direction (load vs store) is
+// still inferred from the operand address spaces.
+def FlyROCDL_TdmMode : FlyROCDL_I32EnumAttr<"TdmMode", "gfx1250 TDM copy mode", [
+  I32EnumAttrCase<"Tiled",  0, "tiled">,
+  I32EnumAttrCase<"Gather", 1, "gather">
+]>;
+
 include "flydsl/Dialect/FlyROCDL/IR/CopyAtom.td"
 include "flydsl/Dialect/FlyROCDL/IR/MmaAtom.td"
 

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

@@ -62,6 +62,11 @@ def FlyROCDL_CopyOpLdsReadTranspose : FlyROCDL_CopyOp<"CopyOpCDNA4LdsReadTranspo
 //     dims 0..rank-2 (innermost stride is assumed 1). Unset falls back to the
 //     tile memref's static layout stride.
 //   `imm_offset` (default 0): i64 byte offset added to base (K-loop tile bump).
+//
+// `mode` selects the descriptor packing: `tiled` moves a contiguous N-D tile;
+// `gather` (rank 2 only) reads a set of rows selected by an `indices` memref
+// operand on the copy_atom_call, with the row-index width taken from that operand's
+// element type (i16 or i32) and the row count from its static extent.
 //===----------------------------------------------------------------------===//
 
 // TDM moves the whole N-D tile in one DMA, so it overrides `getCopyRank` (== rank)
@@ -79,12 +84,14 @@ def FlyROCDL_CopyOpGFX1250TDM : FlyROCDL_StatefulCopyOp<"CopyOpGFX1250TDM", "gfx
     // Descriptor config bits (GROUP1 sgpr0): atomic_barrier_enable [18] sets the
     // HW auto-barrier bit; early_timeout [21] is a multicast-load GL1 knob.
     "bool":$atomicBarrier,
-    "bool":$earlyTimeout
+    "bool":$earlyTimeout,
+    // Copy mode: tiled (contiguous N-D tile) or gather (row gather/scatter).
+    EnumParameter<FlyROCDL_TdmMode>:$mode
   );
   let assemblyFormat = [{
     `<` `rank` `=` $rank `,` `warps` `=` $numWarps `,` `pad` `=` $padInterval `,` $padAmount
     `,` `cache` `=` $cacheModifier `,` `barrier` `=` $atomicBarrier
-    `,` `timeout` `=` $earlyTimeout `>`
+    `,` `timeout` `=` $earlyTimeout `,` `mode` `=` $mode `>`
   }];
   let genVerifyDecl = 1;
 }

lib/Bindings/Python/FlyExtension.cpp

@@ -777,6 +777,11 @@ struct PyCopyAtomType : PyConcreteType<PyCopyAtomType> {
       return wrap(self.toCppType().getCopyOp());
     });
     c.def_prop_ro("val_bits", [](PyCopyAtomType &self) { return self.toCppType().getValBits(); });
+    // Number of memref dims the atom transfers per copy_atom_call (whole-tile atoms
+    // like TDM return N); the DSL gather/scatter routing keys on this.
+    c.def_prop_ro("copy_rank", [](PyCopyAtomType &self) -> unsigned {
+      return cast<CopyOpTypeInterface>(self.toCppType().getCopyOp()).getCopyRank();
+    });
     c.def_prop_ro("thr_layout", [](PyCopyAtomType &self) -> MlirType {
       return wrap(LayoutType::get(cast<LayoutAttr>(self.toCppType().getThrLayout())));
     });

lib/Bindings/Python/FlyROCDLExtension.cpp

@@ -186,20 +186,22 @@ struct PyCopyOpGFX1250TDMType : PyConcreteType<PyCopyOpGFX1250TDMType> {
     c.def_static(
         "get",
         [](int32_t rank, int32_t numWarps, int32_t padInterval, int32_t padAmount,
-           int32_t cacheModifier, bool atomicBarrier, bool earlyTimeout,
+           int32_t cacheModifier, bool atomicBarrier, bool earlyTimeout, int32_t mode,
            DefaultingPyMlirContext context) {
           MLIRContext *ctx = unwrap(context.get()->get());
           return PyCopyOpGFX1250TDMType(
               context->getRef(),
               wrap(CopyOpGFX1250TDMType::get(ctx, rank, numWarps, padInterval, padAmount,
-                                             cacheModifier, atomicBarrier, earlyTimeout)));
+                                             cacheModifier, atomicBarrier, earlyTimeout,
+                                             static_cast<TdmMode>(mode))));
         },
         "rank"_a, "num_warps"_a, "pad_interval"_a = 0, "pad_amount"_a = 0, "cache_modifier"_a = 0,
-        "atomic_barrier"_a = false, "early_timeout"_a = false, nb::kw_only(),
+        "atomic_barrier"_a = false, "early_timeout"_a = false, "mode"_a = 0, nb::kw_only(),
         "context"_a = nb::none(),
         "Create a CopyOpGFX1250TDMType (N-D TDM Global<->LDS copy) with tensor rank (1-5), "
-        "warp count, optional LDS padding (interval/amount in elements), cache modifier, and "
-        "the descriptor atomic_barrier / early_timeout config bits (default false)");
+        "warp count, optional LDS padding (interval/amount in elements), cache modifier, "
+        "the descriptor atomic_barrier / early_timeout config bits (default false), and "
+        "mode (0 = tiled, 1 = gather)");
   }
 };
 

lib/Conversion/FlyToROCDL/FlyToROCDL.cpp

@@ -534,6 +534,7 @@ class CopyAtomCallLowering : public OpConversionPattern<CopyAtomCall> {
     Value copyAtomVal = adaptor.getCopyAtom();
     Value src = adaptor.getSrc();
     Value dst = adaptor.getDst();
+    Value indices = adaptor.getIndices();
     Value pred = adaptor.getPred();
 
     auto srcMemTy = dyn_cast<fly::MemRefType>(op.getSrc().getType());
@@ -546,6 +547,14 @@ class CopyAtomCallLowering : public OpConversionPattern<CopyAtomCall> {
 
     Location loc = op.getLoc();
 
+    // Optional gather/scatter row-index operand; null when absent.
+    Type indicesMemTy = nullptr;
+    if (indices) {
+      indicesMemTy = dyn_cast<fly::MemRefType>(op.getIndices().getType());
+      if (!indicesMemTy)
+        return rewriter.notifyMatchFailure(op, "indices is not a MemRef type");
+    }
+
     Type predMemTy = nullptr;
     if (pred) {
       predMemTy = dyn_cast<fly::MemRefType>(op.getPred().getType());
@@ -554,12 +563,12 @@ class CopyAtomCallLowering : public OpConversionPattern<CopyAtomCall> {
     }
 
     if (pred) {
-      if (failed(copyAtom.emitAtomCall(rewriter, loc, copyAtomType, srcMemTy, dstMemTy, predMemTy,
-                                       copyAtomVal, src, dst, pred)))
+      if (failed(copyAtom.emitAtomCall(rewriter, loc, copyAtomType, srcMemTy, dstMemTy, indicesMemTy,
+                                       predMemTy, copyAtomVal, src, dst, indices, pred)))
         return failure();
     } else {
-      if (failed(copyAtom.emitAtomCall(rewriter, loc, copyAtomType, srcMemTy, dstMemTy, copyAtomVal,
-                                       src, dst)))
+      if (failed(copyAtom.emitAtomCall(rewriter, loc, copyAtomType, srcMemTy, dstMemTy, indicesMemTy,
+                                       copyAtomVal, src, dst, indices)))
         return failure();
     }
     rewriter.eraseOp(op);

lib/Dialect/Fly/IR/FlyTypeDefs.cpp

@@ -450,18 +450,20 @@ Attribute CopyAtomType::getThrValLayoutRef() {
 }
 
 LogicalResult CopyAtomType::emitAtomCall(OpBuilder &builder, Location loc, Type copyAtomTy,
-                                         Type srcMemTy, Type dstMemTy, Value atomVal, Value src,
-                                         Value dst) const {
+                                         Type srcMemTy, Type dstMemTy, Type indicesMemTy,
+                                         Value atomVal, Value src, Value dst, Value indices) const {
   return cast<CopyOpTypeInterface>(getCopyOp())
-      .emitAtomCall(builder, loc, copyAtomTy, srcMemTy, dstMemTy, atomVal, src, dst);
+      .emitAtomCall(builder, loc, copyAtomTy, srcMemTy, dstMemTy, indicesMemTy, atomVal, src, dst,
+                    indices);
 }
 
 LogicalResult CopyAtomType::emitAtomCall(OpBuilder &builder, Location loc, Type copyAtomTy,
-                                         Type srcMemTy, Type dstMemTy, Type predMemTy,
-                                         Value atomVal, Value src, Value dst, Value pred) const {
+                                         Type srcMemTy, Type dstMemTy, Type indicesMemTy,
+                                         Type predMemTy, Value atomVal, Value src, Value dst,
+                                         Value indices, Value pred) const {
   return cast<CopyOpTypeInterface>(getCopyOp())
-      .emitAtomCall(builder, loc, copyAtomTy, srcMemTy, dstMemTy, predMemTy, atomVal, src, dst,
-                    pred);
+      .emitAtomCall(builder, loc, copyAtomTy, srcMemTy, dstMemTy, indicesMemTy, predMemTy, atomVal,
+                    src, dst, indices, pred);
 }
 
 FailureOr<Value> CopyAtomType::emitAtomCallSSA(OpBuilder &builder, Location loc, Type resultTy,