[mlir] Add coro intrinsics operations to LLVM dialect

[ezhulenev]

Summary

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

[ezhulenev]

remove dump input

[ftynse]

Have you considered putting these in a separate dialect that just uses LLVMOpBase.td, similarly to what we have for NVVM or AVX512? My general take on it is that intrinsics listed in https://llvm.org/docs/LangRef.html#intrinsic-functions can stay in the "main" dialect, but "target-specific" intrinsics are better off in separate dialects. I understand there are pros and cons of doing this, so I won't block this commit only because of it.

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

98–109

I'd rather do something like def LLVM_I8PointerType : LLVM_PointerTo<LLVM_i8>; and use it everywhere. Otherwise, this looks like there is an "opaque pointer" type different from "pointer to i8", which is not the case.

[ezhulenev]

Add LLVM_i?Ptr type constraint

[ezhulenev]

In D95143#2515166, @ftynse wrote:

Have you considered putting these in a separate dialect that just uses LLVMOpBase.td, similarly to what we have for NVVM or AVX512? My general take on it is that intrinsics listed in https://llvm.org/docs/LangRef.html#intrinsic-functions can stay in the "main" dialect, but "target-specific" intrinsics are better off in separate dialects. I understand there are pros and cons of doing this, so I won't block this commit only because of it.

Coro intrinsics are actually defined in the "main" LLVM intrinsic definition file, so I think it makes sense to keep them on LLVMOpBase:
https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/IR/Intrinsics.td#L1184

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

98–109

Added a LLVM_i8Ptr constraint similar to simple LLVM_i8

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

 // LLVM dialect type constraints.
 //===----------------------------------------------------------------------===//
 
 // LLVM dialect type.
 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>()">,
   "LLVM integer type">;
 
 // Type constraints accepting LLVM integer type of a specific width.
 class LLVM_IntBase<int width> :
     Type<And<[

 def LLVM_AnyFloat : Type<
   CPred<"::mlir::LLVM::isCompatibleFloatingPointType($_self)">,
   "floating point LLVM type">;
 
 // Type constraint accepting any LLVM pointer type.
 def LLVM_AnyPointer : Type<CPred<"$_self.isa<::mlir::LLVM::LLVMPointerType>()">,
                           "LLVM pointer type">;
 
 // Type constraint accepting LLVM pointer type with an additional constraint
 // on the element type.
 class LLVM_PointerTo<Type pointee> : Type<
   And<[LLVM_AnyPointer.predicate,
        SubstLeaves<
          "$_self",
          "$_self.cast<::mlir::LLVM::LLVMPointerType>().getElementType()",
          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">;
 
 // Type constraint accepting opaque LLVM structure type.
 def LLVM_OpaqueStruct : Type<
   And<[LLVM_AnyStruct.predicate,
        CPred<"$_self.cast<::mlir::LLVM::LLVMStructType>().isOpaque()">]>>;

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

def LLVM_SMulWithOverflowOp
   let arguments = (ins LLVM_Type, LLVM_Type);
 }
 def LLVM_UMulWithOverflowOp
     : LLVM_IntrOp<"umul.with.overflow", [0], [], [], 2> {
   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.
 //
 
 def LLVM_vector_reduce_add : LLVM_VectorReduction<"add">;
 def LLVM_vector_reduce_and : LLVM_VectorReduction<"and">;
 def LLVM_vector_reduce_mul : LLVM_VectorReduction<"mul">;
 def LLVM_vector_reduce_fmax : LLVM_VectorReduction<"fmax">;
 def LLVM_vector_reduce_fmin : LLVM_VectorReduction<"fmin">;

mlir/test/Target/llvmir-intrinsics.mlir

 llvm.func @umul_with_overflow_test(%arg0: i32, %arg1: i32, %arg2: vector<8xi32>, %arg3: vector<8xi32>) {
   // CHECK: call { i32, i1 } @llvm.umul.with.overflow.i32
   "llvm.intr.umul.with.overflow"(%arg0, %arg1) : (i32, i32) -> !llvm.struct<(i32, i1)>
   // CHECK: call { <8 x i32>, <8 x i1> } @llvm.umul.with.overflow.v8i32
   "llvm.intr.umul.with.overflow"(%arg2, %arg3) : (vector<8xi32>, vector<8xi32>) -> !llvm.struct<(vector<8xi32>, vector<8xi1>)>
   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
 // CHECK-DAG: declare float @llvm.fmuladd.f32(float, float, float)
 // CHECK-DAG: declare <8 x float> @llvm.fmuladd.v8f32(<8 x float>, <8 x float>, <8 x float>) #0
 // CHECK-DAG: declare void @llvm.prefetch.p0i8(i8* nocapture readonly, i32 immarg, i32 immarg, i32)
 // CHECK-DAG: declare float @llvm.exp.f32(float)
 // CHECK-DAG: declare <8 x float> @llvm.exp.v8f32(<8 x float>) #0

 // CHECK-DAG: declare { i32, i1 } @llvm.uadd.with.overflow.i32(i32, i32)
 // CHECK-DAG: declare { <8 x i32>, <8 x i1> } @llvm.uadd.with.overflow.v8i32(<8 x i32>, <8 x i32>) #0
 // CHECK-DAG: declare { i32, i1 } @llvm.ssub.with.overflow.i32(i32, i32)
 // CHECK-DAG: declare { <8 x i32>, <8 x i1> } @llvm.ssub.with.overflow.v8i32(<8 x i32>, <8 x i32>) #0
 // CHECK-DAG: declare { i32, i1 } @llvm.usub.with.overflow.i32(i32, i32)
 // 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*)