init sdpa op and flash attention pass

Author: yifeizh2

include/gc/Dialect/Arith/Utils/EasyBuild.h

@@ -264,7 +264,7 @@ inline EBFloatPoint operator-(const EBFloatPoint &a) {
 }
 
 #define DEF_EASYBUILD_CMP_OPERATOR(OP, OPCLASS, TYPE, PRED)                    \
-  EBUnsigned operator OP(const TYPE &a, const TYPE &b) {                       \
+  inline EBUnsigned operator OP(const TYPE &a, const TYPE &b) {                       \
     return OperatorHandlers::handleCmp<OPCLASS>(a, b, PRED);                   \
   }                                                                            \
   template <typename T> EBUnsigned operator OP(const TYPE &a, T b) {           \

include/gc/Dialect/Linalgx/LinalgxStructuredOps.td

@@ -23,6 +23,9 @@ include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
 
+class Linalgx_Op<string mnemonic, list<Trait> traits = []> :
+    Op<LinalgxDialect, mnemonic, traits>;
+
 // Base Tablegen class for Linalg ops.
 // Linalg ops that correspond to library calls operate on ShapedType as their
 // first operands. These may be optionally followed by non-view operands
@@ -312,4 +315,27 @@ def Linalgx_MultiBatchMatmulOp : LinalgxStructuredBase_Op<"multi_batch_matmul",
   }];
 }
 
+def Linalgx_ScaledDotProductAttentionOp
+ : Linalgx_Op<"scaled_dot_product_attention", 
+      [AttrSizedOperandSegments,
+      DeclareOpInterfaceMethods<AggregatedOpInterface, ["decomposeOperation"]>]> {
+  let summary = "Attention structure.";
+  let description = [{
+    Q, K, V, attention_mask.
+    Output = SoftMax(Q @ K.transpose(-2, -1) + attention_mask) @ V.
+  }];
+  let arguments = (ins
+      Variadic<TensorOrMemref>:$inputs,
+      Variadic<TensorOrMemref>:$outputs);
+  let results = (outs Variadic<TensorOrMemref>:$results);
+
+  let hasVerifier = 1;
+  let assemblyFormat = [{
+    attr-dict
+    `ins` `(` $inputs `:` type($inputs) `)`
+    `outs` `(` $outputs `:` type($outputs) `)`
+    (`->` type($results)^)?
+  }];
+}
+
 #endif // LINALGX_STRUCTURED_OPS

include/gc/IR/EasyBuildSCF.h

@@ -150,7 +150,7 @@ inline int IfIterator::operator*() const {
 
 } // namespace impl
 
-impl::IfSimulator makeIfRange(const EasyBuilder &s, Operation *op) {
+inline impl::IfSimulator makeIfRange(const EasyBuilder &s, Operation *op) {
   return impl::IfSimulator{s.builder, op};
 }
 

include/gc/Transforms/Passes.td

@@ -34,7 +34,18 @@ def DeepTileContractionNamedOp
   ];
 }
 
-def GCCPUPipeline : Pass<"gc-cpu-pipeline"> {
+def FlashAttentionConversion
+    : Pass<"flash-attention-conversion", "func::FuncOp"> {
+  let summary = "Flash Attention Conversion";
+  let description =
+      [{The pass converts MHA to flash attention implementation.}];
+  let dependentDialects = [
+    "func::FuncDialect", "linalg::LinalgDialect", "scf::SCFDialect",
+    "tensor::TensorDialect"
+  ];
+}
+
+def GCCPUPipeline: Pass<"gc-cpu-pipeline"> {
   let summary = "All-in-one pipeline for GC for CPU";
   let dependentDialects = [
     "onednn_graph::OneDNNGraphDialect", "tensor::TensorDialect",

lib/gc/Dialect/Linalgx/LinalgxOps.cpp

@@ -9,6 +9,7 @@
 #include "gc/Dialect/Linalgx/LinalgxOps.h"
 #include "gc/Dialect/Linalgx/LinalgxDialect.h"
 #include "mlir/IR/OpImplementation.h"
+#include <utility>
 
 //===----------------------------------------------------------------------===//
 // Builder helper from mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -608,6 +609,80 @@ void MultiBatchMatmulOp::getEffects(
   getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()));
 }
 
+//===----------------------------------------------------------------------===//
+// ScaledDotProductAttentionOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult ScaledDotProductAttentionOp::verify() { return success(); }
+
+/// This method converts ScaledDotProductAttention into the following
+/// sequence of operations:
+/// output = softmax(ins[0] @ transpose(ins[1]) * scale + ins[3]) @ ins[2]
+FailureOr<SmallVector<Value>>
+ScaledDotProductAttentionOp::decomposeOperation(OpBuilder &b) {
+  OpBuilder::InsertionGuard guard(b);
+  b.setInsertionPoint(*this);
+  Location loc = getLoc();
+  Value query = getInputs()[0], key = getInputs()[1], value = getInputs()[2],
+        mask = getInputs()[3];
+  auto dtype = cast<RankedTensorType>(query.getType()).getElementType();
+  auto shape = cast<RankedTensorType>(query.getType()).getShape();
+  float rsqrt_head = 1 / sqrt(shape[3]);
+
+  SmallVector<int64_t> permutation{0, 1, 3, 2};
+  SmallVector<int64_t> transposeShape{shape[0], shape[1], shape[3], shape[2]};
+  auto transposeOut = b.create<tensor::EmptyOp>(loc, transposeShape, dtype);
+  auto transpose = b.create<linalg::TransposeOp>(
+      /*location=*/loc,
+      /*inputs=*/key,
+      /*outputs=*/transposeOut,
+      /*permutation=*/permutation);
+
+  SmallVector<int64_t> matmulQKShape{shape[0], shape[1], shape[2], shape[2]};
+  auto matmulQKOut = b.create<tensor::EmptyOp>(loc, matmulQKShape, dtype);
+  auto matmulQK = b.create<linalgx::MultiBatchMatmulOp>(
+      /*location=*/loc, matmulQKOut.getResult().getType(),
+      /*inputs=*/ValueRange{query, transpose->getResult(0)},
+      /*outputs=*/ValueRange{matmulQKOut.getResult()});
+
+  auto mulOut = b.create<tensor::EmptyOp>(loc, matmulQKShape, dtype);
+  // Broadcast the initial value to the output tensor before convolving.
+  SmallVector<AffineMap, 4> indexingMaps;
+  indexingMaps.push_back(b.getMultiDimIdentityMap(4));
+  indexingMaps.push_back(b.getMultiDimIdentityMap(4));
+  auto mul = b.create<linalg::GenericOp>(
+      /*location=*/loc, matmulQKOut.getResult().getType(),
+      /*inputs=*/ValueRange{matmulQK->getResult(0)},
+      /*outputs=*/ValueRange{mulOut.getResult()}, indexingMaps,
+      SmallVector<utils::IteratorType>(4, utils::IteratorType::parallel),
+      [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
+        Value constant = b.create<arith::ConstantOp>(
+            loc, nestedBuilder.getFloatAttr(dtype, rsqrt_head));
+        Value added =
+            nestedBuilder.create<arith::MulFOp>(loc, args[0], constant);
+        nestedBuilder.create<linalg::YieldOp>(nestedLoc, added);
+      });
+
+  auto addOut = b.create<tensor::EmptyOp>(loc, matmulQKShape, dtype);
+  auto add = b.create<linalg::AddOp>(
+      /*location=*/loc, addOut.getResult().getType(),
+      /*inputs=*/ValueRange{mul->getResult(0), mask},
+      /*outputs=*/ValueRange{addOut.getResult()});
+
+  auto softmaxOut = b.create<tensor::EmptyOp>(loc, matmulQKShape, dtype);
+  auto softmax = b.create<linalg::SoftmaxOp>(
+      /*location=*/loc, softmaxOut.getResult().getType(),
+      /*inputs=*/add->getResult(0),
+      /*outputs=*/softmaxOut.getResult(), 3);
+
+  auto matmulVOut = b.create<tensor::EmptyOp>(loc, shape, dtype);
+  auto matmulV = b.create<linalgx::MultiBatchMatmulOp>(
+      /*location=*/loc, matmulVOut.getResult().getType(),
+      /*inputs=*/ValueRange{softmax->getResult(0), value},
+      /*outputs=*/ValueRange{matmulVOut.getResult()});
+  return SmallVector<Value>{matmulV.getResults()[0]};
+}
+
 /////// Operations corresponding to library calls defined with Tablegen ////////
 
 #define GET_OP_CLASSES

lib/gc/Transforms/CMakeLists.txt

@@ -11,6 +11,7 @@ add_mlir_library(GCPasses
   OneDNNGraphToLinalg.cpp
   Pipeline.cpp
   DeepTileContractionNamedOp.cpp
+  FlashAttentionConversion.cpp
   Tiling.cpp
 
   ADDITIONAL_HEADER_DIRS