Use pipelining for attention

Author: dhernandez0

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

@@ -297,6 +297,12 @@ def Rock_AttnPerfConfig : Rock_Attr<"AttnPerfConfig", [RockTuningParamAttrInterf
       + Twine(getOutputSwizzle()) + ","
       + Twine(getForceUnroll())).toVector(perfStr);
     }
+    AttnPerfConfigAttr withScheduleVersion(int64_t newScheduleVersion) const {
+    return AttnPerfConfigAttr::get(
+      getContext(), getMPerBlockG0(), getMPerBlockG1(), getNPerBlockG0(),
+      getKpackPerBlock(), getMPerWave(), getMnPerXdl(), getKpack(),
+      getSplitKFactor(), newScheduleVersion, getOutputSwizzle(), getForceUnroll());
+  }
   }];
 
   let builders = [

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

@@ -1372,10 +1372,10 @@ def Rock_BlockwiseGemmAccelOp
           I32Attr:$inNPerThread, UnitAttr:$rotateMWithK, UnitAttr:$rotateNWithK,
           UnitAttr:$loadAfromLDS, UnitAttr:$loadBfromLDS,
           UnitAttr:$splitKAcrossThreadsFirstA,
-          UnitAttr:$splitKAcrossThreadsFirstB, UnitAttr:$directToLDS,
-          UnitAttr:$ldsLayoutMxK, UnitAttr:$ldsLayoutNxK,
-          MemRefOf<LdsBufferTypes>:$bufferA, MemRefOf<LdsBufferTypes>:$bufferB,
-          MemRefOf<AccelResTypes>:$matrixC,
+          UnitAttr:$splitKAcrossThreadsFirstB, UnitAttr:$directToLDSA,
+          UnitAttr:$directToLDSB, UnitAttr:$ldsLayoutMxK,
+          UnitAttr:$ldsLayoutNxK, MemRefOf<LdsBufferTypes>:$bufferA,
+          MemRefOf<LdsBufferTypes>:$bufferB, MemRefOf<AccelResTypes>:$matrixC,
           Optional<MemRefOf<LdsBufferTypes>>:$bufferScaleA,
           Optional<MemRefOf<LdsBufferTypes>>:$bufferScaleB,
           TypeAttr:$elementTypeA, TypeAttr:$elementTypeB,
@@ -1414,9 +1414,10 @@ def Rock_BlockwiseLoadTileOp
           TypeAttr:$elementTypeA, TypeAttr:$elementTypeB,
           TypeAttr:$elementTypeALoad, TypeAttr:$elementTypeBLoad,
           UnitAttr:$rotateWithK, UnitAttr:$swapThreadIterSubDims,
-          UnitAttr:$LDSLayoutDxK, Variadic<Index>:$sourceIndices, I64Attr:$G,
-          I64Attr:$M, I64Attr:$N, OptionalAttr<Rock_GemmFeaturesAttr>:$features,
-          I32Attr:$blockSize, RockAccelTuningParamAttrInterface:$params)> {
+          UnitAttr:$LDSLayoutDxK, UnitAttr:$splitKAcrossThreadsFirst,
+          Variadic<Index>:$sourceIndices, I64Attr:$G, I64Attr:$M, I64Attr:$N,
+          OptionalAttr<Rock_GemmFeaturesAttr>:$features, I32Attr:$blockSize,
+          RockAccelTuningParamAttrInterface:$params)> {
   let summary =
       "Blockwise load tile from global memory to LDS and/or registers";
   let description = [{

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

@@ -5,6 +5,7 @@
 #include "mlir/Dialect/Rock/IR/Rock.h"
 #include "mlir/Dialect/Rock/IR/RockGemmGemmWrapperInterface.h"
 #include "mlir/Dialect/Rock/IR/RockGemmWrapperInterface.h"
+#include "mlir/Dialect/Rock/IR/RockTypes.h"
 #include "mlir/Dialect/Rock/Passes.h"
 #include "mlir/Dialect/Rock/Tuning/GridwiseGemmParams.h"
 #include "mlir/Dialect/Rock/Tuning/UtilityParams.h"
@@ -19,6 +20,7 @@
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/LogicalResult.h"
 #include "llvm/Support/raw_ostream.h"
+#include <optional>
 
 namespace mlir {
 namespace rock {
@@ -51,6 +53,41 @@ struct AffixTuningParameters
 };
 } // anonymous namespace
 
+static FailureOr<std::optional<int64_t>> getScheduleVersion(func::FuncOp funcOp,
+                                                            Operation *op) {
+  auto scheduleVersionAttrName = rock::ScheduleVersionAttr::getMnemonic();
+
+  std::optional<int64_t> scheduleVersion = std::nullopt;
+  bool hasPerfConfig = op->hasAttrOfType<StringAttr>("perf_config");
+  if (funcOp->hasAttrOfType<rock::ScheduleVersionAttr>(
+          scheduleVersionAttrName) &&
+      hasPerfConfig) {
+    return op->emitError(
+        "kernel has both perf_config and schedule_version attribute "
+        "set. Please modify schedule version directly inside "
+        "perf_config and remove schedule_version\n");
+  }
+  if (funcOp->hasAttrOfType<rock::ScheduleVersionAttr>(
+          scheduleVersionAttrName)) {
+    scheduleVersion = dyn_cast<rock::ScheduleVersionAttr>(
+                          funcOp->removeAttr(scheduleVersionAttrName))
+                          .getScheduleVersion();
+  } else if (!hasPerfConfig) {
+    // set default schedule
+    scheduleVersion = static_cast<int64_t>(GemmLoadTileType::Default);
+  }
+
+  // check scheduleVersion is valid
+  if (scheduleVersion.has_value()) {
+    std::optional<GemmLoadTileType> maybeLoadType =
+        rock::symbolizeGemmLoadTileType(scheduleVersion.value());
+    if (!maybeLoadType.has_value())
+      return op->emitOpError("schedule version value is incorrect");
+  }
+
+  return scheduleVersion;
+}
+
 void AffixTuningParameters::runOnOperation() {
   func::FuncOp func = getOperation();
   // currently, in rocMLIR we only support one Fusion Root per function.
@@ -131,30 +168,18 @@ void AffixTuningParameters::setUtilityKernelSizes(Value arg, T utilityOp) {
 void AffixTuningParameters::affixTuningParametersImpl(
     RockGemmWrapperInterface op) {
   OpBuilder b(op.getContext());
-  auto scheduleVersionAttrName = rock::ScheduleVersionAttr::getMnemonic();
   auto funcParent = op->getParentOfType<func::FuncOp>();
   std::string perfConfig;
-  if (funcParent->hasAttrOfType<rock::ScheduleVersionAttr>(
-          scheduleVersionAttrName) &&
-      op->hasAttrOfType<StringAttr>("perf_config")) {
-    op->emitError("kernel has both perf_config and schedule_version attribute "
-                  "set. Please modify schedule version directly inside "
-                  "perf_config and remove schedule_version\n");
-    signalPassFailure();
-    return;
-  }
   if (auto perfConfigAttr =
           op->template getAttrOfType<StringAttr>("perf_config")) {
     perfConfig = perfConfigAttr.getValue().str();
   }
-  // by default rocMLIR selects GEMM Schedule V1
-  auto scheduleVersion = 1;
-  if (funcParent->hasAttrOfType<rock::ScheduleVersionAttr>(
-          scheduleVersionAttrName)) {
-    scheduleVersion = dyn_cast<rock::ScheduleVersionAttr>(
-                          funcParent->removeAttr(scheduleVersionAttrName))
-                          .getScheduleVersion();
-  }
+  FailureOr<std::optional<int64_t>> maybeScheduleVersion =
+      getScheduleVersion(funcParent, op);
+  if (failed(maybeScheduleVersion))
+    return signalPassFailure();
+
+  std::optional<int64_t> scheduleVersion = maybeScheduleVersion.value();
 
   GemmFeatures features = rock::getFeatures(op);
   if (isAccel(features)) {
@@ -165,9 +190,9 @@ void AffixTuningParameters::affixTuningParametersImpl(
     // update schedule version to what is provided by the user if and only if
     // user hasn't provided perfConfig, otherwise just keep whatever is inside
     // perfConfig
-    if (!op->hasAttrOfType<StringAttr>("perf_config")) {
-      validParams.gemmScheduleVersion = scheduleVersion;
-    }
+    if (scheduleVersion.has_value())
+      validParams.gemmScheduleVersion = scheduleVersion.value();
+
     if (failed(status)) {
       // Try again if allowed.
       if (fallBackNoConfig) {
@@ -233,9 +258,8 @@ void AffixTuningParameters::affixTuningParametersImpl(
     // update schedule version to what is provided by the user if and only if
     // user hasn't provided perfConfig, otherwise just keep whatever was
     // obtained from perfConfig
-    if (!op->hasAttrOfType<StringAttr>("perf_config")) {
-      validParams.gemmScheduleVersion = scheduleVersion;
-    }
+    if (scheduleVersion.has_value())
+      validParams.gemmScheduleVersion = scheduleVersion.value();
 
     Attribute gemmParams = populateParams.getGemmParamsAttr(b, validParams);
     op.setGemmParamsAttr(gemmParams);
@@ -289,6 +313,13 @@ void AffixTuningParameters::affixTuningParametersImpl(
                  "with matrix accelerator extentions");
     return signalPassFailure();
   }
+  auto funcParent = op->getParentOfType<func::FuncOp>();
+  FailureOr<std::optional<int64_t>> maybeScheduleVersion =
+      getScheduleVersion(funcParent, op);
+  if (failed(maybeScheduleVersion))
+    return signalPassFailure();
+
+  std::optional<int64_t> scheduleVersion = maybeScheduleVersion.value();
 
   Attribute params0 = op.getGemm0Params().value_or(nullptr);
   // set a default one if params is not provided
@@ -305,6 +336,11 @@ void AffixTuningParameters::affixTuningParametersImpl(
     op.emitError("perf config string has an incorrect format.");
     return signalPassFailure();
   }
+
+  if (scheduleVersion.has_value())
+    attnPerfConfig =
+        attnPerfConfig.withScheduleVersion(scheduleVersion.value());
+
   GemmFeatures features = rock::getFeatures(op);
   RockAccelTuningParamAttrInterface accelParams0;
   if (bitEnumContainsAny(features, GemmFeatures::mfma)) {

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

@@ -466,24 +466,26 @@ struct BlockwiseGemmAccelRewritePattern
     // considered a temporary hack until we have a proper way of "searching"
     // through different schedules (either heuristically or automatically)
 
-    bool directToLDS = op.getDirectToLDS();
+    bool directToLDSA = op.getDirectToLDSA();
+    bool directToLDSB = op.getDirectToLDSB();
     Value wrappedLDSBufferForLoadA, wrappedLDSBufferForLoadB;
     if (loadAFromLDS) {
       wrappedLDSBufferForLoadA = accelEmitterPtr->wrapLDSBufferForLoad(
           b, loc, op.getMatrixA(), op.getBlockSize(), op.getInMPerThread(), "m",
-          op.getRotateMWithK(), directToLDS, op.getLdsLayoutMxK(),
+          op.getRotateMWithK(), directToLDSA, op.getLdsLayoutMxK(),
           op.getSplitKAcrossThreadsFirstA());
     }
     if (loadBFromLDS) {
       wrappedLDSBufferForLoadB = accelEmitterPtr->wrapLDSBufferForLoad(
           b, loc, op.getMatrixB(), op.getBlockSize(), op.getInNPerThread(), "n",
-          op.getRotateNWithK(), directToLDS, op.getLdsLayoutNxK(),
+          op.getRotateNWithK(), directToLDSB, op.getLdsLayoutNxK(),
           op.getSplitKAcrossThreadsFirstB());
     }
 
     auto loadBuffer = [&](Value buffer, Value wrappedLDSBufferForLoad,
                           Value loopVar, Type argType, int64_t repeats,
-                          bool loadFromLDS, bool isA) -> Value {
+                          bool loadFromLDS, bool directToLDS,
+                          bool isA) -> Value {
       Value inputBuffer = buffer;
       SmallVector<int64_t> shape;
       if (directToLDS) {
@@ -545,7 +547,7 @@ struct BlockwiseGemmAccelRewritePattern
 
       Value bufferA = adaptor.getBufferA();
       bufferA = loadBuffer(bufferA, wrappedLDSBufferForLoadA, i, argTypeA,
-                           mRepeats, loadAFromLDS, true);
+                           mRepeats, loadAFromLDS, directToLDSA, true);
       Value viewA =
           accelEmitterPtr->generateThreadwiseViewBufferA(b, loc, bufferA);
 
@@ -557,7 +559,7 @@ struct BlockwiseGemmAccelRewritePattern
 
         Value bufferB = adaptor.getBufferB();
         bufferB = loadBuffer(bufferB, wrappedLDSBufferForLoadB, j, argTypeB,
-                             nRepeats, loadBFromLDS, false);
+                             nRepeats, loadBFromLDS, directToLDSB, false);
         Value viewB =
             accelEmitterPtr->generateThreadwiseViewBufferB(b, loc, bufferB);
 

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

@@ -76,13 +76,14 @@ class LoweringBlockwiseLoadTileOp final
       const std::unique_ptr<rock::accel::AccelEmitter> &accelEmitterPtr,
       Value tid, StringRef dName, Value ldsView, Value regs, int64_t blockSize,
       int64_t inDPerThread, bool rotateDWithK, bool forceUnroll,
-      bool directToLDS, bool ldsLayoutDxK) const {
+      bool directToLDS, bool ldsLayoutDxK,
+      bool doSplitKAcrossThreadsFirst) const {
 
     // wrapLDSBufferForLoad is reading a single set of Ks into private memory
     // A/B[m/n, 0:kBasePerThread]
     Value ldsViewForLoad = accelEmitterPtr->wrapLDSBufferForLoad(
         b, loc, ldsView, blockSize, inDPerThread, dName, rotateDWithK,
-        directToLDS, ldsLayoutDxK);
+        directToLDS, ldsLayoutDxK, doSplitKAcrossThreadsFirst);
 
     // We enhance the transformation from wrapLDSBufferForLoad using a builder
     // that, given a single index, splits it into "m"("n") and "k" and lets
@@ -158,6 +159,7 @@ class LoweringBlockwiseLoadTileOp final
     bool doRotateWithK = op.getRotateWithK();
     bool doSwapThreadIterSubDims = op.getSwapThreadIterSubDims();
     bool ldsLayoutDxK = op.getLDSLayoutDxK();
+    bool doSplitKAcrossThreadsFirst = op.getSplitKAcrossThreadsFirst();
     LDSLayoutConfigDim ldsLayoutConfig{doRotateWithK, doSwapThreadIterSubDims,
                                        ldsLayoutDxK};
 
@@ -406,7 +408,8 @@ class LoweringBlockwiseLoadTileOp final
           generateReadLoop(loc, b, accelEmitterPtr, tid, dName, ldsViewForGemm,
                            destRegisters, blockSize, copyDPerThread,
                            ldsLayoutConfig.doRotateWithK, forceUnroll,
-                           directToLDS, ldsLayoutConfig.ldsLayoutDxK);
+                           directToLDS, ldsLayoutConfig.ldsLayoutDxK,
+                           doSplitKAcrossThreadsFirst);
           if (stageLDSReadNew)
             rock::YieldOp::create(b, loc);
         }