[Codegen][AMDGPU] Add UseGlobalTransposeLoad promotion attr and operand promotion

Introduces IREEGPU_UseGlobalTransposeLoad, a promotion attribute for RDNA4
(gfx1200+) that drives matmul operands to be loaded via global_load_tr.

The attribute implements both IREEGPU_PromotionAttr and
IREECodegen_LoweringConfigAttrInterface, with tiling sizes [N=8, K=1]
(vectorSize x 1) derived from the element type.

A new transposePromoteOperand path in GPUPromoteMatmulOperands creates a
linalg.generic copy with K-inner thread mapping:
  - input  map: reads B[K, N] (K-outer, N-inner in memory)
  - output map: writes alloc[N, K] (N-outer, K-inner, for contiguous K writes)

This K-inner tiling aligns with global_load_tr's 8x8 wave-level transpose
semantics: 8 consecutive lanes each read 8 contiguous N-elements, and the
hardware transposes so each lane holds a K-direction slice.

The copy op is tagged with UseGlobalTransposeLoadAttr as its lowering config
so the ROCDLLoadToTransposeLoad pass (PR 2) can recognise and lower it.

amdgpu.fat_raw_buffer_cast is stripped before creating the copy because
global_load_tr requires a flat global pointer, not a fat buffer descriptor.

Part of: #24454

Co-authored-by: Claude Sonnet 4 (1M context) <noreply@anthropic.com>
Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com>

Author: nirvedhmeshram

compiler/src/iree/compiler/Codegen/Common/GPU/GPUPromoteMatmulOperands.cpp

@@ -15,10 +15,12 @@
 #include "iree/compiler/Codegen/Utils/LinalgOpInfo.h"
 #include "iree/compiler/Codegen/Utils/Utils.h"
 #include "llvm/ADT/STLExtras.h"
+#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Interfaces/FunctionInterfaces.h"
@@ -118,6 +120,119 @@ void promoteResult(OpBuilder &builder, Operation *op, Value valToMakeShared) {
   });
 }
 
+/// Traces through tensor.extract_slice ops to find the
+/// iree_codegen.load_from_buffer feeding this value. Returns the slice chain
+/// (outermost first) and the load op, or failure if the pattern isn't matched.
+static FailureOr<std::pair<SmallVector<tensor::ExtractSliceOp>,
+                           IREE::Codegen::LoadFromBufferOp>>
+findLoadFromBuffer(Value v) {
+  SmallVector<tensor::ExtractSliceOp> slices;
+  while (auto sliceOp = v.getDefiningOp<tensor::ExtractSliceOp>()) {
+    slices.push_back(sliceOp);
+    v = sliceOp.getSource();
+  }
+  auto loadOp = v.getDefiningOp<IREE::Codegen::LoadFromBufferOp>();
+  if (!loadOp) {
+    return failure();
+  }
+  return std::make_pair(slices, loadOp);
+}
+
+/// Promotes an operand for global_load_tr by copying to shared memory.
+/// Strips fat_raw_buffer so the copy reads from a flat global pointer
+/// (required by global_load_tr), then creates a linalg.generic that
+/// iterates (K-outer, N-inner) so vectorization produces vector<1x8> reads
+/// — 8 contiguous N-direction elements per lane. This is the access pattern
+/// that global_load_tr_b128 requires.
+/// The matmul's indexing map is updated to read from the transposed [N,K]
+/// shared memory layout.
+Value transposePromoteOperand(OpBuilder &builder, Operation *op,
+                              unsigned index) {
+  OpOperand &operand = op->getOpOperand(index);
+  Location loc = op->getLoc();
+
+  // Strip fat_raw_buffer if present so global_load_tr can use a flat pointer.
+  Value sourceValue = operand.get();
+  auto maybeLoad = findLoadFromBuffer(sourceValue);
+  if (succeeded(maybeLoad)) {
+    auto &[slices, loadOp] = *maybeLoad;
+    if (auto castOp =
+            loadOp.getBuffer().getDefiningOp<amdgpu::FatRawBufferCastOp>()) {
+      OpBuilder::InsertionGuard g(builder);
+      builder.setInsertionPointAfter(loadOp);
+      auto flatMemrefType = cast<MemRefType>(castOp.getSource().getType());
+      auto flatTensorType = RankedTensorType::get(
+          flatMemrefType.getShape(), flatMemrefType.getElementType());
+      Value flatLoad = IREE::Codegen::LoadFromBufferOp::create(
+          builder, loadOp.getLoc(), flatTensorType, castOp.getSource());
+      sourceValue = flatLoad;
+      for (auto sliceOp : llvm::reverse(slices)) {
+        builder.setInsertionPointAfter(sliceOp);
+        sourceValue = tensor::ExtractSliceOp::create(
+            builder, sliceOp.getLoc(), sliceOp.getResultType(), sourceValue,
+            sliceOp.getMixedOffsets(), sliceOp.getMixedSizes(),
+            sliceOp.getMixedStrides());
+      }
+    }
+  }
+
+  auto tensorType = cast<RankedTensorType>(sourceValue.getType());
+  MLIRContext *ctx = op->getContext();
+
+  // Create the transposed output buffer [N, K].
+  SmallVector<OpFoldResult> mixedSizes =
+      tensor::getMixedSizes(builder, loc, sourceValue);
+  SmallVector<OpFoldResult> transposedSizes(mixedSizes.rbegin(),
+                                            mixedSizes.rend());
+  Value empty = tensor::EmptyOp::create(builder, loc, transposedSizes,
+                                        tensorType.getElementType());
+
+  // linalg.generic with (d0=K outer, d1=N inner):
+  //   input  map: (d0, d1) -> (d0, d1)  reads src[k, n]
+  //   output map: (d0, d1) -> (d1, d0)  writes dst[n, k]
+  // With N as the inner (vectorized) dimension, each thread reads
+  // vector<1x8> (8 contiguous N elements at fixed K) — the correct
+  // access pattern for global_load_tr_b128.
+  // Loop iteration order: (d0=N outer, d1=K inner).
+  // With K as the inner (fast-varying per-lane) dimension,
+  // UseGlobalTransposeLoadAttr's tiling [N=vectorSize, K=1] maps K to
+  // linear_dim_0 (fast thread dim): 8 consecutive lanes get 8 consecutive
+  // K rows, which is the correct wave-level setup for global_load_tr.
+  // Each lane reads vector<1x8> (8 contiguous N from its K row). The tag
+  // UseGlobalTransposeLoadAttr drives the thread tiling level sizes.
+  //   input  map (d0=N, d1=K) -> (d1, d0)  reads B[K, N]
+  //   output map (d0=N, d1=K) -> (d0, d1)  writes alloc[N, K]
+  AffineExpr d0 = builder.getAffineDimExpr(0); // N (outer)
+  AffineExpr d1 = builder.getAffineDimExpr(1); // K (inner)
+  AffineMap inputMap = AffineMap::get(2, 0, {d1, d0}, ctx);
+  AffineMap outputMap = AffineMap::get(2, 0, {d0, d1}, ctx);
+  SmallVector<utils::IteratorType> iterTypes(2, utils::IteratorType::parallel);
+
+  auto copyOp = linalg::GenericOp::create(
+      builder, loc, empty.getType(), sourceValue, empty,
+      ArrayRef<AffineMap>{inputMap, outputMap}, iterTypes,
+      [](OpBuilder &b, Location l, ValueRange args) {
+        linalg::YieldOp::create(b, l, args[0]);
+      });
+  // Use UseGlobalTransposeLoadAttr as the lowering config so the tiling pass
+  // produces K-inner thread assignment via getStaticTilingLevelSizes.
+  setLoweringConfig(copyOp, IREE::GPU::UseGlobalTransposeLoadAttr::get(ctx));
+
+  // Update the matmul's indexing map for this operand by reversing its
+  // results to reflect the [N, K] shared memory layout.
+  if (auto genericOp = dyn_cast<linalg::GenericOp>(op)) {
+    SmallVector<AffineMap> maps(genericOp.getIndexingMapsArray());
+    AffineMap oldMap = maps[index];
+    SmallVector<AffineExpr> results(oldMap.getResults().rbegin(),
+                                    oldMap.getResults().rend());
+    maps[index] = AffineMap::get(oldMap.getNumDims(), oldMap.getNumSymbols(),
+                                 results, ctx);
+    genericOp.setIndexingMapsAttr(builder.getAffineMapArrayAttr(maps));
+  }
+
+  return copyOp.getResult(0);
+}
+
 void promoteOperand(OpBuilder &builder, Operation *op, unsigned index,
                     IREE::GPU::PromotionAttr promotionAttr) {
   auto dpsOp = dyn_cast<DestinationStyleOpInterface>(op);
@@ -133,9 +248,14 @@ void promoteOperand(OpBuilder &builder, Operation *op, unsigned index,
     // TODO(qedawkins): Move result promotion to attribute interface.
     return promoteResult(builder, op, op->getResult(index));
   }
-  OpOperand &operand = op->getOpOperand(index);
 
-  Value replacement = promotionAttr.promoteOperand(builder, operand);
+  Value replacement;
+  if (isa<IREE::GPU::UseGlobalTransposeLoadAttr>(promotionAttr)) {
+    replacement = transposePromoteOperand(builder, op, index);
+  } else {
+    OpOperand &operand = op->getOpOperand(index);
+    replacement = promotionAttr.promoteOperand(builder, operand);
+  }
   op->setOperand(index, replacement);
 }
 

compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel

@@ -59,6 +59,7 @@ iree_lit_test_suite(
             "gpu_pad_operands.mlir",
             "gpu_pipeline.mlir",
             "gpu_promote_matmul_operands.mlir",
+            "gpu_promote_matmul_operands_global_transpose.mlir",
             "gpu_promotion_analysis.mlir",
             "gpu_reorder_workgroups.mlir",
             "gpu_reorder_workgroups_static.mlir",

compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt

@@ -54,6 +54,7 @@ iree_lit_test_suite(
     "gpu_pad_operands.mlir"
     "gpu_pipeline.mlir"
     "gpu_promote_matmul_operands.mlir"
+    "gpu_promote_matmul_operands_global_transpose.mlir"
     "gpu_promotion_analysis.mlir"
     "gpu_reorder_workgroups.mlir"
     "gpu_reorder_workgroups_static.mlir"

compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_promote_matmul_operands_global_transpose.mlir

@@ -0,0 +1,91 @@
+// RUN: iree-opt %s --split-input-file \
+// RUN:   --pass-pipeline="builtin.module(func.func(iree-codegen-gpu-promote-matmul-operands),canonicalize)" \
+// RUN:   | FileCheck %s
+
+// Verify UseGlobalTransposeLoad promotion:
+//   - Creates a linalg.generic copy with K-inner thread mapping
+//     (input map reads B[K,N], output map writes alloc[N,K])
+//   - Tags the copy with #iree_gpu.use_global_transpose_load lowering config
+//   - The matmul RHS is updated to use the promoted (transposed) buffer
+
+// -----
+
+#lowering_config = #iree_gpu.lowering_config<{
+  promote_operands = [1],
+  promotion_types = [#iree_gpu.use_global_transpose_load]}>
+
+// CHECK-LABEL: func.func @transpose_promote_rhs
+// CHECK-SAME:    %[[LHS:[A-Za-z0-9]+]]: tensor<32x64xbf16>
+// CHECK-SAME:    %[[RHS:[A-Za-z0-9]+]]: tensor<64x128xbf16>
+func.func @transpose_promote_rhs(%lhs: tensor<32x64xbf16>,
+                                  %rhs: tensor<64x128xbf16>) -> tensor<32x128xbf16> {
+  %cst = arith.constant 0.0 : bf16
+  %empty = tensor.empty() : tensor<32x128xbf16>
+  %fill = linalg.fill ins(%cst : bf16) outs(%empty : tensor<32x128xbf16>) -> tensor<32x128xbf16>
+  %mm = linalg.generic {
+    indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
+                     affine_map<(d0, d1, d2) -> (d2, d1)>,
+                     affine_map<(d0, d1, d2) -> (d0, d1)>],
+    iterator_types = ["parallel", "parallel", "reduction"],
+    lowering_config = #lowering_config}
+    ins(%lhs, %rhs : tensor<32x64xbf16>, tensor<64x128xbf16>)
+    outs(%fill : tensor<32x128xbf16>) {
+  ^bb0(%in0: bf16, %in1: bf16, %out0: bf16):
+    %mul = arith.mulf %in0, %in1 : bf16
+    %add = arith.addf %out0, %mul : bf16
+    linalg.yield %add : bf16
+  } -> tensor<32x128xbf16>
+  return %mm : tensor<32x128xbf16>
+}
+
+// The copy linalg.generic has:
+//   input  map: (d0, d1) -> (d1, d0)   reads B[K, N] (K-outer, N-inner)
+//   output map: (d0, d1) -> (d0, d1)   writes alloc[N, K] (N-outer, K-inner)
+// This gives K-inner thread assignment so global_load_tr's 8x8 wave transpose
+// aligns with the thread→lane mapping.
+//
+// The copy has the transposed output shape [N=128, K=64] and K-inner config.
+// CHECK:       tensor.empty() : tensor<128x64xbf16>
+// CHECK:       linalg.generic
+// CHECK-SAME:    ins({{.*}} : tensor<64x128xbf16>)
+// CHECK-SAME:    outs({{.*}} : tensor<128x64xbf16>)
+// CHECK-SAME:    lowering_config = #iree_gpu.use_global_transpose_load
+//
+// The matmul's RHS input is the promoted [N=128, K=64] buffer.
+// CHECK:       linalg.generic
+// CHECK-SAME:    ins({{.*}}, {{.*}} : tensor<32x64xbf16>, tensor<128x64xbf16>)
+
+// -----
+
+// LHS promotion with UseGlobalTransposeLoad (transposedLhs case).
+#lowering_config_lhs = #iree_gpu.lowering_config<{
+  promote_operands = [0],
+  promotion_types = [#iree_gpu.use_global_transpose_load]}>
+
+// CHECK-LABEL: func.func @transpose_promote_lhs
+// CHECK-SAME:    %[[LHS:[A-Za-z0-9]+]]: tensor<64x32xbf16>
+func.func @transpose_promote_lhs(%lhs: tensor<64x32xbf16>,
+                                  %rhs: tensor<64x128xbf16>) -> tensor<32x128xbf16> {
+  %cst = arith.constant 0.0 : bf16
+  %empty = tensor.empty() : tensor<32x128xbf16>
+  %fill = linalg.fill ins(%cst : bf16) outs(%empty : tensor<32x128xbf16>) -> tensor<32x128xbf16>
+  // transposedLhs: LHS is K-outer (K, M) instead of (M, K)
+  %mm = linalg.generic {
+    indexing_maps = [affine_map<(d0, d1, d2) -> (d2, d0)>,
+                     affine_map<(d0, d1, d2) -> (d2, d1)>,
+                     affine_map<(d0, d1, d2) -> (d0, d1)>],
+    iterator_types = ["parallel", "parallel", "reduction"],
+    lowering_config = #lowering_config_lhs}
+    ins(%lhs, %rhs : tensor<64x32xbf16>, tensor<64x128xbf16>)
+    outs(%fill : tensor<32x128xbf16>) {
+  ^bb0(%in0: bf16, %in1: bf16, %out0: bf16):
+    %mul = arith.mulf %in0, %in1 : bf16
+    %add = arith.addf %out0, %mul : bf16
+    linalg.yield %add : bf16
+  } -> tensor<32x128xbf16>
+  return %mm : tensor<32x128xbf16>
+}
+
+// CHECK:       linalg.generic
+// CHECK-SAME:    ins({{.*}} : tensor<64x32xbf16>)
+// CHECK-SAME:    lowering_config = #iree_gpu.use_global_transpose_load

compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/DerivedConfigUtils.cpp

@@ -221,4 +221,44 @@ SmallVector<int64_t> globalLoadDMATileSizes(Operation *op) {
   return tileSizes;
 }
 
+SmallVector<int64_t> globalTransposeLoadTileSizes(Operation *op) {
+  auto funcOp = op->getParentOfType<FunctionOpInterface>();
+  std::optional<SmallVector<int64_t>> workgroupSize = getWorkgroupSize(funcOp);
+  if (!workgroupSize) {
+    return {};
+  }
+  auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
+  if (!linalgOp) {
+    return {};
+  }
+
+  SmallVector<int64_t> loopRanges = linalgOp.getStaticLoopRanges();
+  if (loopRanges.size() != 2) {
+    return {};
+  }
+
+  int64_t targetSubgroupSize = getGPUTargetAttr(op).getPreferredSubgroupSize();
+  // Vector size: 8 bf16 elements = 128 bits per lane per global_load_tr call.
+  int64_t elemBits = getElementTypeOrSelf(linalgOp->getResultTypes()[0])
+                         .getIntOrFloatBitWidth();
+  int64_t vectorSize = 128 / elemBits; // 8 for bf16/f16, 16 for i8
+
+  int64_t numThreads = llvm::product_of(*workgroupSize);
+
+  // Tile: [d0=N (outer, step=vectorSize), d1=K (inner, step=1)].
+  // The K axis (d1, step=1) maps to linear_dim_0 (fast thread dim), giving
+  // K-inner wave assignment needed by global_load_tr.
+  int64_t kRange = loopRanges[1]; // K dimension (inner of linalg.generic)
+  int64_t nRange = loopRanges[0]; // N dimension (outer of linalg.generic)
+
+  // Each thread handles vectorSize N-values (vectorized) and 1 K-value.
+  // Total tasks: (nRange / vectorSize) * kRange = numThreads * numIterations.
+  (void)nRange;
+  (void)kRange;
+  (void)numThreads;
+  (void)targetSubgroupSize;
+
+  return {vectorSize, 1}; // [N_step=vectorSize, K_step=1]
+}
+
 } // namespace mlir::iree_compiler::IREE::GPU

compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/DerivedConfigUtils.h

@@ -15,6 +15,13 @@ namespace mlir::iree_compiler::IREE::GPU {
 SmallVector<int64_t> deriveThreadTileSizes(Operation *op);
 SmallVector<int64_t> globalLoadDMATileSizes(Operation *op);
 
+/// Returns thread-level tile sizes [N=vectorSize, K=1] for the
+/// global_transpose_load copy. With N as the outer step and K=1 as the inner,
+/// the K axis maps to linear_dim_0 (fast thread dim), giving K-inner wave
+/// assignment. 8 consecutive lanes then get 8 consecutive K rows, which is
+/// required for global_load_tr's wave-level 8x8 cross-lane transpose.
+SmallVector<int64_t> globalTransposeLoadTileSizes(Operation *op);
+
 } // namespace mlir::iree_compiler::IREE::GPU
 
 namespace mlir::iree_compiler {

compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp

@@ -3190,6 +3190,33 @@ bool UseGlobalLoadDMAAttr::hasTilingLevel(unsigned level) const {
   return level == llvm::to_underlying(GPU::TilingLevel::Thread);
 }
 
+//===----------------------------------------------------------------------===//
+// UseGlobalTransposeLoadAttr - LoweringConfigAttrInterface
+//===----------------------------------------------------------------------===//
+
+SmallVector<int64_t>
+UseGlobalTransposeLoadAttr::getStaticTilingLevelSizes(unsigned level,
+                                                      Operation *op) const {
+  if (level == llvm::to_underlying(GPU::TilingLevel::Thread)) {
+    return globalTransposeLoadTileSizes(op);
+  }
+  return {};
+}
+
+SmallVector<OpFoldResult>
+UseGlobalTransposeLoadAttr::getTilingLevelSizes(OpBuilder &b, unsigned level,
+                                                Operation *op) const {
+  if (level == llvm::to_underlying(GPU::TilingLevel::Thread)) {
+    SmallVector<int64_t> sizes = globalTransposeLoadTileSizes(op);
+    return getAsIndexOpFoldResult(b.getContext(), sizes);
+  }
+  return {};
+}
+
+bool UseGlobalTransposeLoadAttr::hasTilingLevel(unsigned level) const {
+  return level == llvm::to_underlying(GPU::TilingLevel::Thread);
+}
+
 //===----------------------------------------------------------------------===//
 // PromoteWithCacheSwizzleAttr
 //===----------------------------------------------------------------------===//

compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td

@@ -97,6 +97,36 @@ def IREEGPU_UseGlobalLoadDma :
   let parameters = (ins);
 }
 
+def IREEGPU_UseGlobalTransposeLoad :
+    AttrDef<IREEGPU_Dialect, "UseGlobalTransposeLoad", [
+      DeclareAttrInterfaceMethods<IREECodegen_LoweringConfigAttrInterface, [
+        "getStaticTilingLevelSizes",
+        "getTilingLevelSizes",
+        "hasTilingLevel",
+      ]>,
+      DeclareAttrInterfaceMethods<IREEGPU_PromotionAttr>
+    ]> {
+  let mnemonic = "use_global_transpose_load";
+  let summary = [{
+    Promote an operand using the global memory transpose load instruction.
+  }];
+  let description = [{
+    Promotion attribute indicating that the operand should be loaded from global
+    memory using the `amdgpu.global_transpose_load` instruction, which loads
+    and transposes a matrix tile in a single subgroup operation.
+
+    This is applicable for gfx1200+ targets when the operand's memory layout
+    requires a transpose relative to what the MMA intrinsic expects:
+    - LHS when `transposedLhs` (K is not the innermost dimension)
+    - RHS when `!transposedRhs` (N is not the innermost dimension)
+
+    Supported element types for gfx1200+: i8, f8E5M2FNUZ, f8E4M3FNUZ,
+    f8E5M2, f8E4M3FN, f16, bf16, i16.
+  }];
+  let assemblyFormat = "";
+  let parameters = (ins);
+}
+
 def IREEGPU_PromoteWithCacheSwizzle :
     AttrDef<IREEGPU_Dialect, "PromoteWithCacheSwizzle", [
       DeclareAttrInterfaceMethods<IREEGPU_PromotionAttr, [