diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td @@ -52,6 +52,12 @@ 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 @@ pointee.predicate>]>, "LLVM pointer to " # pointee.summary>; +// Type constraints accepting LLVM pointer type to integer of a specific width. +class LLVM_IntPtrBase : Type< + LLVM_PointerTo>.predicate, + "LLVM pointer to " # LLVM_IntBase.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()">, "LLVM structure type">; diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td @@ -1226,6 +1226,57 @@ 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. // diff --git a/mlir/test/Target/llvmir-intrinsics.mlir b/mlir/test/Target/llvmir-intrinsics.mlir --- a/mlir/test/Target/llvmir-intrinsics.mlir +++ b/mlir/test/Target/llvmir-intrinsics.mlir @@ -347,6 +347,64 @@ llvm.return } +// CHECK-LABEL: @coro_id +llvm.func @coro_id(%arg0: i32, %arg1: !llvm.ptr) { + // 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) { + // CHECK: call i8* @llvm.coro.begin + llvm.intr.coro.begin %arg0, %arg1 : !llvm.ptr + 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) { + // 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, %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) { + // CHECK: call i8* @llvm.coro.free + %0 = llvm.intr.coro.free %arg0, %arg1 : !llvm.ptr + llvm.return +} + +// CHECK-LABEL: @coro_resume +llvm.func @coro_resume(%arg0: !llvm.ptr) { + // 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 @@ // 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*)