Introduce perfConfig V3 with param to select different schedule (#1767)

This PR adds place holder for selecting different GEMM schedule as tuning param.

Also adds param for output swizzle which may be used later.

Author: umangyadav

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

@@ -225,6 +225,9 @@ def Rock_GeneralGemmParamsAttr : Rock_Attr<"GeneralGemmParams", [RockTuningParam
     - mPerThread: The number of values of m to process as a unit on each thread
     - nPerThread: The number of values of n to process as a unit on each thread
     - kpack: The number of values of k to pack contiguously into the shared buffer
+    - splitKFactor:  Split-k factor for the Split-k GEMM algorithm
+    - scheduleVersion: Param to select GEMM schedule
+    - outputSwizzle: Whether to enable/disable output swizzle or use heuristics 
   }];
   let parameters = (ins
     "uint32_t":$blockSize,
@@ -235,18 +238,23 @@ def Rock_GeneralGemmParamsAttr : Rock_Attr<"GeneralGemmParams", [RockTuningParam
     "int64_t":$mPerThread,
     "int64_t":$nPerThread,
     "int64_t":$kpack,
-    "int64_t":$splitKFactor
+    "int64_t":$splitKFactor,
+    "int64_t":$scheduleVersion,
+    "int64_t":$outputSwizzle
   );
 
   let extraClassDeclaration = [{
     void getPerfConfigStr(::llvm::SmallVectorImpl<char> &perfStr) {
-      ("v2:" + Twine(getBlockSize()) + ","
+      ("v3:" + Twine(getBlockSize()) + ","
       + Twine(getMPerBlock()) + ","
       + Twine(getNPerBlock()) + ","
       + Twine(getKPerBlock()) + ","
       + Twine(getMPerThread()) + ","
       + Twine(getNPerThread()) + ","
-      + Twine(getSplitKFactor())).toVector(perfStr);
+      + Twine(getSplitKFactor()) + ","
+      + Twine(getScheduleVersion()) + ","
+      + Twine(getOutputSwizzle())) 
+      .toVector(perfStr);
     }
     bool getForceUnroll() { return true; }
   }];
@@ -286,6 +294,8 @@ def Rock_AttnPerfConfig : Rock_Attr<"AttnPerfConfig", [RockTuningParamAttrInterf
     }
 
     int64_t getSplitKFactor() { return 1; }
+    int64_t getScheduleVersion() { return 1;}
+    int64_t getOutputSwizzle() { return 2; }
   }];
 
   let builders = [
@@ -310,18 +320,22 @@ def Rock_XdlopsGemmParamsAttr : Rock_Attr<"XdlopsGemmParams", [RockTuningParamAt
     "int64_t":$mPerWave,
     "int64_t":$mnPerXdl,
     "int64_t":$splitKFactor,
+    "int64_t":$scheduleVersion,
+    "int64_t":$outputSwizzle,
     "bool":$forceUnroll
   );
 
   let extraClassDeclaration = [{
     void getPerfConfigStr(::llvm::SmallVectorImpl<char> &perfStr) {
-        ("v2:" + Twine(getMPerBlock()) + ","
+        ("v3:" + Twine(getMPerBlock()) + ","
         + Twine(getNPerBlock()) + ","
         + Twine(getKpackPerBlock()) + ","
         + Twine(getMPerWave()) + ","
         + Twine(getMnPerXdl()) + ","
         + Twine(getKpack()) + ","
         + Twine(getSplitKFactor()) + ","
+        + Twine(getScheduleVersion()) + ","
+        + Twine(getOutputSwizzle()) + ","
         + Twine(getForceUnroll()) + ","
           + "1") /* *ThreadCopyMore* */
         .toVector(perfStr);
@@ -347,18 +361,22 @@ def Rock_XdlopsGemmDerivedParamsAttr : Rock_Attr<"XdlopsGemmDerivedParams", [Roc
     "int64_t":$nPerWave,
     "int64_t":$mnPerXdl,
     "int64_t":$splitKFactor,
+    "int64_t":$scheduleVersion,
+    "int64_t":$outputSwizzle,
     "bool":$forceUnroll
   );
 
   let extraClassDeclaration = [{
     void getPerfConfigStr(::llvm::SmallVectorImpl<char> &perfStr) {
-        ("v2:" + Twine(getMPerBlock()) + ","
+        ("v3:" + Twine(getMPerBlock()) + ","
         + Twine(getNPerBlock()) + ","
         + Twine(getKpackPerBlock()) + ","
         + Twine(getMPerWave()) + ","
         + Twine(getNPerWave()) + ","
         + Twine(getKpack()) + ","
         + Twine(getSplitKFactor()) + ","
+        + Twine(getScheduleVersion()) + ","
+        + Twine(getOutputSwizzle()) + ","
         + Twine(getForceUnroll()) + ","
           + "1") /* *ThreadCopyMore* */
         .toVector(perfStr);
@@ -386,6 +404,8 @@ def Rock_XdlopsGemmDerivedParamsAttr : Rock_Attr<"XdlopsGemmDerivedParams", [Roc
         nPerWave,
         mnPerXdl,
         params.getSplitKFactor(),
+        params.getScheduleVersion(),
+        params.getOutputSwizzle(),
         params.getForceUnroll()
       );
     }]>
@@ -409,18 +429,22 @@ def Rock_WmmaGemmParamsAttr : Rock_Attr<"WmmaGemmParams", [RockTuningParamAttrIn
     "int64_t":$mPerWave,
     "int64_t":$nPerWave,
     "int64_t":$splitKFactor,
+    "int64_t":$scheduleVersion,
+    "int64_t":$outputSwizzle,
     "bool":$forceUnroll
   );
 
   let extraClassDeclaration = [{
     void getPerfConfigStr(SmallVectorImpl<char> &perfStr) {
-        ("v2:" +  Twine(getMPerBlock()) + ","
+        ("v3:" +  Twine(getMPerBlock()) + ","
         + Twine(getNPerBlock()) + ","
         + Twine(getKpackPerBlock()) + ","
         + Twine(getMPerWave()) + ","
         + Twine(getNPerWave()) + ","
         + Twine(getKpack()) + ","
         + Twine(getSplitKFactor()) + ","
+        + Twine(getScheduleVersion()) + ","
+        + Twine(getOutputSwizzle()) + ","
         + Twine(getForceUnroll()) + ","
           + "1") /* *ThreadCopyMore* */
         .toVector(perfStr);

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

@@ -69,7 +69,29 @@ def RockTuningParamAttrInterface : AttrInterface<"RockTuningParamAttrInterface">
         /*methodBody=*/"",
         /*defaultImplementation=*/""
       >,
-
+    InterfaceMethod<
+        /*desc=*/[{
+          Returns value 0 or 1 to decide if rocMLIR should enable output swizzle or not. 
+          Return value 2 means rocMLIR should use heuristics to make this decision
+        }],
+        /*retType=*/"int64_t",
+        /*methodName=*/"getOutputSwizzle",
+        /*args=*/(ins),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/""
+      >,
+    InterfaceMethod<
+        /*desc=*/[{
+          Returns version of the schedule used by the underlying GEMM algorithm. 
+          Return value of 1 means it should use 3 stage schedule with II=2. 
+          Return value of 2 means it should 4 stage schedule with II=1. 
+        }],
+        /*retType=*/"int64_t",
+        /*methodName=*/"getScheduleVersion",
+        /*args=*/(ins),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/""
+      >,
     // TODO: more methods here as needed
   ];
 

mlir/include/mlir/Dialect/Rock/Tuning/GridwiseGemmParams.h

@@ -101,24 +101,30 @@ struct InitParamsNonAccel : InitParams, Serializable<InitParamsNonAccel> {
   int64_t gemmNPerThread;
   uint32_t blockSize;
   int64_t splitKFactor;
+  int64_t gemmScheduleVersion;
+  int64_t outputSwizzle;
 
   constexpr InitParamsNonAccel(uint32_t bSize, int64_t mPerBlock,
                                int64_t nPerBlock, int64_t kPerBlock,
                                int64_t mPerThread, int64_t nPerThread,
-                               int64_t splitKFactor)
+                               int64_t splitKFactor, int64_t scheduleVersion,
+                               int64_t outputSwizzle)
       : InitParams{mPerBlock, nPerBlock, kPerBlock}, gemmMPerThread(mPerThread),
         gemmNPerThread(nPerThread), blockSize(bSize),
-        splitKFactor(splitKFactor) {}
+        splitKFactor(splitKFactor), gemmScheduleVersion(scheduleVersion),
+        outputSwizzle(outputSwizzle) {}
 
   constexpr InitParamsNonAccel()
-      : InitParamsNonAccel(0U, 0LL, 0LL, 0LL, 0LL, 0LL, 1LL) {}
+      : InitParamsNonAccel(0U, 0LL, 0LL, 0LL, 0LL, 0LL, 1LL, 1LL, 2LL) {}
 
   InitParamsNonAccel(GeneralGemmParamsAttr attr)
       : InitParams{attr.getMPerBlock(), attr.getNPerBlock(),
                    attr.getKPerBlock()},
         gemmMPerThread(attr.getMPerThread()),
         gemmNPerThread(attr.getNPerThread()), blockSize(attr.getBlockSize()),
-        splitKFactor(attr.getSplitKFactor()){};
+        splitKFactor(attr.getSplitKFactor()),
+        gemmScheduleVersion(attr.getScheduleVersion()),
+        outputSwizzle(attr.getOutputSwizzle()){};
 
   int64_t getKPack() { return 1; }
 
@@ -130,32 +136,41 @@ struct InitParamsNonAccel : InitParams, Serializable<InitParamsNonAccel> {
     f(self.gemmKPerBlock);
     f(self.gemmMPerThread);
     f(self.gemmNPerThread);
-    if (self.version != Version::V1)
+    if (self.version >= Version::V2)
       f(self.splitKFactor);
+    if (self.version >= Version::V3) {
+      f(self.gemmScheduleVersion);
+      f(self.outputSwizzle);
+    }
   }
 };
 
 struct InitParamsAccel : InitParams, Serializable<InitParamsAccel> {
   constexpr InitParamsAccel(int64_t mPerBlock, int64_t nPerBlock,
                             int64_t kPerBlock, int64_t mPerWave,
                             int64_t nPerWaveOrMnPerXdl, int64_t kPack,
-                            int64_t splitKFactor, bool aThreadCopyMoreGemmK,
+                            int64_t splitKFactor, int64_t scheduleVersion,
+                            int64_t outputSwizzle, bool aThreadCopyMoreGemmK,
                             bool bThreadCopyMoreGemmKPack)
       : InitParams{mPerBlock, nPerBlock, kPerBlock}, gemmMPerWave(mPerWave),
         gemmNPerWaveOrMnPerXdl(nPerWaveOrMnPerXdl), gemmKPack(kPack),
-        splitKFactor(splitKFactor),
+        splitKFactor(splitKFactor), gemmScheduleVersion(scheduleVersion),
+        outputSwizzle(outputSwizzle),
         gemmAThreadCopyMoreGemmK(aThreadCopyMoreGemmK),
         gemmBThreadCopyMoreGemmKPack(bThreadCopyMoreGemmKPack) {}
 
   constexpr InitParamsAccel()
-      : InitParamsAccel(0LL, 0LL, 0LL, 0LL, 0LL, 0LL, 1LL, false, false) {}
+      : InitParamsAccel(0LL, 0LL, 0LL, 0LL, 0LL, 0LL, 1LL, 1LL, 2LL, false,
+                        false) {}
 
   InitParamsAccel(XdlopsGemmParamsAttr attr)
       : InitParams{attr.getMPerBlock(), attr.getNPerBlock(),
                    attr.getKpackPerBlock()},
         gemmMPerWave(attr.getMPerWave()),
         gemmNPerWaveOrMnPerXdl(attr.getMnPerXdl()), gemmKPack(attr.getKpack()),
         splitKFactor(attr.getSplitKFactor()),
+        gemmScheduleVersion(attr.getScheduleVersion()),
+        outputSwizzle(attr.getOutputSwizzle()),
         gemmAThreadCopyMoreGemmK(attr.getForceUnroll()),
         gemmBThreadCopyMoreGemmKPack(false){};
 
@@ -165,6 +180,8 @@ struct InitParamsAccel : InitParams, Serializable<InitParamsAccel> {
         gemmMPerWave(attr.getMPerWave()),
         gemmNPerWaveOrMnPerXdl(attr.getNPerWave()), gemmKPack(attr.getKpack()),
         splitKFactor(attr.getSplitKFactor()),
+        gemmScheduleVersion(attr.getScheduleVersion()),
+        outputSwizzle(attr.getOutputSwizzle()),
         gemmAThreadCopyMoreGemmK(attr.getForceUnroll()),
         gemmBThreadCopyMoreGemmKPack(false){};
 
@@ -174,6 +191,8 @@ struct InitParamsAccel : InitParams, Serializable<InitParamsAccel> {
   int64_t gemmNPerWaveOrMnPerXdl;
   int64_t gemmKPack;
   int64_t splitKFactor;
+  int64_t gemmScheduleVersion;
+  int64_t outputSwizzle;
   bool gemmAThreadCopyMoreGemmK;
   bool gemmBThreadCopyMoreGemmKPack;
 
@@ -185,9 +204,13 @@ struct InitParamsAccel : InitParams, Serializable<InitParamsAccel> {
     f(self.gemmMPerWave);
     f(self.gemmNPerWaveOrMnPerXdl);
     f(self.gemmKPack);
-    if (self.version != Version::V1) {
+    if (self.version >= Version::V2) {
       f(self.splitKFactor);
     }
+    if (self.version >= Version::V3) {
+      f(self.gemmScheduleVersion);
+      f(self.outputSwizzle);
+    }
     f(self.gemmAThreadCopyMoreGemmK);
     f(self.gemmBThreadCopyMoreGemmKPack);
   }
@@ -206,14 +229,15 @@ class BasePopulateParams {
 private:
   struct InitParamData {
     InitParamType paramSet;
-    size_t original_pos;
-    int64_t padding_amount;
+    size_t originalPos;
+    int64_t paddingAmount;
 
     bool operator<(const InitParamData &rhs) const {
-      if (this->padding_amount < rhs.padding_amount) {
+      if (this->paddingAmount < rhs.paddingAmount) {
         return true;
-      } else if (this->padding_amount == rhs.padding_amount) {
-        return (this->original_pos < rhs.original_pos);
+      }
+      if (this->paddingAmount == rhs.paddingAmount) {
+        return (this->originalPos < rhs.originalPos);
       }
       return false;
     }
@@ -227,10 +251,10 @@ class BasePopulateParams {
       InitParamType paramSet = initParams[pos];
       InitParamData paramData;
       paramData.paramSet = paramSet;
-      paramData.original_pos = pos;
-      paramData.padding_amount = calculatePaddingAmount(paramSet, gemmSize);
-      assert(paramData.original_pos >= 0);
-      assert(paramData.padding_amount >= 0);
+      paramData.originalPos = pos;
+      paramData.paddingAmount = calculatePaddingAmount(paramSet, gemmSize);
+      assert(paramData.originalPos >= 0);
+      assert(paramData.paddingAmount >= 0);
       res.push_back(paramData);
     }
     return res;