Specialize problem string for attention and gemm+gemm

Author: dhernandez0

mlir/lib/Dialect/Rock/Tuning/RockTuningImpl.cpp

@@ -582,22 +582,27 @@ getTuningProblemStr(RockGemmGemmWrapperInterface gemmGemmOp,
   ArrayRef<int64_t> vShape = cast<MemRefType>(gemmGemmOp.getCType()).getShape();
   int64_t g = qShape[0];
 
+  bool isAttention = isa<AttentionOp>(gemmGemmOp);
+
   Type elemTypeQ = cast<MemRefType>(gemmGemmOp.getAType()).getElementType();
   problemOS << "-t ";
   if (elemTypeQ.isF32()) {
     problemOS << "f32" << sep;
-  } else if (elemTypeQ.isF16()) {
+  } else if (elemTypeQ.isF16() && isAttention) {
     problemOS << "f16" << sep;
-  } else if (elemTypeQ.isBF16()) {
+  } else if (elemTypeQ.isBF16() && isAttention) {
     problemOS << "bf16" << sep;
-  } else if (elemTypeQ.isInteger(8)) {
+  } else if (elemTypeQ.isInteger(8) && isAttention) {
     problemOS << "i8" << sep;
   } else {
     return gemmGemmOp.emitError("invalid type:") << elemTypeQ << "\n";
   }
 
   // TransQ
-  problemOS << "-transQ ";
+  if (isAttention)
+    problemOS << "-transQ ";
+  else
+    problemOS << "-transA ";
   if (gemmGemmOp.getTransposedA()) {
     seqLenQ = qShape[2];
     headDimQK = qShape[1];
@@ -609,7 +614,10 @@ getTuningProblemStr(RockGemmGemmWrapperInterface gemmGemmOp,
   }
 
   // TransK
-  problemOS << "-transK ";
+  if (isAttention)
+    problemOS << "-transK ";
+  else
+    problemOS << "-transB ";
   if (gemmGemmOp.getTransposedB()) {
     seqLenK = kShape[1];
     problemOS << "true" << sep;
@@ -619,7 +627,10 @@ getTuningProblemStr(RockGemmGemmWrapperInterface gemmGemmOp,
   }
 
   // TransV
-  problemOS << "-transV ";
+  if (isAttention)
+    problemOS << "-transV ";
+  else
+    problemOS << "-transC ";
   if (gemmGemmOp.getTransposedC()) {
     headDimV = vShape[1];
     problemOS << "true" << sep;
@@ -636,10 +647,17 @@ getTuningProblemStr(RockGemmGemmWrapperInterface gemmGemmOp,
     problemOS << "false" << sep;
 
   problemOS << "-g " << g << sep;
-  problemOS << "-seq_len_q " << seqLenQ << sep;
-  problemOS << "-seq_len_k " << seqLenK << sep;
-  problemOS << "-head_dim_qk " << headDimQK << sep;
-  problemOS << "-head_dim_v " << headDimV;
+  if (isAttention) {
+    problemOS << "-seq_len_q " << seqLenQ << sep;
+    problemOS << "-seq_len_k " << seqLenK << sep;
+    problemOS << "-head_dim_qk " << headDimQK << sep;
+    problemOS << "-head_dim_v " << headDimV;
+  } else {
+    problemOS << "-m " << seqLenQ << sep;
+    problemOS << "-n " << seqLenK << sep;
+    problemOS << "-k " << headDimQK << sep;
+    problemOS << "-gemmO " << headDimV;
+  }
   return success();
 }