[mlir] Add coro intrinsics operations to LLVM dialect

This PR only has coro intrinsics needed for the Async to LLVM lowering. Will add other intrinsics as needed in the followup PRs.

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D95143

2 files changed, 67 insertions, 0 deletions

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td
index 0ef223c4b023..21984ae2914c 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td
@@ -52,6 +52,12 @@ def LLVM_Type : DialectType<LLVM_Dialect,
                             CPred<"::mlir::LLVM::isCompatibleType($_self)">,
                             "LLVM dialect-compatible type">;
 
+// Type constraint accepting LLVM token type.
+def LLVM_TokenType : Type<
+  CPred<"$_self.isa<::mlir::LLVM::LLVMTokenType>()">,
+  "LLVM token type">,
+  BuildableType<"::mlir::LLVM::LLVMTokenType::get($_builder.getContext())">;
+
 // Type constraint accepting LLVM integer types.
 def LLVM_AnyInteger : Type<
   CPred<"$_self.isa<::mlir::IntegerType>()">,
@@ -99,6 +105,16 @@ class LLVM_PointerTo<Type pointee> : Type<
          pointee.predicate>]>,
   "LLVM pointer to " # pointee.summary>;
 
+// Type constraints accepting LLVM pointer type to integer of a specific width.
+class LLVM_IntPtrBase<int width> : Type<
+  LLVM_PointerTo<LLVM_IntBase<width>>.predicate,
+  "LLVM pointer to " # LLVM_IntBase<width>.summary>,
+  BuildableType<"::mlir::LLVM::LLVMPointerType::get("
+                "::mlir::IntegerType::get($_builder.getContext(), "
+                # width #"))">;
+
+def LLVM_i8Ptr : LLVM_IntPtrBase<8>;
+
 // Type constraint accepting any LLVM structure type.
 def LLVM_AnyStruct : Type<CPred<"$_self.isa<::mlir::LLVM::LLVMStructType>()">,
                          "LLVM structure type">;
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index cb2eede3040e..908080ace058 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -1227,6 +1227,57 @@ def LLVM_UMulWithOverflowOp
 }
 
 //
+// Coroutine intrinsics.
+//
+
+def LLVM_CoroIdOp : LLVM_IntrOp<"coro.id", [], [], [], 1> {
+  let arguments = (ins LLVM_i32:$align,
+                       LLVM_i8Ptr:$promise,
+                       LLVM_i8Ptr:$coroaddr,
+                       LLVM_i8Ptr:$fnaddrs);
+  let assemblyFormat = "$align `,` $promise `,` $coroaddr `,` $fnaddrs"
+    " attr-dict `:` type($res)";
+}
+
+def LLVM_CoroBeginOp : LLVM_IntrOp<"coro.begin", [], [], [], 1> {
+  let arguments = (ins LLVM_TokenType:$token,
+                       LLVM_i8Ptr:$mem);
+  let assemblyFormat = "$token `,` $mem attr-dict `:` type($res)";
+}
+
+def LLVM_CoroSizeOp : LLVM_IntrOp<"coro.size", [0], [], [], 1> {
+  let assemblyFormat = "attr-dict `:` type($res)";
+}
+
+def LLVM_CoroSaveOp : LLVM_IntrOp<"coro.save", [], [], [], 1> {
+  let arguments = (ins LLVM_i8Ptr:$handle);
+  let assemblyFormat = "$handle attr-dict `:` type($res)";
+}
+
+def LLVM_CoroSuspendOp : LLVM_IntrOp<"coro.suspend", [], [], [], 1> {
+  let arguments = (ins LLVM_TokenType:$save,
+                       LLVM_i1:$final);
+  let assemblyFormat = "$save `,` $final attr-dict `:` type($res)";
+}
+
+def LLVM_CoroEndOp : LLVM_IntrOp<"coro.end", [], [], [], 1> {
+  let arguments = (ins LLVM_i8Ptr:$handle,
+                       LLVM_i1:$unwind);
+  let assemblyFormat = "$handle `,` $unwind attr-dict `:` type($res)";
+}
+
+def LLVM_CoroFreeOp : LLVM_IntrOp<"coro.free", [], [], [], 1> {
+  let arguments = (ins LLVM_TokenType:$id,
+                       LLVM_i8Ptr:$handle);
+  let assemblyFormat = "$id `,` $handle attr-dict `:` type($res)";
+}
+
+def LLVM_CoroResumeOp : LLVM_IntrOp<"coro.resume", [], [], [], 0> {
+  let arguments = (ins LLVM_i8Ptr:$handle);
+  let assemblyFormat = "$handle attr-dict";
+}
+
+//
 // Vector Reductions.
 //