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 @@ -1542,6 +1542,10 @@ let assemblyFormat = "attr-dict `:` type($res)"; } +def LLVM_CoroAlignOp : LLVM_IntrOp<"coro.align", [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)"; diff --git a/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp b/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp --- a/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp +++ b/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp @@ -335,27 +335,31 @@ // Get coroutine frame size: @llvm.coro.size.i64. Value coroSize = rewriter.create(loc, rewriter.getI64Type()); - // The coroutine lowering doesn't properly account for alignment of the - // frame, so align everything to 64 bytes which ought to be enough for - // everyone. https://llvm.org/PR53148 - constexpr int64_t coroAlign = 64; + // Get coroutine frame alignment: @llvm.coro.align.i64. + Value coroAlign = + rewriter.create(loc, rewriter.getI64Type()); + + // Round up the size to be multiple of the alignment. Since aligned_alloc + // requires the size parameter be an integral multiple of the alignment + // parameter. auto makeConstant = [&](uint64_t c) { return rewriter.create( op->getLoc(), rewriter.getI64Type(), rewriter.getI64IntegerAttr(c)); }; - // Round up the size to the alignment. This is a requirement of - // aligned_alloc. - coroSize = rewriter.create(op->getLoc(), coroSize, - makeConstant(coroAlign - 1)); - coroSize = rewriter.create(op->getLoc(), coroSize, - makeConstant(-coroAlign)); + coroSize = rewriter.create(op->getLoc(), coroSize, coroAlign); + coroSize = + rewriter.create(op->getLoc(), coroSize, makeConstant(1)); + Value NegCoroAlign = + rewriter.create(op->getLoc(), makeConstant(0), coroAlign); + coroSize = + rewriter.create(op->getLoc(), coroSize, NegCoroAlign); // Allocate memory for the coroutine frame. auto allocFuncOp = LLVM::lookupOrCreateAlignedAllocFn( op->getParentOfType(), rewriter.getI64Type()); auto coroAlloc = rewriter.create( loc, i8Ptr, SymbolRefAttr::get(allocFuncOp), - ValueRange{makeConstant(coroAlign), coroSize}); + ValueRange{coroAlign, coroSize}); // Begin a coroutine: @llvm.coro.begin. auto coroId = CoroBeginOpAdaptor(adaptor.getOperands()).id(); diff --git a/mlir/test/Conversion/AsyncToLLVM/convert-coro-to-llvm.mlir b/mlir/test/Conversion/AsyncToLLVM/convert-coro-to-llvm.mlir --- a/mlir/test/Conversion/AsyncToLLVM/convert-coro-to-llvm.mlir +++ b/mlir/test/Conversion/AsyncToLLVM/convert-coro-to-llvm.mlir @@ -14,12 +14,14 @@ // CHECK: %[[ID:.*]] = llvm.intr.coro.id %0 = async.coro.id // CHECK: %[[SIZE:.*]] = llvm.intr.coro.size : i64 - // CHECK: %[[C63:.*]] = llvm.mlir.constant(63 : i64) : i64 - // CHECK: %[[SIZE2:.*]] = llvm.add %[[SIZE]], %[[C63]] : i64 - // CHECK: %[[CN64:.*]] = llvm.mlir.constant(-64 : i64) : i64 - // CHECK: %[[SIZE3:.*]] = llvm.and %[[SIZE2]], %[[CN64]] : i64 - // CHECK: %[[ALIGN:.*]] = llvm.mlir.constant(64 : i64) : i64 - // CHECK: %[[ALLOC:.*]] = llvm.call @aligned_alloc(%[[ALIGN]], %[[SIZE3]]) + // CHECK: %[[ALIGN:.*]] = llvm.intr.coro.align : i64 + // CHECK: %[[SIZE_PLUS_ALIGN:.*]] = llvm.add %[[SIZE]], %[[ALIGN]] : i64 + // CHECK: %[[C1:.*]] = llvm.mlir.constant(1 : i64) : i64 + // CHECK: %[[SIZE_PLUS_ALIGN_MINUS_ONE:.*]] = llvm.sub %[[SIZE_PLUS_ALIGN]], %[[C1]] : i64 + // CHECK: %[[C0:.*]] = llvm.mlir.constant(0 : i64) : i64 + // CHECK: %[[NEGATED_ALIGN:.*]] = llvm.sub %[[C0]], %[[ALIGN]] : i64 + // CHECK: %[[ROUNDED_SIZE:.*]] = llvm.and %[[SIZE_PLUS_ALIGN_MINUS_ONE]], %[[NEGATED_ALIGN]] : i64 + // CHECK: %[[ALLOC:.*]] = llvm.call @aligned_alloc(%[[ALIGN]], %[[ROUNDED_SIZE]]) // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin %[[ID]], %[[ALLOC]] %1 = async.coro.begin %0 return diff --git a/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir b/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir --- a/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir +++ b/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir @@ -446,6 +446,15 @@ llvm.return } +// CHECK-LABEL: @coro_align +llvm.func @coro_align() { + // CHECK: call i64 @llvm.coro.align.i64 + %0 = llvm.intr.coro.align : i64 + // CHECK: call i32 @llvm.coro.align.i32 + %1 = llvm.intr.coro.align : i32 + llvm.return +} + // CHECK-LABEL: @coro_save llvm.func @coro_save(%arg0: !llvm.ptr) { // CHECK: call token @llvm.coro.save