feat(compiler/runtime): Introduce a RT dialect

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)

compiler/include/zamalang/Dialect/RT/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(IR)

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)

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

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

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

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

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

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

compiler/lib/Dialect/CMakeLists.txt

@@ -2,3 +2,4 @@ add_subdirectory(HLFHELinalg)
 add_subdirectory(HLFHE)
 add_subdirectory(MidLFHE)
 add_subdirectory(LowLFHE)
+add_subdirectory(RT)

compiler/lib/Dialect/RT/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(IR)

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)

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);
+  }
+}

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"

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>();

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)

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"