Addressing PR comments

Author: dhernandez0

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

@@ -93,7 +93,7 @@ def RockAccelTuningParamAttrInterface : AttrInterface<"RockAccelTuningParamAttrI
       >,
     InterfaceMethod<
         /*desc=*/[{
-          Return param passed to the backend compiler: waves_per_eu.
+          Return waves_per_eu, this is a hint to the backend compiler to tune the number of wavefronts that are capable of fitting within the resources of an EU.
         }],
         /*retType=*/"int64_t",
         /*methodName=*/"getWavesPerEU",

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

@@ -464,9 +464,25 @@ def Rock_ScheduleVersionAttr : Rock_Attr<"ScheduleVersion"> {
   }];
 }
 
-// It is a temporary attribute
 def Rock_EnableSplitKForTuning : Rock_Attr<"EnableSplitKForTuning"> {
   let mnemonic = "enable_splitk_for_tuning";
+  let description = [{
+    Whether we tune for split-k. If unset, split-k=1.
+  }];
+}
+
+def Rock_WavesPerEU : Rock_Attr<"WavesPerEU"> {
+  let mnemonic = "waves_per_eu";
+  let description = [{
+    This is a hint to the backend compiler to tune the number of wavefronts that are capable of fitting within the resources of an EU.
+  }];
+}
+
+def Rock_OutputSwizzle : Rock_Attr<"OutputSwizzle"> {
+  let mnemonic = "output_swizzle";
+  let description = [{
+    Whether we run the output swizzle pass. 0 -> disabled, 1 -> enabled, 2 -> heuristic.
+  }];
 }
 
 def Rock_PrefillAttr : Rock_Attr<"Prefill"> {

mlir/lib/Conversion/RockToGPU/RockToGPU.cpp

@@ -204,9 +204,9 @@ void LowerRockOpsToGPUPass::runOnOperation() {
     gridSize = cast<IntegerAttr>(gridSizeAttr).getInt();
     gpuFunc.setKnownGridSizeAttr(b.getDenseI32ArrayAttr({gridSize, 1, 1}));
 
-    auto wavesPerEUAttr = theFunc->getAttr("waves_per_eu");
+    auto wavesPerEUAttr = theFunc->getAttr(rock::WavesPerEUAttr::getMnemonic());
     if (wavesPerEUAttr) {
-      gpuFunc->setAttr("waves_per_eu", wavesPerEUAttr);
+      gpuFunc->setAttr(rock::WavesPerEUAttr::getMnemonic(), wavesPerEUAttr);
     }
 
     FailureOr<StringAttr> maybeArch = rock::getArch(theFunc);
@@ -397,9 +397,9 @@ void LowerRockOpsToGPUPass::runOnOperation() {
                        b.getI32IntegerAttr(ldsUsage));
     }
     // if waves_per_eu is set, use it
-    if (gpuFunc->hasAttrOfType<IntegerAttr>("waves_per_eu")) {
+    if (gpuFunc->hasAttrOfType<IntegerAttr>(rock::WavesPerEUAttr::getMnemonic())) {
       int64_t wavesPerEU =
-          gpuFunc->getAttrOfType<IntegerAttr>("waves_per_eu").getInt();
+          gpuFunc->getAttrOfType<IntegerAttr>(rock::WavesPerEUAttr::getMnemonic()).getInt();
       // zero means, use heuristic
       if (wavesPerEU != 0) {
         gpuFunc->setAttr("rocdl.waves_per_eu", b.getI32IntegerAttr(wavesPerEU));

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

@@ -1981,7 +1981,7 @@ struct GridwiseAttentionAccelRewritePattern
     rock::accel::AccelEmitterParams accelParamsGemm1 =
         accelEmitterPtrGemm1->getParams();
 
-    // some params are needed in future passes, add them to func as attributes
+    // wavesPerEU is needed in RockToGPU pass and OutputSwizzle for the OutputSwizzle pass. We add them as func attributes.
     assert(gemm0TuningParams.getWavesPerEU() ==
            gemm1TuningParams.getWavesPerEU());
     assert(gemm0TuningParams.getOutputSwizzle() ==
@@ -1991,8 +1991,8 @@ struct GridwiseAttentionAccelRewritePattern
     IntegerAttr outputSwizzleAttr =
         rewriter.getI64IntegerAttr(gemm0TuningParams.getOutputSwizzle());
     func::FuncOp funcOp = cast<func::FuncOp>(op->getParentOp());
-    funcOp->setAttr("waves_per_eu", wavesPerEUAttr);
-    funcOp->setAttr("output_swizzle", outputSwizzleAttr);
+    funcOp->setAttr(rock::WavesPerEUAttr::getMnemonic(), wavesPerEUAttr);
+    funcOp->setAttr(rock::OutputSwizzleAttr::getMnemonic(), outputSwizzleAttr);
 
     // Get current workgroup ID.
     auto bid = WorkgroupIdOp::create(rewriter, loc, rewriter.getIndexType());
@@ -3037,14 +3037,14 @@ struct GridwiseGemmAccelRewritePattern
          gridGroupSize},
         arch);
 
-    // some params are needed in future passes, add them to func as attributes
+    // wavesPerEU is needed in RockToGPU pass and OutputSwizzle for the OutputSwizzle pass. We add them as func attributes.
     IntegerAttr wavesPerEUAttr =
         b.getI64IntegerAttr(tuningParams.getWavesPerEU());
     IntegerAttr outputSwizzleAttr =
         b.getI64IntegerAttr(tuningParams.getOutputSwizzle());
     func::FuncOp funcOp = cast<func::FuncOp>(op->getParentOp());
-    funcOp->setAttr("waves_per_eu", wavesPerEUAttr);
-    funcOp->setAttr("output_swizzle", outputSwizzleAttr);
+    funcOp->setAttr(rock::WavesPerEUAttr::getMnemonic(), wavesPerEUAttr);
+    funcOp->setAttr(rock::OutputSwizzleAttr::getMnemonic(), outputSwizzleAttr);
 
     LDSLayoutConfigDim ldsLayoutConfigA = getLDSLayoutConfigDim(
         elementTypeA, kpack, maybeVecDimInfoA.value(), directToLDS);

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

@@ -416,10 +416,10 @@ void RockOutputSwizzlePass::runOnOperation() {
 
   bool useHeuristic = true;
   bool enableOutputSwizzle = false;
-  if (func->hasAttrOfType<IntegerAttr>("output_swizzle")) {
+  if (func->hasAttrOfType<IntegerAttr>(rock::OutputSwizzleAttr::getMnemonic())) {
     // 0 -> disabled, 1 -> enabled, 2 -> heuristic
     int64_t outputSwizzle =
-        func->getAttrOfType<IntegerAttr>("output_swizzle").getInt();
+        func->getAttrOfType<IntegerAttr>(rock::OutputSwizzleAttr::getMnemonic()).getInt();
     useHeuristic = outputSwizzle == 2;
     enableOutputSwizzle = outputSwizzle == 1;
   }