Add prefetch functionality for gfx1250

Author: dhernandez0

mlir/include/mlir/Dialect/Rock/IR/AmdArchDb.h

@@ -52,6 +52,7 @@ struct AmdArchInfo {
 };
 
 AmdArchInfo lookupArchInfo(StringRef arch);
+bool isGlobalPrefetchSupported(StringRef arch);
 } // namespace rock
 } // namespace mlir
 

mlir/include/mlir/Dialect/Rock/IR/RockOps.td

@@ -1017,6 +1017,26 @@ def Rock_InBoundsStoreOp
   let hasVerifier = 1;
 }
 
+// global_prefetch
+def Rock_GlobalPrefetchOp
+    : Rock_Op<"global_prefetch">,
+      Arguments<(
+          ins Arg<MemRefOf<SupportedMemoryElems>, "source memory">:$source,
+          Variadic<Index>:$sourceCoord)> {
+  let summary = "Prefetch from global memory";
+  let description = [{
+    `global_prefetch` prefetches an item from the provided memref
+    (applying no coordinate transformations), starting at `sourceCoord`. 
+
+    Note that it's ok if `sourceCoord` are out of bounds because we use Speculative Prefetch.
+  }];
+  let assemblyFormat = [{
+    $source `[` $sourceCoord `]` attr-dict
+    `:` type($source)
+  }];
+  let hasVerifier = 1;
+}
+
 // global_load
 def Rock_GlobalLoadOp
     : Rock_Op<
@@ -1188,7 +1208,7 @@ def Rock_ThreadwiseReadIntoOp
     L is the length of `%dest`, V is the maximum vectorization computed
     for MAPS.
 
-    The input to extraViews ; (the transforms on %source) must have the form
+    The input to extraViews: (the transforms on %source) must have the form
     (extraIdx0, ... , extraIdxN, iteration_number)
 
     Primarily, extraIndices would be used to pass in tid and bid. This would need
@@ -1235,6 +1255,59 @@ def Rock_ThreadwiseReadIntoOp
   let hasVerifier = 1;
 }
 
+// threadwise_prefetch
+def Rock_ThreadwisePrefetchOp
+    : Rock_Op<"threadwise_prefetch", [DeclareOpInterfaceMethods<
+                                         RockAcceptingViewOpInterface>]>,
+      Arguments<(ins Arg<MemRefOf<SupportedMemoryElems>, "source view">:$source,
+          TransformMapArrayAttr:$extraViews, Variadic<Index>:$extraIndices,
+          UnitAttr:$forceUnroll, UnitAttr:$useIndexDiffs)> {
+  let summary = "Prefetch values from transformed source";
+
+  let description = [{
+    A high-level representation of a global prefetch loop that
+    accounts for coordinate transformations.
+
+    If `%source = rock.transform #transform_mapN %buffer`
+    (with the one transformation representing an entire sequence),
+    the operation
+
+    ```mlir
+    rock.threadwise_prefetch [#transform_mapM](%source)
+    ```
+
+    will lower to
+    ```mlir
+    %bid = rock.workgroup_id
+    %tid = rock.workitem_id
+    rock.transforming_for
+      (%args, ...) = MAPS(%bid, %tid, %c0)
+      (%_, %_, %i) = [](%c0, %c0, %c0)
+      bounds = [1, 1, L], strides = [1, 1, 1] {
+        rock.global_prefetch %buffer[%args]
+    }
+    ```
+
+    where MAPS is `[#transform_mapM, #transform_mapN]`,
+    L is the length of the upper view last dimension (number of elements).
+
+    The input to extraViews ; (the transforms on %source) must have the form
+    (extraIdx0, ... , extraIdxN, iteration_number)
+
+    Primarily, extraIndices would be used to pass in tid and bid. This would need
+    to have a matching view in [extraViews]source. The extraIndices could be used
+    to integrate loop induction vars that is outside of the op.
+    If extraIndices are used, the [extraViews]source must have the form
+    (extraIdx0, ... , extraIdxN, iteration_number).
+  }];
+
+  let assemblyFormat = [{
+    attr-dict $extraViews `(` $source `)` (`[` $extraIndices^ `]`)?
+    `:` type($source)
+  }];
+  let hasVerifier = 1;
+}
+
 // threadwise_write_all
 def Rock_ThreadwiseWriteAllOp
     : Rock_Op<"threadwise_write_all",

mlir/lib/Dialect/Rock/IR/AmdArchDb.cpp

@@ -112,7 +112,18 @@ static constexpr AmdArchInfo
               /*maxSharedMemPerWG*/ 65536, /*numEUPerCU=*/4, /*minNumCU=*/12,
               /*hasFp8ConversionInstrs=*/false,
               /*hasOcpFp8ConversionInstrs=*/true, /*hasScaledGemm=*/false,
-              /*maxNumXCC=*/1);
+              /*maxNumXCC=*/1),
+    // TODO: update with right information
+    gfx1250Info(GemmFeatures::dot | GemmFeatures::atomic_add |
+                    GemmFeatures::atomic_fmax_f32 | GemmFeatures::wmma |
+                    GemmFeatures::atomic_add_f16 |
+                    GemmFeatures::atomic_add_bf16,
+                /*waveSize=*/32, /*maxWavesPerEU*/ 16, /*totalSGPRPerEU*/ 800,
+                /*totalVGPRPerEU*/ 1536, /*totalSharedMemPerCU*/ 131072,
+                /*maxSharedMemPerWG*/ 65536, /*numEUPerCU=*/4, /*minNumCU=*/12,
+                /*hasFp8ConversionInstrs=*/false,
+                /*hasOcpFp8ConversionInstrs=*/true, /*hasScaledGemm=*/false,
+                /*maxNumXCC=*/1);
 
 static std::tuple<StringRef, unsigned> parseArchString(StringRef arch) {
   std::tuple<StringRef, unsigned> ret("", 0);
@@ -367,7 +378,9 @@ AmdArchInfo mlir::rock::lookupArchInfo(StringRef arch) {
     return rdna3Info;
   }
   if (major == "gfx12") {
-    return rdna4Info;
+    return llvm::StringSwitch<AmdArchInfo>(minor)
+        .Case("50", gfx1250Info)
+        .Default(rdna4Info);
   }
   auto msg = "Unsupported architecture: " + arch.str();
   llvm_unreachable(msg.c_str());
@@ -413,3 +426,7 @@ GemmFeatures mlir::rock::AmdArchInfo::getDefaultFeatures(Type dataType) {
   }
   return theseFeatures;
 }
+
+bool mlir::rock::isGlobalPrefetchSupported(StringRef arch) {
+  return arch == "gfx1250";
+}

mlir/lib/Dialect/Rock/IR/RockDialect.cpp

@@ -1835,15 +1835,13 @@ LogicalResult IndexDiffUpdateOp::verify() {
   return success();
 }
 
-template <typename Load>
-static LogicalResult verifyGlobalLoad(Load op) {
+template <typename LoadOrPrefetch>
+static LogicalResult verifyGlobalLoadAndPrefetch(LoadOrPrefetch op) {
   MemRefType sourceType = op.getSource().getType();
   size_t nDims = sourceType.getRank();
 
   if (op.getSourceCoord().size() != nDims)
-    return op.emitOpError("Expected " + Twine(nDims) + " coordinates for load");
-  if (op.getCanReadOffEnd() && nDims != 1)
-    return op.emitOpError("can only have one dimension in canReadOffEnd loads");
+    return op.emitOpError("Expected " + Twine(nDims) + " coordinates");
   Attribute memSpaceAttr = sourceType.getMemorySpace();
   auto gpuMemSpaceAttr = dyn_cast_or_null<gpu::AddressSpaceAttr>(memSpaceAttr);
   if (memSpaceAttr && (!gpuMemSpaceAttr ||
@@ -1852,6 +1850,27 @@ static LogicalResult verifyGlobalLoad(Load op) {
   return success();
 }
 
+template <typename Load>
+static LogicalResult verifyGlobalLoad(Load op) {
+  if (failed(verifyGlobalLoadAndPrefetch(op)))
+    return failure();
+
+  MemRefType sourceType = op.getSource().getType();
+  size_t nDims = sourceType.getRank();
+
+  if (op.getCanReadOffEnd() && nDims != 1)
+    return op.emitOpError("can only have one dimension in canReadOffEnd loads");
+  return success();
+}
+
+//===-----------------------------------------------------===//
+// GlobalPrefetchOp
+//===-----------------------------------------------------===//
+
+LogicalResult GlobalPrefetchOp::verify() {
+  return verifyGlobalLoadAndPrefetch(*this);
+}
+
 //===-----------------------------------------------------===//
 // GlobalLoadOp
 //===-----------------------------------------------------===//
@@ -1956,6 +1975,70 @@ LogicalResult InBoundsStoreOp::verify() {
   return success();
 }
 
+//===-----------------------------------------------------===//
+// ThreadwisePrefetchOp
+//===-----------------------------------------------------===//
+
+SmallPtrSet<OpOperand *, 2> ThreadwisePrefetchOp::getAcceptingViewOperands() {
+  auto operands = getOperation()->getOpOperands();
+  return {operands.begin()};
+}
+
+std::optional<OperandRange>
+ThreadwisePrefetchOp::getExtraIndices(OpOperand &operand) {
+  if (!getAcceptingViewOperands().contains(&operand)) {
+    return std::nullopt;
+  }
+  // Only one operand supports view
+  return getExtraIndices();
+}
+
+Operation *
+ThreadwisePrefetchOp::cloneWithExtraIndices(OpBuilder &builder,
+                                            OpOperand &operand, Value view,
+                                            ArrayRef<Value> newExtraIndices) {
+  if (!getAcceptingViewOperands().contains(&operand)) {
+    return getOperation();
+  }
+
+  // Only one operand supports view
+  auto newOp = ThreadwisePrefetchOp::create(
+      builder, getLoc(), view, getExtraViews(), newExtraIndices,
+      getForceUnroll(), getUseIndexDiffs());
+  return newOp.getOperation();
+}
+
+LogicalResult ThreadwisePrefetchOp::verify() {
+  MemRefType srcType = getSource().getType();
+  Attribute srcMemSpaceAttr = srcType.getMemorySpace();
+  auto gpuSrcMemSpaceAttr =
+      dyn_cast_or_null<gpu::AddressSpaceAttr>(srcMemSpaceAttr);
+  if (srcMemSpaceAttr &&
+      (!gpuSrcMemSpaceAttr ||
+       gpuSrcMemSpaceAttr.getValue() != gpu::AddressSpace::Global))
+    return emitOpError("prefetching only works for global");
+
+  // we are checking below if extra indices match the upper view bounds.
+  // we should expect zero extra indices if we are prefetching a scalar.
+  // And upperBounds.size() + 1 otherwise.
+  ArrayAttr extraViews = getExtraViews();
+  ArrayRef<int64_t> inputShape;
+  if (extraViews.empty())
+    inputShape = srcType.getShape();
+  else
+    inputShape = cast<TransformMapAttr>(extraViews[0]).getUpperBounds();
+
+  size_t extraIdxCount = getExtraIndices().size();
+  if (inputShape.empty()) {
+    if (extraIdxCount != 0)
+      return emitOpError("read from a scalar value cannot have coordinates");
+  } else if (inputShape.size() != extraIdxCount + 1) {
+    return emitOpError("source view must be extraIndices + 1");
+  }
+
+  return success();
+}
+
 //===-----------------------------------------------------===//
 // ThreadwiseReadIntoOp
 //===-----------------------------------------------------===//

mlir/lib/Dialect/Rock/Transforms/AlignTiling.cpp

@@ -478,6 +478,8 @@ traceToNonViewReaders(Operation *op, Value parentVal,
     if (copyOp.getSource() == parentVal) {
       nonViewReaders.push_back(op);
     }
+  } else if (isa<ThreadwisePrefetchOp>(op)) {
+    // ignore because ThreadwisePrefetchOp is not a reader.
   } else {
     return op->emitError() << "Found an unsupported operator that needs to "
                               "be added reader checks \n"

mlir/lib/Dialect/Rock/Transforms/BlockwiseLoadTileToThreadwise.cpp

@@ -18,9 +18,11 @@
 
 #include "GridLayoutEmitter.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Rock/IR/AccelEmitter.h"
+#include "mlir/Dialect/Rock/IR/AmdArchDb.h"
 #include "mlir/Dialect/Rock/IR/GetRockInfo.h"
 #include "mlir/Dialect/Rock/IR/Rock.h"
 #include "mlir/Dialect/Rock/IR/RockGemmGemmWrapperInterface.h"
@@ -263,6 +265,22 @@ class LoweringBlockwiseLoadTileOp final
                                    /*dynamicValidities=*/ValueRange{},
                                    /*extraViews=*/b.getArrayAttr({}),
                                    /*extraIndices=*/indices, forceUnroll, true);
+
+      if (rock::isGlobalPrefetchSupported(arch)) {
+        // add one to k_loop to prefetch next iteration
+        SmallVector<Value> indicesNext(indices.begin(), indices.end());
+        Value one = b.createOrFold<arith::ConstantIndexOp>(loc, 1);
+        indicesNext[0] =
+            arith::AddIOp::create(b, loc, indicesNext[0], one).getResult();
+
+        // it's acceptable if the indices are out of bounds because we use
+        // GLOBAL_PREFETCH_B8 with Speculative Prefetch. See llvm.prefetch
+        // documentation in AMDGPUUsage.rst
+        rock::ThreadwisePrefetchOp::create(b, loc, wrappedSource,
+                                           /*extraViews=*/b.getArrayAttr({}),
+                                           /*extraIndices=*/indicesNext,
+                                           forceUnroll, true);
+      }
       if (stageGlobalReadNew)
         rock::YieldOp::create(b, loc);
     }

mlir/lib/Dialect/Rock/Transforms/SugarToLoops.cpp

@@ -1164,6 +1164,26 @@ Value selectDataIf4b(Location loc, PatternRewriter &b,
   return b.createOrFold<vector::ExtractOp>(loc, loadedVec, lsb);
 }
 
+struct GlobalPrefetchRewritePattern
+    : public OpRewritePattern<GlobalPrefetchOp> {
+  using OpRewritePattern<GlobalPrefetchOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(GlobalPrefetchOp op,
+                                PatternRewriter &b) const override {
+    Value source = op.getSource();
+
+    source = asGlobal(b, source);
+
+    // it's acceptable if the indices are out of bounds because we use
+    // GLOBAL_PREFETCH_B8 with Speculative Prefetch. See llvm.prefetch
+    // documentation in AMDGPUUsage.rst localityHint=3 is translated to memory
+    // scope SCOPE_SE.
+    b.replaceOpWithNewOp<memref::PrefetchOp>(
+        op, source, op.getSourceCoord(), /*isWrite=*/false, /*localityHint=*/3,
+        /*isDataCache=*/true);
+    return success();
+  }
+};
+
 struct GlobalLoadRewritePattern : public OpRewritePattern<GlobalLoadOp> {
   using OpRewritePattern<GlobalLoadOp>::OpRewritePattern;
   LogicalResult matchAndRewrite(GlobalLoadOp op,
@@ -1620,8 +1640,8 @@ void RockSugarToLoopsPass::runOnOperation() {
   RewritePatternSet patterns(ctx);
   patterns.add<ExtractSliceRewritePattern, InsertSliceRewritePattern,
                GlobalLoadRewritePattern, GlobalLoadToLDSRewritePattern,
-               GlobalStoreRewritePattern, InBoundsLoadRewritePattern,
-               InBoundsStoreRewritePattern>(ctx);
+               GlobalPrefetchRewritePattern, GlobalStoreRewritePattern,
+               InBoundsLoadRewritePattern, InBoundsStoreRewritePattern>(ctx);
   if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
     signalPassFailure();