[Mlir-commits] [mlir] 964dc36 - [AsyncToLLVM] aligned_alloc requires the size to be a multiple of aignment, so round up

[Benjamin Kramer]

Author: Benjamin Kramer
Date: 2022-01-17T21:48:00+01:00
New Revision: 964dc368e7c72ad5b7ab995c97920c4deb624631

URL: https://github.com/llvm/llvm-project/commit/964dc368e7c72ad5b7ab995c97920c4deb624631
DIFF: https://github.com/llvm/llvm-project/commit/964dc368e7c72ad5b7ab995c97920c4deb624631.diff

LOG: [AsyncToLLVM] aligned_alloc requires the size to be a multiple of aignment, so round up

Fixes a crash with debug malloc.

Added: 
    

Modified: 
    mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
    mlir/test/Conversion/AsyncToLLVM/convert-coro-to-llvm.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp b/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
index 7d504709082de..c332f4da5722f 100644
--- a/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
+++ b/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
@@ -333,20 +333,29 @@ class CoroBeginOpConversion : public OpConversionPattern<CoroBeginOp> {
     auto loc = op->getLoc();
 
     // Get coroutine frame size: @llvm.coro.size.i64.
-    auto coroSize =
+    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
-    auto coroAlign = rewriter.create<LLVM::ConstantOp>(
-        op->getLoc(), rewriter.getI64Type(), rewriter.getI64IntegerAttr(64));
+    constexpr int64_t coroAlign = 64;
+    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));
 
     // 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{coroAlign, coroSize.getResult()});
+        ValueRange{makeConstant(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 1377854bf3ab4..85623772249f9 100644
--- a/mlir/test/Conversion/AsyncToLLVM/convert-coro-to-llvm.mlir
+++ b/mlir/test/Conversion/AsyncToLLVM/convert-coro-to-llvm.mlir
@@ -14,8 +14,12 @@ 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]], %[[SIZE]])
+  // CHECK: %[[ALLOC:.*]] = llvm.call @aligned_alloc(%[[ALIGN]], %[[SIZE3]])
   // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin %[[ID]], %[[ALLOC]]
   %1 = async.coro.begin %0
   return