[Mlir-commits] [mlir] db59654 - [MLIR][LLVM] Generate LLVM lifetime intrinsics while inlining.
Tobias Gysi
llvmlistbot at llvm.org
Mon Jan 30 05:37:35 PST 2023
Author: Johannes de Fine Licht
Date: 2023-01-30T14:33:12+01:00
New Revision: db59654e236780ea713eb1015334dc4ee555f576
URL: https://github.com/llvm/llvm-project/commit/db59654e236780ea713eb1015334dc4ee555f576
DIFF: https://github.com/llvm/llvm-project/commit/db59654e236780ea713eb1015334dc4ee555f576.diff
LOG: [MLIR][LLVM] Generate LLVM lifetime intrinsics while inlining.
Extend `LLVMInlinerInterface` to inline lifetime intrinsics for
`LLVM::AllocaOp` operations, and to insert new lifetime intrinsics when
an alloca is moved to the entry block that restrict its scope to where
the call was before inlining.
Depends on D142436
Reviewed By: gysit
Differential Revision: https://reviews.llvm.org/D142701
Added:
Modified:
mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
mlir/test/Dialect/LLVMIR/inlining.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 017a59b0b2066..1940bcda7b619 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -2848,8 +2848,25 @@ struct LLVMOpAsmDialectInterface : public OpAsmDialectInterface {
// DialectInlinerInterface
//===----------------------------------------------------------------------===//
+/// Check whether the given alloca is an input to a lifetime intrinsic,
+/// optionally passing through one or more casts on the way. This is not
+/// transitive through block arguments.
+static bool hasLifetimeMarkers(LLVM::AllocaOp allocaOp) {
+ SmallVector<Operation *> stack(allocaOp->getUsers().begin(),
+ allocaOp->getUsers().end());
+ while (!stack.empty()) {
+ Operation *op = stack.pop_back_val();
+ if (isa<LLVM::LifetimeStartOp, LLVM::LifetimeEndOp>(op))
+ return true;
+ if (isa<LLVM::BitcastOp>(op))
+ stack.append(op->getUsers().begin(), op->getUsers().end());
+ }
+ return false;
+}
+
/// Move all alloca operations with a constant size in the former entry block of
-/// the newly inlined callee into the entry block of the caller.
+/// the newly inlined callee into the entry block of the caller, and insert
+/// lifetime intrinsics that limit their scope to the inlined blocks.
static void moveConstantAllocasToEntryBlock(
iterator_range<Region::iterator> inlinedBlocks) {
Block *calleeEntryBlock = &(*inlinedBlocks.begin());
@@ -2857,22 +2874,52 @@ static void moveConstantAllocasToEntryBlock(
if (calleeEntryBlock == callerEntryBlock)
// Nothing to do.
return;
- SmallVector<std::pair<LLVM::AllocaOp, IntegerAttr>> allocasToMove;
+ SmallVector<std::tuple<LLVM::AllocaOp, IntegerAttr, bool>> allocasToMove;
+ bool shouldInsertLifetimes = false;
// Conservatively only move alloca operations that are part of the entry block
// and do not inspect nested regions, since they may execute conditionally or
// have other unknown semantics.
for (auto allocaOp : calleeEntryBlock->getOps<LLVM::AllocaOp>()) {
IntegerAttr arraySize;
- if (matchPattern(allocaOp.getArraySize(), m_Constant(&arraySize)))
- allocasToMove.emplace_back(allocaOp, arraySize);
+ if (!matchPattern(allocaOp.getArraySize(), m_Constant(&arraySize)))
+ continue;
+ bool shouldInsertLifetime =
+ arraySize.getValue() != 0 && !hasLifetimeMarkers(allocaOp);
+ shouldInsertLifetimes |= shouldInsertLifetime;
+ allocasToMove.emplace_back(allocaOp, arraySize, shouldInsertLifetime);
}
+ if (allocasToMove.empty())
+ return;
OpBuilder builder(callerEntryBlock, callerEntryBlock->begin());
- for (auto &[allocaOp, arraySize] : allocasToMove) {
+ for (auto &[allocaOp, arraySize, shouldInsertLifetime] : allocasToMove) {
auto newConstant = builder.create<LLVM::ConstantOp>(
allocaOp->getLoc(), allocaOp.getArraySize().getType(), arraySize);
+ // Insert a lifetime start intrinsic where the alloca was before moving it.
+ if (shouldInsertLifetime) {
+ OpBuilder::InsertionGuard insertionGuard(builder);
+ builder.setInsertionPoint(allocaOp);
+ builder.create<LLVM::LifetimeStartOp>(
+ allocaOp.getLoc(), arraySize.getValue().getLimitedValue(),
+ allocaOp.getResult());
+ }
allocaOp->moveAfter(newConstant);
allocaOp.getArraySizeMutable().assign(newConstant.getResult());
}
+ if (!shouldInsertLifetimes)
+ return;
+ // Insert a lifetime end intrinsic before each return in the callee function.
+ for (Block &block : inlinedBlocks) {
+ if (!block.getTerminator()->hasTrait<OpTrait::ReturnLike>())
+ continue;
+ builder.setInsertionPoint(block.getTerminator());
+ for (auto &[allocaOp, arraySize, shouldInsertLifetime] : allocasToMove) {
+ if (!shouldInsertLifetime)
+ continue;
+ builder.create<LLVM::LifetimeEndOp>(
+ allocaOp.getLoc(), arraySize.getValue().getLimitedValue(),
+ allocaOp.getResult());
+ }
+ }
}
namespace {
@@ -2912,7 +2959,8 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
return false;
return true;
})
- .Case<LLVM::CallOp, LLVM::AllocaOp>([](auto) { return true; })
+ .Case<LLVM::CallOp, LLVM::AllocaOp, LLVM::LifetimeStartOp,
+ LLVM::LifetimeEndOp>([](auto) { return true; })
.Default([](auto) { return false; });
}
diff --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir
index 5e5c93b973017..65632434b1349 100644
--- a/mlir/test/Dialect/LLVMIR/inlining.mlir
+++ b/mlir/test/Dialect/LLVMIR/inlining.mlir
@@ -231,7 +231,9 @@ llvm.func @test_inline(%cond : i1, %size : i32) -> f32 {
// CHECK: ^{{.+}}:
^bb1:
// CHECK-NOT: llvm.call @static_alloca
+ // CHECK: llvm.intr.lifetime.start
%0 = llvm.call @static_alloca() : () -> f32
+ // CHECK: llvm.intr.lifetime.end
// CHECK: llvm.br
llvm.br ^bb3(%0: f32)
// CHECK: ^{{.+}}:
@@ -275,3 +277,79 @@ llvm.func @test_inline(%cond : i1) -> f32 {
%0 = llvm.call @static_alloca_not_in_entry(%cond) : (i1) -> f32
llvm.return %0 : f32
}
+
+// -----
+
+llvm.func @static_alloca(%cond: i1) -> f32 {
+ %0 = llvm.mlir.constant(4 : i32) : i32
+ %1 = llvm.alloca %0 x f32 : (i32) -> !llvm.ptr
+ llvm.cond_br %cond, ^bb1, ^bb2
+^bb1:
+ %2 = llvm.load %1 : !llvm.ptr -> f32
+ llvm.return %2 : f32
+^bb2:
+ %3 = llvm.mlir.constant(3.14192 : f32) : f32
+ llvm.return %3 : f32
+}
+
+// CHECK-LABEL: llvm.func @test_inline
+llvm.func @test_inline(%cond0 : i1, %cond1 : i1, %funcArg : f32) -> f32 {
+ // CHECK-NOT: llvm.cond_br
+ // CHECK: %[[PTR:.+]] = llvm.alloca
+ // CHECK: llvm.cond_br %{{.+}}, ^[[BB1:.+]], ^{{.+}}
+ llvm.cond_br %cond0, ^bb1, ^bb2
+ // CHECK: ^[[BB1]]
+^bb1:
+ // Make sure the lifetime begin intrinsic has been inserted where the call
+ // used to be, even though the alloca has been moved to the entry block.
+ // CHECK-NEXT: llvm.intr.lifetime.start 4, %[[PTR]]
+ %0 = llvm.call @static_alloca(%cond1) : (i1) -> f32
+ // CHECK: llvm.cond_br %{{.+}}, ^[[BB2:.+]], ^[[BB3:.+]]
+ llvm.br ^bb3(%0: f32)
+ // Make sure the lifetime end intrinsic has been inserted at both former
+ // return sites of the callee.
+ // CHECK: ^[[BB2]]:
+ // CHECK-NEXT: llvm.load
+ // CHECK-NEXT: llvm.intr.lifetime.end 4, %[[PTR]]
+ // CHECK: ^[[BB3]]:
+ // CHECK-NEXT: llvm.intr.lifetime.end 4, %[[PTR]]
+^bb2:
+ llvm.br ^bb3(%funcArg: f32)
+^bb3(%blockArg: f32):
+ llvm.return %blockArg : f32
+}
+
+// -----
+
+llvm.func @alloca_with_lifetime(%cond: i1) -> f32 {
+ %0 = llvm.mlir.constant(4 : i32) : i32
+ %1 = llvm.alloca %0 x f32 : (i32) -> !llvm.ptr
+ llvm.intr.lifetime.start 4, %1 : !llvm.ptr
+ %2 = llvm.load %1 : !llvm.ptr -> f32
+ llvm.intr.lifetime.end 4, %1 : !llvm.ptr
+ %3 = llvm.fadd %2, %2 : f32
+ llvm.return %3 : f32
+}
+
+// CHECK-LABEL: llvm.func @test_inline
+llvm.func @test_inline(%cond0 : i1, %cond1 : i1, %funcArg : f32) -> f32 {
+ // CHECK-NOT: llvm.cond_br
+ // CHECK: %[[PTR:.+]] = llvm.alloca
+ // CHECK: llvm.cond_br %{{.+}}, ^[[BB1:.+]], ^{{.+}}
+ llvm.cond_br %cond0, ^bb1, ^bb2
+ // CHECK: ^[[BB1]]
+^bb1:
+ // Make sure the original lifetime intrinsic has been preserved, rather than
+ // inserting a new one with a larger scope.
+ // CHECK: llvm.intr.lifetime.start 4, %[[PTR]]
+ // CHECK-NEXT: llvm.load %[[PTR]]
+ // CHECK-NEXT: llvm.intr.lifetime.end 4, %[[PTR]]
+ // CHECK: llvm.fadd
+ // CHECK-NOT: llvm.intr.lifetime.end
+ %0 = llvm.call @alloca_with_lifetime(%cond1) : (i1) -> f32
+ llvm.br ^bb3(%0: f32)
+^bb2:
+ llvm.br ^bb3(%funcArg: f32)
+^bb3(%blockArg: f32):
+ llvm.return %blockArg : f32
+}
More information about the Mlir-commits
mailing list