[Transform] Introduce passes for transforming from linalg to microkernel dialect (#266)

Contains passes that lower linalg.batch_reduce_matmul & linalgx.batch_reduce_matmul_vnni into microkernel dialect:
* ConvertLinalgToMicrokernel: convert linalg Op into frontend microkernel Op
* ExpandMicrokernel: expand frontend microkernel Op into detailed microkernel Ops
* ConvertMicrokernelToDnnlFunc: convert microkernel Ops into dnnl runtime invocation and related tests.

Author: Haixin Huang

CMakeLists.txt

@@ -55,7 +55,6 @@ if(GC_ENABLE_LEGACY)
   add_subdirectory(legacy/core)
 endif()
 
-
 if (GC_ENABLE_IMEX)
   # normalize the value for lit config
   set(GC_ENABLE_IMEX ON)

include/gc/Dialect/Microkernel/MicrokernelOps.h

@@ -9,12 +9,14 @@
 #ifndef GC_DIALECTS_MICROKERNELOPS_H
 #define GC_DIALECTS_MICROKERNELOPS_H
 
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/DestinationStyleOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 #include "gc/Dialect/Microkernel/MicrokernelDialect.h"

include/gc/Dialect/Microkernel/MicrokernelOps.td

@@ -11,14 +11,79 @@
 
 include "MicrokernelDialect.td"
 include "gc/Dialect/Microkernel/MicrokernelEnum.td"
+include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td"
+include "mlir/Interfaces/DestinationStyleOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
+class StaticTensorRankOf<list<Type> allowedTypes, list<int> ranks> :
+     Type<And<[TensorOf<allowedTypes>.predicate,
+              HasAnyRankOfPred<ranks>, HasStaticShapePred]>, 
+              !interleave(!foreach(rank, ranks, rank # "D"), "/") # " static " # 
+              TensorOf<allowedTypes>.summary, "::mlir::TensorType">;
+
 class StaticMemRefRankOf<list<Type> allowedTypes, list<int> ranks> :
      Type<And<[MemRefOf<allowedTypes>.predicate,
               HasAnyRankOfPred<ranks>, HasStaticShapePred]>, 
               !interleave(!foreach(rank, ranks, rank # "D"), "/") # " static " # 
               MemRefOf<allowedTypes>.summary, "::mlir::MemRefType">;
 
+def BrgemmTensor : StaticTensorRankOf<[F32, BF16, SI32, SI8, UI8], [2, 3, 4]>;
+
+def BrgemmTensorOrMemRef : AnyTypeOf<[StaticTensorRankOf<[F32, BF16, SI32, SI8, UI8], [2, 3, 4]>,
+                                      StaticMemRefRankOf<[F32, BF16, SI32, SI8, UI8], [2, 3, 4]>]>;
+
+def Microkernel_BrgemmOp : Microkernel_Op<"brgemm", 
+	[DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+         BufferizableOpInterface,
+         DestinationStyleOpInterface]> {
+  let summary = "Abstract Op that execute brgemm kernel on tensors.";
+  let description = [{
+    The operation has the following arguments: 
+    1) Tensors or MemRefs of operand A/B;
+    2) The batch dims and leading dims of operand A/B;
+    And has the following outputs:
+    1) Tensor of operand C;
+  }];
+
+  let arguments = (ins Variadic<BrgemmTensorOrMemRef>:$inputs,
+                       BrgemmTensorOrMemRef:$init, 
+    		       ConfinedAttr<DenseI64ArrayAttr,
+                           [DenseArrayNonNegative<DenseI64ArrayAttr>]>:$batchDims,
+    		       ConfinedAttr<DenseI64ArrayAttr,
+                           [DenseArrayNonNegative<DenseI64ArrayAttr>]>:$leadingDims, 
+                       TypedArrayAttrBase<Microkernel_BrgemmFlags, "brgemm flags">:$flags);
+  let results = (outs Variadic<BrgemmTensor>:$output);
+
+  let extraClassDeclaration = [{
+    Value getOperandA() { return getInputs()[0]; }
+    Value getOperandB() { return getInputs()[1]; }
+    Value getOperandC() { return getInit(); }
+    
+    int64_t getBatchDimA() { return getBatchDims()[0]; }
+    int64_t getLeadingDimA() { return getLeadingDims()[0]; }
+    
+    int64_t getBatchDimB() { return getBatchDims()[1]; }
+    int64_t getLeadingDimB() { return getLeadingDims()[1]; }
+
+    MutableOperandRange getDpsInitsMutable() { return getInitMutable(); }
+
+    bool bufferizesToMemoryRead(OpOperand &,
+                                const bufferization::AnalysisState &);
+    bool bufferizesToMemoryWrite(OpOperand &,
+                                 const bufferization::AnalysisState &);
+    bool bufferizesToElementwiseAccess(const bufferization::AnalysisState &,
+                                       ArrayRef<OpOperand *>);
+    bufferization::AliasingValueList getAliasingValues(OpOperand &opOperand,
+                                      const bufferization::AnalysisState &state);
+    LogicalResult bufferize(RewriterBase &,
+                            const bufferization::BufferizationOptions &);
+  }];
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+  let hasFolder = 1;
+}
+
 def Microkernel_BrgemmDispatchOp : Microkernel_Op<"brgemm.dispatch", [Pure]> {
   let summary = "JIT the brgemm microkernel given the parameters";
   let description = [{
@@ -80,7 +145,7 @@ def Microkernel_BrgemmEpilogueOp : Microkernel_Op<"brgemm.epilogue"> {
  */
 def BrgemmMemRefOrI64 : AnyTypeOf<[StaticMemRefRankOf<[F32, BF16, SI32, SI8, UI8], [2, 3, 4]>, I64]>;
 
-def Microkernel_BrgemmOp : Microkernel_Op<"brgemm"> {
+def Microkernel_BrgemmExecuteOp : Microkernel_Op<"brgemm.execute"> {
   let summary = "execute the JITed brgemm kernel.";
   let description = [{
     The operation has the following arguments: 

include/gc/Transforms/CMakeLists.txt

@@ -11,3 +11,5 @@ mlir_tablegen(Passes.capi.h.inc -gen-pass-capi-header --prefix GraphCompiler)
 mlir_tablegen(Passes.capi.cpp.inc -gen-pass-capi-impl --prefix GraphCompiler)
 add_public_tablegen_target(GraphCompilerPassIncGen)
 add_mlir_doc(Passes GraphCompilerPasses ./ -gen-pass-doc)
+
+add_subdirectory(Microkernel)

include/gc/Transforms/Microkernel/BrgemmRuntimeUtils.h

@@ -0,0 +1,53 @@
+//===-- BrgemmRuntimeUtils.h - Utils for Brgemm Runtime ---------*- C++ -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef GC_MICROKERNEL_BRGEMMRUNTIMEUTILS_H
+#define GC_MICROKERNEL_BRGEMMRUNTIMEUTILS_H
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "oneapi/dnnl/dnnl_types.h"
+
+namespace mlir::microkernel {
+
+// these strings contain symbols for BRGEMM interfaces used in mlir pass
+static const std::string DNNL_BRGEMM_DISPATCH_NAME = "dnnl_brgemm_dispatch";
+static const std::string DNNL_BRGEMM_TILECFG_NAME = "dnnl_brgemm_tileconfig";
+static const std::string DNNL_BRGEMM_TILERELEASE_NAME =
+    "dnnl_brgemm_tilerelease";
+static const std::string DNNL_BRGEMM_EXECUTE_NAME = "dnnl_brgemm_execute";
+
+static inline int64_t getDnnlDataTypeVal(RewriterBase &rewriter,
+                                         Attribute attr) {
+  auto context = rewriter.getContext();
+  auto tattr = dyn_cast_or_null<TypeAttr>(attr);
+  assert(tattr);
+  if (tattr == TypeAttr::get(FloatType::getF32(context))) {
+    return static_cast<int64_t>(dnnl_f32);
+  } else if (tattr == TypeAttr::get(FloatType::getF64(context))) {
+    return static_cast<int64_t>(dnnl_f64);
+  } else if (tattr == TypeAttr::get(FloatType::getBF16(context))) {
+    return static_cast<int64_t>(dnnl_bf16);
+  } else if (tattr == TypeAttr::get(FloatType::getF16(context))) {
+    return static_cast<int64_t>(dnnl_f16);
+  } else if (tattr == TypeAttr::get(
+                          IntegerType::get(context, 32, IntegerType::Signed))) {
+    return static_cast<int64_t>(dnnl_s32);
+  } else if (tattr ==
+             TypeAttr::get(IntegerType::get(context, 8, IntegerType::Signed))) {
+    return static_cast<int64_t>(dnnl_s8);
+  } else if (tattr == TypeAttr::get(IntegerType::get(context, 8,
+                                                     IntegerType::Unsigned))) {
+    return static_cast<int64_t>(dnnl_u8);
+  }
+  return static_cast<int64_t>(dnnl_data_type_undef);
+}
+
+}; // namespace mlir::microkernel
+
+#endif // GC_MICROKERNEL_BRGEMMRUNTIMEUTILS_H

include/gc/Transforms/Microkernel/CMakeLists.txt

@@ -0,0 +1,6 @@
+set(LLVM_TARGET_DEFINITIONS MicrokernelPasses.td)
+mlir_tablegen(MicrokernelPasses.h.inc --gen-pass-decls -name Microkernel)
+mlir_tablegen(MicrokernelPasses.capi.h.inc -gen-pass-capi-header --prefix Microkernel)
+mlir_tablegen(MicrokernelPasses.capi.cpp.inc -gen-pass-capi-impl --prefix Microkernel)
+add_public_tablegen_target(MLIRMicrokernelPassesIncGen)
+add_mlir_doc(MicrokernelPasses GraphCompilerMicrokernelPasses ./ -gen-pass-doc)

include/gc/Transforms/Microkernel/MicrokernelPasses.h

@@ -0,0 +1,28 @@
+//===- MicrokernelPasses.h - Graph Compiler microkerenl passes --*- C++ -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef GC_MICROKERNELPASSES_H
+#define GC_MICROKERNELPASSES_H
+
+#include "gc/Dialect/Linalgx/LinalgxDialect.h"
+#include "gc/Dialect/Microkernel/MicrokernelDialect.h"
+#include "gc/Dialect/Microkernel/MicrokernelOps.h"
+#include "mlir/Pass/Pass.h"
+#include <memory>
+
+namespace mlir {
+namespace microkernel {
+#define GEN_PASS_DECL
+#include "gc/Transforms/Microkernel/MicrokernelPasses.h.inc"
+
+#define GEN_PASS_REGISTRATION
+#include "gc/Transforms/Microkernel/MicrokernelPasses.h.inc"
+} // namespace microkernel
+} // namespace mlir
+
+#endif // GC_MICROKERNELPASSES_H

include/gc/Transforms/Microkernel/MicrokernelPasses.td

@@ -0,0 +1,79 @@
+//===-- MicrokernelPasses.td - microkernel passes ----------*- tablegen -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef GC_DIALECT_MICROKERNELPASSES
+#define GC_DIALECT_MICROKERNELPASSES
+
+include "mlir/Pass/PassBase.td"
+
+def ConvertLinalgToMicrokernel: Pass<"convert-linalg-to-microkernel", "::mlir::func::FuncOp"> {
+  let summary = "Lower eligible linalg ops to microkernels";
+  let description = [{
+    Convert eligible linalg ops to microkernel dialects based on pattern matching.
+    For example:
+    ```
+    scf.forall {
+      linalg.fill ins(...) outs(...) -> tensor<...>
+      linalg.batch_reduce_matmul ins(...) outs(...) -> tensor<...>
+    }
+    ```
+    Will be changed into
+    ```
+    scf.forall {
+      linalg.fill ins(...) outs(...) -> tensor<...>
+      microkernel.brgemm ins(...) outs(...) -> tensor<...>
+    }
+    ```
+  }];
+  let dependentDialects = ["func::FuncDialect",
+                           "tensor::TensorDialect",
+                           "memref::MemRefDialect",
+                           "linalg::LinalgDialect",
+                           "linalgx::LinalgxDialect",
+                           "microkernel::MicrokernelDialect"];
+}
+
+def ExpandMicrokernel: Pass<"expand-microkernel", "::mlir::func::FuncOp"> {
+  let summary = "Expand abstract microkernels into detailed execution phases";
+  let description = [{
+    Expand abstract microkernels into detailed execution phases
+    For example:
+    ```
+    scf.forall {
+      linalg.fill ins(...) outs(...) -> tensor<...>
+      microkernel.brgemm ins(...) outs(...) -> tensor<...>
+    }
+    ```
+    Will be changed into
+    ```
+    scf.forall {
+      linalg.fill ins(...) outs(...) -> tensor<...>
+      %0 = microkernel.brgemm.dispatch(...)
+      microkernel.brgemm.prologue(%0)
+      microkernel.brgemm.execute(%0, ...)
+      microkernel.brgemm.epilogue(%0)
+    }
+    ```
+  }];
+  let dependentDialects = ["func::FuncDialect",
+                           "memref::MemRefDialect",
+                           "microkernel::MicrokernelDialect"];
+}
+
+def ConvertMicrokernelToDnnlFunc: Pass<"convert-microkernel-to-dnnl-func", "::mlir::ModuleOp"> {
+  let summary = "Lower microkernel dialects to dnnl func call";
+  let description = [{
+    Convert microkernel dialects to runtime function call to oneDNN library.
+  }];
+  let dependentDialects = ["func::FuncDialect",
+                           "memref::MemRefDialect",
+                           "LLVM::LLVMDialect",
+                           "microkernel::MicrokernelDialect"];
+}
+
+#endif // GC_DIALECT_MICROKERNELPASSES

include/gc/Transforms/Utils/StructuredOpMatcher.h

@@ -217,6 +217,16 @@ template <typename T> struct EqualsTo {
 };
 template <typename T> EqualsTo(T) -> EqualsTo<T>;
 
+// Callable object to check if the input is greater than or equal to specified
+// `value`.
+struct GreaterThanOrEqualTo {
+  GreaterThanOrEqualTo() = delete;
+  explicit GreaterThanOrEqualTo(size_t value) : value(value){};
+  const size_t value;
+
+  bool operator()(size_t value) const { return value >= this->value; }
+};
+
 // Callable object to validate number of init operands for `op`.
 struct NumDpsInits {
   NumDpsInits() = delete;

include/gc/Transforms/Utils/ValueUtils.h

@@ -27,8 +27,7 @@ FailureOr<SmallVector<int64_t>> getStaticStrides(Value val);
 
 // Return the offset and ptr for `val`. Assert if `val`
 // is not a memref.
-std::pair<Value, Value> getPtrAndOffset(OpBuilder &builder, Value val,
-                                        Location loc);
+std::pair<Value, Value> getPtrAndOffset(OpBuilder &builder, Value operand);
 
 } // namespace utils
 } // namespace mlir