[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

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">;

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

@@ -1226,6 +1226,57 @@ def LLVM_UMulWithOverflowOp
   let arguments = (ins LLVM_Type, LLVM_Type);
 }
 
+//
+// 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.
 //

mlir/test/Target/llvmir-intrinsics.mlir

@@ -347,6 +347,64 @@ llvm.func @umul_with_overflow_test(%arg0: i32, %arg1: i32, %arg2: vector<8xi32>,
   llvm.return
 }
 
+// CHECK-LABEL: @coro_id
+llvm.func @coro_id(%arg0: i32, %arg1: !llvm.ptr<i8>) {
+  // CHECK: call token @llvm.coro.id
+  llvm.intr.coro.id %arg0, %arg1, %arg1, %arg1 : !llvm.token
+  llvm.return
+}
+
+// CHECK-LABEL: @coro_begin
+llvm.func @coro_begin(%arg0: !llvm.token, %arg1: !llvm.ptr<i8>) {
+  // CHECK: call i8* @llvm.coro.begin
+  llvm.intr.coro.begin %arg0, %arg1 : !llvm.ptr<i8>
+  llvm.return
+}
+
+// CHECK-LABEL: @coro_size
+llvm.func @coro_size() {
+  // CHECK: call i64 @llvm.coro.size.i64
+  %0 = llvm.intr.coro.size : i64
+  // CHECK: call i32 @llvm.coro.size.i32
+  %1 = llvm.intr.coro.size : i32
+  llvm.return
+}
+
+// CHECK-LABEL: @coro_save
+llvm.func @coro_save(%arg0: !llvm.ptr<i8>) {
+  // CHECK: call token @llvm.coro.save
+  %0 = llvm.intr.coro.save %arg0 : !llvm.token
+  llvm.return
+}
+
+// CHECK-LABEL: @coro_suspend
+llvm.func @coro_suspend(%arg0: !llvm.token, %arg1 : i1) {
+  // CHECK: call i8 @llvm.coro.suspend
+  %0 = llvm.intr.coro.suspend %arg0, %arg1 : !llvm.i8
+  llvm.return
+}
+
+// CHECK-LABEL: @coro_end
+llvm.func @coro_end(%arg0: !llvm.ptr<i8>, %arg1 : i1) {
+  // CHECK: call i1 @llvm.coro.end
+  %0 = llvm.intr.coro.end %arg0, %arg1 : !llvm.i1
+  llvm.return
+}
+
+// CHECK-LABEL: @coro_free
+llvm.func @coro_free(%arg0: !llvm.token, %arg1 : !llvm.ptr<i8>) {
+  // CHECK: call i8* @llvm.coro.free
+  %0 = llvm.intr.coro.free %arg0, %arg1 : !llvm.ptr<i8>
+  llvm.return
+}
+
+// CHECK-LABEL: @coro_resume
+llvm.func @coro_resume(%arg0: !llvm.ptr<i8>) {
+  // CHECK: call void @llvm.coro.resume
+  llvm.intr.coro.resume %arg0
+  llvm.return
+}
+
 // Check that intrinsics are declared with appropriate types.
 // CHECK-DAG: declare float @llvm.fma.f32(float, float, float)
 // CHECK-DAG: declare <8 x float> @llvm.fma.v8f32(<8 x float>, <8 x float>, <8 x float>) #0
@@ -393,3 +451,12 @@ llvm.func @umul_with_overflow_test(%arg0: i32, %arg1: i32, %arg2: vector<8xi32>,
 // CHECK-DAG: declare { <8 x i32>, <8 x i1> } @llvm.usub.with.overflow.v8i32(<8 x i32>, <8 x i32>) #0
 // CHECK-DAG: declare { i32, i1 } @llvm.umul.with.overflow.i32(i32, i32)
 // CHECK-DAG: declare { <8 x i32>, <8 x i1> } @llvm.umul.with.overflow.v8i32(<8 x i32>, <8 x i32>) #0
+// CHECK-DAG: declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*)
+// CHECK-DAG: declare i8* @llvm.coro.begin(token, i8* writeonly)
+// CHECK-DAG: declare i64 @llvm.coro.size.i64()
+// CHECK-DAG: declare i32 @llvm.coro.size.i32()
+// CHECK-DAG: declare token @llvm.coro.save(i8*)
+// CHECK-DAG: declare i8 @llvm.coro.suspend(token, i1)
+// CHECK-DAG: declare i1 @llvm.coro.end(i8*, i1)
+// CHECK-DAG: declare i8* @llvm.coro.free(token, i8* nocapture readonly)
+// CHECK-DAG: declare void @llvm.coro.resume(i8*)