From d793c87cdda030679d57c8d56c29277cab71a4b7 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Wed, 11 Nov 2020 16:27:57 +0800
Subject: [PATCH] refactor(mlir/dialect): redefine mgb dialect

GitOrigin-RevId: ac0653474655a2b5b89002e888592415fd472190
---
 CMakeLists.txt                                     | 21 +++++++
 src/jit/impl/mlir/ir/CMakeLists.txt                |  8 +++
 .../include/megbrain/jit/mlir/ir/mgb_dialect.td    | 65 ++++++++++++++++++++++
 3 files changed, 94 insertions(+)
 create mode 100644 src/jit/include/megbrain/jit/mlir/ir/mgb_dialect.td

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 9ae0418a..1072cd97 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -750,6 +750,27 @@ target_include_directories(mgb_opr_param_defs
 add_dependencies(mgb_opr_param_defs _mgb_opr_param_defs)
 install(TARGETS mgb_opr_param_defs EXPORT ${MGE_EXPORT_TARGETS})
 
+if(MGE_WITH_JIT_MLIR)
+    # generate param_defs.td
+    set(MGE_GENFILE_DIR ${PROJECT_BINARY_DIR}/src/genfiles)
+    set(OPR_PARAM_DEFS_SRCS ${MGE_GENFILE_DIR}/opr_param_defs.py)
+    set(OPR_PARAM_DEFS_SCRIPT ${PROJECT_SOURCE_DIR}/dnn/scripts/gen_tablegen.py)
+    set(OPR_PARAM_DEFS_OUT ${MGE_GENFILE_DIR}/param_defs.td)
+    file(COPY ${PROJECT_SOURCE_DIR}/dnn/scripts/opr_param_defs.py DESTINATION ${MGE_GENFILE_DIR})
+    file(READ ${PROJECT_SOURCE_DIR}/tools/param_defs/mgb_opr_param_defs.py CONTENTS)
+    file(APPEND ${OPR_PARAM_DEFS_SRCS} ${CONTENTS})
+    add_custom_target(param_defs_tblgen
+        COMMAND ${PYTHON_EXECUTABLE} ${OPR_PARAM_DEFS_SCRIPT} ${OPR_PARAM_DEFS_SRCS} ${OPR_PARAM_DEFS_OUT}
+        DEPENDS ${OPR_PARAM_DEFS_SRCS} ${OPR_PARAM_DEFS_SCRIPT}
+        VERBATIM
+    )
+    # mlir tblgen sources
+    set(MGE_IR_DIR ${PROJECT_SOURCE_DIR}/src/core/include/megbrain/ir)
+    set(MGE_IR_INCLUDE_DIRS ${MLIR_LLVM_INCLUDE_DIR} ${MGE_GENFILE_DIR} ${MGE_IR_DIR})
+    list(TRANSFORM MGE_IR_INCLUDE_DIRS PREPEND "-I")
+    file(GLOB_RECURSE MGE_IR_TDS ${MGE_IR_DIR}/*.td)
+endif()
+
 if(MGE_WITH_DISTRIBUTED)
     add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/MegRay)
 endif()
diff --git a/src/jit/impl/mlir/ir/CMakeLists.txt b/src/jit/impl/mlir/ir/CMakeLists.txt
index 5c89c7c3..bd340af6 100644
--- a/src/jit/impl/mlir/ir/CMakeLists.txt
+++ b/src/jit/impl/mlir/ir/CMakeLists.txt
@@ -29,3 +29,11 @@ external_tablegen_library(
         -gen-op-decls  include/megbrain/jit/mlir/ir/ops.h.inc
         -gen-op-defs   include/megbrain/jit/mlir/ir/ops.cpp.inc
 )
+
+# mgb_dialect
+set(MGB_DIALECT_TD ${PROJECT_SOURCE_DIR}/src/jit/include/megbrain/jit/mlir/ir/mgb_dialect.td)
+set(LLVM_TARGET_DEFINITIONS ${MGB_DIALECT_TD})
+tablegen(MLIR mgb_dialect.h.inc ${MGE_IR_INCLUDE_DIRS} "--gen-op-decls")
+tablegen(MLIR mgb_dialect.cpp.inc ${MGE_IR_INCLUDE_DIRS} "--gen-op-defs")
+add_custom_target(mgb_dialect DEPENDS mgb_dialect.h.inc mgb_dialect.cpp.inc ${MGB_DIALECT_TD} ${MGE_IR_TDS})
+add_dependencies(mgb_dialect param_defs_tblgen)
diff --git a/src/jit/include/megbrain/jit/mlir/ir/mgb_dialect.td b/src/jit/include/megbrain/jit/mlir/ir/mgb_dialect.td
new file mode 100644
index 00000000..7ae37ffb
--- /dev/null
+++ b/src/jit/include/megbrain/jit/mlir/ir/mgb_dialect.td
@@ -0,0 +1,65 @@
+/**
+ * \file src/jit/include/megbrain/jit/mlir/ir/mgb_dialect.td
+ * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
+ *
+ * Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ * implied.
+ */
+
+#ifndef MGB_DIALECT
+#define MGB_DIALECT
+
+include "ops.td"
+
+class GenericOp<string mnemonic, list<OpTrait> traits = []> :
+  Op<Mgb_Dialect, mnemonic, traits>;
+
+def ReturnOp : GenericOp<"return",
+    [NoSideEffect, HasParent<"FuncOp">, Terminator]> {
+  let summary = "return operation";
+  let description = [{
+    The "return" operation represents a return operation within a function.
+    The operation takes an no tensor operand and produces no results.
+  }];
+
+  // The return operation takes an optional input operand to return. This
+  // value must match the return type of the enclosing function.
+  let arguments = (ins);
+
+  // The return operation only emits the input in the format if it is present.
+  let assemblyFormat = "attr-dict";
+}
+
+def ConstantScalarOp: GenericOp<"sconst", [NoSideEffect]> {
+  let summary = "scalar constant";
+  let arguments = (ins AnyAttr:$value);
+  let results = (outs F32:$result);
+
+  let builders = [OpBuilder<
+    "Builder* builder, OperationState& result, float value", [{
+      result.addAttribute("value", builder->getF32FloatAttr(value));
+      result.addTypes(builder->getF32Type());
+    }]
+  >];
+
+  let extraClassDeclaration = [{
+    Attribute getValue() { return getAttr("value"); }
+    FloatAttr getFloatAttr() { return getAttrOfType<FloatAttr>("value"); }
+  }];
+
+}
+
+def AssignOp :  GenericOp<"assign", []> {
+  let summary = "assign op";
+  let description = [{
+    assign rhs to lhs without results
+  }];
+
+  let arguments = (ins AnyMemRef:$lhs, AnyMemRef:$rhs);
+}
+
+#endif // MGB_DIALECT