Merge branch 'develop' into upstream_merge_55

Author: dorde-antic

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

@@ -220,7 +220,8 @@ def Rock_AttentionOp
           Optional<TensorOrMemRefOf<[F32, F16, BF16]>>:$lse,
           UnitAttr:$qTransposed, UnitAttr:$kTransposed, UnitAttr:$vTransposed,
           UnitAttr:$oTransposed, UnitAttr:$causal, StrAttr:$arch,
-          Rock_GemmFeaturesAttr:$features, OptionalAttr<I32Attr>:$numCU,
+          Rock_GemmFeaturesAttr:$features, OptionalAttr<TypeAttr>:$softmaxType,
+          OptionalAttr<I32Attr>:$numCU,
           OptionalAttr<RockTuningParamAttrInterface>:$params0,
           OptionalAttr<RockTuningParamAttrInterface>:$params1,
           I32Attr:$firstGemmIdx)>,
@@ -568,6 +569,7 @@ def Rock_GridwiseAttentionAccelOp
           I32Attr:$gridSize, UnitAttr:$disableQBypassLDS,
           OptionalAttr<IndexAttr>:$prePadG0M,
           OptionalAttr<IndexAttr>:$prePadG0N,
+          OptionalAttr<TypeAttr>:$softmaxType,
           RockAccelTuningParamAttrInterface:$params0,
           RockAccelTuningParamAttrInterface:$params1, I32Attr:$firstGemmIdx,
           DefaultValuedOptionalAttr<BoolAttr, "true">:$enableSoftmax)> {

mlir/include/mlir/Dialect/Rock/utility/builderUtils.h

@@ -32,8 +32,16 @@ Value createZeroConstantOp(OpBuilder &b, Location loc, Type type);
 Value createTypeConversionOp(OpBuilder &b, Location loc, Value source,
                              Type destType);
 
+// Utility function to perform cast
+// and copy to another memref using a Linalg Generic.
 void createTypeConversionLaGeneric(PatternRewriter &rewriter, Location loc,
                                    Value src, Value dst);
+
+// Utility function to perform cast
+// and copy to another memref using a vector store. This flattens the vectors.
+void createTypeConversionFlatAndStore(PatternRewriter &rewriter, Location loc,
+                                      Value src, Value dst);
+
 /// Utility function to collapse an multi-dimensional memref to 1D.
 Value createCollapseShapeOp(OpBuilder &b, Location loc, Value source);
 

mlir/include/mlir/Dialect/Rock/utility/loweringUtils.h

@@ -196,6 +196,9 @@ AffineMap getIdxReversalMap(OpBuilder &b);
 // helper to create ReassociationIndices for flattening
 ReassociationIndices getReassociationForFlattening(ShapedType srcTp);
 
+// helper to obtain a flattened memref
+Value getFlattenedMemref(OpBuilder &b, Value nonFlatMemRef);
+
 /// Construct a `memref.view` operation that interprets the buffer `buffer`,
 /// whose elements are bytes, as a buffer of `type`.
 TypedValue<MemRefType> viewBufferAs(OpBuilder &b, Value buffer, Type type);

mlir/lib/Conversion/TosaToRock/TosaToRock.cpp

@@ -1843,7 +1843,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
         /*kTransposed=*/nullptr,
         /*vTransposed=*/nullptr,
         /*oTransposed=*/nullptr, causalAttr, arch,
-        rewriter.getAttr<rock::GemmFeaturesAttr>(features), numCUAttr,
+        rewriter.getAttr<rock::GemmFeaturesAttr>(features),
+        /*softmaxType=*/nullptr, numCUAttr,
         /*params0=*/nullptr, /*params1=*/nullptr,
         /*firstGemmIdx=*/rewriter.getI32IntegerAttr(0));
 

mlir/lib/Dialect/Rock/IR/RockDialect.cpp

@@ -1876,6 +1876,10 @@ LogicalResult GridwiseAttentionAccelOp::verify() {
     return emitError("LSE only works for attention.");
   }
 
+  if (!getEnableSoftmax() && getSoftmaxType()) {
+    return emitError("Setting softmax type only works for attention.");
+  }
+
   int64_t linalgOpCount = 0;
   getPreSoftmaxBody().walk([&](linalg::GenericOp genOp) { linalgOpCount++; });
   if (linalgOpCount > 1) {

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

@@ -129,12 +129,11 @@ computeGridSizeAttentionGemmElmtGemm(ConversionPatternRewriter &rw, Op op,
   return success();
 }
 
-static LogicalResult
-commonAttentionGemmElmtGemm(ConversionPatternRewriter &rw,
-                            RockGemmGemmWrapperInterface op, Value a, Value b,
-                            Value c, Value out, Value lse, Value currentSeqLen,
-                            UnitAttr causal, ValueRange elementwiseInputs,
-                            Region &preSecondOpRegion, bool enableSoftmax) {
+static LogicalResult commonAttentionGemmElmtGemm(
+    ConversionPatternRewriter &rw, RockGemmGemmWrapperInterface op, Value a,
+    Value b, Value c, Value out, Value lse, Value currentSeqLen,
+    UnitAttr causal, ValueRange elementwiseInputs, Region &preSecondOpRegion,
+    bool enableSoftmax, TypeAttr softmaxType) {
   Location loc = op->getLoc();
 
   if (!isa<MemRefType>(op.getAType()))
@@ -218,8 +217,8 @@ commonAttentionGemmElmtGemm(ConversionPatternRewriter &rw,
       rw.getStringAttr(op.getArch()),
       rw.getAttr<rock::GemmFeaturesAttr>(op.getGemmFeatures()), blockSizeAttr,
       gridSizeAttr,
-      /*disableQBypassLDS=*/nullptr, prePadG0MAttr, prePadG0NAttr, params0,
-      params1, rw.getI32IntegerAttr(op.getFirstGemmIndex()),
+      /*disableQBypassLDS=*/nullptr, prePadG0MAttr, prePadG0NAttr, softmaxType,
+      params0, params1, rw.getI32IntegerAttr(op.getFirstGemmIndex()),
       rw.getBoolAttr(enableSoftmax));
   bool linalgOpFound = false;
   preSecondOpRegion.walk(
@@ -584,7 +583,7 @@ AttentionRewritePattern::matchAndRewrite(AttentionOp op,
       adaptor.getOut(), adaptor.getLse(), adaptor.getCurrentSeqLen(),
       adaptor.getCausalAttr(), adaptor.getPreSoftmaxElemWiseInputs(),
       op.getPreSoftmaxBody(),
-      /*enableSoftmax=*/true);
+      /*enableSoftmax=*/true, op.getSoftmaxTypeAttr());
 }
 
 LogicalResult GemmElementwiseGemmRewritePattern::matchAndRewrite(
@@ -595,7 +594,7 @@ LogicalResult GemmElementwiseGemmRewritePattern::matchAndRewrite(
       /*lse=*/nullptr,
       /*currentSeqLen=*/nullptr, /*causal=*/nullptr,
       adaptor.getElemwiseInputs(), op.getPreSecondGemmBody(),
-      /*enableSoftmax=*/false);
+      /*enableSoftmax=*/false, /*softmaxType=*/nullptr);
 }
 
 void RockGemmToGridwisePass::runOnOperation() {