From dbbe0109086d8e813fbabf3114da737086250ff9 Mon Sep 17 00:00:00 2001
From: Chuanqi Xu <yedeng.yd@linux.alibaba.com>
Date: Tue, 25 Jan 2022 19:24:49 +0800
Subject: [PATCH] [MLIR] [AsyncToLLVM] Use llvm.coro.align intrinsic

Use llvm.coro.align to align coroutine frame properly.

Reviewed By: bkramer

Differential Revision: https://reviews.llvm.org/D117978
---
 mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td        |  4 ++++
 mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp    | 26 +++++++++++++---------
 .../AsyncToLLVM/convert-coro-to-llvm.mlir          | 14 +++++++-----
 mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir     |  9 ++++++++
 4 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index 58d8cc0..946e674 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -1542,6 +1542,10 @@ def LLVM_CoroSizeOp : LLVM_IntrOp<"coro.size", [0], [], [], 1> {
   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
index f0342f3..5e2bc50 100644
--- a/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
+++ b/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
@@ -335,27 +335,31 @@ public:
     // Get coroutine frame size: @llvm.coro.size.i64.
     Value coroSize =
         rewriter.create<LLVM::CoroSizeOp>(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<LLVM::CoroAlignOp>(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<LLVM::ConstantOp>(
           op->getLoc(), rewriter.getI64Type(), rewriter.getI64IntegerAttr(c));
     };
-    // Round up the size to the alignment. This is a requirement of
-    // aligned_alloc.
-    coroSize = rewriter.create<LLVM::AddOp>(op->getLoc(), coroSize,
-                                            makeConstant(coroAlign - 1));
-    coroSize = rewriter.create<LLVM::AndOp>(op->getLoc(), coroSize,
-                                            makeConstant(-coroAlign));
+    coroSize = rewriter.create<LLVM::AddOp>(op->getLoc(), coroSize, coroAlign);
+    coroSize =
+        rewriter.create<LLVM::SubOp>(op->getLoc(), coroSize, makeConstant(1));
+    Value NegCoroAlign =
+        rewriter.create<LLVM::SubOp>(op->getLoc(), makeConstant(0), coroAlign);
+    coroSize =
+        rewriter.create<LLVM::AndOp>(op->getLoc(), coroSize, NegCoroAlign);
 
     // Allocate memory for the coroutine frame.
     auto allocFuncOp = LLVM::lookupOrCreateAlignedAllocFn(
         op->getParentOfType<ModuleOp>(), rewriter.getI64Type());
     auto coroAlloc = rewriter.create<LLVM::CallOp>(
         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
index 8562377..f84c596 100644
--- a/mlir/test/Conversion/AsyncToLLVM/convert-coro-to-llvm.mlir
+++ b/mlir/test/Conversion/AsyncToLLVM/convert-coro-to-llvm.mlir
@@ -14,12 +14,14 @@ func @coro_begin() {
   // 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
index 752e9e9..f2f83d7 100644
--- a/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
@@ -446,6 +446,15 @@ llvm.func @coro_size() {
   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<i8>) {
   // CHECK: call token @llvm.coro.save
-- 
2.7.4