From db683f4a0ed3a3cd34dac427e5214365dd5414e7 Mon Sep 17 00:00:00 2001
From: Quentin Bourgerie <bourgerie.quentin@gmail.com>
Date: Wed, 29 Sep 2021 15:56:08 +0200
Subject: [PATCH] feat(compiler/runtime): Introduce a RT dialect

---
 .../include/zamalang/Dialect/CMakeLists.txt   |  1 +
 .../zamalang/Dialect/RT/CMakeLists.txt        |  1 +
 .../zamalang/Dialect/RT/IR/CMakeLists.txt     |  9 +++
 .../zamalang/Dialect/RT/IR/RTDialect.h        | 13 ++++
 .../zamalang/Dialect/RT/IR/RTDialect.td       | 15 +++++
 .../include/zamalang/Dialect/RT/IR/RTOps.h    | 14 ++++
 .../include/zamalang/Dialect/RT/IR/RTOps.td   | 67 +++++++++++++++++++
 .../include/zamalang/Dialect/RT/IR/RTTypes.h  | 12 ++++
 .../include/zamalang/Dialect/RT/IR/RTTypes.td | 52 ++++++++++++++
 compiler/lib/Dialect/CMakeLists.txt           |  1 +
 compiler/lib/Dialect/RT/CMakeLists.txt        |  1 +
 compiler/lib/Dialect/RT/IR/CMakeLists.txt     | 14 ++++
 compiler/lib/Dialect/RT/IR/RTDialect.cpp      | 38 +++++++++++
 compiler/lib/Dialect/RT/IR/RTOps.cpp          |  8 +++
 compiler/lib/Support/CompilerEngine.cpp       |  2 +
 compiler/src/CMakeLists.txt                   |  3 +-
 compiler/src/main.cpp                         |  1 +
 17 files changed, 250 insertions(+), 2 deletions(-)
 create mode 100644 compiler/include/zamalang/Dialect/RT/CMakeLists.txt
 create mode 100644 compiler/include/zamalang/Dialect/RT/IR/CMakeLists.txt
 create mode 100644 compiler/include/zamalang/Dialect/RT/IR/RTDialect.h
 create mode 100644 compiler/include/zamalang/Dialect/RT/IR/RTDialect.td
 create mode 100644 compiler/include/zamalang/Dialect/RT/IR/RTOps.h
 create mode 100644 compiler/include/zamalang/Dialect/RT/IR/RTOps.td
 create mode 100644 compiler/include/zamalang/Dialect/RT/IR/RTTypes.h
 create mode 100644 compiler/include/zamalang/Dialect/RT/IR/RTTypes.td
 create mode 100644 compiler/lib/Dialect/RT/CMakeLists.txt
 create mode 100644 compiler/lib/Dialect/RT/IR/CMakeLists.txt
 create mode 100644 compiler/lib/Dialect/RT/IR/RTDialect.cpp
 create mode 100644 compiler/lib/Dialect/RT/IR/RTOps.cpp

diff --git a/compiler/include/zamalang/Dialect/CMakeLists.txt b/compiler/include/zamalang/Dialect/CMakeLists.txt
index 249b735ae..0fa987877 100644
--- a/compiler/include/zamalang/Dialect/CMakeLists.txt
+++ b/compiler/include/zamalang/Dialect/CMakeLists.txt
@@ -2,3 +2,4 @@ add_subdirectory(HLFHE)
 add_subdirectory(HLFHELinalg)
 add_subdirectory(MidLFHE)
 add_subdirectory(LowLFHE)
+add_subdirectory(RT)
diff --git a/compiler/include/zamalang/Dialect/RT/CMakeLists.txt b/compiler/include/zamalang/Dialect/RT/CMakeLists.txt
new file mode 100644
index 000000000..f33061b2d
--- /dev/null
+++ b/compiler/include/zamalang/Dialect/RT/CMakeLists.txt
@@ -0,0 +1 @@
+add_subdirectory(IR)
diff --git a/compiler/include/zamalang/Dialect/RT/IR/CMakeLists.txt b/compiler/include/zamalang/Dialect/RT/IR/CMakeLists.txt
new file mode 100644
index 000000000..0e4c8e6e8
--- /dev/null
+++ b/compiler/include/zamalang/Dialect/RT/IR/CMakeLists.txt
@@ -0,0 +1,9 @@
+set(LLVM_TARGET_DEFINITIONS RTOps.td)
+mlir_tablegen(RTOps.h.inc -gen-op-decls)
+mlir_tablegen(RTOps.cpp.inc -gen-op-defs)
+mlir_tablegen(RTOpsTypes.h.inc -gen-typedef-decls -typedefs-dialect=RT)
+mlir_tablegen(RTOpsTypes.cpp.inc -gen-typedef-defs -typedefs-dialect=RT)
+mlir_tablegen(RTOpsDialect.h.inc -gen-dialect-decls -dialect=RT)
+mlir_tablegen(RTOpsDialect.cpp.inc -gen-dialect-defs -dialect=RT)
+add_public_tablegen_target(MLIRRTOpsIncGen)
+add_dependencies(mlir-headers MLIRRTOpsIncGen)
diff --git a/compiler/include/zamalang/Dialect/RT/IR/RTDialect.h b/compiler/include/zamalang/Dialect/RT/IR/RTDialect.h
new file mode 100644
index 000000000..c908a6306
--- /dev/null
+++ b/compiler/include/zamalang/Dialect/RT/IR/RTDialect.h
@@ -0,0 +1,13 @@
+#ifndef ZAMALANG_DIALECT_RT_IR_RTDIALECT_H
+#define ZAMALANG_DIALECT_RT_IR_RTDIALECT_H
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Dialect.h"
+
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Dialect.h"
+
+#include "zamalang/Dialect/RT/IR/RTOpsDialect.h.inc"
+
+#endif
diff --git a/compiler/include/zamalang/Dialect/RT/IR/RTDialect.td b/compiler/include/zamalang/Dialect/RT/IR/RTDialect.td
new file mode 100644
index 000000000..73318e448
--- /dev/null
+++ b/compiler/include/zamalang/Dialect/RT/IR/RTDialect.td
@@ -0,0 +1,15 @@
+#ifndef ZAMALANG_DIALECT_RT_IR_RT_DIALECT
+#define ZAMALANG_DIALECT_RT_IR_RT_DIALECT
+
+include "mlir/IR/OpBase.td"
+
+def RT_Dialect : Dialect {
+  let name = "RT";
+  let summary = "Runtime dialect";
+  let description = [{
+    A dialect for representation the abstraction needed for the runtime.
+  }];
+  let cppNamespace = "::mlir::zamalang::RT";
+}
+
+#endif
diff --git a/compiler/include/zamalang/Dialect/RT/IR/RTOps.h b/compiler/include/zamalang/Dialect/RT/IR/RTOps.h
new file mode 100644
index 000000000..9dd2def01
--- /dev/null
+++ b/compiler/include/zamalang/Dialect/RT/IR/RTOps.h
@@ -0,0 +1,14 @@
+#ifndef ZAMALANG_DIALECT_RT_IR_RTOPS_H
+#define ZAMALANG_DIALECT_RT_IR_RTOPS_H
+
+#include <mlir/IR/BuiltinOps.h>
+#include <mlir/IR/BuiltinTypes.h>
+#include <mlir/Interfaces/ControlFlowInterfaces.h>
+#include <mlir/Interfaces/SideEffectInterfaces.h>
+
+#include "zamalang/Dialect/RT/IR/RTTypes.h"
+
+#define GET_OP_CLASSES
+#include "zamalang/Dialect/RT/IR/RTOps.h.inc"
+
+#endif
diff --git a/compiler/include/zamalang/Dialect/RT/IR/RTOps.td b/compiler/include/zamalang/Dialect/RT/IR/RTOps.td
new file mode 100644
index 000000000..3a3f2066e
--- /dev/null
+++ b/compiler/include/zamalang/Dialect/RT/IR/RTOps.td
@@ -0,0 +1,67 @@
+#ifndef ZAMALANG_DIALECT_RT_IR_RT_OPS
+#define ZAMALANG_DIALECT_RT_IR_RT_OPS
+
+include "mlir/Interfaces/ControlFlowInterfaces.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+include "zamalang/Dialect/RT/IR/RTDialect.td"
+include "zamalang/Dialect/RT/IR/RTTypes.td"
+
+class RT_Op<string mnemonic, list<OpTrait> traits = []> :
+    Op<RT_Dialect, mnemonic, traits>;
+
+def DataflowTaskOp : RT_Op<"dataflow_task", [SingleBlockImplicitTerminator<"DataflowYieldOp">]> {
+    let arguments = (ins Variadic<AnyType>: $inputs);
+    let results = (outs Variadic<AnyType>:$outputs);
+
+    let regions = (region AnyRegion:$body);
+
+    let summary = "Dataflow task operation";
+    let description = [{
+`RT.dataflow_task` allows to specify a task that will be concurrently executed when their operands are ready.
+
+Example:
+
+```mlir
+func @test(%0 : i64): (i64, i64) {
+    // Execute right now as %0 is ready.
+    %1, %2 = "RT.dataflow_task"(%0) ({
+        %a = addi %0, %0 : i64
+        %b = muli %0, %0 : i64
+        "RT.dataflow_yield"(%a, %b) : (i64, i64) -> i64
+    }) : (i64, i64) -> (i64, i64)
+    // Concurrently execute both tasks below when the task above is completed.
+    %3 = "RT.dataflow_task"(%1) ({
+        %c = constant 1 : %i64
+        %a = addi %1, %c : i64
+        "RT.dataflow_yield"(%a) : (i64, i64) -> i64
+    }) : (i64, i64) -> (i64, i64)
+    %4 = "RT.dataflow_task"(%2) ({
+        %c = constant 2 : %i64
+        %a = addi %2, %c : i64
+        "RT.dataflow_yield"(%a) : (i64, i64) -> i64
+    }) : (i64, i64) -> (i64, i64)
+    return %3, %4 : (i64, i64)
+}
+    }];
+}
+
+def DataflowYieldOp : RT_Op<"dataflow_yield", [ReturnLike, Terminator]> {
+    let arguments = (ins Variadic<AnyType>: $values);
+
+    let summary = "Dataflow yield operation";
+    let description = [{
+`RT.dataflow_yield` is a special terminator operation for blocks inside the region
+in `RT.dataflow_task`. It allows to specify the returns values of a `RT.dataflow_task`.
+
+Example:
+
+```mlir
+%0 = constant 1 : i64
+%1 = constant 2 : i64
+"RT.dataflow_yield" %0, %1 : i64, i64
+```
+  }];
+}
+
+#endif
diff --git a/compiler/include/zamalang/Dialect/RT/IR/RTTypes.h b/compiler/include/zamalang/Dialect/RT/IR/RTTypes.h
new file mode 100644
index 000000000..55fd15833
--- /dev/null
+++ b/compiler/include/zamalang/Dialect/RT/IR/RTTypes.h
@@ -0,0 +1,12 @@
+#ifndef ZAMALANG_DIALECT_RT_IR_RTTYPES_H
+#define ZAMALANG_DIALECT_RT_IR_RTTYPES_H
+
+#include "llvm/ADT/TypeSwitch.h"
+#include <mlir/IR/BuiltinOps.h>
+#include <mlir/IR/BuiltinTypes.h>
+#include <mlir/IR/DialectImplementation.h>
+
+#define GET_TYPEDEF_CLASSES
+#include "zamalang/Dialect/RT/IR/RTOpsTypes.h.inc"
+
+#endif
diff --git a/compiler/include/zamalang/Dialect/RT/IR/RTTypes.td b/compiler/include/zamalang/Dialect/RT/IR/RTTypes.td
new file mode 100644
index 000000000..13cb887e4
--- /dev/null
+++ b/compiler/include/zamalang/Dialect/RT/IR/RTTypes.td
@@ -0,0 +1,52 @@
+#ifndef ZAMALANG_DIALECT_HLFHE_IR_HLFHE_TYPES
+#define ZAMALANG_DIALECT_HLFHE_IR_HLFHE_TYPES
+
+include "zamalang/Dialect/RT/IR/RTDialect.td"
+include "mlir/IR/BuiltinTypes.td"
+
+class RT_Type<string name, list<Trait> traits = []> :
+  TypeDef<RT_Dialect, name, traits> { }
+
+def RT_Future : RT_Type<"Future"> {
+  let mnemonic = "future";
+
+  let summary = "future with a parameterized element type";
+
+  let description = [{
+    The value of a `!RT.future` type represents the result of an asynchronous operation.
+
+    Examples:
+
+    ```mlir
+    !RT.future<i64>
+    ```
+  }];
+
+  let parameters = (ins "Type":$elementType);
+
+  let builders = [
+    TypeBuilderWithInferredContext<(ins "Type":$elementType), [{
+      return $_get(elementType.getContext(), elementType);
+    }]>
+  ];
+
+  let printer = [{
+    $_printer << "future<";
+    $_printer.printType(getElementType());
+    $_printer << ">";
+  }];
+
+  let parser = [{
+    if ($_parser.parseLess())
+      return Type();
+    Type elementType;
+    if ($_parser.parseType(elementType))
+      return Type();
+    if ($_parser.parseGreater())
+      return Type();
+    return get($_ctxt, elementType);
+  }];
+  //let genVerifyDecl = 1;
+}
+
+#endif
diff --git a/compiler/lib/Dialect/CMakeLists.txt b/compiler/lib/Dialect/CMakeLists.txt
index ef6fdf7df..3f199807c 100644
--- a/compiler/lib/Dialect/CMakeLists.txt
+++ b/compiler/lib/Dialect/CMakeLists.txt
@@ -2,3 +2,4 @@ add_subdirectory(HLFHELinalg)
 add_subdirectory(HLFHE)
 add_subdirectory(MidLFHE)
 add_subdirectory(LowLFHE)
+add_subdirectory(RT)
diff --git a/compiler/lib/Dialect/RT/CMakeLists.txt b/compiler/lib/Dialect/RT/CMakeLists.txt
new file mode 100644
index 000000000..f33061b2d
--- /dev/null
+++ b/compiler/lib/Dialect/RT/CMakeLists.txt
@@ -0,0 +1 @@
+add_subdirectory(IR)
diff --git a/compiler/lib/Dialect/RT/IR/CMakeLists.txt b/compiler/lib/Dialect/RT/IR/CMakeLists.txt
new file mode 100644
index 000000000..41bf8fb46
--- /dev/null
+++ b/compiler/lib/Dialect/RT/IR/CMakeLists.txt
@@ -0,0 +1,14 @@
+add_mlir_dialect_library(RTDialect
+  RTDialect.cpp
+  RTOps.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${PROJECT_SOURCE_DIR}/include/zamalang/Dialect/RT
+
+  DEPENDS
+  MLIRRTOpsIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR)
+
+target_link_libraries(RTDialect PUBLIC MLIRIR)
diff --git a/compiler/lib/Dialect/RT/IR/RTDialect.cpp b/compiler/lib/Dialect/RT/IR/RTDialect.cpp
new file mode 100644
index 000000000..ef6d6ae19
--- /dev/null
+++ b/compiler/lib/Dialect/RT/IR/RTDialect.cpp
@@ -0,0 +1,38 @@
+#include "zamalang/Dialect/RT/IR/RTDialect.h"
+#include "zamalang/Dialect/RT/IR/RTOps.h"
+#include "zamalang/Dialect/RT/IR/RTTypes.h"
+
+#define GET_TYPEDEF_CLASSES
+#include "zamalang/Dialect/RT/IR/RTOpsTypes.cpp.inc"
+
+#include "zamalang/Dialect/RT/IR/RTOpsDialect.cpp.inc"
+
+using namespace mlir::zamalang::RT;
+
+void RTDialect::initialize() {
+  addOperations<
+#define GET_OP_LIST
+#include "zamalang/Dialect/RT/IR/RTOps.cpp.inc"
+      >();
+
+  addTypes<
+#define GET_TYPEDEF_LIST
+#include "zamalang/Dialect/RT/IR/RTOpsTypes.cpp.inc"
+      >();
+}
+
+::mlir::Type RTDialect::parseType(::mlir::DialectAsmParser &parser) const {
+  mlir::Type type;
+  if (parser.parseOptionalKeyword("future").succeeded()) {
+    generatedTypeParser(this->getContext(), parser, "future", type);
+    return type;
+  }
+  return type;
+}
+
+void RTDialect::printType(::mlir::Type type,
+                          ::mlir::DialectAsmPrinter &printer) const {
+  if (generatedTypePrinter(type, printer).failed()) {
+    printer.printType(type);
+  }
+}
\ No newline at end of file
diff --git a/compiler/lib/Dialect/RT/IR/RTOps.cpp b/compiler/lib/Dialect/RT/IR/RTOps.cpp
new file mode 100644
index 000000000..e8b59fc5b
--- /dev/null
+++ b/compiler/lib/Dialect/RT/IR/RTOps.cpp
@@ -0,0 +1,8 @@
+#include "mlir/IR/Region.h"
+#include "mlir/IR/TypeUtilities.h"
+
+#include "zamalang/Dialect/RT/IR/RTOps.h"
+#include "zamalang/Dialect/RT/IR/RTTypes.h"
+
+#define GET_OP_CLASSES
+#include "zamalang/Dialect/RT/IR/RTOps.cpp.inc"
diff --git a/compiler/lib/Support/CompilerEngine.cpp b/compiler/lib/Support/CompilerEngine.cpp
index d5869ca48..4f0e65c40 100644
--- a/compiler/lib/Support/CompilerEngine.cpp
+++ b/compiler/lib/Support/CompilerEngine.cpp
@@ -14,6 +14,7 @@
 #include <zamalang/Dialect/HLFHELinalg/IR/HLFHELinalgDialect.h>
 #include <zamalang/Dialect/LowLFHE/IR/LowLFHEDialect.h>
 #include <zamalang/Dialect/MidLFHE/IR/MidLFHEDialect.h>
+#include <zamalang/Dialect/RT/IR/RTDialect.h>
 #include <zamalang/Support/CompilerEngine.h>
 #include <zamalang/Support/Error.h>
 #include <zamalang/Support/Jit.h>
@@ -43,6 +44,7 @@ mlir::MLIRContext *CompilationContext::getMLIRContext() {
   if (this->mlirContext == nullptr) {
     this->mlirContext = new mlir::MLIRContext();
 
+    this->mlirContext->getOrLoadDialect<mlir::zamalang::RT::RTDialect>();
     this->mlirContext->getOrLoadDialect<mlir::zamalang::HLFHE::HLFHEDialect>();
     this->mlirContext
         ->getOrLoadDialect<mlir::zamalang::MidLFHE::MidLFHEDialect>();
diff --git a/compiler/src/CMakeLists.txt b/compiler/src/CMakeLists.txt
index 617919b1d..f40d6aaba 100644
--- a/compiler/src/CMakeLists.txt
+++ b/compiler/src/CMakeLists.txt
@@ -5,7 +5,6 @@ llvm_update_compile_flags(zamacompiler)
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
 
-
 if(ZAMALANG_PARALLEL_EXECUTION_ENABLED)
   target_link_libraries(zamacompiler
     PRIVATE
@@ -20,6 +19,7 @@ if(ZAMALANG_PARALLEL_EXECUTION_ENABLED)
     MLIRIR
     MLIRLLVMIR
     MLIRLLVMToLLVMIRTranslation
+    RTDialect
 
     ZamalangSupport
 
@@ -47,5 +47,4 @@ else()
     )
 endif()
 
-
 mlir_check_all_link_libraries(zamacompiler)
diff --git a/compiler/src/main.cpp b/compiler/src/main.cpp
index ceeae6880..2e3ff53ee 100644
--- a/compiler/src/main.cpp
+++ b/compiler/src/main.cpp
@@ -23,6 +23,7 @@
 #include "zamalang/Dialect/LowLFHE/IR/LowLFHETypes.h"
 #include "zamalang/Dialect/MidLFHE/IR/MidLFHEDialect.h"
 #include "zamalang/Dialect/MidLFHE/IR/MidLFHETypes.h"
+#include "zamalang/Dialect/RT/IR/RTDialect.h"
 #include "zamalang/Support/Error.h"
 #include "zamalang/Support/JitCompilerEngine.h"
 #include "zamalang/Support/KeySet.h"