diff --git a/examples/06-cdna5_tensor_copy.py b/examples/06-cdna5_tensor_copy.py new file mode 100644 index 000000000..df6cc1921 --- /dev/null +++ b/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}") diff --git a/include/flydsl/Conversion/FlyToROCDL/Passes.td b/include/flydsl/Conversion/FlyToROCDL/Passes.td index e23eb7b35..66218921e 100644 --- a/include/flydsl/Conversion/FlyToROCDL/Passes.td +++ b/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" - ]; -} diff --git a/include/flydsl/Dialect/Fly/IR/FlyInterfaces.td b/include/flydsl/Dialect/Fly/IR/FlyInterfaces.td index 414c8dab2..5a789dc88 100644 --- a/include/flydsl/Dialect/Fly/IR/FlyInterfaces.td +++ b/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.", diff --git a/include/flydsl/Dialect/Fly/IR/FlyOps.td b/include/flydsl/Dialect/Fly/IR/FlyOps.td index c759840c4..d605a0150 100644 --- a/include/flydsl/Dialect/Fly/IR/FlyOps.td +++ b/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:$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]> { @@ -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:$pred); + let arguments = (ins Fly_CopyAtom:$copyAtom, Fly_TensorLikeType:$src, + Fly_TensorLikeType:$dst, Optional:$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:$pred); + let arguments = (ins AnyType:$copyAtom, Fly_TensorLikeType:$src, Fly_TensorLikeType:$dst, + Optional:$pred); let assemblyFormat = "`(` $copyAtom `,` $src `,` $dst (`,` $pred^)? `)` attr-dict `:` functional-type(operands, results)"; } def Fly_GemmOp : Fly_Op<"gemm"> { diff --git a/include/flydsl/Dialect/FlyROCDL/CMakeLists.txt b/include/flydsl/Dialect/FlyROCDL/CMakeLists.txt index f33061b2d..9f57627c3 100644 --- a/include/flydsl/Dialect/FlyROCDL/CMakeLists.txt +++ b/include/flydsl/Dialect/FlyROCDL/CMakeLists.txt @@ -1 +1,2 @@ add_subdirectory(IR) +add_subdirectory(Transforms) diff --git a/include/flydsl/Dialect/FlyROCDL/IR/Atom.td b/include/flydsl/Dialect/FlyROCDL/IR/Atom.td index 1daac72b5..b0fec02fe 100644 --- a/include/flydsl/Dialect/FlyROCDL/IR/Atom.td +++ b/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; diff --git a/include/flydsl/Dialect/FlyROCDL/IR/CopyAtom.td b/include/flydsl/Dialect/FlyROCDL/IR/CopyAtom.td index 3ca656ec0..768cb19a2 100644 --- a/include/flydsl/Dialect/FlyROCDL/IR/CopyAtom.td +++ b/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 + : FlyROCDL_StatefulCopyOp + ]> { + 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 diff --git a/include/flydsl/Dialect/FlyROCDL/IR/Ops.td b/include/flydsl/Dialect/FlyROCDL/IR/Ops.td index ec3719fc1..2b7450657 100644 --- a/include/flydsl/Dialect/FlyROCDL/IR/Ops.td +++ b/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 + : FlyROCDL_Op]> { + dag tdmArgs = (ins + AnyType:$tensor, + AnyType:$smemLayout, + AnyType:$tiler, + DefaultValuedOptionalAttr:$numWarps, + DefaultValuedOptionalAttr:$initBoundaryCheck, + DefaultValuedOptionalAttr:$cacheModifier, + DefaultValuedOptionalAttr:$atomicBarrier, + OptionalAttr:$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 diff --git a/include/flydsl/Dialect/FlyROCDL/Transforms/CMakeLists.txt b/include/flydsl/Dialect/FlyROCDL/Transforms/CMakeLists.txt new file mode 100644 index 000000000..33bccad56 --- /dev/null +++ b/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) diff --git a/include/flydsl/Dialect/FlyROCDL/Transforms/Passes.h b/include/flydsl/Dialect/FlyROCDL/Transforms/Passes.h new file mode 100644 index 000000000..24536684c --- /dev/null +++ b/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 diff --git a/include/flydsl/Dialect/FlyROCDL/Transforms/Passes.td b/include/flydsl/Dialect/FlyROCDL/Transforms/Passes.td new file mode 100644 index 000000000..aeb0378ca --- /dev/null +++ b/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 diff --git a/include/flydsl/Dialect/FlyROCDL/Utils/TdmAtomBuilder.h b/include/flydsl/Dialect/FlyROCDL/Utils/TdmAtomBuilder.h new file mode 100644 index 000000000..9c8f7d429 --- /dev/null +++ b/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 +FailureOr deriveTdmAtom(AdaptorT adaptor, tdm::Request &request, Type &dataType, + function_ref emitError); + +extern template FailureOr +deriveTdmAtom(MakeTiledTdmLoadAtomOpAdaptor, tdm::Request &, Type &, + function_ref); +extern template FailureOr +deriveTdmAtom(MakeTiledTdmStoreAtomOpAdaptor, tdm::Request &, + Type &, function_ref); + +FailureOr tdmLoadOpType(MLIRContext *ctx, const tdm::Geometry &geometry, Type dataType, + fly::IntTupleAttr tensor2tdm, bool atomicBarrier, + int32_t cacheModifier, function_ref emitError); +FailureOr tdmStoreOpType(MLIRContext *ctx, const tdm::Geometry &geometry, Type dataType, + fly::IntTupleAttr tensor2tdm, bool atomicBarrier, + int32_t cacheModifier, function_ref emitError); + +FailureOr tdmPartitionLayout(Type atomType, Type stensorType, Type gtensorType, + int32_t numWarps, + function_ref emitError); + +} // namespace mlir::fly_rocdl + +#endif // FLYDSL_DIALECT_FLYROCDL_UTILS_TDMATOMBUILDER_H diff --git a/include/flydsl/Dialect/FlyROCDL/Utils/TdmGeometry.h b/include/flydsl/Dialect/FlyROCDL/Utils/TdmGeometry.h new file mode 100644 index 000000000..83be78a3b --- /dev/null +++ b/include/flydsl/Dialect/FlyROCDL/Utils/TdmGeometry.h @@ -0,0 +1,183 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors + +#ifndef FLYDSL_DIALECT_FLYROCDL_UTILS_TDMGEOMETRY_H +#define FLYDSL_DIALECT_FLYROCDL_UTILS_TDMGEOMETRY_H + +#include "mlir/IR/Diagnostics.h" +#include "mlir/Support/LLVM.h" +#include "llvm/ADT/SmallVector.h" + +#include +#include + +#include "flydsl/Dialect/Fly/IR/FlyDialect.h" + +namespace mlir::fly_rocdl::tdm { + +// TDM descriptors hold at most five dims. Beyond that the trailing modes are packed into +// the last one. +constexpr int32_t kMaxRank = 5; + +// `tensor_dimN_stride` is a 48-bit field, in elements of data_size. +constexpr uint64_t kMaxTensorStride = (uint64_t{1} << 48) - 1; + +// `iterate_count` is a 16-bit field encoded as "value minus one". +constexpr int32_t kMaxIterateCount = 256; + +// Iteration is paid for out of the descriptor's own slots: `tensor_dim2_stride` becomes +// `global_addr_increment`, `tensor_dim3` becomes `lds_addr_increment`, and `tile_dim3` +// becomes `iterate_count`. Losing dim 2's *stride* is what sets the bound -- a dim the +// descriptor cannot step is not a dim -- so an iterating descriptor holds two. The third +// axis is the one the iteration itself walks, which is why iteration covers a 3D tensor +// with a 2D descriptor. +constexpr int32_t kMaxIterateRank = 2; + +/// One descriptor scalar — an extent or a stride. +/// +/// Either a compile-time constant, or a reference to a mode of the global tensor whose +/// value only exists at run time. The reference is what lets type inference and the +/// expansion share one derivation: inference needs the geometry's *shape* and reads it +/// off the static half, while the expansion needs the SSA value and looks it up by +/// `mode` in the layout operand. `divisor` is the ratio applied to the looked-up value +/// when the descriptor counts in a wider unit than the tensor's element, 1 otherwise. +struct Scalar { + int32_t value = 0; ///< valid when `isStatic` + bool isStatic = true; + int32_t mode = -1; ///< global mode (depth-first leaf index) when dynamic + bool fromShape = false; ///< dynamic: read the mode's extent rather than its stride + int32_t divisor = 1; + + static Scalar getStatic(int32_t v) { return Scalar{v, true, -1, false, 1}; } + static Scalar getDynamic(int32_t mode, bool fromShape) { + return Scalar{0, false, mode, fromShape, 1}; + } +}; + +/// One descriptor dim. +struct Dim { + int32_t box = 1; ///< tile extent along this dim + SmallVector modes; ///< global modes feeding it, depth-first leaf indices + Scalar stride; ///< global stride, in descriptor elements + Scalar tensorDim; ///< global extent, for the hardware boundary check + /// Whether this dim's bound is a faithful per-mode rectangular bound. A dim fed by one + /// mode always is. The rank-5 packing's is not — its extent comes from a gcd recurrence + /// over modes that are unrelated in memory. + bool clampable = true; +}; + +/// The derived descriptor geometry, in descriptor dim order (innermost first). +struct Geometry { + SmallVector dims; + int32_t padInterval = 0; + int32_t padAmount = 0; + int32_t ratio = 1; ///< tensor elements per descriptor element + int32_t iterCount = 1; ///< descriptor replays, 1 when it does not iterate + Scalar iterStride; ///< global step between replays, valid when iterCount > 1 + int32_t iterMode = -1; ///< the global mode hardware iteration absorbed + + /// Per global mode, its geometry *after* the recast — what the descriptor counts in. + /// Kept because the coordinate tensor's basis scales are read back off it. + SmallVector modeExtent; + SmallVector modeStride; + /// The single stride-1 global mode, or -1. Under a recast this is the mode whose + /// coordinate has to be divided by `ratio`. + int32_t contiguousMode = -1; + + int32_t rank() const { return static_cast(dims.size()); } +}; + +struct Request { + fly::LayoutAttr gLayout; ///< the global tensor's layout + fly::LayoutAttr smemLayout; ///< the LDS tile layout, fully static + /// The value map, from tile values to global modes: `makeValueMap` of the tiler. + fly::LayoutAttr valueMap; + int32_t elemBits = 0; + int32_t internalBits = 0; + int32_t numWarps = 1; + /// Whether the atom starts out bounding what it can. One flag for the whole tensor: + /// which *individual* modes clamp is per-call state, so the builder only says whether + /// to switch the checkable ones on. A mode with no bound to give is left off either + /// way — this is "clamp what can be clamped", never a request that can be refused. + bool initBoundaryCheck = true; +}; + +/// Derive the descriptor geometry: invert the LDS tile layout to find the largest +/// contiguous run, trace that run back through the global tensor's modes, and read the +/// per-dim extents and strides off the result. +FailureOr derive(const Request &request, function_ref emitError); + +/// Split an LDS tile layout into `(padInterval, padAmount, compactStride)`: the descriptor +/// carries the skip as a pad field and the geometry runs on the tile with the skip taken +/// back out. Exposed because `tdm_partition` cuts the tile by its *compact* order too. +FailureOr> +analyzeLdsTile(fly::LayoutAttr smemLayout, function_ref emitError); + +/// The value map: `identity(gshape) . tiler`, with the tiler right-padded with 1s to the +/// tensor's profile. +FailureOr makeValueMap(fly::IntTupleAttr gshape, fly::TileAttr tiler, + function_ref emitError); + +/// Per global mode, the descriptor axis its boundary-check flag lands on — the atom type's +/// `tensor2tdm`. Shaped like `gshape`, so an `boundary_check` tuple can be checked against it +/// for profile. A leaf is `scale E axis`, or `0` for a mode with no bound to give. +/// +/// Axis indices are in *tensor dim order* (the reverse of descriptor order), matching the +/// atom's `tileShape` and its state slots. +FailureOr makeTensor2Tdm(const Geometry &geometry, fly::IntTupleAttr gshape, + function_ref emitError); + +/// `tensor2tdm` read the other way round: per global mode, the descriptor axis index +/// its flag lands on, flattened into mode order; `-1` for a mode with no bound to give. +/// +/// The basis keeps the tensor's nesting so an `boundary_check` tuple can be profile-checked against +/// it, but every consumer works mode by mode. A leaf's *scale* says how the mode's +/// coordinate maps onto its axis, which the extent — measured on that axis — has already +/// accounted for, so only the axis index survives here. +void boundaryCheckAxes(fly::IntTupleAttr tensor2tdm, SmallVectorImpl &axes); + +/// The coordinate tensor's layout: per global mode, the descriptor axis it moves along, as +/// a basis stride. Slicing it in logical order yields an origin already expressed in +/// descriptor axes. +FailureOr coordLayout(const Geometry &geometry, fly::IntTupleAttr gshape, + function_ref emitError); + +/// The atom's initial `boundary_check` state, shaped like the tensor: `enable` on every +/// mode that has a bound to give, and off on every mode that has none. +/// +/// A mode with nothing to clamp — size-1, stride-0, not spanned by the box, or sharing the +/// rank-5 packing's dim — simply comes out off. The builder's flag is "clamp what can be +/// clamped" and so has no way to fail; naming such a mode is only an error when a call +/// site names it, which is `fly.atom.set_value "boundary_check"`. +FailureOr initialBoundaryCheck(const Geometry &geometry, + fly::IntTupleAttr gshape, bool enable, + function_ref emitError); + +/// The tile shape the atom type carries, in tensor dim order. +SmallVector tileShape(const Geometry &geometry); + +/// The layout `tdm_partition` cuts both the LDS tile and the coordinate tile by, so the +/// two keep describing the same elements. +/// +/// The result is `((ATOM), (ITER))`, or `((ATOM), (WARP), (ITER))` when the workgroup's +/// warps split the tile: mode 0 is one atom call's worth of values, the middle mode (when +/// present) is the warp the caller slices out, and the last counts the calls. Composing a +/// tensor with it is what the caller still does in IR; everything up to it is static and +/// is folded here, which is why this hands back a layout rather than a tensor. +/// +/// `smemLayout` decides the split. Inverting it gives the order the tile is laid out in +/// LDS; that order is cut into `numWarps` equal chunks and the first ATOM values of a +/// chunk are what one descriptor fills. The padding is divided out first, because a pad is +/// a hole in the addresses and not in the values. +/// +/// `atomValBits` / `atomValShape` come from the copy atom: how many values one *call* +/// moves, counted in the atom's own unit, which a recast makes wider than the LDS tile's +/// element. `ldsElemBits` is that element, and the two are reconciled in bits. +FailureOr partitionLayout(fly::IntTupleAttr atomValShape, int32_t atomValBits, + fly::LayoutAttr smemLayout, int32_t ldsElemBits, + fly::IntTupleAttr coordShape, int32_t numWarps, + function_ref emitError); + +} // namespace mlir::fly_rocdl::tdm + +#endif // FLYDSL_DIALECT_FLYROCDL_UTILS_TDMGEOMETRY_H diff --git a/lib/Bindings/Python/FlyROCDLExtension.cpp b/lib/Bindings/Python/FlyROCDLExtension.cpp index 6263830bd..07e27fffe 100644 --- a/lib/Bindings/Python/FlyROCDLExtension.cpp +++ b/lib/Bindings/Python/FlyROCDLExtension.cpp @@ -5,11 +5,18 @@ #include "mlir/IR/MLIRContext.h" #include "mlir/IR/Value.h" +#include "mlir/IR/Diagnostics.h" + #include "flydsl/Dialect/Fly/IR/FlyDialect.h" #include "flydsl/Dialect/FlyROCDL/IR/Dialect.h" +#include "flydsl/Dialect/FlyROCDL/Utils/TdmAtomBuilder.h" #include "BindingUtils.h" +#include +#include +#include + namespace nb = nanobind; using namespace nb::literals; using namespace ::mlir::fly; @@ -241,6 +248,75 @@ struct PyCopyOpCDNA4LdsReadTransposeType : PyConcreteType { + FLYDSL_REGISTER_TYPE_BINDING(CopyOpCDNA5TensorLoadType, "CopyOpCDNA5TensorLoadType"); + + static void bindDerived(ClassTy &c) { + c.def_static( + "get", + [](std::vector tileShape, PyType &dataType, PyType &tensor2tdm, bool atomicBarrier, + int32_t cacheModifier, int32_t iterCount, int32_t padInterval, int32_t padAmount, + DefaultingPyMlirContext context) { + MLIRContext *ctx = unwrap(context.get()->get()); + return PyCopyOpCDNA5TensorLoadType( + context->getRef(), + wrap(CopyOpCDNA5TensorLoadType::get( + ctx, tileShape, unwrap(dataType), + ::mlir::dyn_cast(unwrap(tensor2tdm)).getAttr(), atomicBarrier, + cacheModifier, iterCount, padInterval, padAmount))); + }, + "tile_shape"_a, "data_type"_a, "tensor2tdm"_a, "atomic_barrier"_a = false, + "cache_modifier"_a = 0, "iter_count"_a = 1, "pad_interval"_a = 0, "pad_amount"_a = 0, + nb::kw_only(), "context"_a = nb::none(), + "Create a CopyOpCDNA5TensorLoadType (N-D TDM Global->LDS DMA) with the static per-dim tile " + "shape (tensor dim order, len = rank 1-5)"); + } +}; + +struct PyCopyOpCDNA5TensorStoreType : PyConcreteType { + FLYDSL_REGISTER_TYPE_BINDING(CopyOpCDNA5TensorStoreType, "CopyOpCDNA5TensorStoreType"); + + static void bindDerived(ClassTy &c) { + c.def_static( + "get", + [](std::vector tileShape, PyType &dataType, PyType &tensor2tdm, bool atomicBarrier, + int32_t cacheModifier, int32_t iterCount, DefaultingPyMlirContext context) { + MLIRContext *ctx = unwrap(context.get()->get()); + return PyCopyOpCDNA5TensorStoreType( + context->getRef(), wrap(CopyOpCDNA5TensorStoreType::get( + ctx, tileShape, unwrap(dataType), + ::mlir::dyn_cast(unwrap(tensor2tdm)).getAttr(), + atomicBarrier, cacheModifier, iterCount))); + }, + "tile_shape"_a, "data_type"_a, "tensor2tdm"_a, "atomic_barrier"_a = false, + "cache_modifier"_a = 0, "iter_count"_a = 1, nb::kw_only(), "context"_a = nb::none(), + "Create a CopyOpCDNA5TensorStoreType (N-D TDM LDS->Global DMA). Same parameters as " + "the load minus the padding: the store instruction has no de-padding, it drains LDS " + "with the packed tile stride"); + } +}; + +PyType tdm_partition_layout(PyType &atomType, PyType &stensorType, PyType >ensorType, + int32_t numWarps) { + ::mlir::MLIRContext *ctx = unwrap(atomType).getContext(); + std::string error; + ::mlir::ScopedDiagnosticHandler handler(ctx, + [&](::mlir::Diagnostic &diag) -> ::mlir::LogicalResult { + if (!error.empty()) + error += "; "; + error += diag.str(); + return ::mlir::success(); + }); + auto emitError = [&]() -> ::mlir::InFlightDiagnostic { + return ::mlir::emitError(::mlir::UnknownLoc::get(ctx)) << "tdm_partition_layout: "; + }; + ::mlir::FailureOr layout = ::mlir::fly_rocdl::tdmPartitionLayout( + unwrap(atomType), unwrap(stensorType), unwrap(gtensorType), numWarps, emitError); + if (::mlir::failed(layout)) + throw std::invalid_argument(error.empty() ? "tdm_partition_layout failed" : error); + return PyType(atomType.getContext(), wrap(*layout)); +} + } // namespace fly_rocdl } // namespace MLIR_BINDINGS_PYTHON_DOMAIN } // namespace python @@ -261,5 +337,13 @@ NB_MODULE(_mlirDialectsFlyROCDL, m) { ::mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::fly_rocdl::PyCopyOpCDNA3BufferAtomicType::bind(m); ::mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::fly_rocdl::PyCopyOpGFX1250TDMType::bind(m); ::mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::fly_rocdl::PyCopyOpCDNA4LdsReadTransposeType::bind(m); + ::mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::fly_rocdl::PyCopyOpCDNA5TensorLoadType::bind(m); + ::mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::fly_rocdl::PyCopyOpCDNA5TensorStoreType::bind(m); // clang-format on + + m.def("tdm_partition_layout", + &::mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::fly_rocdl::tdm_partition_layout, + "atom_type"_a, "stensor_type"_a, "gtensor_type"_a, "num_warps"_a, + "Create a tdm_partition_layout with the given atom type, stensor type, gtensor type, and " + "num warps"); } diff --git a/lib/CAPI/Dialect/FlyROCDL/FlyROCDLDialect.cpp b/lib/CAPI/Dialect/FlyROCDL/FlyROCDLDialect.cpp index d835245eb..c2a3096f7 100644 --- a/lib/CAPI/Dialect/FlyROCDL/FlyROCDLDialect.cpp +++ b/lib/CAPI/Dialect/FlyROCDL/FlyROCDLDialect.cpp @@ -5,13 +5,13 @@ #include "flydsl/Conversion/Passes.h" #include "flydsl/Dialect/FlyROCDL/IR/Dialect.h" +#include "flydsl/Dialect/FlyROCDL/Transforms/Passes.h" #include "mlir/CAPI/IR.h" #include "mlir/CAPI/Registration.h" MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(FlyROCDL, fly_rocdl, mlir::fly_rocdl::FlyROCDLDialect) void mlirRegisterFlyToROCDLConversionPass(void) { mlir::registerFlyToROCDLConversionPass(); } -void mlirRegisterFlyROCDLClusterAttrPass(void) { mlir::registerFlyROCDLClusterAttrPass(); } void flydsl_register_rocdl_dialects(MlirDialectRegistry registry) { unwrap(registry)->insert(); @@ -19,5 +19,5 @@ void flydsl_register_rocdl_dialects(MlirDialectRegistry registry) { void flydsl_register_rocdl_passes(void) { mlirRegisterFlyToROCDLConversionPass(); - mlirRegisterFlyROCDLClusterAttrPass(); + mlir::fly_rocdl::registerFlyROCDLPasses(); } diff --git a/lib/Conversion/FlyToROCDL/FlyToROCDL.cpp b/lib/Conversion/FlyToROCDL/FlyToROCDL.cpp index 22babc5c4..8bc2704d2 100644 --- a/lib/Conversion/FlyToROCDL/FlyToROCDL.cpp +++ b/lib/Conversion/FlyToROCDL/FlyToROCDL.cpp @@ -26,7 +26,6 @@ namespace mlir { #define GEN_PASS_DEF_FLYTOROCDLCONVERSIONPASS -#define GEN_PASS_DEF_FLYROCDLCLUSTERATTRPASS #include "flydsl/Conversion/FlyToROCDL/Passes.h.inc" } // namespace mlir @@ -359,16 +358,10 @@ class MakeViewOpLowering : public OpConversionPattern { LogicalResult matchAndRewrite(MakeViewOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - if (isa(op.getResult().getType())) { - if (!op.getResult().use_empty()) - return rewriter.notifyMatchFailure(op, "coord_tensor result should have no uses"); - rewriter.eraseOp(op); - return success(); - } else { - Value base = adaptor.getIter(); - rewriter.replaceOp(op, base); - return success(); - } + // A view's runtime value is its iterator: a pointer for a memref, a coordinate + // for a coordinate tensor. + rewriter.replaceOp(op, adaptor.getIter()); + return success(); } }; @@ -478,8 +471,12 @@ class MakeCopyAtomOpLowering : public OpConversionPattern { auto statefulOp = dyn_cast(copyAtomTy.getCopyOp()); if (statefulOp) { - Value state = statefulOp.getDefaultState(rewriter, op.getLoc()); + Value state = statefulOp.getAtomState(rewriter, op.getLoc(), adaptor.getArgs()); + if (!state) + return failure(); rewriter.replaceOp(op, state); + } else if (!adaptor.getArgs().empty()) { + return rewriter.notifyMatchFailure(op, "stateless copy atom takes no construction arguments"); } else { rewriter.replaceOpWithNewOp(op, convertedTy); } @@ -500,6 +497,8 @@ class MakeMmaAtomOpLowering : public OpConversionPattern { auto statefulOp = dyn_cast(mmaAtomTy.getMmaOp()); if (statefulOp) { Value state = statefulOp.getDefaultState(rewriter, op.getLoc()); + if (!state) + return failure(); rewriter.replaceOp(op, state); } else { rewriter.replaceOpWithNewOp(op, convertedTy); @@ -580,12 +579,15 @@ class CopyAtomCallLowering : public OpConversionPattern { Value dst = adaptor.getDst(); Value pred = adaptor.getPred(); - auto srcMemTy = dyn_cast(op.getSrc().getType()); - auto dstMemTy = dyn_cast(op.getDst().getType()); - - if (!srcMemTy || !dstMemTy) - return rewriter.notifyMatchFailure(op, "expected MemRef types on original op"); - if (srcMemTy.getElemTy() != dstMemTy.getElemTy()) + Type srcTy = op.getSrc().getType(); + Type dstTy = op.getDst().getType(); + auto srcMemTy = dyn_cast(srcTy); + auto dstMemTy = dyn_cast(dstTy); + if (!srcMemTy && !isa(srcTy)) + return rewriter.notifyMatchFailure(op, "src is neither a MemRef nor a coord tensor"); + if (!dstMemTy && !isa(dstTy)) + return rewriter.notifyMatchFailure(op, "dst is neither a MemRef nor a coord tensor"); + if (srcMemTy && dstMemTy && srcMemTy.getElemTy() != dstMemTy.getElemTy()) return rewriter.notifyMatchFailure(op, "src/dst element types mismatch"); Location loc = op.getLoc(); @@ -598,12 +600,12 @@ class CopyAtomCallLowering : public OpConversionPattern { } if (pred) { - if (failed(copyAtom.emitAtomCall(rewriter, loc, copyAtomType, srcMemTy, dstMemTy, predMemTy, + if (failed(copyAtom.emitAtomCall(rewriter, loc, copyAtomType, srcTy, dstTy, predMemTy, copyAtomVal, src, dst, pred))) return failure(); } else { - if (failed(copyAtom.emitAtomCall(rewriter, loc, copyAtomType, srcMemTy, dstMemTy, copyAtomVal, - src, dst))) + if (failed(copyAtom.emitAtomCall(rewriter, loc, copyAtomType, srcTy, dstTy, copyAtomVal, src, + dst))) return failure(); } rewriter.eraseOp(op); @@ -793,6 +795,9 @@ class FlyTypeConverter : public TypeConverter { unsigned as = mapAttrToLLVMAddressSpace(flyMemRefTy.getAddressSpace()); return LLVM::LLVMPointerType::get(flyMemRefTy.getContext(), as); }); + addConversion([&](fly::CoordTensorType coordTy) -> Type { + return fly::IntTupleType::get(coordTy.getBase()); + }); addConversion([&](fly::PointerType flyPtrTy) -> Type { if (isTargetAddressSpace(flyPtrTy.getAddressSpace())) return BufferFatPtr::getType(flyPtrTy.getContext()); @@ -919,44 +924,4 @@ class FlyToROCDLConversionPass } }; -// --------------------------------------------------------------------------- -// FlyROCDLClusterAttrPass — inject amdgpu-cluster-dims into llvm.func -// passthrough. Run inside gpu.module() AFTER convert-gpu-to-rocdl. -// -// The upstream ROCDL dialect does not translate `rocdl.cluster_dims` to the -// LLVM IR function attribute `amdgpu-cluster-dims`. This pass bridges the -// gap by converting the discardable attribute that `GPUFuncOpLowering` -// copied from gpu.func into an LLVM passthrough entry that the LLVM IR -// emitter honours. -// --------------------------------------------------------------------------- -class FlyROCDLClusterAttrPass - : public mlir::impl::FlyROCDLClusterAttrPassBase { -public: - using mlir::impl::FlyROCDLClusterAttrPassBase< - FlyROCDLClusterAttrPass>::FlyROCDLClusterAttrPassBase; - - void runOnOperation() override { - getOperation()->walk([&](LLVM::LLVMFuncOp func) { - auto clusterAttr = func->getAttrOfType("rocdl.cluster_dims"); - if (!clusterAttr) - return; - - MLIRContext *ctx = func.getContext(); - - // Build the new passthrough entry: ["amdgpu-cluster-dims", "2,2,1"]. - auto key = StringAttr::get(ctx, "amdgpu-cluster-dims"); - auto entry = ArrayAttr::get(ctx, {key, clusterAttr}); - - // Append to existing passthrough list (if any). - SmallVector passthroughAttrs; - if (auto existing = func.getPassthroughAttr()) - passthroughAttrs.append(existing.begin(), existing.end()); - passthroughAttrs.push_back(entry); - - func.setPassthroughAttr(ArrayAttr::get(ctx, passthroughAttrs)); - func->removeAttr("rocdl.cluster_dims"); - }); - } -}; - } // namespace diff --git a/lib/Dialect/Fly/Transforms/ConvertAtomCallToSSAForm.cpp b/lib/Dialect/Fly/Transforms/ConvertAtomCallToSSAForm.cpp index 6b9c81cf0..174b4ec6a 100644 --- a/lib/Dialect/Fly/Transforms/ConvertAtomCallToSSAForm.cpp +++ b/lib/Dialect/Fly/Transforms/ConvertAtomCallToSSAForm.cpp @@ -23,7 +23,12 @@ namespace fly { namespace { -bool isEligibleToPromote(fly::MemRefType memRefTy) { +// A copy operand may be a coordinate tensor rather than a memref (it names a +// position, not storage), and such an operand is never register-promotable. +bool isEligibleToPromote(Type ty) { + auto memRefTy = dyn_cast(ty); + if (!memRefTy) + return false; if (!isGenericAddressSpace(memRefTy.getAddressSpace())) return false; auto layoutAttr = dyn_cast(memRefTy.getLayout()); @@ -49,9 +54,7 @@ class FlyConvertAtomCallToSSAFormPass SmallVector mmaOpsToConvert; moduleOp->walk([&](CopyAtomCall op) { - auto srcTy = cast(op.getSrc().getType()); - auto dstTy = cast(op.getDst().getType()); - if (isEligibleToPromote(srcTy) || isEligibleToPromote(dstTy)) + if (isEligibleToPromote(op.getSrc().getType()) || isEligibleToPromote(op.getDst().getType())) copyOpsToConvert.push_back(op); }); @@ -68,10 +71,10 @@ class FlyConvertAtomCallToSSAFormPass OpBuilder builder(moduleOp->getContext()); for (CopyAtomCall copyOp : copyOpsToConvert) { - auto srcTy = cast(copyOp.getSrc().getType()); - auto dstTy = cast(copyOp.getDst().getType()); - bool srcEligible = isEligibleToPromote(srcTy); - bool dstEligible = isEligibleToPromote(dstTy); + auto srcTy = dyn_cast(copyOp.getSrc().getType()); + auto dstTy = dyn_cast(copyOp.getDst().getType()); + bool srcEligible = isEligibleToPromote(copyOp.getSrc().getType()); + bool dstEligible = isEligibleToPromote(copyOp.getDst().getType()); builder.setInsertionPoint(copyOp); Location loc = copyOp.getLoc(); diff --git a/lib/Dialect/Fly/Transforms/LayoutLowering.cpp b/lib/Dialect/Fly/Transforms/LayoutLowering.cpp index 44c7b063b..c7bed9903 100644 --- a/lib/Dialect/Fly/Transforms/LayoutLowering.cpp +++ b/lib/Dialect/Fly/Transforms/LayoutLowering.cpp @@ -2170,8 +2170,6 @@ class ExpandCopyOpLowering : public OpRewritePattern { Value dst = op.getDst(); Value pred = op.getPred(); - auto srcMemRefTy = cast(src.getType()); - auto dstMemRefTy = cast(dst.getType()); auto predMemRefTy = pred ? cast(pred.getType()) : nullptr; std::function getLayoutAttr = [&](Attribute attr) -> LayoutAttr { @@ -2179,9 +2177,20 @@ class ExpandCopyOpLowering : public OpRewritePattern { return layout; return getLayoutAttr(cast(attr).getOuter()); }; + // An operand is a memref or a coordinate tensor -- a TMA-style atom is addressed by + // a coordinate, and only its layout is needed to decide how the tile decomposes. + auto tensorLikeLayout = [&](Type ty) -> LayoutAttr { + if (auto memref = dyn_cast(ty)) + return getLayoutAttr(memref.getLayout()); + if (auto coord = dyn_cast(ty)) + return getLayoutAttr(coord.getLayout()); + return nullptr; + }; - LayoutAttr srcLayoutAttr = getLayoutAttr(srcMemRefTy.getLayout()); - LayoutAttr dstLayoutAttr = getLayoutAttr(dstMemRefTy.getLayout()); + LayoutAttr srcLayoutAttr = tensorLikeLayout(src.getType()); + LayoutAttr dstLayoutAttr = tensorLikeLayout(dst.getType()); + if (!srcLayoutAttr || !dstLayoutAttr) + return rewriter.notifyMatchFailure(op, "src/dst are not tensor-like"); LayoutAttr predLayoutAttr = nullptr; if (pred) predLayoutAttr = getLayoutAttr(predMemRefTy.getLayout()); @@ -2192,13 +2201,27 @@ class ExpandCopyOpLowering : public OpRewritePattern { if (srcRank != dstRank) return rewriter.notifyMatchFailure(op, "src/dst ranks mismatch"); - // A whole-tile copy atom (e.g. the gfx1250 TDM DMA) moves the entire N-D tile - // in one call and reads its geometry from the operand memref layout, so emit a - // single call on the tile instead of decomposing it per element. The atom's own - // emitAtomCall verifies the operand rank matches the tile. Detected via a - // boundary-safe type trait rather than a concrete cross-dialect cast. + // tests `size(src) == NumValSrc` *before* it peels or loops: a call is + // issued the moment the operand holds exactly one atom's worth of values. + // Without that test the decomposition below keeps descending until the + // shape is a leaf, which is only accidentally right -- it is wrong for + // every atom whose values are a nested tuple. + // TODO: need a better solution for this. + LayoutBuilder layoutBuilder(ctx); + auto sizeOf = [&](LayoutAttr layout) { + return intTupleProduct(layoutBuilder, layout.getShape()).getLeafAsInt(); + }; + auto numValOf = [&](Attribute thrVal) { return sizeOf(cast(thrVal).at(1)); }; if (auto copyAtomTy = dyn_cast(copyAtomVal.getType())) { - if (copyAtomTy.getCopyOp().hasTrait()) { + IntAttr srcVals = sizeOf(srcLayoutAttr); + IntAttr dstVals = sizeOf(dstLayoutAttr); + IntAttr atomSrcVals = numValOf(copyAtomTy.getThrValLayoutSrc()); + IntAttr atomDstVals = numValOf(copyAtomTy.getThrValLayoutDst()); + bool singleAtomCall = srcVals.isStatic() && dstVals.isStatic() && atomSrcVals.isStatic() && + atomDstVals.isStatic() && + srcVals.getValue() == atomSrcVals.getValue() && + dstVals.getValue() == atomDstVals.getValue(); + if (singleAtomCall || copyAtomTy.getCopyOp().hasTrait()) { CopyAtomCall::create(rewriter, loc, copyAtomVal, src, dst, pred); rewriter.eraseOp(op); return success(); diff --git a/lib/Dialect/Fly/Utils/IntTupleUtils.cpp b/lib/Dialect/Fly/Utils/IntTupleUtils.cpp index 3360406a3..0f8d7f641 100644 --- a/lib/Dialect/Fly/Utils/IntTupleUtils.cpp +++ b/lib/Dialect/Fly/Utils/IntTupleUtils.cpp @@ -782,8 +782,24 @@ IntTupleValueAdaptor intTupleBasis2Tuple(const IntTupleBuilder modes = basisAttr.getModes(); + assert(!modes.empty() && "modes must not be empty"); + + IntAttr zero = IntAttr::getStatic(attr.getContext(), 0); + IntTupleValueAdaptor result = + basisAttr.getValue().isStatic() + ? builder.materializeConstantLeaf(basisAttr.getValue()) + : IntTupleValueAdaptor{basis.getValue(), IntTupleAttr::get(basisAttr.getValue())}; + for (auto it = modes.rbegin(); it != modes.rend(); ++it) { + IntTupleBuilder::ElemCollector elements; + for (int32_t i = 0; i < *it; ++i) { + elements.push_back(builder.materializeConstantLeaf(zero)); + } + elements.push_back(result); + result = builder.makeTuple(elements); + } + return result; } static IntTupleAttr intTupleMakeBasisTupleLikeImpl(MLIRContext *ctx, IntTupleAttr profile, diff --git a/lib/Dialect/FlyROCDL/CDNA5/CopyAtom.cpp b/lib/Dialect/FlyROCDL/CDNA5/CopyAtom.cpp new file mode 100644 index 000000000..4397f246c --- /dev/null +++ b/lib/Dialect/FlyROCDL/CDNA5/CopyAtom.cpp @@ -0,0 +1,1082 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/LLVMIR/ROCDLDialect.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Vector/IR/VectorOps.h" +#include "mlir/IR/BuiltinTypes.h" + +#include + +#include "llvm/Support/MathExtras.h" + +#include "flydsl/Dialect/Fly/IR/FlyDialect.h" +#include "flydsl/Dialect/Fly/Utils/IntTupleUtils.h" +#include "flydsl/Dialect/Fly/Utils/LayoutUtils.h" +#include "flydsl/Dialect/Fly/Utils/ThrValLayoutMacro.h.inc" +#include "flydsl/Dialect/FlyROCDL/IR/Dialect.h" +#include "flydsl/Dialect/FlyROCDL/Utils/TdmGeometry.h" + +using namespace mlir; +using namespace mlir::fly; + +namespace mlir::fly_rocdl { + +//===----------------------------------------------------------------------===// +// CopyOpCDNA5TensorLoadType / CopyOpCDNA5TensorStoreType — N-D TDM whole-tile DMA +// (rank 1-5), one type per hardware instruction (TENSOR_LOAD_TO_LDS / +// TENSOR_STORE_FROM_LDS), addressed by a tile coordinate. +// +// Everything about the *global tensor* is fixed for the atom's life and arrives +// as construction arguments of `fly.make_copy_atom`: the base pointer, the +// per-dim stride, and every dim's extent, all measured from the tensor origin. +// Everything about the *tile* is a coordinate, and it arrives as the +// `!fly.coord_tensor` operand of the copy: the kernel tiles and slices that +// tensor with ordinary layout algebra, which folds the tile's position into its +// type, and the operand's runtime value is whatever of that position was +// genuinely dynamic. +// +// Why a coordinate and not an address. TDM's `global_addr` is the address of the +// tile, not of the tensor, so the hardware measures `tensor_dim` from the tile start +// as well: moving the tile moves the base *and* shrinks the in-bounds window. Those +// two must agree — a base that has advanced past a clamp that has not silently reads +// outside the tensor — so both are derived here from the one coordinate, and there is +// no way to set either directly. +// +// Which dims clamp is atom state (`boundary_check`, an int_tuple with a leaf per *global tensor +// mode*, all-clamp by default), not a type parameter, because the cost is per dim *and* +// per call: a clamping dim spends a subtract and a max on its extent, while a dim left +// alone passes that extent straight through as `tensor_dim` and spends nothing. A +// tiled loop is usually ragged in only some dims, and often only on some iterations — +// the K-loop that needs its K bound checked on the last trip and nowhere else is the +// shape this exists for. The tuple is set as a unit, which is what makes its rank +// checkable; a static leaf folds the select and kills the losing side's arithmetic, so +// the state slots only cost anything when a flag is genuinely dynamic. +// +// The flags are indexed by the tensor's modes and not the descriptor's, because the +// descriptor's are not the caller's to know: it takes its order from the LDS tile's +// majorness, and packs the tail of a tensor with more than five modes. `tensor2tdm` is the +// translation from one to the other, and it is a type parameter because a call site far from the +// builder must still be able to spell `set_value("boundary_check", ...)` in the tensor's terms. It +// is shaped like the tensor rather than flattened, so the `boundary_check` tuple can be +// checked against it for profile and not merely for leaf count. +//===----------------------------------------------------------------------===// + +namespace { + +constexpr unsigned kMaxTdmRank = 5; +constexpr uint64_t kMaxTensorStride = (uint64_t{1} << 48) - 1; +constexpr int32_t kMaxAtomicBarrierAddress = 0x7FFF8; +constexpr int32_t kMaxIterateCount = 256; + +// Filler for the extent slots past the atom's rank, which the lowering never reads -- +// a value that would be obviously wrong in a descriptor, so a stray read shows up. +constexpr int32_t kUnusedExtent = -1; + +// LDS. `atomic_barrier_addr` is a pointer into it rather than an integer, so the barrier +// the transfer arrives on is named the way the kernel already holds it -- what an +// allocator hands back -- instead of being flattened to an address at the call site. +// It says only *which* barrier: whether there is one at all is the atom's type, so no +// pointer value is reserved to mean "none" and offset 0 is a barrier like any other. +constexpr unsigned kSharedAddrSpace = ROCDL::ROCDLDialect::kSharedMemoryAddressSpace; + +// TDM padding descriptor bitfield: +// encoded_interval = log2(interval_dw) - 1 -> bits [24:22] (3 bits) +// encoded_amount = amount_dw - 1 -> bits [31:25] (7 bits) +// where *_dw = pad_*_elems * elem_bits / 32 +struct PadEncoding { + int32_t interval = 0; + int32_t amount = 0; + bool enable = false; +}; + +// Encoding for active padding, a disabled encoding when none is requested, or +// failure() when it cannot be represented (not dword-aligned, dword interval not a +// power of two, or a field out of range) — failing here avoids silently emitting a +// wrong descriptor. +FailureOr computePadEncoding(int32_t padIntervalElems, int32_t padAmountElems, + int32_t elemBits) { + PadEncoding e; + if (padIntervalElems <= 0 || padAmountElems <= 0) + return e; // disabled + if ((padIntervalElems * elemBits) % 32 != 0 || (padAmountElems * elemBits) % 32 != 0) + return failure(); + int32_t intervalDw = padIntervalElems * elemBits / 32; + int32_t amountDw = padAmountElems * elemBits / 32; + if (intervalDw <= 0 || amountDw <= 0 || (intervalDw & (intervalDw - 1)) != 0) + return failure(); + int32_t encInterval = llvm::Log2_32(static_cast(intervalDw)) - 1; + int32_t encAmount = amountDw - 1; + if (encInterval < 0 || encInterval > 0x7 || encAmount < 0 || encAmount > 0x7F) + return failure(); + e.interval = encInterval; + e.amount = encAmount; + e.enable = true; + return e; +} + +Value i32Const(OpBuilder &b, Location loc, int32_t v) { + return arith::ConstantIntOp::create(b, loc, v, 32); +} + +//===----------------------------------------------------------------------===// +// Atom-state struct layout +//===----------------------------------------------------------------------===// + +constexpr unsigned kNoSlot = ~0u; + +// The state-struct layout is a per-type property, so the load's MCAST fields simply +// do not exist on the store type instead of sitting there as slots the hardware +// ignores. The geometry slots (extent, boundary_check) are present on both. +struct TdmSlots { + unsigned workgroupMask; // i32, load only (kNoSlot on store) + unsigned earlyTimeout; // i32, load only (kNoSlot on store) + unsigned atomicBarrierAddr; // !llvm.ptr<3>, which barrier; *whether* is the type + unsigned basePtr; // !llvm.ptr<1>, the tensor origin + unsigned stride0; // stride_i at stride0 + i, i < kMaxTdmRank - 1 (i64) + unsigned extent0; // extent_i at extent0 + i, i < kMaxTdmRank (i32) + unsigned boundaryCheck0; // boundary_check_i at boundaryCheck0 + i, i < kMaxTdmRank (i1) + unsigned iterStride; // i64, global step between descriptor replays + unsigned numSlots; + // Whether `atomicBarrierAddr` is a field a call may write. The *index* stays valid + // either way, so the struct layout is the same for both and does not depend on a type + // parameter; only `tdmFieldIndex` hides it, which is what makes `set_value` fail. + bool hasAtomicBarrier = true; +}; + +// {mask, early_timeout, atomic_barrier_addr, base, stride_0..3, extent_0..4, boundary_check_0..4, +// iter} +constexpr TdmSlots kLoadSlots = {0, 1, 2, 3, 4, 8, 13, 18, 19}; +// {atomic_barrier_addr, base, stride_0..3, extent_0..4, boundary_check_0..4, iter} +constexpr TdmSlots kStoreSlots = {kNoSlot, kNoSlot, 0, 1, 2, 6, 11, 16, 17}; + +// The `atomic_barrier_addr` field belongs to an atom whose *type* enables the HW +// auto-barrier and to no other: on one that does not, `atom.set_value` finds no slot +// and fails to legalize instead of writing an address nothing arrives on. The slot +// index itself stays reserved so a type parameter does not move the fields around it; +// a slot no call writes dies with the scalarized struct, like an unread extent. +TdmSlots tdmSlots(const TdmSlots &base, bool hasAtomicBarrier) { + TdmSlots slots = base; + slots.hasAtomicBarrier = hasAtomicBarrier; + return slots; +} + +// A field the store type does not carry maps to no slot, so `atom.set_value` on it +// fails to legalize instead of writing a descriptor bit the engine ignores. +std::optional optSlot(unsigned slot) { + if (slot == kNoSlot) + return std::nullopt; + return slot; +} + +std::optional tdmFieldIndex(const TdmSlots &slots, AtomStateField field) { + switch (field) { + case AtomStateField::WorkgroupMask: + return optSlot(slots.workgroupMask); + case AtomStateField::EarlyTimeout: + return optSlot(slots.earlyTimeout); + case AtomStateField::AtomicBarrierAddr: + return slots.hasAtomicBarrier ? optSlot(slots.atomicBarrierAddr) : std::nullopt; + default: + return std::nullopt; + } +} + +Type tdmConvertedType(MLIRContext *ctx, const TdmSlots &slots) { + auto i32 = IntegerType::get(ctx, 32); + auto i64 = IntegerType::get(ctx, 64); + SmallVector fields(slots.numSlots, i32); + fields[slots.basePtr] = LLVM::LLVMPointerType::get(ctx, /*global*/ 1); + // Unconditional, so the two atoms have one layout: an atom whose type carries no + // barrier still has the slot, it is simply never written and dies with the scalarized + // struct. Making it conditional would make `getConvertedType` disagree with the state + // the builder hands back. + fields[slots.atomicBarrierAddr] = LLVM::LLVMPointerType::get(ctx, kSharedAddrSpace); + for (unsigned i = 0; i + 1 < kMaxTdmRank; ++i) + fields[slots.stride0 + i] = i64; + for (unsigned i = 0; i < kMaxTdmRank; ++i) + fields[slots.boundaryCheck0 + i] = IntegerType::get(ctx, 1); + fields[slots.iterStride] = i64; + return LLVM::LLVMStructType::getLiteral(ctx, fields); +} + +// Construction arguments of `fly.make_copy_atom`, in order: +// base pointer, stride_0..stride_{rank-2} (i64), extent_0..extent_{rank-1} (i32) — +// all in tensor dim order, all measured from the tensor origin. The innermost stride +// is 1 by construction and is not passed. Every dim passes its extent whether or not +// it currently clamps: `boundary_check` is per-call state, so a later call site may switch its +// dim's clamping on and needs the bound already in the atom. An extent no call reads +// dies with the rest of the scalarized state struct, so this costs nothing. +Value tdmInitialState(OpBuilder &builder, Location loc, const TdmSlots &slots, unsigned rank, + IntTupleAttr tensor2tdm, bool iterates, ValueRange args) { + unsigned expected = 1 + (rank - 1) + rank + (iterates ? 1 : 0); + if (args.size() != expected) { + mlir::emitError(loc) << "cdna5 TDM: expected " << expected << " construction arguments (base + " + << (rank - 1) << " strides + " << rank << " extents" + << (iterates ? " + 1 iteration stride" : "") << "), got " << args.size(); + return nullptr; + } + if (iterates && !args.back().getType().isInteger(64)) { + mlir::emitError(loc) << "cdna5 TDM: the iteration stride must be i64, got " + << args.back().getType(); + return nullptr; + } + if (!isa(args[0].getType())) { + mlir::emitError(loc) << "cdna5 TDM: the base must be a pointer, got " << args[0].getType(); + return nullptr; + } + for (unsigned i = 0; i + 1 < rank; ++i) + if (!args[1 + i].getType().isInteger(64)) { + mlir::emitError(loc) << "cdna5 TDM: stride_" << i << " must be i64, got " + << args[1 + i].getType(); + return nullptr; + } + for (unsigned i = 0; i < rank; ++i) + if (!args[rank + i].getType().isInteger(32)) { + mlir::emitError(loc) << "cdna5 TDM: extent_" << i << " must be i32, got " + << args[rank + i].getType(); + return nullptr; + } + + auto structTy = cast(tdmConvertedType(builder.getContext(), slots)); + Value state = LLVM::UndefOp::create(builder, loc, structTy); + auto set = [&](unsigned slot, Value v) { + state = LLVM::InsertValueOp::create(builder, loc, state, v, ArrayRef{slot}); + }; + if (slots.workgroupMask != kNoSlot) + set(slots.workgroupMask, i32Const(builder, loc, 0)); + if (slots.earlyTimeout != kNoSlot) + set(slots.earlyTimeout, i32Const(builder, loc, 0)); + set(slots.atomicBarrierAddr, + LLVM::ZeroOp::create(builder, loc, + LLVM::LLVMPointerType::get(builder.getContext(), kSharedAddrSpace))); + set(slots.basePtr, args[0]); + // Slots past the atom's rank are never read by the lowering; they only exist so + // the struct layout is rank-independent. + for (unsigned i = 0; i + 1 < kMaxTdmRank; ++i) + set(slots.stride0 + i, + i + 1 < rank ? args[1 + i] : arith::ConstantIntOp::create(builder, loc, 0, 64)); + for (unsigned i = 0; i < kMaxTdmRank; ++i) + set(slots.extent0 + i, i < rank ? args[rank + i] : i32Const(builder, loc, kUnusedExtent)); + // A dim starts clamping when some tensor mode can put a bound on it. The builder + // normally overwrites this immediately with the caller's `boundary_check`; this is what an atom + // built without one does, and "clamp what can be clamped" is the safe default. + SmallVector axes; + tdm::boundaryCheckAxes(tensor2tdm, axes); + for (unsigned i = 0; i < kMaxTdmRank; ++i) { + bool bounded = llvm::is_contained(axes, static_cast(i)); + set(slots.boundaryCheck0 + i, arith::ConstantIntOp::create(builder, loc, bounded, 1)); + } + set(slots.iterStride, iterates ? args.back() : arith::ConstantIntOp::create(builder, loc, 0, 64)); + return state; +} + +// Per-dim clamping is TDM-private, so it does not go through the shared AtomStateField +// enum: it is set as a whole, as one int_tuple — a nonzero leaf clamps that descriptor +// dim to the extent baked in at construction, zero lets the tile run its full size there. +// Setting it as one aggregate rather than a leaf at a time is what makes its rank +// checkable. +// +// A caller writes `"boundary_check"` in the *tensor's* modes, because the descriptor's are not +// theirs to know. `fly-rocdl-expand-ops` translates that through `tensor2tdm` into +// the `"boundary_check_axes"` this reads: one leaf per descriptor axis, already OR-merged where +// several modes shared one. So nothing about modes survives to here, and a `"boundary_check"` that +// reaches this point is a pipeline that skipped the pass rather than something to +// translate a second time. +// +// A static leaf lives in the tuple's *type*, so it reaches the descriptor math as a +// constant and the select it feeds folds away; a dynamic leaf is an i32/i64 operand of +// `make_int_tuple` and is compared against zero to reach the i1 the slot holds. +LogicalResult tdmUnpackBoundaryCheckAxes(OpBuilder &builder, Location loc, Value flags, + unsigned rank, SmallVectorImpl &out) { + auto tupleTy = dyn_cast(flags.getType()); + if (!tupleTy) + return mlir::emitError(loc) << "cdna5 TDM: \"boundary_check_axes\" must be an int_tuple, got " + << flags.getType(); + IntTupleAttr attr = tupleTy.getAttr(); + if (attr.isLeaf() || attr.rank() != static_cast(rank)) + return mlir::emitError(loc) << "cdna5 TDM: \"boundary_check_axes\" is " << attr + << ", expected the flat " << rank + << "-leaf tuple `fly-rocdl-expand-ops` produces, one leaf per " + "descriptor axis"; + + auto tupleOp = flags.getDefiningOp(); + OperandRange dyn = tupleOp ? tupleOp.getDyncElems() : OperandRange(nullptr, 0); + auto dynIt = dyn.begin(); + out.clear(); + for (int32_t axis = 0; axis < attr.rank(); ++axis) { + IntTupleAttr leaf = attr.at(axis); + if (!leaf.isLeaf()) + return mlir::emitError(loc) << "cdna5 TDM: \"boundary_check_axes\" leaf " << axis + << " is nested"; + IntAttr value = leaf.extractIntFromLeaf(); + if (value.isStatic()) { + out.push_back(arith::ConstantIntOp::create(builder, loc, value.getValue() != 0, 1)); + continue; + } + if (!tupleOp || dynIt == dyn.end()) + return mlir::emitError(loc) << "cdna5 TDM: \"boundary_check_axes\" leaf " << axis + << " is dynamic but the tuple is not normal form"; + Value v = *dynIt++; + Value zero = arith::ConstantIntOp::create(builder, loc, v.getType(), 0); + out.push_back(arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::ne, v, zero)); + } + return success(); +} + +Value tdmSetAtomState(OpBuilder &builder, Location loc, const TdmSlots &slots, unsigned rank, + Value atomStruct, Attribute fieldAttr, Value fieldValue) { + auto fieldStr = dyn_cast(fieldAttr); + if (!fieldStr) + return nullptr; + auto insert = [&](unsigned slot, Value v) { + atomStruct = LLVM::InsertValueOp::create(builder, loc, atomStruct, v, ArrayRef{slot}); + }; + + if (fieldStr.getValue() == "boundary_check") { + mlir::emitError(loc) + << "cdna5 TDM: \"boundary_check\" is written in the global tensor's modes and has " + "to be translated by `fly-rocdl-expand-ops` into \"boundary_check_axes\" before " + "it can be lowered"; + return nullptr; + } + if (fieldStr.getValue() == "boundary_check_axes") { + SmallVector boundaryCheck; + if (failed(tdmUnpackBoundaryCheckAxes(builder, loc, fieldValue, rank, boundaryCheck))) + return nullptr; + for (auto [i, v] : llvm::enumerate(boundaryCheck)) + insert(slots.boundaryCheck0 + i, v); + return atomStruct; + } + + auto field = symbolizeAtomStateField(fieldStr.getValue()); + if (!field) + return nullptr; + std::optional idx = tdmFieldIndex(slots, *field); + if (!idx) + return nullptr; + if (*field == AtomStateField::AtomicBarrierAddr) { + auto ptrTy = dyn_cast(fieldValue.getType()); + if (!ptrTy || ptrTy.getAddressSpace() != kSharedAddrSpace) { + mlir::emitError(loc) << "cdna5 TDM: \"atomic_barrier_addr\" is the barrier itself, so it " + "takes a shared-memory pointer, got " + << fieldValue.getType(); + return nullptr; + } + } + insert(*idx, fieldValue); + return atomStruct; +} + +// Only the width of the descriptor's data type reaches the hardware (`data_size`); +// the type itself is there so the atom says which element it moves. The verifier has +// already rejected anything without a width. +int32_t tdmElemBits(Type dataType) { + return static_cast(dataType.getIntOrFloatBitWidth()); +} + +// One TDM call is issued by a single thread and moves the whole tile, so the atom is +// (1 thread) x (product(tileShape) * elemBits bits). +int32_t tdmNumBits(ArrayRef tileShape, int32_t elemBits, int32_t iterCount) { + int64_t numElems = iterCount; + for (int32_t d : tileShape) + numElems *= d; + return static_cast(numElems * elemBits); +} + +// Spelled out instead of using the FxLayout/FxShape macros: those expand to a bare +// `getContext()` call and so only work inside a type member function. +Attribute tdmThrBitLayout(MLIRContext *ctx, ArrayRef tileShape, int32_t elemBits, + int32_t iterCount) { + Attribute one = IntTupleAttr::getLeafStatic(ctx, 1); + Attribute bits = IntTupleAttr::getLeafStatic(ctx, tdmNumBits(tileShape, elemBits, iterCount)); + return LayoutAttr::get(IntTupleAttr::get(ArrayAttr::get(ctx, {one, bits})), + IntTupleAttr::get(ArrayAttr::get(ctx, {one, one}))); +} + +// Shared verifier. `padInterval` / `padAmount` are always zero on the store type, +// which has no padding parameters at all. +LogicalResult tdmVerify(function_ref emitError, ArrayRef tileShape, + Type dataType, IntTupleAttr tensor2tdm, int32_t iterCount, + int32_t padInterval, int32_t padAmount) { + int32_t rank = static_cast(tileShape.size()); + if (rank < 1 || rank > static_cast(kMaxTdmRank)) + return emitError() << "TDM rank must be in [1, " << kMaxTdmRank << "], got " << rank; + IntTupleBuilder tupleBuilder(tensor2tdm.getContext()); + SmallVector mapLeaves; + intTupleFlattenToVector(tupleBuilder, tensor2tdm, mapLeaves); + // A tensor always has at least as many modes as the descriptor has dims: dims are + // built from modes, and several modes can share one, never the other way round. + if (static_cast(mapLeaves.size()) < rank) + return emitError() << "TDM tensor2tdm has " << mapLeaves.size() + << " modes but the descriptor has " << rank << " dims"; + // Only a static map is representable: a mode either lands on a known axis at a known + // scale, or it has no bound to give. A mode with a dynamic stride never shares a dim + // with another, so it either owns its axis at scale 1 or reaches here as a `0` leaf. + for (auto [mode, leaf] : llvm::enumerate(mapLeaves)) { + if (leaf.isLeafBasis()) { + BasisAttr basis = leaf.getLeafAsBasis(); + if (basis.getModes().size() != 1) + return emitError() << "TDM tensor2tdm mode " << mode + << " must name one descriptor dim, got " << leaf; + int32_t dim = basis.getModes().front(); + if (dim < 0 || dim >= rank) + return emitError() << "TDM tensor2tdm mode " << mode << " lands on dim " << dim + << ", which is not a descriptor dim in [0, " << rank << ")"; + if (!basis.getValue().isStatic()) + return emitError() << "TDM tensor2tdm mode " << mode << " must have a static scale, got " + << leaf; + } else if (!leaf.isLeafStaticValue(0)) { + return emitError() << "TDM tensor2tdm mode " << mode << " must be a basis stride or 0, got " + << leaf; + } + } + for (int32_t d : tileShape) + if (d < 1) + return emitError() << "TDM tile shape dims must be >= 1, got " << d; + // A width is all the descriptor takes from the data type, so a type without one + // (a memref, a tuple) has nothing to give it -- and `getIntOrFloatBitWidth` asserts + // rather than answering, so this guard comes first. + if (!dataType.isIntOrFloat()) + return emitError() << "TDM dataType must be an integer or float type, got " << dataType; + int32_t elemBits = tdmElemBits(dataType); + // data_size is exactly two bits: 0/1/2/3 encode 1/2/4/8 bytes. + if (elemBits != 8 && elemBits != 16 && elemBits != 32 && elemBits != 64) + return emitError() << "TDM element width must be one of 8, 16, 32 or 64 bits, got " << dataType; + if ((padInterval == 0) != (padAmount == 0)) + return emitError() << "padInterval and padAmount must both be zero or both non-zero"; + if (padInterval != 0) { + if (padInterval < 0 || padAmount < 0) + return emitError() << "padInterval and padAmount must be non-negative, got " << padInterval + << ", " << padAmount; + // interval_in_dwords is a power of two iff padInterval (in elements) is, since + // element bits are a power of two; the exact dword/bitfield check needs the + // element type and runs at lowering. + if ((padInterval & (padInterval - 1)) != 0) + return emitError() << "padInterval must be a power of two (in elements), got " << padInterval; + // The hardware pad counter is free-running over the whole transfer -- it is seeded + // once, before the dim loops, and only reset where the pad fires -- so the interval + // is not tied to tile_dim0 and may span several rows. What it cannot straddle is a + // *call* boundary: each instruction re-seeds the counter, so an interval that does + // not divide one call's element count would put the second call's holes in the + // wrong place. (Descriptor iteration re-seeds it too, which is why padding and + // iterCount > 1 are refused together below.) + int64_t tileElems = 1; + for (int32_t dim : tileShape) + tileElems *= dim; + if (tileElems % padInterval != 0) + return emitError() << "padInterval must divide the descriptor's tile (" << tileElems + << " elements), so every call starts on a pad boundary, got " + << padInterval; + if ((static_cast(tileShape.back()) * elemBits) % 32 != 0) + return emitError() << "padded TDM tile_dim0 must span a whole number of dwords, got " + << tileShape.back() << " elements at " << elemBits << " bits"; + } + // `iterate_count` is a 16-bit field encoded as value-minus-one, and iteration is paid + // for out of GROUP2's own slots -- see `tdm::kMaxIterateRank` for which ones. + if (iterCount < 1 || iterCount > kMaxIterateCount) + return emitError() << "TDM iterCount must be in [1, " << kMaxIterateCount << "], got " + << iterCount; + if (iterCount > 1 && rank > tdm::kMaxIterateRank) + return emitError() << "TDM descriptor iteration takes dim 2's stride for its own, so it " + "needs a descriptor of at most " + << tdm::kMaxIterateRank << " dims, got " << rank; + if (iterCount > 1 && padInterval != 0) + return emitError() << "TDM descriptor iteration and LDS padding are not modelled together"; + int64_t totalBits = elemBits; + for (int32_t dim : tileShape) { + if (totalBits > std::numeric_limits::max() / dim) + return emitError() << "TDM tile bit count exceeds the layout integer range"; + totalBits *= dim; + } + return success(); +} + +// Rebuild a tile coordinate from a lowered coordinate tensor. +// +// A coordinate tensor lowers to its coordinate the way a memref lowers to its +// pointer (see `MakeViewOpLowering`), and an `IntTuple` keeps its static leaves in +// its type and only its dynamic ones as SSA operands. So the coordinate is read from +// both halves: constants for the leaves the layout algebra already folded, and the +// `make_int_tuple` operands for the ones a kernel computed. +LogicalResult unpackCoord(OpBuilder &builder, Location loc, IntTupleAttr baseAttr, Value coordValue, + SmallVectorImpl &out) { + IntTupleBuilder tupleBuilder(builder.getContext()); + SmallVector leaves; + intTupleFlattenToVector(tupleBuilder, baseAttr, leaves); + + // Only the dynamic leaves are carried by the value, and only a normal-form tuple + // exposes them in order. + auto tupleOp = coordValue.getDefiningOp(); + OperandRange dyn = tupleOp ? tupleOp.getDyncElems() : OperandRange(nullptr, 0); + auto dynIt = dyn.begin(); + auto dynEnd = dyn.end(); + Type i32Ty = builder.getI32Type(); + for (auto [dim, leaf] : llvm::enumerate(leaves)) { + auto intAttr = leaf.extractIntFromLeaf(); + if (intAttr.isStatic()) { + out.push_back(arith::ConstantIntOp::create(builder, loc, i32Ty, intAttr.getValue())); + continue; + } + if (!tupleOp || dynIt == dynEnd) + return mlir::emitError(loc) << "cdna5 TDM: coordinate leaf " << dim + << " is dynamic but the coordinate tensor is not in normal form"; + Value v = *dynIt++; + if (!v.getType().isInteger(32)) + return mlir::emitError(loc) << "cdna5 TDM: coordinate leaf " << dim + << " must be i32 (a tile index in elements), got " << v.getType(); + out.push_back(v); + } + return success(); +} + +//===----------------------------------------------------------------------===// +// Descriptor construction + instruction emission (shared by both directions) +//===----------------------------------------------------------------------===// + +// The compile-time half of the descriptor, read off the atom type. +struct TdmStatic { + ArrayRef tileShape; + int32_t elemBits; + int32_t padInterval; // always 0 on the store path + int32_t padAmount; // always 0 on the store path + int32_t cacheModifier; + int32_t iterCount; // 1 = no descriptor iteration + bool isLoad; +}; + +// Read off the atom type, once per direction: the two `emitAtomCall` overloads of a type +// want the same config, and the direction is the only thing that varies between them. +TdmStatic tdmConfig(CopyOpCDNA5TensorLoadType ty) { + return {ty.getTileShape(), tdmElemBits(ty.getDataType()), ty.getPadInterval(), + ty.getPadAmount(), ty.getCacheModifier(), ty.getIterCount(), + /*isLoad=*/true}; +} + +TdmStatic tdmConfig(CopyOpCDNA5TensorStoreType ty) { + return {ty.getTileShape(), tdmElemBits(ty.getDataType()), /*padInterval=*/0, + /*padAmount=*/0, ty.getCacheModifier(), ty.getIterCount(), + /*isLoad=*/false}; +} + +// Emits one TENSOR_LOAD_TO_LDS / TENSOR_STORE_FROM_LDS. When `pred` is non-null the +// call is wrapped in an `scf.if`. +LogicalResult emitTdmAtomCall(OpBuilder &builder, Location loc, const TdmStatic &cfg, + const TdmSlots &slots, Type srcTyArg, Type dstTyArg, Value atomVal, + Value src, Value dst, Type predMemTyArg = nullptr, + Value pred = nullptr) { + // The direction is a property of the atom type (one type per opcode). The global + // side is the coordinate tensor — it has no pointer, the atom holds the base — and + // the LDS side is a shared memref. + Type glbTy = cfg.isLoad ? srcTyArg : dstTyArg; + Type ldsTyArg = cfg.isLoad ? dstTyArg : srcTyArg; + Value ldsPtr = cfg.isLoad ? dst : src; + auto ldsMemTy = dyn_cast(ldsTyArg); + if (!isa(glbTy) || !ldsMemTy || + !isGenericAddressSpace(ldsMemTy.getAddressSpace())) + return mlir::emitError(loc) << "cdna5 TDM: " << (cfg.isLoad ? "tensor_load" : "tensor_store") + << " needs a " + << (cfg.isLoad ? "coord-tensor source and a shared destination" + : "shared source and a coord-tensor destination") + << ", got " << srcTyArg << " -> " << dstTyArg; + + OpBuilder::InsertionGuard guard(builder); + if (pred) { + auto predMemTy = cast(predMemTyArg); + Value predVal = LLVM::LoadOp::create(builder, loc, predMemTy.getElemTy(), pred); + auto ifOp = scf::IfOp::create(builder, loc, TypeRange{}, predVal, /*withElse=*/false); + builder.setInsertionPointToStart(&ifOp.getThenRegion().front()); + } + + SmallVector tileShape(cfg.tileShape.begin(), cfg.tileShape.end()); + int32_t rank = static_cast(tileShape.size()); + auto coordTy = cast(glbTy); + + // `elemBits` is the *descriptor's* unit, which a recast may + // have widened past the tensor's own element. The LDS operand keeps the tensor's + // element type, so the two are checked against each other in bits rather than by + // width: what has to agree is how much data one call moves, not how it is counted. + int32_t elemBits = cfg.elemBits; + int32_t ldsElemBits = static_cast(ldsMemTy.getElemTy().getIntOrFloatBitWidth()); + int32_t elemBytes = elemBits / 8; // verified byte-granular, power-of-two + int32_t dataSizeCode = llvm::Log2_32(static_cast(elemBytes)); + + Type i32Ty = builder.getI32Type(); + Type i64Ty = builder.getI64Type(); + Value zeroC = i32Const(builder, loc, 0); + Value zero64 = arith::ConstantIntOp::create(builder, loc, 0, 64); + Value c16 = i32Const(builder, loc, 16); + Value mask16 = i32Const(builder, loc, 0xFFFF); + auto slotField = [&](unsigned slot) { + return LLVM::ExtractValueOp::create(builder, loc, atomVal, ArrayRef{slot}); + }; + + // Tensor geometry, from the tensor origin, baked in at construction. + SmallVector strideElems(rank); + for (int32_t i = 0; i < rank - 1; ++i) + strideElems[i] = slotField(slots.stride0 + i); + strideElems[rank - 1] = arith::ConstantIntOp::create(builder, loc, 1, 64); // innermost contiguous + for (int32_t i = 0; i < rank - 1; ++i) { + Value inRange = + arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::ule, strideElems[i], + arith::ConstantIntOp::create(builder, loc, kMaxTensorStride, 64)); + LLVM::AssumeOp::create(builder, loc, inRange); + } + + // The tile coordinate comes from the coordinate-tensor operand, whose runtime value + // is that coordinate: static leaves live in its type and dynamic ones are the + // operands of the `make_int_tuple` that built it. A static leaf therefore costs a + // constant the descriptor math folds away, and a dynamic one costs exactly the + // integer the kernel already had. + Value glbCoord = cfg.isLoad ? src : dst; + SmallVector coord; + if (failed(unpackCoord(builder, loc, coordTy.getBase(), glbCoord, coord))) + return failure(); + if (static_cast(coord.size()) != rank) + return mlir::emitError(loc) << "cdna5 TDM: the coordinate has " << coord.size() + << " leaves but the atom's tile is rank " << rank; + + // global_addr = base + elem_bytes * sum_i coord_i * stride_i. A coord left at zero + // folds its whole term away, so a tile that never moves along a dim costs nothing. + Value glbAddr = LLVM::PtrToIntOp::create(builder, loc, i64Ty, slotField(slots.basePtr)); + Value elemOff = zero64; + for (int32_t i = 0; i < rank; ++i) { + Value coordNonNegative = + arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::sge, coord[i], zeroC); + LLVM::AssumeOp::create(builder, loc, coordNonNegative); + Value c64 = arith::ExtUIOp::create(builder, loc, i64Ty, coord[i]); + elemOff = arith::AddIOp::create(builder, loc, elemOff, + arith::MulIOp::create(builder, loc, c64, strideElems[i])); + } + glbAddr = arith::AddIOp::create( + builder, loc, glbAddr, + arith::MulIOp::create(builder, loc, elemOff, + arith::ConstantIntOp::create(builder, loc, elemBytes, 64))); + + // The bounds move with the tile, from the same coordinate. `boundary_check_i` selects between + // the moving bound and the tensor's own extent; it is atom state rather than a + // compile-time flag, so a call site can turn a dim's clamping off (or back on) + // without a second atom. The usual case is a constant, and then the select folds and + // the losing side's arithmetic dies -- the knob costs nothing until it varies. + SmallVector tensorDim(rank); + for (int32_t i = 0; i < rank; ++i) { + Value extent = slotField(slots.extent0 + i); + // Clamping off is the caller asserting the tile lies inside the tensor, i.e. + // `coord_i + tile_i <= extent_i`. The un-shifted extent is then already a bound the + // tile cannot reach, so the off side needs no subtract and no separate "never out of + // bounds" sentinel: it reuses an SGPR that is live anyway, and stays a length the + // hardware reads as positive -- where a saturated 0xFFFFFFFF would be an empty window + // to a signed read. + // + // On side: `max(extent - coord, 0)` rather than the equivalent `coord < extent ? ... : 0`. + // Extents and coordinates are non-negative, so the two agree, but the compare form is + // the `usub.sat` idiom and AMDGPU only has that on the VALU -- which would drag a + // uniform descriptor field onto the vector path and back through readfirstlane. + Value rem = arith::SubIOp::create(builder, loc, extent, coord[i]); + Value clamped = arith::MaxSIOp::create(builder, loc, rem, zeroC); + tensorDim[i] = + arith::SelectOp::create(builder, loc, slotField(slots.boundaryCheck0 + i), clamped, extent); + } + + Value ldsAddr = LLVM::PtrToIntOp::create(builder, loc, i32Ty, ldsPtr); + + auto ldsLayout = dyn_cast(ldsMemTy.getLayout()); + if (!ldsLayout || !ldsLayout.isStatic()) + return mlir::emitError(loc) << "cdna5 TDM: the LDS operand needs a fully static layout"; + LayoutBuilder layoutBuilder(builder.getContext()); + int64_t ldsCapacity = layoutCosize(layoutBuilder, ldsLayout).getLeafAsInt().getValue(); + int64_t tileElems = 1; + for (int32_t dim : tileShape) + tileElems *= dim; + // Every replay stacks another box into LDS, so the footprint counts them all. + int64_t requiredLds = tileElems * cfg.iterCount; + if (cfg.padAmount) { + int64_t rows = tileElems / cfg.padInterval; + requiredLds += (rows - 1) * cfg.padAmount; + } + // Both sides in bits, since a recast leaves the two counting in different units. + if (ldsCapacity * ldsElemBits < requiredLds * elemBits) + return mlir::emitError(loc) << "cdna5 TDM: LDS view holds " << ldsCapacity * ldsElemBits + << " bits, less than the descriptor footprint of " + << requiredLds * elemBits; + if (tileElems * cfg.iterCount * elemBits % ldsElemBits != 0) + return mlir::emitError(loc) << "cdna5 TDM: the descriptor moves " + << tileElems * cfg.iterCount * elemBits + << " bits, which is not a whole number of the LDS operand's " + << ldsElemBits << "-bit elements"; + + // GROUP0 (vector<4xi32>): count, lds_addr, glb_lo, glb_hi | type. `count = 1` is + // "valid tensor" (0 would be a NULL Tensor that moves nothing). The global address + // is split from the full i64, so a tile coordinate walking past 4 GiB carries into + // glb_hi automatically. + Value g0s2 = LLVM::TruncOp::create(builder, loc, i32Ty, glbAddr); + Value glbHiRaw = LLVM::LShrOp::create(builder, loc, glbAddr, + arith::ConstantIntOp::create(builder, loc, 32, 64)); + // No mask on the high word: bits [30:25] are reserved and a canonical AMDGPU VA + // (48-bit) never reaches them, so an `and` here only costs a SALU op per issue site + // inside the K loop -- `s_or`/`s_bitset1` instead of `s_and`+`s_or`. + Value g0s3 = + arith::OrIOp::create(builder, loc, LLVM::TruncOp::create(builder, loc, i32Ty, glbHiRaw), + i32Const(builder, loc, /*type field [31:30]=2*/ 1 << 31)); + Value dgroup0 = vector::FromElementsOp::create( + builder, loc, VectorType::get({4}, i32Ty), + ValueRange{i32Const(builder, loc, /*count=*/1), ldsAddr, g0s2, g0s3}); + + // Padding describes the padded LDS tile the DMA engine fills, which only the load + // direction has: the store type carries no padding parameters, so this always + // encodes "disabled" there. + FailureOr padOr = computePadEncoding(cfg.padInterval, cfg.padAmount, elemBits); + if (failed(padOr)) + return mlir::emitError(loc) + << "cdna5 TDM: padding (interval=" << cfg.padInterval << ", amount=" << cfg.padAmount + << " elements at " << elemBits + << "-bit) is not encodable — the dword interval must be a power of two and the encoded " + "fields must fit the descriptor bitfield"; + PadEncoding pad = *padOr; + + // Descriptor dims are innermost-first: descriptor dim j maps to tensor dim + // (rank-1-j). LDS padding is carried by the pad bitfield, never by widening + // tile_dim, so the global transfer extent stays the true tile size either way. + auto descTensorDim = [&](int32_t j) -> Value { return tensorDim[rank - 1 - j]; }; + auto descTileDim = [&](int32_t j) -> int32_t { return tileShape[rank - 1 - j]; }; + // 48-bit stride slots: descriptor stride k = stride of tensor dim (rank-2-k). + auto descStrideLo32 = [&](int32_t k) -> Value { + return LLVM::TruncOp::create(builder, loc, i32Ty, strideElems[rank - 2 - k]); + }; + auto descStrideHi16 = [&](int32_t k) -> Value { + Value hi = LLVM::LShrOp::create(builder, loc, strideElems[rank - 2 - k], + arith::ConstantIntOp::create(builder, loc, 32, 64)); + return arith::AndIOp::create(builder, loc, LLVM::TruncOp::create(builder, loc, i32Ty, hi), + mask16); + }; + auto lo16 = [&](Value v) { return arith::AndIOp::create(builder, loc, v, mask16); }; + auto hi16 = [&](Value v) { + return arith::AndIOp::create(builder, loc, arith::ShRUIOp::create(builder, loc, v, c16), + mask16); + }; + auto shl16 = [&](Value v) { return arith::ShLIOp::create(builder, loc, v, c16); }; + auto orr = [&](Value a, Value b) { return arith::OrIOp::create(builder, loc, a, b); }; + + // Whether the transfer arrives on a barrier is the atom's *type* and nothing else, so + // config bit [18] is a constant here and the state says only which barrier -- the + // pointer the kernel already holds, flattened to the atomic_barrier_address[18:3] the + // descriptor wants (GROUP1 bits 47:32, i.e. sgpr1 [15:0]). An LDS pointer is 32-bit, so + // that flattening is a bitcast and not a truncation. + // + // The pointer is never read as an enable, which is what lets LDS offset 0 be a barrier + // like any other. It also means an atom whose type asks for a barrier and whose state + // was never given one arrives on offset 0 rather than nowhere: naming the barrier is + // the caller's to do, exactly as the base pointer is. + Value barrierEnableBit = zeroC, barrierAddrField = zeroC; + if (slots.hasAtomicBarrier) { + Value barrier = + LLVM::PtrToIntOp::create(builder, loc, i32Ty, slotField(slots.atomicBarrierAddr)); + Value barrierInRange = arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::ule, barrier, + i32Const(builder, loc, kMaxAtomicBarrierAddress)); + Value barrierAligned = arith::CmpIOp::create( + builder, loc, arith::CmpIPredicate::eq, + arith::AndIOp::create(builder, loc, barrier, i32Const(builder, loc, 7)), zeroC); + LLVM::AssumeOp::create(builder, loc, + arith::AndIOp::create(builder, loc, barrierInRange, barrierAligned)); + barrierEnableBit = i32Const(builder, loc, 1 << 18); + barrierAddrField = + lo16(arith::ShRUIOp::create(builder, loc, barrier, i32Const(builder, loc, 3))); + } + + // MCAST mask and its early-timeout companion exist on the load only. + Value maskValue = slots.workgroupMask == kNoSlot ? zeroC : slotField(slots.workgroupMask); + if (slots.workgroupMask != kNoSlot) { + Value maskInRange = arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::ule, maskValue, + i32Const(builder, loc, 0xFFFF)); + LLVM::AssumeOp::create(builder, loc, maskInRange); + } + Value maskLow = lo16(maskValue); + Value earlyTimeoutBit = zeroC; + if (slots.earlyTimeout != kNoSlot) { + Value early = slotField(slots.earlyTimeout); + Value earlyIsBool = arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::ule, early, + i32Const(builder, loc, 1)); + Value earlyOff = arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::eq, early, zeroC); + Value hasMask = arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::ne, maskValue, zeroC); + LLVM::AssumeOp::create( + builder, loc, + arith::AndIOp::create(builder, loc, earlyIsBool, + arith::OrIOp::create(builder, loc, earlyOff, hasMask))); + earlyTimeoutBit = arith::ShLIOp::create( + builder, loc, arith::AndIOp::create(builder, loc, early, i32Const(builder, loc, 1)), + i32Const(builder, loc, 21)); + } + + // GROUP1: config | mask, tensor_dim0/1 | barrier addr, tile_dim0/1/2, stride0/1. + int32_t g1s0Upper = (dataSizeCode << 16) | ((cfg.iterCount > 1 ? 1 : 0) << 19) | + ((pad.enable ? 1 : 0) << 20) | (pad.interval << 22) | (pad.amount << 25); + Value g1s0 = + orr(orr(orr(i32Const(builder, loc, g1s0Upper), maskLow), barrierEnableBit), earlyTimeoutBit); + + Value td0 = descTensorDim(0); + Value td1 = rank >= 2 ? descTensorDim(1) : zeroC; + int32_t tile0 = descTileDim(0); + int32_t tile1 = rank >= 2 ? descTileDim(1) : 0; + int32_t tile2 = rank >= 3 ? descTileDim(2) : 0; + Value g1s1 = orr(shl16(lo16(td0)), barrierAddrField); // tensor_dim0 lo16 | barrier addr + Value g1s2 = orr(hi16(td0), shl16(lo16(td1))); // dim0 hi16 | dim1 lo16 + auto upper16 = [](int32_t value) { + return static_cast(static_cast(value) << 16); + }; + Value g1s3 = orr(hi16(td1), i32Const(builder, loc, upper16(tile0))); // dim1 hi16 | tile0 + int32_t g1s4c = (tile1 & 0xFFFF) | upper16(tile2); + Value g1s5 = zeroC, g1s6 = zeroC, g1s7 = zeroC; + if (rank >= 2) { + g1s5 = descStrideLo32(0); + g1s6 = descStrideHi16(0); + if (rank >= 3) { + g1s6 = orr(g1s6, shl16(lo16(descStrideLo32(1)))); + g1s7 = orr(hi16(descStrideLo32(1)), shl16(descStrideHi16(1))); + } + } + Value dgroup1 = vector::FromElementsOp::create( + builder, loc, VectorType::get({8}, i32Ty), + ValueRange{g1s0, g1s1, g1s2, g1s3, i32Const(builder, loc, g1s4c), g1s5, g1s6, g1s7}); + + // GROUP2 (rank>=3): tensor_dim2, tensor_dim3, stride2, tile_dim3 -- or, when the + // descriptor iterates, the same three slots redefined as lds_addr_increment, + // global_addr_increment and iterate_count. That is why iteration costs the fourth dim. + Value g2s0 = zeroC, g2s1 = zeroC, g2s2 = zeroC, g2s3 = zeroC; + if (rank >= 3) + g2s0 = descTensorDim(2); + if (cfg.iterCount > 1) { + // Each replay lands the next box directly after the previous one in LDS, so the LDS + // step is one whole box; the global step is the residual axis's stride, baked in as + // a construction argument because a tensor stride can be dynamic. + int64_t boxElems = 1; + for (int32_t d : tileShape) + boxElems *= d; + Value glbInc = slotField(slots.iterStride); + g2s1 = i32Const(builder, loc, static_cast(boxElems)); + g2s2 = LLVM::TruncOp::create(builder, loc, i32Ty, glbInc); + Value glbIncHi = arith::AndIOp::create( + builder, loc, + LLVM::TruncOp::create( + builder, loc, i32Ty, + LLVM::LShrOp::create(builder, loc, glbInc, + arith::ConstantIntOp::create(builder, loc, 32, 64))), + mask16); + g2s3 = orr(glbIncHi, i32Const(builder, loc, upper16(cfg.iterCount - 1))); + } else if (rank >= 4) { + g2s1 = descTensorDim(3); + g2s2 = descStrideLo32(2); + g2s3 = orr(descStrideHi16(2), i32Const(builder, loc, upper16(descTileDim(3)))); + } + Value dg2 = vector::FromElementsOp::create(builder, loc, VectorType::get({4}, i32Ty), + ValueRange{g2s0, g2s1, g2s2, g2s3}); + + // GROUP3 (rank==5): stride3, tensor_dim4, tile_dim4. + Value g3s0 = zeroC, g3s1 = zeroC, g3s2 = zeroC, g3s3 = zeroC; + if (rank == 5) { + Value td4 = descTensorDim(4); + g3s0 = descStrideLo32(3); + g3s1 = orr(descStrideHi16(3), shl16(lo16(td4))); + g3s2 = orr(hi16(td4), i32Const(builder, loc, upper16(descTileDim(4)))); + } + Value dg3 = vector::FromElementsOp::create(builder, loc, VectorType::get({4}, i32Ty), + ValueRange{g3s0, g3s1, g3s2, g3s3}); + + Value dg4 = vector::FromElementsOp::create( + builder, loc, VectorType::get({8}, i32Ty), + ValueRange{zeroC, zeroC, zeroC, zeroC, zeroC, zeroC, zeroC, zeroC}); + + // The ROCDL intrinsic takes the cache policy as an attribute, which is why + // `cacheModifier` has to stay a compile-time atom-type parameter. + auto cachePolicy = builder.getI32IntegerAttr(static_cast(cfg.cacheModifier)); + ArrayAttr noAliasScopes; + if (cfg.isLoad) + ROCDL::TensorLoadToLDSOp::create(builder, loc, dgroup0, dgroup1, dg2, dg3, dg4, cachePolicy, + noAliasScopes, noAliasScopes, noAliasScopes); + else + ROCDL::TensorStoreFromLDSOp::create(builder, loc, dgroup0, dgroup1, dg2, dg3, dg4, cachePolicy, + noAliasScopes, noAliasScopes, noAliasScopes); + + return success(); +} + +} // namespace + +//===----------------------------------------------------------------------===// +// CopyOpCDNA5TensorLoadType — TENSOR_LOAD_TO_LDS (global -> LDS) +//===----------------------------------------------------------------------===// + +// Static, so it answers for the direction rather than for one atom: whether a given +// atom carries `atomic_barrier_addr` is a type parameter, and `setAtomState` reads it. +std::optional CopyOpCDNA5TensorLoadType::getFieldIndex(AtomStateField field) { + return tdmFieldIndex(kLoadSlots, field); +} + +Type CopyOpCDNA5TensorLoadType::getConvertedType(MLIRContext *ctx) const { + return tdmConvertedType(ctx, kLoadSlots); +} + +// The tensor geometry is not optional for this atom, so there is no resting state +// to hand back: an atom without its base pointer could only lower to a descriptor +// pointing at nothing. `getAtomState` is the way in. +Value CopyOpCDNA5TensorLoadType::getDefaultState(OpBuilder &builder, Location loc) const { + mlir::emitError(loc) << "cdna5 TDM: this atom is built from a tensor and needs its " + "construction arguments (base, strides, extents)"; + return nullptr; +} + +Value CopyOpCDNA5TensorLoadType::getAtomState(OpBuilder &builder, Location loc, + ValueRange args) const { + return tdmInitialState(builder, loc, tdmSlots(kLoadSlots, getAtomicBarrier()), + getTileShape().size(), getTensor2tdm(), getIterCount() > 1, args); +} + +Value CopyOpCDNA5TensorLoadType::setAtomState(OpBuilder &builder, Location loc, Value atomStruct, + Attribute fieldAttr, Value fieldValue) const { + return tdmSetAtomState(builder, loc, tdmSlots(kLoadSlots, getAtomicBarrier()), + getTileShape().size(), atomStruct, fieldAttr, fieldValue); +} + +Attribute CopyOpCDNA5TensorLoadType::getThrLayout() const { return FxLayout(FxC(1), FxC(1)); } + +Attribute CopyOpCDNA5TensorLoadType::getThrBitLayoutSrc() const { + return tdmThrBitLayout(getContext(), getTileShape(), tdmElemBits(getDataType()), getIterCount()); +} +Attribute CopyOpCDNA5TensorLoadType::getThrBitLayoutDst() const { return getThrBitLayoutSrc(); } +Attribute CopyOpCDNA5TensorLoadType::getThrBitLayoutRef() const { return getThrBitLayoutSrc(); } + +LogicalResult CopyOpCDNA5TensorLoadType::verify(function_ref emitError, + ArrayRef tileShape, Type dataType, + IntTupleAttr tensor2tdm, bool atomicBarrier, + int32_t cacheModifier, int32_t iterCount, + int32_t padInterval, int32_t padAmount) { + return tdmVerify(emitError, tileShape, dataType, tensor2tdm, iterCount, padInterval, padAmount); +} + +FailureOr CopyOpCDNA5TensorLoadType::emitAtomCallSSA(OpBuilder &builder, Location loc, + Type resultTy, Type copyAtomTyArg, + Type srcTyArg, Type dstTyArg, + Value atomVal, Value src, + Value dst) const { + if (failed(emitAtomCall(builder, loc, copyAtomTyArg, srcTyArg, dstTyArg, atomVal, src, dst))) + return failure(); + return Value{}; +} + +FailureOr CopyOpCDNA5TensorLoadType::emitAtomCallSSA( + OpBuilder &builder, Location loc, Type resultTy, Type copyAtomTyArg, Type srcTyArg, + Type dstTyArg, Type predTyArg, Value atomVal, Value src, Value dst, Value pred) const { + if (failed(emitAtomCall(builder, loc, copyAtomTyArg, srcTyArg, dstTyArg, predTyArg, atomVal, src, + dst, pred))) + return failure(); + return Value{}; +} + +LogicalResult CopyOpCDNA5TensorLoadType::emitAtomCall(OpBuilder &builder, Location loc, + Type copyAtomTyArg, Type srcTyArg, + Type dstTyArg, Value atomVal, Value src, + Value dst) const { + TdmStatic cfg = tdmConfig(*this); + return emitTdmAtomCall(builder, loc, cfg, tdmSlots(kLoadSlots, getAtomicBarrier()), srcTyArg, + dstTyArg, atomVal, src, dst); +} + +LogicalResult CopyOpCDNA5TensorLoadType::emitAtomCall(OpBuilder &builder, Location loc, + Type copyAtomTyArg, Type srcTyArg, + Type dstTyArg, Type predMemTyArg, + Value atomVal, Value src, Value dst, + Value pred) const { + TdmStatic cfg = tdmConfig(*this); + return emitTdmAtomCall(builder, loc, cfg, tdmSlots(kLoadSlots, getAtomicBarrier()), srcTyArg, + dstTyArg, atomVal, src, dst, predMemTyArg, pred); +} + +//===----------------------------------------------------------------------===// +// CopyOpCDNA5TensorStoreType — TENSOR_STORE_FROM_LDS (LDS -> global) +//===----------------------------------------------------------------------===// + +std::optional CopyOpCDNA5TensorStoreType::getFieldIndex(AtomStateField field) { + return tdmFieldIndex(kStoreSlots, field); +} + +Type CopyOpCDNA5TensorStoreType::getConvertedType(MLIRContext *ctx) const { + return tdmConvertedType(ctx, kStoreSlots); +} + +// The tensor geometry is not optional for this atom, so there is no resting state +// to hand back: an atom without its base pointer could only lower to a descriptor +// pointing at nothing. `getAtomState` is the way in. +Value CopyOpCDNA5TensorStoreType::getDefaultState(OpBuilder &builder, Location loc) const { + mlir::emitError(loc) << "cdna5 TDM: this atom is built from a tensor and needs its " + "construction arguments (base, strides, extents)"; + return nullptr; +} + +Value CopyOpCDNA5TensorStoreType::getAtomState(OpBuilder &builder, Location loc, + ValueRange args) const { + return tdmInitialState(builder, loc, tdmSlots(kStoreSlots, getAtomicBarrier()), + getTileShape().size(), getTensor2tdm(), getIterCount() > 1, args); +} + +Value CopyOpCDNA5TensorStoreType::setAtomState(OpBuilder &builder, Location loc, Value atomStruct, + Attribute fieldAttr, Value fieldValue) const { + return tdmSetAtomState(builder, loc, tdmSlots(kStoreSlots, getAtomicBarrier()), + getTileShape().size(), atomStruct, fieldAttr, fieldValue); +} + +Attribute CopyOpCDNA5TensorStoreType::getThrLayout() const { return FxLayout(FxC(1), FxC(1)); } + +Attribute CopyOpCDNA5TensorStoreType::getThrBitLayoutSrc() const { + return tdmThrBitLayout(getContext(), getTileShape(), tdmElemBits(getDataType()), getIterCount()); +} +Attribute CopyOpCDNA5TensorStoreType::getThrBitLayoutDst() const { return getThrBitLayoutSrc(); } +Attribute CopyOpCDNA5TensorStoreType::getThrBitLayoutRef() const { return getThrBitLayoutSrc(); } + +LogicalResult CopyOpCDNA5TensorStoreType::verify(function_ref emitError, + ArrayRef tileShape, Type dataType, + IntTupleAttr tensor2tdm, bool atomicBarrier, + int32_t cacheModifier, int32_t iterCount) { + return tdmVerify(emitError, tileShape, dataType, tensor2tdm, iterCount, /*padInterval=*/0, + /*padAmount=*/0); +} + +FailureOr CopyOpCDNA5TensorStoreType::emitAtomCallSSA(OpBuilder &builder, Location loc, + Type resultTy, Type copyAtomTyArg, + Type srcTyArg, Type dstTyArg, + Value atomVal, Value src, + Value dst) const { + if (failed(emitAtomCall(builder, loc, copyAtomTyArg, srcTyArg, dstTyArg, atomVal, src, dst))) + return failure(); + return Value{}; +} + +FailureOr CopyOpCDNA5TensorStoreType::emitAtomCallSSA( + OpBuilder &builder, Location loc, Type resultTy, Type copyAtomTyArg, Type srcTyArg, + Type dstTyArg, Type predTyArg, Value atomVal, Value src, Value dst, Value pred) const { + if (failed(emitAtomCall(builder, loc, copyAtomTyArg, srcTyArg, dstTyArg, predTyArg, atomVal, src, + dst, pred))) + return failure(); + return Value{}; +} + +LogicalResult CopyOpCDNA5TensorStoreType::emitAtomCall(OpBuilder &builder, Location loc, + Type copyAtomTyArg, Type srcTyArg, + Type dstTyArg, Value atomVal, Value src, + Value dst) const { + TdmStatic cfg = tdmConfig(*this); + return emitTdmAtomCall(builder, loc, cfg, tdmSlots(kStoreSlots, getAtomicBarrier()), srcTyArg, + dstTyArg, atomVal, src, dst); +} + +LogicalResult CopyOpCDNA5TensorStoreType::emitAtomCall(OpBuilder &builder, Location loc, + Type copyAtomTyArg, Type srcTyArg, + Type dstTyArg, Type predMemTyArg, + Value atomVal, Value src, Value dst, + Value pred) const { + TdmStatic cfg = tdmConfig(*this); + return emitTdmAtomCall(builder, loc, cfg, tdmSlots(kStoreSlots, getAtomicBarrier()), srcTyArg, + dstTyArg, atomVal, src, dst, predMemTyArg, pred); +} + +} // namespace mlir::fly_rocdl diff --git a/lib/Dialect/FlyROCDL/CDNA5/TdmAtomBuilder.cpp b/lib/Dialect/FlyROCDL/CDNA5/TdmAtomBuilder.cpp new file mode 100644 index 000000000..67fccb061 --- /dev/null +++ b/lib/Dialect/FlyROCDL/CDNA5/TdmAtomBuilder.cpp @@ -0,0 +1,143 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors + +#include "flydsl/Dialect/FlyROCDL/Utils/TdmAtomBuilder.h" + +using namespace mlir; +using namespace mlir::fly; + +namespace mlir::fly_rocdl { + +template +FailureOr deriveTdmAtom(AdaptorT adaptor, tdm::Request &request, Type &dataType, + function_ref emitError) { + auto memRefTy = dyn_cast(adaptor.getTensor().getType()); + if (!memRefTy) + return emitError() << "expected a !fly.memref tensor, got " << adaptor.getTensor().getType(); + + auto gLayout = dyn_cast(memRefTy.getLayout()); + if (!gLayout) + return emitError() << "expected a plain #fly.layout on the tensor, got " + << memRefTy.getLayout(); + + auto smemLayoutTy = dyn_cast(adaptor.getSmemLayout().getType()); + if (!smemLayoutTy) + return emitError() << "expected !fly.layout for the LDS layout"; + + Type elemTy = memRefTy.getElemTy(); + if (!elemTy.isIntOrFloat()) + return emitError() << "the tensor's element type must be an integer or float, got " << elemTy; + // The descriptor's unit, as a type: `internal_type` renames it as well as widening it, + // and the atom is the one place that says which element it moves. + std::optional internalType = adaptor.getInternalType(); + dataType = internalType ? *internalType : elemTy; + if (!dataType.isIntOrFloat()) + return emitError() << "internal_type must be an integer or float type, got " << dataType; + + request.gLayout = gLayout; + request.smemLayout = smemLayoutTy.getAttr(); + // The tiler is entirely static, so like `smemLayout` it is read off its operand's type + // and nothing is emitted for the operand itself. + auto tilerTy = dyn_cast(adaptor.getTiler().getType()); + if (!tilerTy) + return emitError() << "expected !fly.tile for the tiler, got " << adaptor.getTiler().getType(); + + FailureOr valueMap = + tdm::makeValueMap(gLayout.getShape(), tilerTy.getAttr(), emitError); + if (failed(valueMap)) + return failure(); + + request.valueMap = *valueMap; + request.elemBits = static_cast(elemTy.getIntOrFloatBitWidth()); + request.internalBits = static_cast(dataType.getIntOrFloatBitWidth()); + request.numWarps = adaptor.getNumWarps(); + request.initBoundaryCheck = adaptor.getInitBoundaryCheck(); + + if (request.numWarps < 1) + return emitError() << "num_warps must be positive, got " << request.numWarps; + + return tdm::derive(request, emitError); +} + +template FailureOr +deriveTdmAtom(MakeTiledTdmLoadAtomOpAdaptor, tdm::Request &, Type &, + function_ref); +template FailureOr +deriveTdmAtom(MakeTiledTdmStoreAtomOpAdaptor, tdm::Request &, + Type &, function_ref); + +FailureOr tdmPartitionLayout(Type atomType, Type stensorType, Type gtensorType, + int32_t numWarps, + function_ref emitError) { + auto atomTy = dyn_cast(atomType); + if (!atomTy) + return emitError() << "expected a !fly.copy_atom, got " << atomType; + auto atomVal = dyn_cast(atomTy.getThrValLayoutSrc()); + if (!atomVal || atomVal.getShape().isLeaf()) + return emitError() << "the atom has no (thread, value) layout"; + + auto smemTy = dyn_cast(stensorType); + if (!smemTy) + return emitError() << "expected a !fly.memref LDS tile, got " << stensorType; + auto smemLayout = dyn_cast(smemTy.getLayout()); + if (!smemLayout) + return emitError() << "the LDS tile needs a plain layout, got " << smemTy.getLayout(); + if (!smemTy.getElemTy().isIntOrFloat()) + return emitError() << "the LDS tile's element needs a width, got " << smemTy.getElemTy(); + + // The coordinate tile is a `!fly.coord_tensor` in a kernel; a memref stands in for it + // where a test cuts the tile out of an ordinary tensor. Only its shape is read either way. + Attribute coordLayoutAttr; + if (auto coordTy = dyn_cast(gtensorType)) + coordLayoutAttr = coordTy.getLayout(); + else if (auto memTy = dyn_cast(gtensorType)) + coordLayoutAttr = memTy.getLayout(); + else + return emitError() << "expected a !fly.coord_tensor or !fly.memref coordinate tile, got " + << gtensorType; + auto coordLayout = dyn_cast(coordLayoutAttr); + if (!coordLayout) + return emitError() << "the coordinate tile needs a plain layout, got " << coordLayoutAttr; + + FailureOr layout = + tdm::partitionLayout(atomVal.getShape().at(1), atomTy.getValBits(), smemLayout, + static_cast(smemTy.getElemTy().getIntOrFloatBitWidth()), + coordLayout.getShape(), numWarps, emitError); + if (failed(layout)) + return failure(); + return LayoutType::get(*layout); +} + +FailureOr tdmLoadOpType(MLIRContext *ctx, const tdm::Geometry &geometry, Type dataType, + IntTupleAttr tensor2tdm, bool atomicBarrier, int32_t cacheModifier, + function_ref emitError) { + SmallVector shapeStorage = tdm::tileShape(geometry); + ArrayRef shape(shapeStorage); + + Type ty = CopyOpCDNA5TensorLoadType::getChecked(emitError, ctx, shape, dataType, tensor2tdm, + atomicBarrier, cacheModifier, geometry.iterCount, + geometry.padInterval, geometry.padAmount); + if (!ty) + return failure(); + return ty; +} + +FailureOr tdmStoreOpType(MLIRContext *ctx, const tdm::Geometry &geometry, Type dataType, + IntTupleAttr tensor2tdm, bool atomicBarrier, int32_t cacheModifier, + function_ref emitError) { + if (geometry.padAmount) + return emitError() << "a TDM store cannot drain a padded LDS tile (pad interval " + << geometry.padInterval << ", amount " << geometry.padAmount + << "); TENSOR_STORE_FROM_LDS has no de-padding, it walks LDS with the " + "packed tile stride"; + SmallVector shapeStorage = tdm::tileShape(geometry); + ArrayRef shape(shapeStorage); + Type ty = + CopyOpCDNA5TensorStoreType::getChecked(emitError, ctx, shape, dataType, tensor2tdm, + atomicBarrier, cacheModifier, geometry.iterCount); + if (!ty) + return failure(); + return ty; +} + +} // namespace mlir::fly_rocdl diff --git a/lib/Dialect/FlyROCDL/CDNA5/TdmGeometry.cpp b/lib/Dialect/FlyROCDL/CDNA5/TdmGeometry.cpp new file mode 100644 index 000000000..a5fd0b567 --- /dev/null +++ b/lib/Dialect/FlyROCDL/CDNA5/TdmGeometry.cpp @@ -0,0 +1,1014 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors +// +// Descriptor construction for the CDNA5 TDM copy atom. +// +// The question this answers: *given a global tensor and an LDS tile layout, which global +// modes does the DMA box span, in what order, and what are the per-dim extents and +// strides?* +// +// A global mode is identified by its depth-first **leaf index** into the (possibly +// hierarchical) tensor shape. A Fly basis stride carries the *path* instead (`1E1E0`), so +// the two are translated by `flatIndexOfPath`; a hierarchical global tensor therefore +// needs no flattening of its own. +// +// What the LDS tile's padding does to all of this: TDM has no smem swizzle, it has an LDS +// padding field that the DMA engine walks on its own. So the tile is split as +// `padded tile -> (pad field, compact tile)` and the algebra runs on the compact tile. +// Without that split the smem vector would stop at the first padded row and the box would +// collapse to a single row. + +#include "flydsl/Dialect/FlyROCDL/Utils/TdmGeometry.h" + +#include "flydsl/Dialect/Fly/Utils/IntTupleUtils.h" +#include "flydsl/Dialect/Fly/Utils/LayoutUtils.h" + +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/STLExtras.h" + +#include +#include + +using namespace mlir; +using namespace mlir::fly; + +namespace mlir::fly_rocdl::tdm { + +namespace { + +using TupleBuilder = IntTupleBuilder; +using AttrLayoutBuilder = LayoutBuilder; + +//===----------------------------------------------------------------------===// +// Nested int tuples, addressed by depth-first leaf index +//===----------------------------------------------------------------------===// + +void flattenLeaves(IntTupleAttr t, SmallVectorImpl &out) { + intTupleFlattenToVector(TupleBuilder(t.getContext()), t, out); +} + +int32_t leafCount(IntTupleAttr t) { + SmallVector leaves; + flattenLeaves(t, leaves); + return static_cast(leaves.size()); +} + +/// The depth-first leaf index a basis stride's mode path names, or -1 when the path does +/// not address a leaf of `profile`. +int32_t flatIndexOfPath(IntTupleAttr profile, ArrayRef path) { + int32_t base = 0; + IntTupleAttr cur = profile; + for (int32_t step : path) { + if (cur.isLeaf() || step < 0 || step >= cur.rank()) + return -1; + for (int32_t i = 0; i < step; ++i) + base += leafCount(cur.at(i)); + cur = cur.at(step); + } + return cur.isLeaf() ? base : -1; +} + +/// An int tuple's elements are int tuples, never the bare `IntAttr` / `BasisAttr` +/// underneath them; wrap one that is bare, and pass a tuple through. +IntTupleAttr asLeaf(Attribute a) { + if (auto tuple = dyn_cast(a)) + return tuple; + return IntTupleAttr::get(a); +} + +/// Rebuild a tuple shaped like `profile`, taking leaves from `flat` in depth-first order. +/// The leaf at index `splitLeaf` is replaced by `splitValue` instead, which makes the +/// result one level deeper there — that is how a recast's integer division is written +/// into the coordinate tensor's shape. `intTupleUnflatten` takes whatever attribute sits +/// at a profile leaf, so a two-element `splitValue` needs no special case. +IntTupleAttr unflattenLike(MLIRContext *ctx, IntTupleAttr profile, ArrayRef flat, + int32_t splitLeaf = -1, Attribute splitValue = nullptr) { + SmallVector leaves; + for (auto [i, elem] : llvm::enumerate(flat)) + leaves.push_back(static_cast(i) == splitLeaf ? splitValue : Attribute(asLeaf(elem))); + return intTupleUnflatten(TupleBuilder(ctx), IntTupleAttr::get(ArrayAttr::get(ctx, leaves)), + profile); +} + +/// A two-element int tuple, as an attribute usable as a leaf replacement. +Attribute makePair(MLIRContext *ctx, Attribute a, Attribute b) { + return IntTupleAttr::get(ArrayAttr::get(ctx, {asLeaf(a), asLeaf(b)})); +} + +//===----------------------------------------------------------------------===// +// LDS padding: split the tile into a pad field and the compact tile +//===----------------------------------------------------------------------===// + +struct LdsEntry { + int32_t leaf; + int32_t extent; + int32_t step; +}; + +} // namespace + +FailureOr> +analyzeLdsTile(LayoutAttr smemLayout, function_ref emitError) { + MLIRContext *ctx = smemLayout.getContext(); + IntTupleAttr shape = smemLayout.getShape(); + IntTupleAttr stride = smemLayout.getStride(); + + SmallVector shapeLeaves, strideLeaves; + flattenLeaves(shape, shapeLeaves); + flattenLeaves(stride, strideLeaves); + if (shapeLeaves.size() != strideLeaves.size()) + return emitError() << "the LDS tile's shape and stride do not have the same profile"; + + SmallVector entries; + for (auto [i, extentLeaf, stepLeaf] : llvm::enumerate(shapeLeaves, strideLeaves)) { + if (!extentLeaf.isLeafInt() || !extentLeaf.isStatic() || !stepLeaf.isLeafInt() || + !stepLeaf.isStatic()) + return emitError() << "the LDS tile must be fully static, mode " << i << " is not"; + int32_t extent = extentLeaf.getLeafAsInt().getValue(); + int32_t step = stepLeaf.getLeafAsInt().getValue(); + if (extent < 1) + return emitError() << "LDS mode " << i << " has extent " << extent; + if (extent > 1 && step == 0) + return emitError() << "LDS mode " << i << " has extent " << extent + << " and stride 0; TDM walks LDS linearly and cannot broadcast one " + "location over several tile elements"; + entries.push_back({static_cast(i), extent, step}); + } + + // The modes are ordered by stride rather than assumed to run outermost-first, so a + // column-major, permuted or hierarchical tile is as ordinary as a row-major one, and the + // pad may sit at any level rather than only on the innermost row. + SmallVector active; + for (const LdsEntry &e : entries) + if (e.extent > 1) + active.push_back(e); + llvm::stable_sort(active, [](const LdsEntry &a, const LdsEntry &b) { return a.step < b.step; }); + + SmallVector compact(active.size(), 1); + for (size_t k = 1; k < active.size(); ++k) + compact[k] = compact[k - 1] * active[k - 1].extent; + if (!active.empty() && active[0].step != 1) + return emitError() << "the LDS tile must be contiguous along its fastest axis, but its " + "smallest stride is " + << active[0].step << " (mode " << active[0].leaf << ")"; + + int32_t padInterval = 0, padAmount = 0; + size_t split = active.size(); + for (size_t k = 0; k < active.size(); ++k) + if (active[k].step != compact[k]) { + split = k; + break; + } + if (split < active.size()) { + padAmount = active[split].step - compact[split]; + if (padAmount < 0) + return emitError() << "LDS mode " << active[split].leaf << " has stride " + << active[split].step << ", smaller than the " << compact[split] + << " elements inside it; the tile overlaps itself"; + padInterval = compact[split]; + SmallVector padded(compact); + padded[split] = active[split].step; + for (size_t k = split + 1; k < active.size(); ++k) + padded[k] = padded[k - 1] * active[k - 1].extent; + for (size_t k = 0; k < active.size(); ++k) + if (active[k].step != padded[k]) + return emitError() << "the LDS tile's stride (modes sorted by stride) is not a " + "single-pad form; TDM can express one constant skip every " + "pad_interval elements, so mode " + << active[k].leaf << " must have stride " << padded[k] << ", got " + << active[k].step; + } + + // `compactStride` is the same layout with the skip taken back out, shaped like `shape`. + // That is what the geometry algebra runs on. + Attribute one = IntAttr::getStatic(ctx, 1); + SmallVector flatStride(shapeLeaves.size(), one); + for (auto [k, e] : llvm::enumerate(active)) + flatStride[e.leaf] = IntAttr::getStatic(ctx, compact[k]); + return std::make_tuple(padInterval, padAmount, unflattenLike(ctx, shape, flatStride)); +} + +namespace { + +//===----------------------------------------------------------------------===// +// The value map: the tensor's identity layout composed with the tiler +//===----------------------------------------------------------------------===// + +/// `profile` with every extent replaced by 1. +Attribute onesLike(MLIRContext *ctx, IntTupleAttr profile) { + if (profile.isLeaf()) + return IntAttr::getStatic(ctx, 1); + SmallVector elems; + for (int32_t i = 0; i < profile.rank(); ++i) + elems.push_back(onesLike(ctx, profile.at(i))); + return TileAttr::get(ArrayAttr::get(ctx, elems)); +} + +/// Right-pad a tiler with 1s until it has the tensor's profile. +/// +/// Fly's composition keeps every mode a shorter tiler does not reach — at every nesting +/// level, not just the top — so the 1s are written out here: taking one element from a mode +/// is the same thing as not tiling it, and it is what keeps the value map's size equal to +/// the LDS tile's. +/// +/// A scalar against a hierarchical mode is left alone: the tiler splits itself across the +/// sub-modes there, and composition already gets it right. +/// +/// A `Tile`'s modes are bare attributes rather than a uniform tuple — an `IntAttr` for a +/// plain extent, a nested `TileAttr` for a hierarchical one, a `LayoutAttr` for a +/// strided tiler — so this walks `Attribute` and names what it cannot take. +FailureOr padTiler(MLIRContext *ctx, Attribute tiler, IntTupleAttr profile, + int32_t depth, function_ref emitError) { + if (auto extent = dyn_cast(tiler)) { + if (!extent.isStatic()) + return emitError() << "the tiler must be static, mode at nesting depth " << depth + << " is not"; + return Attribute(IntAttr::getStatic(ctx, extent.getValue())); + } + auto tile = dyn_cast(tiler); + if (!tile) + return emitError() << "the tiler's mode at nesting depth " << depth + << " must be an extent or a nested tile, got " << tiler + << "; a strided (layout) tiler mode is not modelled"; + // A one-element `Tile` wraps its mode rather than nesting it. + if (tile.isLeaf()) + return padTiler(ctx, tile.getValue(), profile, depth, emitError); + if (profile.isLeaf()) + return emitError() << "the tiler is deeper than the tensor at nesting depth " << depth + << "; the tensor has a single extent there but the tiler has " << tile.rank() + << " sub-modes"; + if (tile.rank() > profile.rank()) + return emitError() << "the tiler has " << tile.rank() << " modes at nesting depth " << depth + << " but the tensor has " << profile.rank() + << "; a tiler selects from a mode's leading sub-modes"; + SmallVector elems; + for (int32_t i = 0; i < profile.rank(); ++i) { + if (i < tile.rank()) { + FailureOr sub = padTiler(ctx, tile.at(i), profile.at(i), depth + 1, emitError); + if (failed(sub)) + return failure(); + elems.push_back(*sub); + } else { + elems.push_back(onesLike(ctx, profile.at(i))); + } + } + return Attribute(TileAttr::get(ArrayAttr::get(ctx, elems))); +} + +} // namespace + +FailureOr makeValueMap(IntTupleAttr gshape, TileAttr tiler, + function_ref emitError) { + MLIRContext *ctx = gshape.getContext(); + FailureOr padded = padTiler(ctx, tiler, gshape, 0, emitError); + if (failed(padded)) + return failure(); + auto tileAttr = dyn_cast(*padded); + if (!tileAttr) + tileAttr = TileAttr::get(*padded); + + AttrLayoutBuilder layoutBuilder(ctx); + LayoutAttr identity = LayoutAttr::get(gshape, intTupleMakeBasisTupleLike(gshape)); + return layoutComposition(layoutBuilder, identity, tileAttr); +} + +namespace { + +//===----------------------------------------------------------------------===// +// The transfer box: which modes it spans, and how far +//===----------------------------------------------------------------------===// + +/// One entry of the smem run: how many elements it takes, and from which global mode. +struct VectorEntry { + int32_t extent; + int32_t mode; +}; + +/// `coalesce(composition(valueMap, right_inverse(slayout)))`. +/// +/// The vector is the smem run innermost first, truncated at the first mode whose basis +/// coefficient is not 1 (no starting in the middle of a global +/// mode). `nextAfterCut` is the mode the cut stopped on, or -1 when the vector runs to the +/// end — it matters because it is the mode the box is *adjacent to in LDS*, and hardware +/// iteration walks LDS with a constant increment of one box, so that mode and only that +/// mode can be folded into the instruction. +FailureOr> ldsRun(MLIRContext *ctx, LayoutAttr valueMap, + IntTupleAttr gshape, IntTupleAttr smemShape, + IntTupleAttr compactStride, int32_t &nextAfterCut, + function_ref emitError) { + AttrLayoutBuilder layoutBuilder(ctx); + LayoutAttr compactSmem = LayoutAttr::get(smemShape, compactStride); + LayoutAttr invSmem = layoutRightInverse(layoutBuilder, compactSmem); + LayoutAttr sidx2gmode = + layoutCoalesce(layoutBuilder, layoutComposition(layoutBuilder, valueMap, invSmem)); + + SmallVector extents, strides; + flattenLeaves(sidx2gmode.getShape(), extents); + flattenLeaves(sidx2gmode.getStride(), strides); + + SmallVector vector; + nextAfterCut = -1; + for (auto [extentLeaf, strideLeaf] : llvm::zip(extents, strides)) { + if (!extentLeaf.isLeafInt() || !extentLeaf.isStatic()) + return emitError() << "the tile/global vectorization must be static"; + int32_t extent = extentLeaf.getLeafAsInt().getValue(); + int32_t coeff = 0, mode = -1; + if (strideLeaf.isLeafBasis()) { + BasisAttr basis = strideLeaf.getLeafAsBasis(); + if (!basis.getValue().isStatic()) + return emitError() << "a dynamic basis coefficient is not supported"; + coeff = basis.getValue().getValue(); + mode = flatIndexOfPath(gshape, basis.getModes()); + } else if (strideLeaf.isLeafInt() && strideLeaf.isStatic()) { + // A stride-0 / stride-1 constant leaf carries no global mode. + coeff = strideLeaf.getLeafAsInt().getValue(); + } else { + return emitError() << "unsupported stride leaf in the tile/global vectorization"; + } + if (coeff != 1 || mode < 0) { + nextAfterCut = mode; // -1 when the leaf carried no mode at all + break; // stop at the first non-unit basis + } + vector.push_back({extent, mode}); + } + if (vector.empty()) + return emitError() << "no common tile/global vectorization — the LDS tile and the global " + "tile do not share a contiguous innermost run. Does the tiler select " + "out the major global mode?"; + return vector; +} + +/// One warp's share of the smem run: its leading `1 / numWarps`. +/// +/// The whole run belongs to the workgroup; the box only has to be one participant's share +/// of it, and the participants take equal contiguous chunks. Every chunk is then the same +/// box translated along a single global mode, which is what lets one descriptor serve all +/// of them: the caller moves the origin, not the geometry. +FailureOr> splitAcrossWarps(ArrayRef vector, int32_t numWarps, + function_ref emitError) { + if (numWarps == 1) + return SmallVector(vector); + int32_t total = 1; + for (const VectorEntry &e : vector) + total *= e.extent; + if (total % numWarps) + return emitError() << numWarps << " warps cannot split this transfer's " << total + << "-element contiguous run evenly; num_warps has to divide it"; + int32_t want = total / numWarps; + + SmallVector out; + int32_t acc = 1; + for (const VectorEntry &e : vector) { + int32_t room = want / acc; + if (room == 1) + break; + if (e.extent <= room) { + out.push_back(e); + acc *= e.extent; + continue; + } + if (e.extent % room) + return emitError() << "splitting the transfer run " << numWarps << " ways cuts global mode " + << e.mode << " at " << room << " of the " << e.extent + << " elements the tile takes from it, which does not divide it"; + out.push_back({room, e.mode}); + break; + } + if (out.empty()) + out.push_back({1, vector.front().mode}); + return out; +} + +//===----------------------------------------------------------------------===// +// The multi-mode gcd recurrence +//===----------------------------------------------------------------------===// + +struct RawDim { + int32_t box; + SmallVector modes; + Scalar stride; +}; + +/// The extent/stride recurrence for a dim fed by several global modes. +/// +/// A descriptor dim fed by several global modes takes `gcd` of their strides and grows its +/// extent to span them: `g_shape = (s_i - 1) * (d_i / gcd) + ... + 1`. Only defined over +/// static values — a dynamic extent or stride has no build-time gcd — and the rank-5 +/// packing is the only thing that builds such a dim, so a tensor whose tail it has to lump +/// together is refused here unless that tail's geometry is static. +FailureOr> foldModes(ArrayRef modes, ArrayRef modeExtent, + ArrayRef modeStride, + function_ref emitError) { + if (modes.size() == 1) + return std::make_pair(modeExtent[modes[0]], modeStride[modes[0]]); + + // The recurrence below is compile-time `int32_t` arithmetic, so the packing is static + // or it does not happen. That is the companion of the restriction in + // `axisBasisPerMode`: see the TODO there for what lifting both would take. + for (int32_t m : modes) + if (!modeExtent[m].isStatic || !modeStride[m].isStatic) + return emitError() << "a descriptor dim spans several global modes but mode " << m + << " has a dynamic extent or stride; a multi-mode dim needs static " + "geometry"; + int32_t stride = 0; + for (int32_t m : modes) + stride = std::gcd(stride, modeStride[m].value); + if (stride == 0) + return std::make_pair(modeExtent[modes[0]], Scalar::getStatic(0)); + int32_t extent = 1; + for (int32_t m : modes) + extent += (modeExtent[m].value - 1) * (modeStride[m].value / stride); + return std::make_pair(Scalar::getStatic(extent), Scalar::getStatic(stride)); +} + +/// Divide a scalar by the recast ratio. A static value is checked; a dynamic one records +/// the divisor for the expansion, because the caller is asserting that its tensor is laid +/// out in whole internal units and a run-time value cannot be checked against that any more +/// than a static check can verify a dynamic value. +LogicalResult recastDivide(Scalar &s, int32_t ratio, const Twine &what, + function_ref emitError) { + if (!s.isStatic) { + s.divisor *= ratio; + return success(); + } + if (s.value % ratio) + return emitError() << what << " is " << s.value + << ", which is not divisible by the recast ratio " << ratio; + s.value /= ratio; + return success(); +} + +} // namespace + +//===----------------------------------------------------------------------===// +// derive +//===----------------------------------------------------------------------===// + +FailureOr derive(const Request &request, function_ref emitError) { + LayoutAttr gLayout = request.gLayout; + MLIRContext *ctx = gLayout.getContext(); + TupleBuilder tupleBuilder(ctx); + + IntTupleAttr gshape = gLayout.getShape(); + SmallVector gshapeLeaves, gstrideLeaves; + flattenLeaves(gshape, gshapeLeaves); + flattenLeaves(gLayout.getStride(), gstrideLeaves); + if (gshapeLeaves.size() != gstrideLeaves.size()) + return emitError() << "the global tensor's shape and stride do not have the same profile"; + int32_t numModes = static_cast(gshapeLeaves.size()); + + // Per global mode, its extent and stride — static leaves as values, dynamic ones as + // references the expansion resolves against the layout operand. + SmallVector modeExtent(numModes), modeStride(numModes); + for (int32_t i = 0; i < numModes; ++i) { + if (!gshapeLeaves[i].isLeafInt() || !gstrideLeaves[i].isLeafInt()) + return emitError() << "global mode " << i << " must have an integer extent and stride"; + modeExtent[i] = gshapeLeaves[i].isStatic() + ? Scalar::getStatic(gshapeLeaves[i].getLeafAsInt().getValue()) + : Scalar::getDynamic(i, /*fromShape=*/true); + modeStride[i] = gstrideLeaves[i].isStatic() + ? Scalar::getStatic(gstrideLeaves[i].getLeafAsInt().getValue()) + : Scalar::getDynamic(i, /*fromShape=*/false); + } + + // The single contiguous mode, recorded before the recast rewrites the strides. + int32_t contiguousMode = -1; + { + int32_t found = 0; + for (int32_t i = 0; i < numModes; ++i) + if (modeStride[i].isStatic && modeStride[i].value == 1) { + contiguousMode = i; + ++found; + } + if (found != 1) + contiguousMode = -1; + } + + FailureOr> lds = + analyzeLdsTile(request.smemLayout, emitError); + if (failed(lds)) + return failure(); + + int32_t padInterval = std::get<0>(*lds); + int32_t padAmount = std::get<1>(*lds); + IntTupleAttr compactStride = std::get<2>(*lds); + + LayoutAttr valueMap = request.valueMap; + if (!valueMap) + return emitError() << "the request carries no value map; build one from the tiler with " + "`makeValueMap`"; + + // The tiler cuts the global tile that the LDS tile receives, so the two hold the same + // elements. Only the + // sizes are compared, so the tiler may be hierarchical where the LDS tile is flat. + IntTupleAttr smemSizeAttr = intTupleProduct(tupleBuilder, request.smemLayout.getShape()); + IntTupleAttr vSizeAttr = intTupleProduct(tupleBuilder, valueMap.getShape()); + if (!smemSizeAttr.isStatic() || !vSizeAttr.isStatic()) + return emitError() << "the tiler and the LDS tile must both be static"; + + int32_t smemSize = smemSizeAttr.getLeafAsInt().getValue(); + int32_t vSize = vSizeAttr.getLeafAsInt().getValue(); + if (smemSize != vSize) + return emitError() << "the tiler holds " << vSize << " elements but the LDS tile holds " + << smemSize << "; a tiler reshapes the tile, it does not shrink it"; + + // Per global mode, how many elements the tile takes and the step it takes them by. The + // step is the basis coefficient: 1 for the ordinary case where the tile walks a mode + // element by element, and more when it strides — `8:2E0` takes 8 rows two apart. + SmallVector tileExtent(numModes, 1), tileStep(numModes, 1); + { + llvm::DenseSet tiled; + SmallVector vExtents, vStrides; + flattenLeaves(valueMap.getShape(), vExtents); + flattenLeaves(valueMap.getStride(), vStrides); + for (auto [extentLeaf, strideLeaf] : llvm::zip(vExtents, vStrides)) { + if (!strideLeaf.isLeafBasis()) + continue; + BasisAttr basis = strideLeaf.getLeafAsBasis(); + if (!basis.getValue().isStatic()) + return emitError() << "a dynamic basis coefficient is not supported in the tiler"; + int32_t mode = flatIndexOfPath(gshape, basis.getModes()); + if (mode < 0) + continue; + if (!extentLeaf.isLeafInt() || !extentLeaf.isStatic()) + return emitError() << "the tiler must be static"; + int32_t extent = extentLeaf.getLeafAsInt().getValue(); + int32_t coeff = basis.getValue().getValue(); + if (tiled.insert(mode).second) { + tileExtent[mode] = extent; + tileStep[mode] = coeff; + } else { + tileExtent[mode] *= extent; + tileStep[mode] = std::min(tileStep[mode], coeff); + } + } + } + + int32_t nextAfterCut = -1; + FailureOr> vectorOr = ldsRun( + ctx, valueMap, gshape, request.smemLayout.getShape(), compactStride, nextAfterCut, emitError); + if (failed(vectorOr)) + return failure(); + SmallVector vector = *vectorOr; + + // How much of each mode the box covers, counted before the recast rewrites the run in + // wider units — `tileExtent` is in the tensor's own elements and the two are compared + // below. + SmallVector coveredExtent(numModes, 1); + for (const VectorEntry &e : vector) + coveredExtent[e.mode] *= e.extent; + + // Cut before the recast, so the other warps' share turns into + // neither a wider unit nor an iteration — it is simply not this box. + FailureOr> splitOr = + splitAcrossWarps(vector, request.numWarps, emitError); + if (failed(splitOr)) + return failure(); + vector = *splitOr; + + // The recast: view the transfer in a wider unit. The innermost run, + // the contiguous global mode's extent, and every global stride are divided by the ratio. + // This is what lets a sub-byte or awkward element type ride on a TDM data size the + // hardware can encode. + if (request.elemBits <= 0 || request.internalBits % request.elemBits != 0) + return emitError() << "internal width " << request.internalBits + << " is not a multiple of element width " << request.elemBits; + int32_t ratio = request.internalBits / request.elemBits; + if (ratio != 1) { + if (vector.front().extent % ratio) + return emitError() << "the innermost run is " << vector.front().extent + << ", which is not divisible by the recast ratio " << ratio; + vector.front().extent /= ratio; + for (int32_t i = 0; i < numModes; ++i) { + if (modeStride[i].isStatic && modeStride[i].value == 0) + continue; + if (modeStride[i].isStatic && modeStride[i].value == 1) { + // The contiguous mode is the one measured in elements; recast shrinks it. + if (failed(recastDivide(modeExtent[i], ratio, + "the extent of the contiguous global mode " + Twine(i), emitError))) + return failure(); + } else if (failed(recastDivide(modeStride[i], ratio, "the stride of global mode " + Twine(i), + emitError))) { + return failure(); + } + } + // The pad fields ride in the same units as the box (the lowering rebuilds the LDS pitch + // as `tileShape[-1] + padAmount`), so a recast rescales them too. + if (padAmount) { + if (padInterval % ratio || padAmount % ratio) + return emitError() << "LDS padding (" << padInterval << ", " << padAmount + << ") is not divisible by the recast ratio " << ratio + << "; the pad must be a whole number of internal units"; + padInterval /= ratio; + padAmount /= ratio; + } + } + + // The box need not span the whole tile. What the cut above leaves behind is not an + // error: a TiledCopy says how many values one *tiled* + // operation covers, the atom says how many one *call* covers, and the V mode carries the + // difference as more calls. + // + // What TDM adds is one optimization on top: `iterate_enable` replays a single descriptor + // `iterate_count` times, advancing global and LDS by a constant increment each time. That + // folds *one* of those axes back into the instruction. Which one is not a free choice — + // the LDS increment is one whole box, so it has to be the mode the box is adjacent to in + // LDS, which is exactly the mode the cut stopped on. Hardware iteration also + // steps LDS by one box, which is the other warps' space once the box is only a share of + // the tile, so a split gives the axis up. + llvm::DenseSet coveredModes; + for (const VectorEntry &e : vector) + coveredModes.insert(e.mode); + + int32_t iterCount = 1, iterMode = -1; + Scalar iterStride; + if (request.numWarps == 1 && nextAfterCut >= 0 && + tileExtent[nextAfterCut] > coveredExtent[nextAfterCut]) { + int32_t count = tileExtent[nextAfterCut] / coveredExtent[nextAfterCut]; + if (count <= kMaxIterateCount && !padAmount) { + Scalar step = modeStride[nextAfterCut]; + if (!step.isStatic && tileStep[nextAfterCut] != 1) { + // The step would need a run-time multiply the expansion does not model; keep the + // axis in the V mode instead, which is always correct and only costs instructions. + } else { + iterMode = nextAfterCut; + iterCount = count; + iterStride = step.isStatic ? Scalar::getStatic(step.value * tileStep[nextAfterCut]) : step; + } + } + } + + SmallVector dims; + for (const VectorEntry &e : vector) + dims.push_back({e.extent, {e.mode}, modeStride[e.mode]}); + // A tile extent of one still needs a descriptor coordinate when the global tensor can + // move on that mode. Without this filler, changing a batch index leaves global_addr at + // batch zero. + for (int32_t i = 0; i < numModes; ++i) { + if (coveredModes.contains(i)) + continue; + if (modeExtent[i].isStatic && modeExtent[i].value == 1) + continue; + if (modeStride[i].isStatic && modeStride[i].value == 0) + continue; + dims.push_back({1, {i}, modeStride[i]}); + } + + // The lowering hardwires descriptor dim 0's stride to 1, so dim 0 must be the contiguous + // global mode. It can legitimately be missing: `coalesce` drops a size-1 mode from the + // smem vector, so a tile that is one element wide along the contiguous mode loses it. Put + // it back as an extent-1 dim instead of rejecting the copy. + if (!dims.front().stride.isStatic || dims.front().stride.value != 1) { + llvm::DenseSet seen; + for (const RawDim &d : dims) + for (int32_t m : d.modes) + seen.insert(m); + if (contiguousMode >= 0 && !seen.contains(contiguousMode) && tileExtent[contiguousMode] == 1) + dims.insert(dims.begin(), {1, {contiguousMode}, Scalar::getStatic(1)}); + else + return emitError() << "the innermost descriptor dim must be contiguous in global memory " + "(stride 1), got global mode " + << dims.front().modes.front() + << " — the LDS tile's majorness does not match the global tensor's"; + } + + // Every mode the box spans keeps its own descriptor dim: two dims that happen to be + // adjacent in global memory are not folded into one, so a dim's bound is always a + // per-mode rectangular bound and the descriptor says exactly what the tile said. + SmallVector clampable(dims.size(), true); + + // Past five dims the trailing modes are packed into the + // last descriptor dim, whose extent and stride come from a gcd recurrence over modes that + // are unrelated in memory. That is the one dim whose single bound cannot speak for the + // modes under it, so it is marked unclampable and the boundary-check state refuses to switch it + // on. + if (static_cast(dims.size()) > kMaxRank) { + SmallVector head(dims.begin(), dims.begin() + (kMaxRank - 1)); + SmallVector packedModes; + int32_t packedBox = 1; + for (const RawDim &d : llvm::drop_begin(dims, kMaxRank - 1)) { + packedModes.append(d.modes.begin(), d.modes.end()); + packedBox *= d.box; + } + FailureOr> packed = + foldModes(packedModes, modeExtent, modeStride, emitError); + if (failed(packed)) + return failure(); + head.push_back({packedBox, packedModes, packed->second}); + dims = head; + clampable.truncate(kMaxRank - 1); + clampable.push_back(false); + } + + // TDM writes LDS strictly linearly -- the address only ever advances, across all the + // dim loops -- and the pad counter runs with it rather than restarting per row. So a + // skip every N elements is expressible for any N the tile is a whole number of, and + // the interval is free to span several rows. It has to divide the *whole* box because + // each instruction re-seeds that counter: a tile issued as several calls would + // otherwise put the later calls' holes in the wrong place. + if (padAmount) { + int64_t boxElems = 1; + for (const RawDim &d : dims) + boxElems *= d.box; + if (boxElems % padInterval != 0) + return emitError() << "the padded LDS tile skips every " << padInterval + << " elements, which does not divide the descriptor's " << boxElems + << "-element box; one call's holes would fall in the wrong place"; + } + + Geometry geometry; + geometry.padInterval = padInterval; + geometry.padAmount = padAmount; + geometry.ratio = ratio; + geometry.modeExtent = modeExtent; + geometry.modeStride = modeStride; + geometry.contiguousMode = contiguousMode; + for (auto [d, raw] : llvm::enumerate(dims)) { + FailureOr> folded = + foldModes(raw.modes, modeExtent, modeStride, emitError); + if (failed(folded)) + return failure(); + if (folded->second.isStatic && (folded->second.value < 0 || + static_cast(folded->second.value) > kMaxTensorStride)) + return emitError() << "descriptor dim " << d << " stride " << folded->second.value + << " is outside the unsigned 48-bit tensor_dim_stride range"; + Dim dim; + dim.box = raw.box; + dim.modes = raw.modes; + dim.tensorDim = folded->first; + dim.stride = folded->second; + dim.clampable = clampable[d]; + geometry.dims.push_back(dim); + } + + // Iteration is paid for out of the descriptor's own slots (`tdm::kMaxIterateRank`), so a + // descriptor that needs them keeps its residual axis in the V mode instead. Declining is + // not a failure: the copy still moves the whole tile, it just spends one more instruction + // per step. + if (iterCount > 1 && geometry.rank() > kMaxIterateRank) + iterCount = 1; + if (iterCount > 1) { + geometry.iterCount = iterCount; + geometry.iterStride = iterStride; + geometry.iterMode = iterMode; + // The residual axis is stepped by the hardware, not by a `tensor_dim`, so whichever dim + // carries its tile origin has no bound to give. + for (Dim &dim : geometry.dims) + if (llvm::is_contained(dim.modes, iterMode)) + dim.clampable = false; + } + + return geometry; +} + +FailureOr initialBoundaryCheck(const Geometry &geometry, IntTupleAttr gshape, + bool enable, + function_ref emitError) { + MLIRContext *ctx = gshape.getContext(); + FailureOr basis = makeTensor2Tdm(geometry, gshape, emitError); + if (failed(basis)) + return failure(); + + // `makeTensor2Tdm` already answers "does this mode have a bound to give": a basis + // leaf if it does, a `0` if it does not. So the initial state is that map with the + // caller's one flag written onto the modes that can carry it. + SmallVector basisLeaves; + flattenLeaves(*basis, basisLeaves); + SmallVector flat; + for (IntTupleAttr basisLeaf : basisLeaves) + flat.push_back(IntAttr::getStatic(ctx, enable && basisLeaf.isLeafBasis() ? 1 : 0)); + return unflattenLike(ctx, gshape, flat); +} + +//===----------------------------------------------------------------------===// +// The inverse of the mode map: per global mode, the descriptor axis it moves along +//===----------------------------------------------------------------------===// + +namespace { + +/// Per global mode, the descriptor axis it moves along, as a basis leaf (`0` = none). +/// +/// Shared by the coordinate tensor's strides and by the atom's boundary-check mode map, +/// which differ only in whether a dim the descriptor cannot put a +/// single bound on still contributes: a coordinate on such a dim is still a coordinate, a +/// bound on it would not be one. +FailureOr> axisBasisPerMode(MLIRContext *ctx, const Geometry &geometry, + IntTupleAttr gshape, bool clampableOnly, + function_ref emitError) { + SmallVector gshapeLeaves; + flattenLeaves(gshape, gshapeLeaves); + int32_t numModes = static_cast(gshapeLeaves.size()); + if (numModes != static_cast(geometry.modeStride.size())) + return emitError() << "internal: the coordinate tensor's shape has " << numModes + << " modes but the derivation recorded " << geometry.modeStride.size(); + + Attribute zero = IntAttr::getStatic(ctx, 0); + SmallVector byMode(numModes, zero); + + int32_t descRank = geometry.rank(); + for (auto [descDim, dim] : llvm::enumerate(geometry.dims)) { + if (clampableOnly && !dim.clampable) + continue; + // Axis indices are in tensor dim order, the reverse of descriptor order, matching the + // atom's `tileShape` and its state slots. + int32_t axis = descRank - 1 - static_cast(descDim); + for (int32_t mode : dim.modes) { + // A size-1 axis has only coordinate 0 and a stride-0 one is a + // broadcast; neither can move the tile, so neither gets a basis. + const Scalar &extent = geometry.modeExtent[mode]; + const Scalar &stride = geometry.modeStride[mode]; + if (extent.isStatic && extent.value == 1) + continue; + if (stride.isStatic && stride.value == 0) + continue; + if (dim.modes.size() == 1) { + byMode[mode] = BasisAttr::get(IntAttr::getStatic(ctx, 1), axis); + continue; + } + // A dim covering several modes only ever comes out of the rank-5 packing. Each of + // those modes rides the shared axis at its own scale, `mode_stride / dim_stride`, + // and the map has nowhere to put a scale that is not a compile-time integer: a + // `tensor2tdm` leaf is a `BasisAttr`, so the whole 1-to-many mapping has to be + // static or it is not expressible. + // + // TODO(rank>5, dynamic strides): a tensor with more than five modes and dynamic + // strides is refused right here, need to support tensor2tdm with dynamic basis strides. + if (!stride.isStatic || !dim.stride.isStatic || dim.stride.value == 0) + return emitError() << "descriptor dim " << descDim << " covers global modes but mode " + << mode + << " has a dynamic stride; a shared axis needs static strides to " + "scale its coordinate"; + if (stride.value % dim.stride.value) + return emitError() << "global mode " << mode << " has stride " << stride.value + << ", which is not a multiple of the descriptor dim " << descDim + << " stride " << dim.stride.value + << "; its coordinate cannot be expressed on that axis"; + byMode[mode] = BasisAttr::get(IntAttr::getStatic(ctx, stride.value / dim.stride.value), axis); + } + } + return byMode; +} + +} // namespace + +FailureOr makeTensor2Tdm(const Geometry &geometry, IntTupleAttr gshape, + function_ref emitError) { + MLIRContext *ctx = gshape.getContext(); + FailureOr> byMode = + axisBasisPerMode(ctx, geometry, gshape, /*clampableOnly=*/true, emitError); + if (failed(byMode)) + return failure(); + return unflattenLike(ctx, gshape, *byMode); +} + +void boundaryCheckAxes(IntTupleAttr tensor2tdm, SmallVectorImpl &axes) { + SmallVector leaves; + flattenLeaves(tensor2tdm, leaves); + for (IntTupleAttr leaf : leaves) + axes.push_back(leaf.isLeafBasis() ? leaf.getLeafAsBasis().getModes().front() : -1); +} + +FailureOr coordLayout(const Geometry &geometry, IntTupleAttr gshape, + function_ref emitError) { + MLIRContext *ctx = gshape.getContext(); + FailureOr> byMode = + axisBasisPerMode(ctx, geometry, gshape, /*clampableOnly=*/false, emitError); + if (failed(byMode)) + return failure(); + + if (geometry.ratio == 1) + return LayoutAttr::get(gshape, unflattenLike(ctx, gshape, *byMode)); + + // Under a recast the descriptor counts in wider units than the tensor does, so the + // contiguous mode's coordinate has to be divided by the ratio — and a rational basis + // scale is not something a Fly basis can hold (its coefficient is an + // integer). The shape carries it instead: that mode becomes `(ratio, N / ratio)` with + // strides `(0, 1E)`, so a logical coordinate `n` decomposes into + // `(n % ratio, n / ratio)` and only the second half reaches the axis. Integer division, + // done by the layout algebra, invisible to the kernel — which goes on tiling in tensor + // elements. + if (geometry.contiguousMode < 0) + return emitError() << "a recast needs exactly one contiguous global mode to divide the " + "coordinate on, found none"; + int32_t split = geometry.contiguousMode; + const Scalar &recastExtent = geometry.modeExtent[split]; + Attribute extentAttr = recastExtent.isStatic + ? Attribute(IntAttr::getStatic(ctx, recastExtent.value)) + : Attribute(IntAttr::getDynamic(ctx)); + SmallVector gshapeLeaves; + flattenLeaves(gshape, gshapeLeaves); + SmallVector flatShape; + for (IntTupleAttr leaf : gshapeLeaves) + flatShape.push_back(leaf.getValue()); + + IntTupleAttr shape = + unflattenLike(ctx, gshape, flatShape, split, + makePair(ctx, IntAttr::getStatic(ctx, geometry.ratio), extentAttr)); + IntTupleAttr stride = unflattenLike(ctx, gshape, *byMode, split, + makePair(ctx, IntAttr::getStatic(ctx, 0), (*byMode)[split])); + return LayoutAttr::get(shape, stride); +} + +SmallVector tileShape(const Geometry &geometry) { + SmallVector shape; + for (const Dim &dim : llvm::reverse(geometry.dims)) + shape.push_back(dim.box); + return shape; +} + +//===----------------------------------------------------------------------===// +// tma_partition +//===----------------------------------------------------------------------===// + +namespace { + +/// The product of a fully static int tuple's leaves, or failure if any is not. +FailureOr staticSize(IntTupleAttr t, const Twine &what, + function_ref emitError) { + SmallVector leaves; + flattenLeaves(t, leaves); + int64_t n = 1; + for (IntTupleAttr leaf : leaves) { + if (!leaf.isLeafInt() || !leaf.isStatic()) + return emitError() << what << " must be a static int tuple, got a dynamic or basis leaf"; + n *= leaf.getLeafAsInt().getValue(); + } + return n; +} + +} // namespace + +FailureOr partitionLayout(IntTupleAttr atomValShape, int32_t atomValBits, + LayoutAttr smemLayout, int32_t ldsElemBits, + IntTupleAttr coordShape, int32_t numWarps, + function_ref emitError) { + MLIRContext *ctx = smemLayout.getContext(); + AttrLayoutBuilder layoutBuilder(ctx); + + FailureOr vSize = staticSize(smemLayout.getShape(), "the LDS tile's shape", emitError); + if (failed(vSize)) + return failure(); + FailureOr gSize = staticSize(coordShape, "the coordinate tile's shape", emitError); + if (failed(gSize)) + return failure(); + if (*vSize != *gSize) + return emitError() << "the LDS tile holds " << *vSize << " values but the coordinate tile " + << "holds " << *gSize + << "; the two are cut by the same layout, so they must agree"; + if (numWarps < 1) + return emitError() << "the warp layout must have a positive size, got " << numWarps; + if (*vSize % numWarps) + return emitError() << numWarps << " warps do not divide the tile's " << *vSize + << " values evenly"; + int64_t want = *vSize / numWarps; + + // Tensor elements one atom *call* moves. The atom counts in its own `val_bits`, which a + // recast makes wider than the tile's element, while the tile + // counts in its own. TDM's two sides carry the same value layout -- one instruction + // moves one whole box either way -- so which of src / dst is read does not matter. + FailureOr vals = staticSize(atomValShape, "the atom's value mode", emitError); + if (failed(vals)) + return failure(); + int64_t bits = *vals * atomValBits; + if (ldsElemBits <= 0 || bits % ldsElemBits) + return emitError() << "one call moves " << bits + << " bits, which is not a whole number of the LDS tile's " << ldsElemBits + << "-bit elements"; + int64_t numElems = bits / ldsElemBits; + if (numElems <= 0 || want % numElems) + return emitError() << "each of the " << numWarps << " warps takes " << want + << " values but one call moves " << numElems << ", which does not divide it"; + + // The inverse of the compact tile, composed directly: the pad has already been taken + // back out, so it covers the tile exactly and needs no tiling up to the tile's size. + FailureOr> lds = analyzeLdsTile(smemLayout, emitError); + if (failed(lds)) + return failure(); + LayoutAttr invSmem = + layoutRightInverse(layoutBuilder, LayoutAttr::get(smemLayout.getShape(), std::get<2>(*lds))); + + auto flat = [&](int64_t extent, int64_t stride) { + return LayoutAttr::get(IntTupleAttr::get(IntAttr::getStatic(ctx, extent)), + IntTupleAttr::get(IntAttr::getStatic(ctx, stride))); + }; + LayoutAttr layoutV = layoutComposition(layoutBuilder, invSmem, flat(numElems, 1)); + LayoutAttr layoutIter = + layoutComposition(layoutBuilder, invSmem, flat(want / numElems, numElems)); + + SmallVector shapes{layoutV.getShape()}; + SmallVector strides{layoutV.getStride()}; + if (numWarps > 1) { + // `((ATOM), (WARP), (ITER))` before the warp coordinate is sliced out: the + // warps take equal contiguous chunks of the LDS order, so warp `w` starts at + // `w * want` -- the stride of the WARP mode. + LayoutAttr layoutWarp = layoutComposition(layoutBuilder, invSmem, flat(numWarps, want)); + shapes.push_back(layoutWarp.getShape()); + strides.push_back(layoutWarp.getStride()); + } + shapes.push_back(layoutIter.getShape()); + strides.push_back(layoutIter.getStride()); + return LayoutAttr::get(IntTupleAttr::get(ArrayAttr::get(ctx, shapes)), + IntTupleAttr::get(ArrayAttr::get(ctx, strides))); +} + +} // namespace mlir::fly_rocdl::tdm diff --git a/lib/Dialect/FlyROCDL/CMakeLists.txt b/lib/Dialect/FlyROCDL/CMakeLists.txt index d73bf8776..c9be7c709 100644 --- a/lib/Dialect/FlyROCDL/CMakeLists.txt +++ b/lib/Dialect/FlyROCDL/CMakeLists.txt @@ -7,6 +7,11 @@ add_mlir_dialect_library(MLIRFlyROCDLDialect CDNA3/CopyAtom.cpp CDNA4/MmaAtom.cpp CDNA4/CopyAtom.cpp + CDNA5/CopyAtom.cpp + CDNA5/TdmGeometry.cpp + CDNA5/TdmAtomBuilder.cpp + Transforms/ExpandOps.cpp + Transforms/ClusterAttr.cpp GFX1250/MmaAtom.cpp GFX1250/MmaAtomScale.cpp GFX1250/CopyAtom.cpp @@ -15,4 +20,5 @@ add_mlir_dialect_library(MLIRFlyROCDLDialect DEPENDS MLIRFlyROCDLIncGen + FlyROCDLTransformPassIncGen ) diff --git a/lib/Dialect/FlyROCDL/Ops.cpp b/lib/Dialect/FlyROCDL/Ops.cpp index c33db625d..bb6c13765 100644 --- a/lib/Dialect/FlyROCDL/Ops.cpp +++ b/lib/Dialect/FlyROCDL/Ops.cpp @@ -7,6 +7,7 @@ #include "flydsl/Dialect/Fly/IR/FlyDialect.h" #include "flydsl/Dialect/FlyROCDL/IR/Dialect.h" #include "flydsl/Dialect/FlyROCDL/Utils/BufferFatPtr.h" +#include "flydsl/Dialect/FlyROCDL/Utils/TdmAtomBuilder.h" using namespace mlir; using namespace mlir::fly; @@ -28,3 +29,86 @@ LogicalResult GetBufferRsrcOp::inferReturnTypes(MLIRContext *context, inferredReturnTypes.assign({BufferFatPtr::getRsrcPtrType(context)}); return success(); } + +//===----------------------------------------------------------------------===// +// MakeTiledTdmLoadAtomOp / MakeTiledTdmStoreAtomOp +//===----------------------------------------------------------------------===// + +namespace { + +/// Both builders infer their results the same way: one derivation, and only the atom type +/// it lands on is the direction's. +template +LogicalResult inferTdmAtomTypes(MLIRContext *context, std::optional location, + ValueRange operands, DictionaryAttr attributes, + PropertyRef properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + using Adaptor = typename OpT::Adaptor; + constexpr bool isLoad = std::is_same_v; + + Location loc = location.value_or(UnknownLoc::get(context)); + auto emitError = [&]() -> InFlightDiagnostic { + return mlir::emitError(loc) << OpT::getOperationName() << ": "; + }; + + typename OpT::Properties parsed; + PropertyRef effective = properties; + if (attributes && !attributes.empty()) { + if (failed(OpT::setPropertiesFromAttr(parsed, attributes, emitError))) + return failure(); + effective = PropertyRef(TypeID::get(), &parsed); + } + Adaptor adaptor(operands, attributes, effective, regions); + + tdm::Request request; + Type dataType; + FailureOr geometry = deriveTdmAtom(adaptor, request, dataType, emitError); + if (failed(geometry)) + return failure(); + + FailureOr tensor2tdm = + tdm::makeTensor2Tdm(*geometry, request.gLayout.getShape(), emitError); + if (failed(tensor2tdm)) + return failure(); + + FailureOr copyOp = + isLoad ? tdmLoadOpType(context, *geometry, dataType, *tensor2tdm, adaptor.getAtomicBarrier(), + adaptor.getCacheModifier(), emitError) + : tdmStoreOpType(context, *geometry, dataType, *tensor2tdm, adaptor.getAtomicBarrier(), + adaptor.getCacheModifier(), emitError); + if (failed(copyOp)) + return failure(); + + FailureOr coord = tdm::coordLayout(*geometry, request.gLayout.getShape(), emitError); + if (failed(coord)) + return failure(); + + SmallVector zeros(geometry->rank(), IntTupleAttr::get(IntAttr::getStatic(context, 0))); + IntTupleAttr base = IntTupleAttr::get(ArrayAttr::get(context, zeros)); + + inferredReturnTypes.assign( + {CopyAtomType::get(*copyOp, static_cast(dataType.getIntOrFloatBitWidth())), + CoordTensorType::get(base, *coord)}); + return success(); +} + +} // namespace + +LogicalResult MakeTiledTdmLoadAtomOp::inferReturnTypes(MLIRContext *context, + std::optional location, + ValueRange operands, + DictionaryAttr attributes, + PropertyRef properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + return inferTdmAtomTypes(context, location, operands, attributes, + properties, regions, inferredReturnTypes); +} + +LogicalResult +MakeTiledTdmStoreAtomOp::inferReturnTypes(MLIRContext *context, std::optional location, + ValueRange operands, DictionaryAttr attributes, + PropertyRef properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + return inferTdmAtomTypes(context, location, operands, attributes, + properties, regions, inferredReturnTypes); +} diff --git a/lib/Dialect/FlyROCDL/Transforms/ClusterAttr.cpp b/lib/Dialect/FlyROCDL/Transforms/ClusterAttr.cpp new file mode 100644 index 000000000..4f9e166c9 --- /dev/null +++ b/lib/Dialect/FlyROCDL/Transforms/ClusterAttr.cpp @@ -0,0 +1,61 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors +// +// Inject amdgpu-cluster-dims into llvm.func passthrough. Run inside gpu.module() AFTER +// convert-gpu-to-rocdl. +// +// The upstream ROCDL dialect does not translate `rocdl.cluster_dims` to the LLVM IR +// function attribute `amdgpu-cluster-dims`. This pass bridges the gap by converting the +// discardable attribute that `GPUFuncOpLowering` copied from gpu.func into an LLVM +// passthrough entry that the LLVM IR emitter honours. +// +// It rewrites an attribute on ops the GPU-to-ROCDL conversion already produced, so it is a +// transform on this dialect's own annotation rather than part of any Fly conversion. + +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Pass/Pass.h" + +#include "flydsl/Dialect/FlyROCDL/Transforms/Passes.h" + +using namespace mlir; + +namespace mlir { +namespace fly_rocdl { +#define GEN_PASS_DEF_FLYROCDLCLUSTERATTRPASS +#include "flydsl/Dialect/FlyROCDL/Transforms/Passes.h.inc" +} // namespace fly_rocdl +} // namespace mlir + +namespace { + +class FlyROCDLClusterAttrPass + : public mlir::fly_rocdl::impl::FlyROCDLClusterAttrPassBase { +public: + using mlir::fly_rocdl::impl::FlyROCDLClusterAttrPassBase< + FlyROCDLClusterAttrPass>::FlyROCDLClusterAttrPassBase; + + void runOnOperation() override { + getOperation()->walk([&](LLVM::LLVMFuncOp func) { + auto clusterAttr = func->getAttrOfType("rocdl.cluster_dims"); + if (!clusterAttr) + return; + + MLIRContext *ctx = func.getContext(); + + // Build the new passthrough entry: ["amdgpu-cluster-dims", "2,2,1"]. + auto key = StringAttr::get(ctx, "amdgpu-cluster-dims"); + auto entry = ArrayAttr::get(ctx, {key, clusterAttr}); + + // Append to existing passthrough list (if any). + SmallVector passthroughAttrs; + if (auto existing = func.getPassthroughAttr()) + passthroughAttrs.append(existing.begin(), existing.end()); + passthroughAttrs.push_back(entry); + + func.setPassthroughAttr(ArrayAttr::get(ctx, passthroughAttrs)); + func->removeAttr("rocdl.cluster_dims"); + }); + } +}; + +} // namespace diff --git a/lib/Dialect/FlyROCDL/Transforms/ExpandOps.cpp b/lib/Dialect/FlyROCDL/Transforms/ExpandOps.cpp new file mode 100644 index 000000000..7f3de0740 --- /dev/null +++ b/lib/Dialect/FlyROCDL/Transforms/ExpandOps.cpp @@ -0,0 +1,417 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors +// +// Spend the geometry `fly_rocdl.make_tiled_tdm_{load,store}_atom` derived during type +// inference. +// +// The op's *type* already says everything static about the descriptor — the box, the +// mode map, the coordinate strides — because the derivation ran to produce it. What is +// left is the run-time half: the tensor's dynamic extents and strides, which have to be +// opened out of the tensor operand and handed to `fly.make_copy_atom` as the atom's +// construction arguments. So the derivation runs a second time here rather than +// being cached on the op: a side table could disagree with the type, and re-deriving is +// pure attribute algebra that costs nothing. +// +// This pass is also where an `boundary_check` state stops being written in the tensor's language +// and starts being written in the descriptor's. A caller sets one flag per *tensor mode*, because +// the descriptor's dims are not theirs to know — they permute, and they pack the tail. Translating +// that is `tensor2tdm` and nothing else, so it is done once, here, and the result is a flat +// `boundary_check_axes` tuple with one leaf per descriptor axis that every later pass can read +// without knowing what a mode was. + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" + +#include "flydsl/Dialect/Fly/Utils/IntTupleUtils.h" +#include "flydsl/Dialect/FlyROCDL/Transforms/Passes.h" +#include "flydsl/Dialect/FlyROCDL/Utils/TdmAtomBuilder.h" + +using namespace mlir; +using namespace mlir::fly; +using namespace mlir::fly_rocdl; + +namespace mlir { +namespace fly_rocdl { +#define GEN_PASS_DEF_FLYROCDLEXPANDOPSPASS +#include "flydsl/Dialect/FlyROCDL/Transforms/Passes.h.inc" +} // namespace fly_rocdl +} // namespace mlir + +namespace { + +/// Pulls the dynamic leaves out of the tensor's layout, once per side and on demand. +/// +/// A Fly int tuple keeps its static leaves in its type and only its dynamic ones as SSA +/// values, so a mode's run-time extent is the *n*-th result of `fly.get_leaves`, where +/// *n* counts the dynamic leaves before it. A tensor with no dynamic geometry never +/// reaches this and emits nothing — not even the `fly.get_layout` that opens the tensor +/// up, which is why that one is materialized here rather than by the caller. +class LayoutLeaves { +public: + LayoutLeaves(OpBuilder &builder, Location loc, Value tensor, LayoutAttr attr) + : builder(builder), loc(loc), tensor(tensor), attr(attr) {} + + /// The SSA value of global mode `mode`'s extent (`fromShape`) or stride. + Value get(int32_t mode, bool fromShape) { + Side &side = fromShape ? shapeSide : strideSide; + IntTupleAttr tuple = fromShape ? attr.getShape() : attr.getStride(); + if (!side.materialized) { + if (!layout) + layout = GetLayoutOp::create(builder, loc, tensor); + Value tupleVal = fromShape ? Value(GetShapeOp::create(builder, loc, layout)) + : Value(GetStrideOp::create(builder, loc, layout)); + auto leaves = GetLeavesOp::create(builder, loc, tupleVal, /*dynamicOnly=*/true); + side.values.assign(leaves.getResults().begin(), leaves.getResults().end()); + side.materialized = true; + } + // The n-th dynamic leaf, counting depth-first the way `get_leaves` orders them. + int32_t dynIndex = 0; + int32_t seen = 0; + if (!countTo(tuple, mode, seen, dynIndex)) + return nullptr; + if (dynIndex < 0 || dynIndex >= static_cast(side.values.size())) + return nullptr; + return side.values[dynIndex]; + } + +private: + struct Side { + bool materialized = false; + SmallVector values; + }; + + /// Walk to leaf `target`, counting the dynamic leaves passed on the way. Returns false + /// when the target leaf is itself static (it has no SSA value to find). + static bool countTo(IntTupleAttr t, int32_t target, int32_t &seen, int32_t &dynIndex) { + if (t.isLeaf()) { + if (seen == target) { + if (t.isStatic()) + return false; + return true; + } + if (!t.isStatic()) + ++dynIndex; + ++seen; + return false; + } + for (int32_t i = 0; i < t.rank(); ++i) + if (countTo(t.at(i), target, seen, dynIndex)) + return true; + return false; + } + + OpBuilder &builder; + Location loc; + Value tensor; + LayoutAttr attr; + Value layout; + Side shapeSide, strideSide; +}; + +/// Materialize one descriptor scalar as a value of `width` bits. +Value materializeScalar(OpBuilder &builder, Location loc, const tdm::Scalar &scalar, unsigned width, + LayoutLeaves &leaves) { + Type ty = builder.getIntegerType(width); + if (scalar.isStatic) + return arith::ConstantIntOp::create(builder, loc, ty, scalar.value); + + Value leaf = leaves.get(scalar.mode, scalar.fromShape); + if (!leaf) + return nullptr; + unsigned leafWidth = leaf.getType().getIntOrFloatBitWidth(); + if (leafWidth < width) + leaf = arith::ExtUIOp::create(builder, loc, ty, leaf); + else if (leafWidth > width) + leaf = arith::TruncIOp::create(builder, loc, ty, leaf); + if (scalar.divisor != 1) { + // The recast's division, deferred to run time: the caller asserted the tensor is laid + // out in whole internal units, which a value only known now cannot be checked against. + Value divisor = arith::ConstantIntOp::create(builder, loc, ty, scalar.divisor); + leaf = arith::DivUIOp::create(builder, loc, leaf, divisor); + } + return leaf; +} + +/// An all-static int tuple, as a value. +Value staticTuple(OpBuilder &builder, Location loc, IntTupleAttr attr) { + return MakeIntTupleOp::create(builder, loc, IntTupleType::get(attr), ValueRange{}); +} + +/// The coordinate tensor's shape, as a value. +/// +/// Unlike its base and its stride, this one is not all-static: it is the global tensor's +/// own shape — except under a recast, which splits the contiguous mode into +/// `(ratio, extent / ratio)` — so a tensor with run-time extents leaves dynamic leaves in +/// it, and those have to be resolved out of the layout operand exactly like the +/// descriptor's extents are. An operand-less `fly.make_int_tuple` would type-check and +/// then have nothing behind those leaves for a later pass to read. +Value materializeCoordShape(OpBuilder &builder, Location loc, IntTupleAttr shape, + const tdm::Geometry &geometry, LayoutLeaves &leaves) { + // One scalar per leaf, in the order `tdm::coordLayout` laid them out. + SmallVector scalars; + for (int32_t mode = 0; mode < static_cast(geometry.modeExtent.size()); ++mode) { + if (geometry.ratio != 1 && mode == geometry.contiguousMode) + scalars.push_back(tdm::Scalar::getStatic(geometry.ratio)); + scalars.push_back(geometry.modeExtent[mode]); + } + + SmallVector flat; + intTupleFlattenToVector(IntTupleBuilder(shape.getContext()), shape, flat); + if (flat.size() != scalars.size()) + return nullptr; + + SmallVector dynamic; + for (auto [leaf, scalar] : llvm::zip_equal(flat, scalars)) { + if (leaf.isStatic()) + continue; + Value v = materializeScalar(builder, loc, scalar, leaf.extractIntFromLeaf().getWidth(), leaves); + if (!v) + return nullptr; + dynamic.push_back(v); + } + return MakeIntTupleOp::create(builder, loc, IntTupleType::get(shape), dynamic); +} + +/// Rewrite one builder into the ops that build the atom. +/// +/// Driven by a plain walk rather than the greedy driver: the expansion never produces +/// another builder, so there is nothing to iterate to a fixed point, and the greedy +/// driver would additionally delete every other dead `Pure` op in the function — which +/// is not this pass's business, and would erase exactly the IR a caller wants to inspect. +template LogicalResult expandOne(OpT op, IRRewriter &rewriter) { + Location loc = op.getLoc(); + rewriter.setInsertionPoint(op); + auto emitError = [&]() -> InFlightDiagnostic { return op.emitOpError(); }; + + tdm::Request request; + Type dataType; + FailureOr geometry = + deriveTdmAtom(typename OpT::Adaptor(op), request, dataType, emitError); + if (failed(geometry)) + return failure(); + + auto atomTy = dyn_cast(op.getAtom().getType()); + auto coordTy = dyn_cast(op.getCoordTensor().getType()); + if (!atomTy || !coordTy) + return failure(); + + int32_t descRank = geometry->rank(); + LayoutLeaves leaves(rewriter, loc, op.getTensor(), request.gLayout); + + // Construction arguments, in tensor dim order (the reverse of descriptor order, which + // is what the atom's `tileShape` and its state slots use): base pointer, then the + // strides of dims 0..rank-2 (the innermost is 1 by construction and is not passed), + // then every dim's extent. Every dim passes its extent whether or not it currently + // clamps: `boundary_check` is per-call state, so a later call site may switch its dim on and + // needs the bound already in the atom, and an extent no call reads dies with the rest + // of the scalarized state struct. + SmallVector args; + args.push_back(GetIterOp::create(rewriter, loc, op.getTensor())); + for (int32_t i = 0; i < descRank - 1; ++i) { + const tdm::Dim &dim = geometry->dims[descRank - 1 - i]; + Value v = materializeScalar(rewriter, loc, dim.stride, 64, leaves); + if (!v) + return op.emitOpError() << "could not materialize the stride of descriptor dim " + << (descRank - 1 - i); + args.push_back(v); + } + for (int32_t i = 0; i < descRank; ++i) { + const tdm::Dim &dim = geometry->dims[descRank - 1 - i]; + Value v = materializeScalar(rewriter, loc, dim.tensorDim, 32, leaves); + if (!v) + return op.emitOpError() << "could not materialize the extent of descriptor dim " + << (descRank - 1 - i); + args.push_back(v); + } + if (geometry->iterCount > 1) { + Value v = materializeScalar(rewriter, loc, geometry->iterStride, 64, leaves); + if (!v) + return op.emitOpError() << "could not materialize the iteration stride"; + args.push_back(v); + } + + Value atom = MakeCopyAtomOp::create(rewriter, loc, atomTy, args, + static_cast(dataType.getIntOrFloatBitWidth())); + + // The initial `boundary_check` state, shaped like the global tensor: the builder's one + // flag written onto every mode that has a bound to give. A mode with nothing to clamp — + // size-1, stride-0, not spanned by the box, or sharing the rank-5 packing's dim — comes + // out off, so the flag cannot fail here; naming such a mode is only an error when a call + // site names it, which is the `set_value` below. + FailureOr boundaryCheck = tdm::initialBoundaryCheck( + *geometry, request.gLayout.getShape(), request.initBoundaryCheck, emitError); + if (failed(boundaryCheck)) + return failure(); + atom = AtomSetValueOp::create(rewriter, loc, atom.getType(), atom, + rewriter.getStringAttr("boundary_check"), + staticTuple(rewriter, loc, *boundaryCheck)); + + // The coordinate tensor, at the tensor origin: tiling and slicing it is what moves it, + // and that folds into its type. + Value origin = staticTuple(rewriter, loc, coordTy.getBase()); + auto coordLayoutAttr = cast(coordTy.getLayout()); + Value shape = materializeCoordShape(rewriter, loc, coordLayoutAttr.getShape(), *geometry, leaves); + if (!shape) + return op.emitOpError() << "could not materialize the coordinate tensor's shape"; + Value stride = staticTuple(rewriter, loc, coordLayoutAttr.getStride()); + Value layout = + MakeLayoutOp::create(rewriter, loc, LayoutType::get(coordLayoutAttr), shape, stride); + Value coord = MakeViewOp::create(rewriter, loc, coordTy, origin, layout); + + rewriter.replaceOp(op, {atom, coord}); + return success(); +} + +//===----------------------------------------------------------------------===// +// "boundary_check" (tensor modes) -> "boundary_check_axes" (descriptor axes) +//===----------------------------------------------------------------------===// + +/// The CDNA5 TDM atom a `fly.atom.set_value` is setting, or a null pair if it is setting +/// something else's state. +std::pair tdmAtomShape(Value atom) { + auto atomTy = dyn_cast(atom.getType()); + if (!atomTy) + return {}; + if (auto load = dyn_cast(atomTy.getCopyOp())) + return {load.getTensor2tdm(), static_cast(load.getTileShape().size())}; + if (auto store = dyn_cast(atomTy.getCopyOp())) + return {store.getTensor2tdm(), static_cast(store.getTileShape().size())}; + return {}; +} + +/// Translate one tensor-order `boundary_check` tuple into the flat axis-order tuple the lowering +/// reads, or fail with the diagnostic the caller earned. +/// +/// Several modes can land on one axis — a merge or the rank-5 packing puts them there — +/// and the axis clamps if any of them asked for a bound, so this is an OR. A static flag +/// wins it outright, which is why the static leaves are swept first: an axis already +/// pinned on spends no `cmpi` on the dynamic leaves that would only OR into it. Both +/// sweeps still walk every dynamic leaf, because their SSA values are positional in the +/// tuple's `make_int_tuple` and skipping one would shift every later mode's. +FailureOr normalizeBoundaryCheckToAxes(OpBuilder &builder, Location loc, Value flags, + IntTupleAttr tensor2tdm, int32_t descRank, + function_ref emitError) { + auto tupleTy = dyn_cast(flags.getType()); + if (!tupleTy) + return emitError() + << "\"boundary_check\" must be an int_tuple shaped like the global tensor, got " + << flags.getType(); + if (!intTupleIsCongruent(tensor2tdm, tupleTy.getAttr())) + return emitError() << "\"boundary_check\" is " << tupleTy.getAttr() + << ", which is not congruent with the global tensor's " << tensor2tdm; + + MLIRContext *ctx = builder.getContext(); + IntTupleBuilder tupleBuilder(ctx); + SmallVector leaves; + intTupleFlattenToVector(tupleBuilder, tupleTy.getAttr(), leaves); + SmallVector axes; + tdm::boundaryCheckAxes(tensor2tdm, axes); + + auto tupleOp = flags.getDefiningOp(); + OperandRange dyn = tupleOp ? tupleOp.getDyncElems() : OperandRange(nullptr, 0); + + SmallVector pinnedOn(descRank, false); + for (auto [mode, leaf] : llvm::enumerate(leaves)) { + IntAttr value = leaf.extractIntFromLeaf(); + if (!value.isStatic() || value.getValue() == 0) + continue; + if (axes[mode] < 0) + return emitError() << "\"boundary_check\" asks to clamp tensor mode " << mode + << ", which has no bound of its own — it is size-1 or stride-0, or it " + "shares a descriptor dim whose single bound would not say what a " + "per-mode one says"; + pinnedOn[axes[mode]] = true; + } + + SmallVector merged(descRank); + auto dynIt = dyn.begin(); + for (auto [mode, leaf] : llvm::enumerate(leaves)) { + if (leaf.extractIntFromLeaf().isStatic()) + continue; + if (!tupleOp || dynIt == dyn.end()) + return emitError() << "\"boundary_check\" leaf " << mode + << " is dynamic but the tuple is not normal " << "form"; + Value v = *dynIt++; + int32_t axis = axes[mode]; + if (axis < 0) + return emitError() << "\"boundary_check\" leaf " << mode + << " is dynamic but that tensor mode has no bound to switch"; + if (pinnedOn[axis]) + continue; + Value zero = arith::ConstantIntOp::create(builder, loc, v.getType(), 0); + Value flag = arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::ne, v, zero); + merged[axis] = merged[axis] ? arith::OrIOp::create(builder, loc, merged[axis], flag) : flag; + } + + // The tuple's dynamic leaves are integers, so the merged i1 goes back out as one. The + // lowering compares it against zero again, and the pair folds. + Type i32 = builder.getI32Type(); + SmallVector axisLeaves; + SmallVector axisDyn; + for (int32_t axis = 0; axis < descRank; ++axis) { + if (pinnedOn[axis] || !merged[axis]) { + axisLeaves.push_back(IntTupleAttr::get(IntAttr::getStatic(ctx, pinnedOn[axis] ? 1 : 0))); + continue; + } + axisLeaves.push_back(IntTupleAttr::get(IntAttr::getDynamic(ctx, /*width=*/32))); + axisDyn.push_back(arith::ExtUIOp::create(builder, loc, i32, merged[axis])); + } + IntTupleAttr result = IntTupleAttr::get(ArrayAttr::get(ctx, axisLeaves)); + return Value(MakeIntTupleOp::create(builder, loc, IntTupleType::get(result), axisDyn)); +} + +/// Rewrite one `set_value("boundary_check", ...)` on a CDNA5 TDM atom into +/// `set_value("boundary_check_axes", ...)`. +LogicalResult normalizeOneBoundaryCheck(AtomSetValueOp op, IRRewriter &rewriter) { + auto [tensor2tdm, descRank] = tdmAtomShape(op.getAtom()); + if (!tensor2tdm) + return success(); // not ours + rewriter.setInsertionPoint(op); + auto emitError = [&]() -> InFlightDiagnostic { return op.emitOpError(); }; + FailureOr axisFlags = normalizeBoundaryCheckToAxes(rewriter, op.getLoc(), op.getValue(), + tensor2tdm, descRank, emitError); + if (failed(axisFlags)) + return failure(); + rewriter.replaceOpWithNewOp(op, op.getResult().getType(), op.getAtom(), + rewriter.getStringAttr("boundary_check_axes"), + *axisFlags); + return success(); +} + +class FlyROCDLExpandOpsPass + : public mlir::fly_rocdl::impl::FlyROCDLExpandOpsPassBase { +public: + using mlir::fly_rocdl::impl::FlyROCDLExpandOpsPassBase< + FlyROCDLExpandOpsPass>::FlyROCDLExpandOpsPassBase; + + void runOnOperation() override { + IRRewriter rewriter(&getContext()); + SmallVector builders; + getOperation()->walk([&](Operation *op) { + if (isa(op)) + builders.push_back(op); + }); + for (Operation *op : builders) { + LogicalResult expanded = isa(op) + ? expandOne(cast(op), rewriter) + : expandOne(cast(op), rewriter); + if (failed(expanded)) + return signalPassFailure(); + } + + // Second, in its own walk rather than inside the expansion, because the builder's own + // initial `boundary_check` is one of these and translating it twice — once there, once here — + // is how the two would drift apart. + SmallVector boundaryCheckSetters; + getOperation()->walk([&](AtomSetValueOp op) { + if (op.getField() == "boundary_check") + boundaryCheckSetters.push_back(op); + }); + for (AtomSetValueOp op : boundaryCheckSetters) + if (failed(normalizeOneBoundaryCheck(op, rewriter))) + return signalPassFailure(); + } +}; + +} // namespace diff --git a/python/flydsl/compiler/backends/rocm.py b/python/flydsl/compiler/backends/rocm.py index 45f891f0f..13390ca35 100644 --- a/python/flydsl/compiler/backends/rocm.py +++ b/python/flydsl/compiler/backends/rocm.py @@ -80,6 +80,7 @@ def _pipeline_parts(self, *, compile_hints: dict) -> Tuple[List[str], str]: pre_binary_fragments = [ "fly-rewrite-func-signature", "fly-canonicalize", + "fly-rocdl-expand-ops", "fly-layout-lowering", "fly-int-swizzle-simplify", "canonicalize", diff --git a/python/flydsl/expr/primitive.py b/python/flydsl/expr/primitive.py index 884ee8ccf..acb9dee2f 100644 --- a/python/flydsl/expr/primitive.py +++ b/python/flydsl/expr/primitive.py @@ -1035,7 +1035,7 @@ def make_copy_atom(copy_op_type, elem_type): else: raise TypeError(f"make_copy_atom: elem_type must be NumericType, ir.Type, or int, got {type(elem_type)}") copy_atom_ty = CopyAtomType.get(copy_op=copy_op_type, val_bits=val_bits) - return fly.make_copy_atom(copy_atom_ty, val_bits=val_bits) + return fly.make_copy_atom(copy_atom_ty, [], val_bits=val_bits) @dsl_loc_tracing @@ -1107,7 +1107,7 @@ def mma_make_fragment(operand_id, tiled_mma, input, *, stages=None): @dsl_loc_tracing def copy(copy_atom, src, dst, *, pred=None, **kwargs): - return fly.copy(copy_atom.set_value(kwargs), src, dst, pred=pred) + return fly.copy(copy_atom.set_value(kwargs) if kwargs else copy_atom, src, dst, pred=pred) @dsl_loc_tracing diff --git a/python/flydsl/expr/rocdl/cdna5.py b/python/flydsl/expr/rocdl/cdna5.py index fa082c766..a9b17efda 100644 --- a/python/flydsl/expr/rocdl/cdna5.py +++ b/python/flydsl/expr/rocdl/cdna5.py @@ -1,17 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025 FlyDSL Project Contributors -"""gfx1250-specific ROCDL atom builders (MX-scaled WMMA + N-D TDM copy).""" +"""CDNA5 / gfx1250 ROCDL atom builders.""" from ..._mlir import ir from ..._mlir._mlir_libs._mlirDialectsFlyROCDL import MmaOpGFX1250_WMMAScaleType -from ..._mlir.dialects.fly_rocdl import CopyOpGFX1250TDMType +from ..._mlir.dialects import fly_rocdl from ..typing import Int32, Int64, Tensor __all__ = [ "WMMAScale", + "TensorLoad", + "TensorStore", "TDM", "make_tdm_atom", + "make_tiled_tdm_atom", + "tdm_partition", ] @@ -67,6 +71,36 @@ def WMMAScale( ) +class TensorLoad: + """CDNA5 TDM Global -> LDS DMA (``TENSOR_LOAD_TO_LDS``). + + Current atom state: + - `workgroup_mask` (i32): the workgroup mask. + - `early_timeout` (i32): the early timeout mask. + - `atomic_barrier_addr` (shared ptr): *which* LDS barrier this copy arrives on. + *Whether* it arrives on one is the atom's type, so only ``atomic_barrier=True`` + has the field. + - `boundary_check` (int_tuple): per mode boundary check, congruent with the global tensor. + """ + + def __init__(self, cache_modifier=0): + self.cache_modifier = cache_modifier + + +class TensorStore: + """CDNA5 TDM LDS -> Global DMA (``TENSOR_STORE_FROM_LDS``). + + Current atom state: + - `atomic_barrier_addr` (shared ptr): *which* LDS barrier this copy arrives on. + *Whether* it arrives on one is the atom's type, so only ``atomic_barrier=True`` + has the field. + - `boundary_check` (int_tuple): per mode boundary check, congruent with the global tensor. + """ + + def __init__(self, cache_modifier=0): + self.cache_modifier = cache_modifier + + def TDM( rank, num_warps, @@ -91,7 +125,7 @@ def TDM( MCAST ``workgroup_mask`` are runtime atom state set via ``fx.atom.set_value``. :func:`make_tdm_atom` builds the atom and populates the descriptor from a tensor. """ - return CopyOpGFX1250TDMType.get( + return fly_rocdl.CopyOpGFX1250TDMType.get( rank, num_warps, pad_interval, @@ -148,7 +182,7 @@ def make_tdm_atom( if len(strides) != rank: raise ValueError(f"make_tdm_atom: expected {rank} strides, got {len(strides)}") - copy_op = CopyOpGFX1250TDMType.get( + copy_op = fly_rocdl.CopyOpGFX1250TDMType.get( rank, num_warps, pad_interval, @@ -173,3 +207,131 @@ def make_tdm_atom( ) atom = atom_set_value(atom, f"stride_{i}", st) return atom + + +def make_tiled_tdm_atom( + op, + tensor: Tensor, + smem_layout, + tdm_tile, + num_warps=1, + *, + init_boundary_check=True, + atomic_barrier=False, + internal_type=None, +): + """Build a wave-scoped CDNA5 TDM copy atom and its coordinate tensor. + + * ``op`` — a ``TensorLoad(...)`` or ``TensorStore(...)`` instance. + * ``tensor`` — the global tensor. + * ``smem_layout`` — the LDS tile layout. + * ``tdm_tile`` — the tiler: how many elements to take from each global mode. + * ``num_warps`` — how many warps of the workgroup split this tile. The same + number must be handed to :func:`tdm_partition` as the size of its warp + layout. + * ``init_boundary_check`` — The *initial* ``boundary_check`` state. + * ``atomic_barrier`` — whether this atom arrives on the atomic barrier when finished. + * ``internal_type`` — the unit the *descriptor* counts in, which may be + wider than the tensor's element (its width must be a multiple). It + is what lets a sub-byte element ride on a ``data_size`` the hardware can + encode. + + Example: + Loading a 128x64 tile of a row-major ``gA`` into LDS. + + sA_layout = fx.make_layout((128, 64), (64, 1)) + atom, mA = make_tiled_tdm_atom(TensorLoad(), gA, sA_layout, (128, 64)) + + mA = fx.zipped_divide(mA, (128, 64))[None, (bid_x, bid_y)] + sA = fx.Tensor(fx.make_view(smem_ptr, sA_layout)) + + tAsA, tAgA = tdm_partition(atom, warp_coord, warp_layout, sA, mA) + fx.copy(atom, tAgA, tAsA) + + Choosing ``sA_layout``: + The layouts below all hold that same 128x64 tile and differ only in how + it sits in LDS. + + # Plain row-major. No skip, so the atom carries no padding fields. + fx.make_layout((128, 64), (64, 1)) + + # 8 elements of slack after every 64-element row -- the usual bank-conflict + # dodge. The atom picks it up as `padInterval = 64, padAmount = 8`. + fx.make_layout((128, 64), (72, 1)) + + # The same addresses with M split 8x16. + fx.make_layout(((8, 16), 64), ((72, 576), 1)) + + # This pads once every 8 rows (`padInterval = 512, padAmount = 64`) + # instead of once every row. + fx.make_layout(((8, 16), 64), ((64, 576), 1)) + + An LDS tile may also be column-major, but that is a property it has to share + with the tensor: the innermost descriptor dim is the one TDM reads + contiguously from global memory, so a column-major tile wants a column-major + ``gA`` and is refused over the row-major one above. + """ + from ..primitive import make_tile + + if not isinstance(op, (TensorLoad, TensorStore)): + raise TypeError( + f"make_tiled_tdm_atom: first argument must be a TensorLoad() or " f"TensorStore() instance, got {op!r}" + ) + + smem_layout = smem_layout.layout if isinstance(smem_layout, Tensor) else smem_layout + # An `!fly.tile` operand, like `smem_layout`: it is entirely static, so it lives in the + # value's type and the derivation reads it there. + tiler = tdm_tile if isinstance(tdm_tile, ir.Value) else make_tile(*tdm_tile) + + common = dict( + init_boundary_check=ir.BoolAttr.get(init_boundary_check), + num_warps=num_warps, + cache_modifier=op.cache_modifier, + atomic_barrier=bool(atomic_barrier), + internal_type=( + None + if internal_type is None + else (internal_type.ir_type if hasattr(internal_type, "ir_type") else internal_type) + ), + ) + if isinstance(op, TensorLoad): + atom, tdm_tensor = fly_rocdl.make_tiled_tdm_load_atom(tensor, smem_layout, tiler, **common) + else: + atom, tdm_tensor = fly_rocdl.make_tiled_tdm_store_atom(tensor, smem_layout, tiler, **common) + return atom, tdm_tensor + + +def tdm_partition( + atom, + warp_coord, + warp_layout, + stensor, + gtensor, +): + """Cut an LDS tile and a coordinate tile into the calls the atom makes. + + Both tiles come out shaped ``((ATOM), (ITER))`` -- mode 0 is one call's worth of + values and mode 1 counts the calls. + + ``warp_coord`` / ``warp_layout`` say how the warps split the tile: each + issues one instruction over its own share, and the assembled tile belongs to + the whole workgroup. Pass ``0`` and ``make_layout(1)`` when a single warp + does the copy. There is no thread index and no per-lane slice, so within a + warp every lane sees the same partition. + """ + from ..primitive import composition, crd2idx, size + from ..typing import static + + n_warps = size(warp_layout).unpack() + + layout_V = static(fly_rocdl.tdm_partition_layout(atom.type, stensor.type, gtensor.type, n_warps)) + if n_warps == 1: + return composition(stensor, layout_V), composition(gtensor, layout_V) + + # The multicast coordinate is sliced out of the middle mode: the warps take equal + # contiguous chunks of the LDS order, and this one is `warp_coord`'s. + warp_id = crd2idx(warp_coord, warp_layout) + return ( + composition(stensor, layout_V)[None, warp_id, None], + composition(gtensor, layout_V)[None, warp_id, None], + ) diff --git a/tests/mlir/Conversion/tdm_cdna5.mlir b/tests/mlir/Conversion/tdm_cdna5.mlir new file mode 100644 index 000000000..65decd745 --- /dev/null +++ b/tests/mlir/Conversion/tdm_cdna5.mlir @@ -0,0 +1,315 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors +// RUN: %fly-opt %s --split-input-file --fly-canonicalize --fly-rocdl-expand-ops --fly-layout-lowering --canonicalize --fly-convert-atom-call-to-ssa-form --convert-fly-to-rocdl --canonicalize | FileCheck %s +// +// `--fly-convert-atom-call-to-ssa-form` is in the pipeline only to keep it +// honest: it walks every copy_atom_call to decide register promotion, so a +// coordinate-tensor operand must survive it. It has nothing to promote here. +// +// `--fly-rocdl-expand-ops` is here for the `boundary_check` state: a caller writes one flag per +// *tensor mode* and that pass rewrites it through `tensor2tdm` into the flat +// `boundary_check_axes` tuple the lowering reads. These tensors map one mode to one axis, so the +// translation is the identity and what is checked below is unchanged by it -- but the +// lowering now rejects an untranslated `boundary_check`, so the pass has to be in the pipeline. +// +// The trailing `--canonicalize` is load-bearing rather than cosmetic: per-dim +// clamping is atom *state*, so the lowering always emits the select between the +// clamped bound and the untouched extent, and it is the folder walking +// extractvalue back through insertvalue that turns a constant `boundary_check` leaf into no +// arithmetic at all. Checking that here is checking the claim the design rests on. + +// CDNA5 TDM atom: a whole-tile DMA addressed by a coordinate tensor. +// +// The atom bakes the tensor (base pointer, per-dim stride, per-dim extent) as +// construction arguments of `fly.make_copy_atom`; the tile's position arrives as the +// `!fly.coord_tensor` operand, whose runtime value *is* that coordinate -- exactly as a +// memref operand's is its pointer. Which dims clamp is the `boundary_check` state -- one int_tuple +// leaf per descriptor dim, all clamping by default. +// load struct: {mask, early_timeout, atomic_barrier_addr (shared ptr), base (ptr), +// stride_0..3 (i64), extent_0..4 (i32), boundary_check_0..4 (i1), +// iter_stride (i64)} +// store struct: the same minus the two MCAST fields. + +// ----- + +// CHECK-LABEL: @test_cdna5_type +// CHECK-SAME: (%{{.*}}: !llvm.struct<(i32, i32, ptr<3>, ptr<1>, i64, i64, i64, i64, i32, i32, i32, i32, i32, i1, i1, i1, i1, i1, i64)>) +func.func @test_cdna5_type( + %atom: !fly.copy_atom, 0>) { + return +} + +// ----- + +// The store type has no MCAST slots, so its state struct is two i32 fields shorter. + +// CHECK-LABEL: @test_cdna5_store_type +// CHECK-SAME: (%{{.*}}: !llvm.struct<(ptr<3>, ptr<1>, i64, i64, i64, i64, i32, i32, i32, i32, i32, i1, i1, i1, i1, i1, i64)>) +func.func @test_cdna5_store_type( + %atom: !fly.copy_atom, 0>) { + return +} + +// ----- + +// A static tile coordinate. `local_tile` has already folded it into the coord +// tensor's type — origin (384, 128) — so the copy site carries no arithmetic at all; +// layout lowering turns that origin into coord_0 / coord_1 and the descriptor math +// falls out as constants folded against the runtime strides. +// +// Both dims clamp (the default `boundary_check`), so the same coordinate that advances the +// address also shrinks the window: tensor_dim_i = max(extent_i - coord_i, 0). + +// CHECK-LABEL: @test_cdna5_load_static_coord +func.func @test_cdna5_load_static_coord( + %base: !fly.ptr, %s0: i64, %e0: i32, %e1: i32, + %lds: !fly.memref) { + %atom = fly.make_copy_atom(%base, %s0, %e0, %e1 : !fly.ptr, i64, i32, i32) {valBits = 16 : i32} : !fly.copy_atom, 16> + %org = fly.make_coord() : () -> !fly.int_tuple<(384,128)> + %shp = fly.make_int_tuple() : () -> !fly.int_tuple<(128,64)> + %str = fly.make_int_tuple() : () -> !fly.int_tuple<(1E0,1E1)> + %lay = fly.make_layout(%shp, %str) : (!fly.int_tuple<(128,64)>, !fly.int_tuple<(1E0,1E1)>) -> !fly.layout<(128,64):(1E0,1E1)> + %gt = fly.make_view(%org, %lay) : (!fly.int_tuple<(384,128)>, !fly.layout<(128,64):(1E0,1E1)>) -> !fly.coord_tensor<(384,128), (128,64):(1E0,1E1)> + // A fully static origin lives in the operand's *type*, so it materializes as two + // constants and nothing was computed to get them. + // CHECK-DAG: %[[C0:.*]] = arith.constant 384 : i32 + // CHECK-DAG: %[[C1:.*]] = arith.constant 128 : i32 + // The address half: sum_i coord_i * stride_i, scaled to bytes, added to the baked + // base pointer (%arg0). + // CHECK-DAG: %[[BI:.*]] = llvm.ptrtoint %arg0 : !llvm.ptr<1> to i64 + // CHECK-DAG: arith.addi %[[BI]] + // The bounds half: the same two coordinates shrink the in-bounds window. Both + // `boundary_check` leaves are the constructed default, so the outer state selects fold away. + // `maxsi` against zero, not the equivalent compare-and-select: AMDGPU has the max on + // the SALU but only a VALU saturating subtract, and a descriptor field is uniform. + // CHECK-DAG: %[[R0:.*]] = arith.subi %arg2, %[[C0]] : i32 + // CHECK-DAG: arith.maxsi %[[R0]] + // CHECK-DAG: %[[R1:.*]] = arith.subi %arg3, %[[C1]] : i32 + // CHECK-DAG: arith.maxsi %[[R1]] + // CHECK: rocdl.tensor.load.to.lds %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, 0 : vector<4xi32>, vector<8xi32> + fly.copy_atom_call(%atom, %gt, %lds) : (!fly.copy_atom, 16>, !fly.coord_tensor<(384,128), (128,64):(1E0,1E1)>, !fly.memref) -> () + return +} + +// ----- + +// No dim clamps: the tile is guaranteed in range by the tiling, so every `boundary_check` +// leaf is zero and `tensor_dim` is the extent itself, passed through untouched. Not a +// saturated sentinel: the caller has asserted `coord_i + tile_i <= extent_i`, so the +// un-shifted extent is already a bound the tile cannot reach, and it is an SGPR that is +// live anyway. So the dim costs no subtract, no clamp, and no constant to materialize -- +// checking for the *absence* of that arithmetic is the only way to keep it honest. +// Descriptor dim 0 is the tensor's innermost mode, hence %arg3 before %arg2. + +// CHECK-LABEL: @test_cdna5_load_no_boundary_check +// CHECK-NOT: arith.maxsi +// CHECK-NOT: arith.subi +// CHECK-NOT: arith.select +// CHECK-DAG: arith.andi %arg3, %{{.*}} : i32 +// CHECK-DAG: arith.andi %arg2, %{{.*}} : i32 +// CHECK: rocdl.tensor.load.to.lds +func.func @test_cdna5_load_no_boundary_check( + %base: !fly.ptr, %s0: i64, %e0: i32, %e1: i32, + %lds: !fly.memref) { + %off = fly.make_int_tuple() : () -> !fly.int_tuple<(0,0)> + %a0 = fly.make_copy_atom(%base, %s0, %e0, %e1 : !fly.ptr, i64, i32, i32) {valBits = 16 : i32} : !fly.copy_atom, 16> + %atom = fly.atom.set_value(%a0, "boundary_check", %off) : (!fly.copy_atom, 16>, !fly.int_tuple<(0,0)>) -> !fly.copy_atom, 16> + %org = fly.make_coord() : () -> !fly.int_tuple<(384,128)> + %shp = fly.make_int_tuple() : () -> !fly.int_tuple<(128,64)> + %str = fly.make_int_tuple() : () -> !fly.int_tuple<(1E0,1E1)> + %lay = fly.make_layout(%shp, %str) : (!fly.int_tuple<(128,64)>, !fly.int_tuple<(1E0,1E1)>) -> !fly.layout<(128,64):(1E0,1E1)> + %gt = fly.make_view(%org, %lay) : (!fly.int_tuple<(384,128)>, !fly.layout<(128,64):(1E0,1E1)>) -> !fly.coord_tensor<(384,128), (128,64):(1E0,1E1)> + fly.copy_atom_call(%atom, %gt, %lds) : (!fly.copy_atom, 16>, !fly.coord_tensor<(384,128), (128,64):(1E0,1E1)>, !fly.memref) -> () + return +} + +// ----- + +// Mixed: dim 0 clamps, dim 1 does not. The clamp is priced per dim -- a subtract-and-clamp +// on the extent at *every* call -- so it is bought per dim. Dim 1 passes its extent +// through unmodified, which is why exactly one subtract-and-clamp reaches the +// instruction while both dims still read an extent. Because this is state and not a type +// parameter, the *same* atom type serves both this function and the two above. The +// tuple names every dim, including the one it leaves clamping: that is the price of +// checking the rank rather than accepting a leaf at a time. + +// CHECK-LABEL: @test_cdna5_load_mixed_boundary_check +// CHECK-COUNT-1: arith.maxsi +// CHECK-NOT: arith.maxsi +// CHECK-NOT: arith.select +// CHECK: rocdl.tensor.load.to.lds +func.func @test_cdna5_load_mixed_boundary_check( + %base: !fly.ptr, %s0: i64, %e0: i32, %e1: i32, + %lds: !fly.memref) { + %off = fly.make_int_tuple() : () -> !fly.int_tuple<(1,0)> + %a0 = fly.make_copy_atom(%base, %s0, %e0, %e1 : !fly.ptr, i64, i32, i32) {valBits = 16 : i32} : !fly.copy_atom, 16> + %atom = fly.atom.set_value(%a0, "boundary_check", %off) : (!fly.copy_atom, 16>, !fly.int_tuple<(1,0)>) -> !fly.copy_atom, 16> + %org = fly.make_coord() : () -> !fly.int_tuple<(384,128)> + %shp = fly.make_int_tuple() : () -> !fly.int_tuple<(128,64)> + %str = fly.make_int_tuple() : () -> !fly.int_tuple<(1E0,1E1)> + %lay = fly.make_layout(%shp, %str) : (!fly.int_tuple<(128,64)>, !fly.int_tuple<(1E0,1E1)>) -> !fly.layout<(128,64):(1E0,1E1)> + %gt = fly.make_view(%org, %lay) : (!fly.int_tuple<(384,128)>, !fly.layout<(128,64):(1E0,1E1)>) -> !fly.coord_tensor<(384,128), (128,64):(1E0,1E1)> + fly.copy_atom_call(%atom, %gt, %lds) : (!fly.copy_atom, 16>, !fly.coord_tensor<(384,128), (128,64):(1E0,1E1)>, !fly.memref) -> () + return +} + +// ----- + +// A genuinely dynamic `boundary_check` leaf is the case that does not fold: the select survives, +// and so does the arithmetic on both of its sides. This is the honest upper bound on +// what the knob costs -- one v_cndmask -- and it is only paid by a caller that actually +// varies the flag at runtime. An int_tuple's dynamic leaves are i32, so reaching the i1 +// the slot holds also costs the compare; a static leaf pays neither. + +// CHECK-LABEL: @test_cdna5_load_dynamic_boundary_check +// CHECK-DAG: arith.cmpi ne +// CHECK-DAG: arith.select +// CHECK: rocdl.tensor.load.to.lds +func.func @test_cdna5_load_dynamic_boundary_check( + %base: !fly.ptr, %s0: i64, %e0: i32, %e1: i32, %flag: i32, + %lds: !fly.memref) { + %off = fly.make_int_tuple(%flag) : (i32) -> !fly.int_tuple<(1,?)> + %a0 = fly.make_copy_atom(%base, %s0, %e0, %e1 : !fly.ptr, i64, i32, i32) {valBits = 16 : i32} : !fly.copy_atom, 16> + %atom = fly.atom.set_value(%a0, "boundary_check", %off) : (!fly.copy_atom, 16>, !fly.int_tuple<(1,?)>) -> !fly.copy_atom, 16> + %org = fly.make_coord() : () -> !fly.int_tuple<(384,128)> + %shp = fly.make_int_tuple() : () -> !fly.int_tuple<(128,64)> + %str = fly.make_int_tuple() : () -> !fly.int_tuple<(1E0,1E1)> + %lay = fly.make_layout(%shp, %str) : (!fly.int_tuple<(128,64)>, !fly.int_tuple<(1E0,1E1)>) -> !fly.layout<(128,64):(1E0,1E1)> + %gt = fly.make_view(%org, %lay) : (!fly.int_tuple<(384,128)>, !fly.layout<(128,64):(1E0,1E1)>) -> !fly.coord_tensor<(384,128), (128,64):(1E0,1E1)> + fly.copy_atom_call(%atom, %gt, %lds) : (!fly.copy_atom, 16>, !fly.coord_tensor<(384,128), (128,64):(1E0,1E1)>, !fly.memref) -> () + return +} + +// ----- + +// A dynamic tile coordinate (a block index) is the one leaf that survives into the +// IR; everything else about the position stayed in the type. + +// CHECK-LABEL: @test_cdna5_load_dynamic_coord +func.func @test_cdna5_load_dynamic_coord( + %base: !fly.ptr, %s0: i64, %e0: i32, %e1: i32, %m: i32, + %lds: !fly.memref) { + %atom = fly.make_copy_atom(%base, %s0, %e0, %e1 : !fly.ptr, i64, i32, i32) {valBits = 16 : i32} : !fly.copy_atom, 16> + %org = fly.make_coord(%m) : (i32) -> !fly.int_tuple<(?,0)> + %shp = fly.make_int_tuple() : () -> !fly.int_tuple<(128,64)> + %str = fly.make_int_tuple() : () -> !fly.int_tuple<(1E0,1E1)> + %lay = fly.make_layout(%shp, %str) : (!fly.int_tuple<(128,64)>, !fly.int_tuple<(1E0,1E1)>) -> !fly.layout<(128,64):(1E0,1E1)> + %gt = fly.make_view(%org, %lay) : (!fly.int_tuple<(?,0)>, !fly.layout<(128,64):(1E0,1E1)>) -> !fly.coord_tensor<(?,0), (128,64):(1E0,1E1)> + // The dynamic leaf is the block index itself, taken straight off the operand; the + // static one is still a constant from the type. + // CHECK: arith.subi %{{.*}}, %arg4 : i32 + // CHECK: rocdl.tensor.load.to.lds + fly.copy_atom_call(%atom, %gt, %lds) : (!fly.copy_atom, 16>, !fly.coord_tensor<(?,0), (128,64):(1E0,1E1)>, !fly.memref) -> () + return +} + +// ----- + +// The store direction: LDS -> global, with the coordinate tensor on the dst side. + +// CHECK-LABEL: @test_cdna5_store +func.func @test_cdna5_store( + %base: !fly.ptr, %s0: i64, %e0: i32, %e1: i32, + %lds: !fly.memref) { + %atom = fly.make_copy_atom(%base, %s0, %e0, %e1 : !fly.ptr, i64, i32, i32) {valBits = 16 : i32} : !fly.copy_atom, 16> + %org = fly.make_coord() : () -> !fly.int_tuple<(0,0)> + %shp = fly.make_int_tuple() : () -> !fly.int_tuple<(128,64)> + %str = fly.make_int_tuple() : () -> !fly.int_tuple<(1E0,1E1)> + %lay = fly.make_layout(%shp, %str) : (!fly.int_tuple<(128,64)>, !fly.int_tuple<(1E0,1E1)>) -> !fly.layout<(128,64):(1E0,1E1)> + %gt = fly.make_view(%org, %lay) : (!fly.int_tuple<(0,0)>, !fly.layout<(128,64):(1E0,1E1)>) -> !fly.coord_tensor<(0,0), (128,64):(1E0,1E1)> + // CHECK: rocdl.tensor.store.from.lds %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, 0 : vector<4xi32>, vector<8xi32> + fly.copy_atom_call(%atom, %lds, %gt) : (!fly.copy_atom, 16>, !fly.memref, !fly.coord_tensor<(0,0), (128,64):(1E0,1E1)>) -> () + return +} + +// ----- + +// A recast: `elem` is the *descriptor's* unit, and it +// may be wider than the tensor's own element. Here an FP4 tile is moved as bytes -- +// `data_size` is 1/2/4/8 bytes and 4 bits is none of them, so this is the only way to +// describe such a tensor at all. The LDS operand keeps its FP4 element type; what has to +// agree is the bit count, 128*32*8 == 128*64*4. +// CHECK-LABEL: @test_cdna5_load_recast_subbyte +func.func @test_cdna5_load_recast_subbyte( + %base: !fly.ptr, %s0: i64, %e0: i32, %e1: i32, + %lds: !fly.memref) { + %atom = fly.make_copy_atom(%base, %s0, %e0, %e1 : !fly.ptr, i64, i32, i32) {valBits = 8 : i32} : !fly.copy_atom, 8> + %org = fly.make_coord() : () -> !fly.int_tuple<(384,64)> + %shp = fly.make_int_tuple() : () -> !fly.int_tuple<(128,32)> + %str = fly.make_int_tuple() : () -> !fly.int_tuple<(1E0,1E1)> + %lay = fly.make_layout(%shp, %str) : (!fly.int_tuple<(128,32)>, !fly.int_tuple<(1E0,1E1)>) -> !fly.layout<(128,32):(1E0,1E1)> + %gt = fly.make_view(%org, %lay) : (!fly.int_tuple<(384,64)>, !fly.layout<(128,32):(1E0,1E1)>) -> !fly.coord_tensor<(384,64), (128,32):(1E0,1E1)> + // data_size 0 == 1 byte, and the address arithmetic scales the coordinate by that + // byte rather than by the tensor's 4-bit element. + // CHECK: rocdl.tensor.load.to.lds %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, 0 : vector<4xi32>, vector<8xi32> + fly.copy_atom_call(%atom, %gt, %lds) : (!fly.copy_atom, 8>, !fly.coord_tensor<(384,64), (128,32):(1E0,1E1)>, !fly.memref) -> () + return +} + +// ----- + +// Descriptor iteration: one instruction replays the same D# `iterCount` times, stepping +// global by the trailing construction argument and LDS by one whole box. It is the TDM +// spelling of a residual axis the box cannot travel on, and it is paid for out of GROUP2 +// (lds_addr_increment, global_addr_increment, iterate_count), which is why the descriptor +// is left with two dims. GROUP1 bit 19 turns it on. +// CHECK-LABEL: @test_cdna5_load_iterate +func.func @test_cdna5_load_iterate( + %base: !fly.ptr, %s0: i64, %e0: i32, %e1: i32, %istride: i64, + %lds: !fly.memref) { + %atom = fly.make_copy_atom(%base, %s0, %e0, %e1, %istride : !fly.ptr, i64, i32, i32, i64) {valBits = 16 : i32} : !fly.copy_atom, 16> + %org = fly.make_coord() : () -> !fly.int_tuple<(0,0)> + %shp = fly.make_int_tuple() : () -> !fly.int_tuple<(1,64)> + %str = fly.make_int_tuple() : () -> !fly.int_tuple<(1E0,1E1)> + %lay = fly.make_layout(%shp, %str) : (!fly.int_tuple<(1,64)>, !fly.int_tuple<(1E0,1E1)>) -> !fly.layout<(1,64):(1E0,1E1)> + %gt = fly.make_view(%org, %lay) : (!fly.int_tuple<(0,0)>, !fly.layout<(1,64):(1E0,1E1)>) -> !fly.coord_tensor<(0,0), (1,64):(1E0,1E1)> + // GROUP1 word 0 carries data_size (1 << 16 for 2-byte elements) together with + // iterate_enable (1 << 19): 0x80000 | 0x10000 == 589824. + // CHECK-DAG: arith.constant 589824 : i32 + // GROUP2 word 1 is the LDS increment, one box of 64 elements. + // CHECK-DAG: arith.constant 64 : i32 + // GROUP2 word 3's upper half is iterate_count encoded as value-minus-one: 7 << 16. + // CHECK-DAG: arith.constant 458752 : i32 + // CHECK: rocdl.tensor.load.to.lds %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, 0 : vector<4xi32>, vector<8xi32> + fly.copy_atom_call(%atom, %gt, %lds) : (!fly.copy_atom, 16>, !fly.coord_tensor<(0,0), (1,64):(1E0,1E1)>, !fly.memref) -> () + return +} + +// ----- + +// The HW auto-barrier: *whether* the copy arrives on one is a type parameter and *which* +// one is atom state, and the two are kept genuinely separate. The state is the barrier +// itself -- a shared pointer, as the kernel holds it -- and only the descriptor's byte +// address is taken from it, here rather than at the call site. Nothing reads it as an +// enable, so no pointer value is spent on meaning "none" and LDS offset 0 is a barrier +// like any other; the price is that an atom whose type asks for a barrier and whose +// state was never given one arrives on offset 0, which is the caller's to get right, as +// the base pointer is. On an atom whose type does not enable it neither the bit nor the +// address appears, and the `set_value` has no field to write at all, which +// `tdm_cdna5_neg.mlir` checks. + +// CHECK-LABEL: @test_cdna5_atomic_barrier +// The enable is the type, so bit 18 is not conditional on anything and folds into the +// GROUP1 config constant -- 327680, that bit alongside the tile's own bit 16. A constant +// there is the whole claim: it could not fold if the pointer were also read as an enable. +// CHECK-DAG: %[[CFG:.*]] = arith.constant 327680 : i32 +// The barrier arrives as `!llvm.ptr<3>` and is flattened once; an LDS pointer is 32-bit, +// so this is a bitcast and not a truncation. +// CHECK-DAG: %[[B:.*]] = llvm.ptrtoint %arg4 : !llvm.ptr<3> to i32 +// CHECK-DAG: arith.shrui %[[B]], %{{.*}} : i32 +// CHECK: vector.from_elements %[[CFG]], +// CHECK: rocdl.tensor.load.to.lds +func.func @test_cdna5_atomic_barrier( + %base: !fly.ptr, %s0: i64, %e0: i32, %e1: i32, + %bar: !fly.ptr, + %lds: !fly.memref) { + %a0 = fly.make_copy_atom(%base, %s0, %e0, %e1 : !fly.ptr, i64, i32, i32) {valBits = 16 : i32} : !fly.copy_atom, 16> + %atom = fly.atom.set_value(%a0, "atomic_barrier_addr", %bar) : (!fly.copy_atom, 16>, !fly.ptr) -> !fly.copy_atom, 16> + %org = fly.make_coord() : () -> !fly.int_tuple<(0,0)> + %shp = fly.make_int_tuple() : () -> !fly.int_tuple<(128,64)> + %str = fly.make_int_tuple() : () -> !fly.int_tuple<(1E0,1E1)> + %lay = fly.make_layout(%shp, %str) : (!fly.int_tuple<(128,64)>, !fly.int_tuple<(1E0,1E1)>) -> !fly.layout<(128,64):(1E0,1E1)> + %gt = fly.make_view(%org, %lay) : (!fly.int_tuple<(0,0)>, !fly.layout<(128,64):(1E0,1E1)>) -> !fly.coord_tensor<(0,0), (128,64):(1E0,1E1)> + fly.copy_atom_call(%atom, %gt, %lds) : (!fly.copy_atom, 16>, !fly.coord_tensor<(0,0), (128,64):(1E0,1E1)>, !fly.memref) -> () + return +} diff --git a/tests/mlir/Conversion/tdm_cdna5_neg.mlir b/tests/mlir/Conversion/tdm_cdna5_neg.mlir new file mode 100644 index 000000000..f5d237cfd --- /dev/null +++ b/tests/mlir/Conversion/tdm_cdna5_neg.mlir @@ -0,0 +1,22 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors +// RUN: { %fly-opt --split-input-file %s 2>&1 || true; } | FileCheck %s + +// Iteration is paid for out of the descriptor's own GROUP2 slots: `tensor_dim2_stride` +// becomes `global_addr_increment`, `tensor_dim3` becomes `lds_addr_increment`, and +// `tile_dim3` becomes `iterate_count`. A dim whose stride the descriptor no longer holds +// is not a dim, so an iterating descriptor has two -- the third axis is the one the +// iteration itself walks. The builder never reaches this: `foldModes` drops iteration +// rather than refusing the atom. Hand-written IR does, and must be refused. + +// ----- + +// CHECK: TDM descriptor iteration takes dim 2's stride for its own, so it needs a descriptor of at most 2 dims, got 3 +func.func private @bad_cdna5_iterate_rank3( + %a: !fly.copy_atom, 16>) + +// ----- + +// CHECK: TDM descriptor iteration takes dim 2's stride for its own, so it needs a descriptor of at most 2 dims, got 3 +func.func private @bad_cdna5_iterate_rank3_store( + %a: !fly.copy_atom, 16>)