Pipelining for attention (#1990)

Author: dhernandez0

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

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

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/lib/Dialect/Rock/IR/AmdArchDb.cpp

@@ -413,3 +413,8 @@ GemmFeatures mlir::rock::AmdArchInfo::getDefaultFeatures(Type dataType) {
   }
   return theseFeatures;
 }
+
+bool mlir::rock::isDirectToLDSSupported(GemmFeatures features) {
+  return bitEnumContainsAll(features, GemmFeatures::direct_to_lds_128b) ||
+         bitEnumContainsAll(features, GemmFeatures::direct_to_lds_32b);
+}

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.
@@ -128,33 +165,37 @@ void AffixTuningParameters::setUtilityKernelSizes(Value arg, T utilityOp) {
   funcOp->setAttr("grid_size", gridSizeAttr);
 }
 
+static LogicalResult isScheduleVersionSupported(int64_t scheduleVersion,
+                                                GemmFeatures features) {
+  std::optional<GemmLoadTileType> maybeLoadType =
+      rock::symbolizeGemmLoadTileType(scheduleVersion);
+  if (!maybeLoadType.has_value())
+    return failure();
+
+  auto loadType = maybeLoadType.value();
+  bool directToLDS = loadType == GemmLoadTileType::DirectToLDSDefault ||
+                     loadType == GemmLoadTileType::DirectToLDSDoubleBuffer;
+  if (directToLDS && !isDirectToLDSSupported(features))
+    return failure();
+
+  return success();
+}
+
 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 +206,15 @@ 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(isScheduleVersionSupported(validParams.gemmScheduleVersion,
+                                          features))) {
+      op->emitError("schedule version not supported\n");
+      return signalPassFailure();
     }
+
     if (failed(status)) {
       // Try again if allowed.
       if (fallBackNoConfig) {
@@ -233,9 +280,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 +335,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 +358,17 @@ void AffixTuningParameters::affixTuningParametersImpl(
     op.emitError("perf config string has an incorrect format.");
     return signalPassFailure();
   }
+
+  if (scheduleVersion.has_value())
+    attnPerfConfig =
+        attnPerfConfig.withScheduleVersion(scheduleVersion.value());
+
+  if (failed(isScheduleVersionSupported(attnPerfConfig.getScheduleVersion(),
+                                        rock::getFeatures(op)))) {
+    op->emitError("schedule version not supported\n");
+    return signalPassFailure();
+  }
+
   GemmFeatures features = rock::getFeatures(op);
   RockAccelTuningParamAttrInterface accelParams0;
   if (bitEnumContainsAny(features, GemmFeatures::mfma)) {