[Dialect] Introduce microkernel dialect (#114)

* add microkernel dialect

* remove Utils borrowed from TPP

Author: huanghaixin008

include/gc/CMakeLists.txt

@@ -1,2 +1,2 @@
 add_subdirectory(Dialect)
-add_subdirectory(Transforms)
+add_subdirectory(Transforms)

include/gc/Dialect/CMakeLists.txt

@@ -1,4 +1,4 @@
 add_subdirectory(CPURuntime)
 add_subdirectory(OneDNNGraph)
 add_subdirectory(Microkernel)
-add_subdirectory(Linalgx)
+add_subdirectory(Linalgx)

include/gc/Dialect/Microkernel/CMakeLists.txt

@@ -1,3 +1,8 @@
+set(LLVM_TARGET_DEFINITIONS MicrokernelEnum.td)
+mlir_tablegen(MicrokernelEnum.h.inc -gen-enum-decls)
+mlir_tablegen(MicrokernelEnum.cpp.inc -gen-enum-defs)
+add_public_tablegen_target(MLIRMicrokernelAttrDefIncGen)
+
 add_mlir_dialect(MicrokernelOps microkernel)
 add_mlir_doc(MicrokernelOps MicrokernelOps gc/Dialect/Microkernel/ -gen-op-doc)
 add_mlir_doc(MicrokernelDialect MicrokernelDialect gc/Dialect/Microkernel/ -gen-dialect-doc)

include/gc/Dialect/Microkernel/MicrokernelDialect.h

@@ -10,6 +10,7 @@
 #define GC_DIALECTS_MICROKERNELDIALECT_H
 
 #include "mlir/IR/Dialect.h"
+#include "mlir/Interfaces/DestinationStyleOpInterface.h"
 
 #include "gc/Dialect/Microkernel/MicrokernelOpsDialect.h.inc"
 

include/gc/Dialect/Microkernel/MicrokernelDialect.td

@@ -15,15 +15,20 @@ include "mlir/IR/OpBase.td"
 // Microkernel dialect definition.
 //===----------------------------------------------------------------------===//
 
-def MicrokernelDialect : Dialect {
+def Microkernel_Dialect : Dialect {
     let name = "microkernel";
     let summary = "A dialect for microkernel abstraction.";
     let description = [{
-        The dialect wraps the BRGEMM API to set up the HW context etc.
+        This dialect contains wrappers for microkernel primitives like BRGEMM.
     }];
     let cppNamespace = "::mlir::microkernel";
-
-    let useDefaultTypePrinterParser = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// Base microkernel operation definition.
+//===----------------------------------------------------------------------===//
+
+class Microkernel_Op<string mnemonic, list<Trait> traits = []> :
+        Op<Microkernel_Dialect, mnemonic, traits>;
+
 #endif // MICROKERNEL_DIALECT

include/gc/Dialect/Microkernel/MicrokernelEnum.h

@@ -0,0 +1,18 @@
+//===- MicrokernelEnum.h - microkernel dialect enums ------------*- 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_DIALECTS_MICROKERNELENUM_H
+#define GC_DIALECTS_MICROKERNELENUM_H
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/DialectImplementation.h"
+
+#define GET_ATTRDEF_CLASSES
+#include "gc/Dialect/Microkernel/MicrokernelEnum.h.inc"
+
+#endif // GC_DIALECTS_MICROKERNELENUM_H

include/gc/Dialect/Microkernel/MicrokernelEnum.td

@@ -0,0 +1,26 @@
+//===- MicrokernelEnum.td - microkernel dialect enum -------*- 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 MICROKERNEL_ENUM
+#define MICROKERNEL_ENUM
+
+include "mlir/IR/EnumAttr.td"
+include "gc/Dialect/Microkernel/MicrokernelDialect.td"
+
+def Microkernel_BrgemmFlags : I64EnumAttr<
+    "BrgemmFlags", "Flags for indicating optional behaviours of Brgemm",
+    [
+      I64EnumAttrCase<"NONE", 0, "none">,
+      I64EnumAttrCase<"BETA_0", 1, "beta_0">,
+      I64EnumAttrCase<"STRIDE", 2, "stride">,
+      I64EnumAttrCase<"LIST", 4, "list">
+    ]> {
+  let cppNamespace = "::mlir::microkernel";
+}
+
+#endif // MICROKERNEL_ENUM

include/gc/Dialect/Microkernel/MicrokernelOps.h

@@ -9,7 +9,16 @@
 #ifndef GC_DIALECTS_MICROKERNELOPS_H
 #define GC_DIALECTS_MICROKERNELOPS_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/SideEffectInterfaces.h"
+
+#include "gc/Dialect/Microkernel/MicrokernelDialect.h"
+#include "gc/Dialect/Microkernel/MicrokernelEnum.h"
 
 #define GET_OP_CLASSES
 #include "gc/Dialect/Microkernel/MicrokernelOps.h.inc"

include/gc/Dialect/Microkernel/MicrokernelOps.td

@@ -10,5 +10,107 @@
 #define MICROKERNEL_OPS
 
 include "MicrokernelDialect.td"
+include "gc/Dialect/Microkernel/MicrokernelEnum.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
 
-#endif // MICROKERNEL_OPS
+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 Microkernel_BrgemmDispatchOp : Microkernel_Op<"brgemm.dispatch", [Pure]> {
+  let summary = "JIT the brgemm microkernel given the parameters";
+  let description = [{
+    The operation has the following arguments: 1) m, n, k, lda, ldb, ldc, stride_a and stride_b. 
+    Inputs is a dense attribute of I64 elements. 2) flags carry information on
+    the different flags that can be used for brgemm like whether beta == 0 or strided batch. For
+    more details, see: `Microkernel_BrgemmFlags`. 3) data_types of operand A & B.
+    Outpus is the id of JITed kernel.
+  }];
+
+  let arguments = (ins
+    ConfinedAttr<DenseI64ArrayAttr,
+                [DenseArrayNonNegative<DenseI64ArrayAttr>]>:$inputs,
+    TypedArrayAttrBase<Microkernel_BrgemmFlags, "brgemm flags">:$flags,
+    TypedArrayAttrBase<TypeAttr, "brgemm dtypes">:$data_type);
+
+  let results = (outs I64:$results);
+  let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
+}
+
+def Microkernel_BrgemmPrologueOp : Microkernel_Op<"brgemm.prologue"> {
+  let summary = "Prologue before executing the JITed brgemm "
+                "microkernel, and the context is considered core-level";
+  let description = [{
+    The operation has the following arguments: Input is the id of JITed kernel.
+    There is no output.
+  }];
+
+  let arguments = (ins I64:$inputs);
+
+  let assemblyFormat = [{
+    `(` $inputs `)`
+    attr-dict `:` functional-type($inputs, results)
+  }];
+}
+
+def Microkernel_BrgemmEpilogueOp : Microkernel_Op<"brgemm.epilogue"> {
+  let summary = "Epilogue after executing the JITed brgemm microkernel";
+  let description = [{
+    The operation has the following arguments: Input is the id of JITed kernel.
+    There is no output.
+  }];
+
+  let arguments = (ins I64:$inputs);
+
+  let assemblyFormat = [{
+    `(` $inputs `)`
+    attr-dict `:` functional-type($inputs, results)
+  }];
+}
+
+/* A generic input type of Microkernel_BrgemmOp, allowing for `BrgemmMemRef` and I64.
+ * The `BrgemmMemRef` should be a static MemRef, and for each operand its shape should be:
+ *     Operand A: StaticMemRefRankOf<[F32, BF16, SI8, UI8], [3]>;
+ *     Operand B (none-VNNI): StaticMemRefRankOf<[F32], [3]>;
+ *     Operand B (VNNI): StaticMemRefRankOf<[BF16, SI8, UI8], [4]>;
+ *     Operand C: StaticMemRefRankOf<[F32, SI32], [2]>;
+ */
+def BrgemmMemRefOrI64 : AnyTypeOf<[StaticMemRefRankOf<[F32, BF16, SI32, SI8, UI8], [2, 3, 4]>, I64]>;
+
+def Microkernel_BrgemmOp : Microkernel_Op<"brgemm"> {
+  let summary = "execute the JITed brgemm kernel.";
+  let description = [{
+    The operation has the following arguments: 
+    1) For stride mode, id of JITed kernel, MemRef of operand A/B/C, and the batch size;
+    2) For addr mode, plus the length of addr list at the end.
+    There is no output.
+  }];
+
+  let arguments = (ins Variadic<BrgemmMemRefOrI64>:$inputs);
+
+  let assemblyFormat = [{
+    `(` $inputs `)`
+    attr-dict `:` functional-type($inputs, results)
+  }];
+
+  let extraClassDeclaration = [{
+    Value getDispatch() { return getInputs()[0]; }
+
+    Value getOperandA() { return getInputs()[1]; }
+
+    Value getOperandB() { return getInputs()[2]; }
+
+    Value getOutput() { return getInputs()[3]; }
+
+    Value getBatch() { return getInputs()[4]; }
+
+    Value getAddrLen() { return getInputs()[5]; }
+  }];
+
+  let hasVerifier = 1;
+}
+
+#endif // MICROKERNEL_OPS

lib/gc/CMakeLists.txt

@@ -7,4 +7,4 @@ include(functions)
 add_subdirectory(CAPI)
 add_subdirectory(Dialect)
 add_subdirectory(Transforms)
-add_subdirectory(ExecutionEngine)
+add_subdirectory(ExecutionEngine)

lib/gc/Dialect/Microkernel/CMakeLists.txt

@@ -1,6 +1,7 @@
 gc_set_mlir_link_components(MLIR_LINK_COMPONENTS MLIRIR)
 
 add_mlir_dialect_library(MLIRMicrokernel
+    MicrokernelEnum.cpp
     MicrokernelDialect.cpp
     MicrokernelOps.cpp
 
@@ -12,5 +13,6 @@ add_mlir_dialect_library(MLIRMicrokernel
 
     LINK_LIBS PUBLIC
     ${MLIR_LINK_COMPONENTS}
+    GCUtilsIR
 )
-set_property(GLOBAL APPEND PROPERTY GC_DIALECT_LIBS MLIRMicrokernel)
+set_property(GLOBAL APPEND PROPERTY GC_DIALECT_LIBS MLIRMicrokernel)

lib/gc/Dialect/Microkernel/MicrokernelDialect.cpp

@@ -7,11 +7,18 @@
 //===----------------------------------------------------------------------===//
 
 #include "gc/Dialect/Microkernel/MicrokernelDialect.h"
+#include "gc/Dialect/Microkernel/MicrokernelEnum.h"
 #include "gc/Dialect/Microkernel/MicrokernelOps.h"
 
 using namespace mlir;
 using namespace mlir::microkernel;
 
+#include "gc/Dialect/Microkernel/MicrokernelOpsDialect.cpp.inc"
+
+//===----------------------------------------------------------------------===//
+// Microkernel dialect.
+//===----------------------------------------------------------------------===//
+
 void MicrokernelDialect::initialize() {
   addOperations<
 #define GET_OP_LIST

lib/gc/Dialect/Microkernel/MicrokernelEnum.cpp

@@ -0,0 +1,15 @@
+//===-- MicrokernelEnum.cpp - microkernel dialect enum ----------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "gc/Dialect/Microkernel/MicrokernelEnum.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+using namespace mlir;
+using namespace mlir::microkernel;
+
+#include "gc/Dialect/Microkernel/MicrokernelEnum.cpp.inc"