Add Gemm+Elementwise+Gemm support

Author: dhernandez0

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

@@ -49,12 +49,6 @@ class FusionRoot : public TraitBase<ConcreteType, FusionRoot> {};
 } // namespace OpTrait
 } // namespace mlir
 
-// Following ifdef could be used to change
-// the attention operator to be a fused gemm-gemm
-// kernel for debugging purposes. This will also
-// adjust the test harness to verify the same as well
-// #define ROCK_DEBUG_ATTENTION_REMOVE_SOFTMAX
-
 namespace mlir {
 namespace rock {
 //===----------------------------------------------------------------------===//

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

@@ -63,11 +63,12 @@ def KernelTypeConvBwdData : I32EnumAttrCase<"ConvBwdData", 1>;
 def KernelTypeConvBwdWeight : I32EnumAttrCase<"ConvBwdWeight", 2>;
 def KernelTypeGemm : I32EnumAttrCase<"Gemm", 3>;
 def KernelTypeAttention : I32EnumAttrCase<"Attention", 4>;
+def KernelTypeGemmElementwiseGemm : I32EnumAttrCase<"GemmElementwiseGemm", 5>;
 
 def KernelType : Rock_I32Enum<"KernelType", "Any of the possible types of a rock kernel",
   [KernelTypeConv, KernelTypeConvBwdData,
     KernelTypeConvBwdWeight, KernelTypeGemm,
-    KernelTypeAttention]>;
+    KernelTypeAttention, KernelTypeGemmElementwiseGemm]>;
 
 /// TransformType
 def PassThrough : I32EnumAttrCase<"PassThrough", 0>;

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

@@ -236,7 +236,7 @@ def Rock_AttentionOp :
     and the number of compute units (numCu) available. The parameters
     `gridSize`, and `blockSize` are optional as they can be inferred by
     a tuning process or a heuristic, but they must be set before the `attention` is
-    lowered into the `gridwise_attention` stage of the code generation pipeline.
+    lowered into the `gridwise_attention_accel` stage of the code generation pipeline.
 
     `features` specifies what hardware features can be used in the generated code.
   }];
@@ -255,6 +255,54 @@ def Rock_AttentionOp :
   }];
 }
 
+def Rock_GemmElementwiseGemmOp:
+  Rock_Op<"gemm_elementwise_gemm", [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>, RockFusionRoot]>,
+  Arguments<(ins
+    TensorOrMemRefOf<GemmInputTypes>:$a,
+    TensorOrMemRefOf<GemmInputTypes>:$b,
+    TensorOrMemRefOf<GemmInputTypes>:$c,
+    Variadic<AnyTensorOrMemRef>:$elemwiseInputs,
+    TensorOrMemRefOf<GemmAccumulatorTypes>:$out,
+    UnitAttr:$aTransposed,
+    UnitAttr:$bTransposed,
+    UnitAttr:$cTransposed,
+    UnitAttr:$oTransposed,
+    StrAttr:$arch,
+    Rock_GemmFeaturesAttr:$features,
+    OptionalAttr<I32Attr>:$numCU,
+    OptionalAttr<RockTuningParamAttrInterface>:$params0,
+    OptionalAttr<RockTuningParamAttrInterface>:$params1,
+    I32Attr:$firstGemmIdx
+  )>,
+  Results<(outs Optional<TensorOf<[F32, F16, BF16]>>:$result)> {
+  let summary = "GEMM-elementwise-GEMM operation";
+  let description = [{
+    Performs the operation out = (a * b) * c.
+
+    This operation performs fused GEMM-elementwise-GEMM. 
+
+    Those creating a `rock.gemm_elementwise_gemm` must specify the GPU architecture being targetted
+    and the number of compute units (numCu) available. The parameters
+    `gridSize`, and `blockSize` are optional as they can be inferred by
+    a tuning process or a heuristic, but they must be set before the `gemm_elementwise_gemm` is
+    lowered into the `gridwise_attention_accel` stage of the code generation pipeline.
+
+    `features` specifies what hardware features can be used in the generated code.
+  }];
+  let hasVerifier = 1;
+  let regions = (region AnyRegion:$preSecondGemmBody);
+  let assemblyFormat = [{
+    `{` `\n`
+        ` ` `ab` `=` (`tr` $aTransposed^)? $a `*` (`tr` $bTransposed^)? $b `:` type($a) `,` type($b) `\n`
+        (`ab` `=` `elementwise` (`otherIns` `(` $elemwiseInputs^ `:` type($elemwiseInputs) `)`)? $preSecondGemmBody^ `\n`)?
+        (`tr` $oTransposed^)? $out `=` `softmax` `(` `qk` `)` `*` (`tr` $cTransposed^)? $c `:` type($c) `->` type($out) `\n`
+    `}` attr-dict (`->` type($result)^)?
+  }];
+  let extraClassDeclaration = [{
+    ::mlir::OpOperand* getOutArgument() { return &(*this)->getOpOperands().back(); }
+  }];
+}
+
 def Rock_InitKernelOp :
     Rock_Op<"init_kernel", []>,
     Arguments<(ins AnyTensorOrMemRef:$buffer,
@@ -434,12 +482,12 @@ def Rock_GridwiseGemmAccelOp :
 // gridwise_attention_accel
 def Rock_GridwiseAttentionAccelOp :
     Rock_Op<"gridwise_attention_accel", [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>, RockFusionRoot, AttrSizedOperandSegments]>,
-    Arguments<(ins MemRefRankOf<[F32, F16, BF16, I8], [3]>:$queries,
-                   MemRefRankOf<[F32, F16, BF16, I8], [3]>:$keys,
-                   MemRefRankOf<[F32, F16, BF16,], [3]>:$values,
+    Arguments<(ins MemRefRankOf<GemmInputTypes, [3]>:$queries,
+                   MemRefRankOf<GemmInputTypes, [3]>:$keys,
+                   MemRefRankOf<GemmInputTypes, [3]>:$values,
                    Variadic<AnyTensorOrMemRef>:$preSoftmaxElemWiseInputs,
                    Optional<MemRefRankOf<[I32], [1]>>:$currentSeqLen,
-                   MemRefRankOf<[F32, F16, BF16], [3]>:$out,
+                   MemRefRankOf<GemmAccumulatorTypes, [3]>:$out,
                    StrAttr:$arch,
                    Rock_GemmFeaturesAttr:$features,
                    I32Attr:$blockSize,
@@ -449,7 +497,8 @@ def Rock_GridwiseAttentionAccelOp :
                    OptionalAttr<IndexAttr>:$prePadG0N,
                    RockAccelTuningParamAttrInterface:$params0,
                    RockAccelTuningParamAttrInterface:$params1,
-                   I32Attr:$firstGemmIdx)> {
+                   I32Attr:$firstGemmIdx,
+                   DefaultValuedOptionalAttr<BoolAttr, "true">:$enableSoftmax)> {
   let summary = "Gridwise attention accelerated version";
   let description = [{
     The `rock.gridwise_attention_accel` op computes gridwise attention with acceleration.