[Mlir-commits] [mlir] 7387972 - [MLIR][LLVM] Inline LLVM::StackSaveOp and LLVM::StackRestoreOp.
Johannes de Fine Licht
llvmlistbot at llvm.org
Thu Apr 13 01:37:37 PDT 2023
Author: Johannes de Fine Licht
Date: 2023-04-13T08:27:29Z
New Revision: 7387972352dbcbfd230b465b54aff65cdc411576
URL: https://github.com/llvm/llvm-project/commit/7387972352dbcbfd230b465b54aff65cdc411576
DIFF: https://github.com/llvm/llvm-project/commit/7387972352dbcbfd230b465b54aff65cdc411576.diff
LOG: [MLIR][LLVM] Inline LLVM::StackSaveOp and LLVM::StackRestoreOp.
Support LLVM::StackSaveOp and LLVM::StackRestoreOp in the LLVM dialect
inliner in MLIR.
Inserts new LLVM::StackSaveOp and LLVM::StackRestoreOp intrinsics when
dynamic allocas are detected in the inlined blocks. This may result in
multiple saves/restores in the same block if some are already present in
the caller, which is legal IR, but is cleaned up in LLVM. There is not
yet a canonicalization pattern for this on LLVM dialect in MLIR.
Reviewed By: Dinistro
Differential Revision: https://reviews.llvm.org/D148011
Added:
Modified:
mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
mlir/test/Dialect/LLVMIR/inlining.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
index aacde746224b8..d1d452124cefd 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
@@ -37,11 +37,17 @@ static bool hasLifetimeMarkers(LLVM::AllocaOp allocaOp) {
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, and insert
-/// lifetime intrinsics that limit their scope to the inlined blocks.
-static void moveConstantAllocasToEntryBlock(
- iterator_range<Region::iterator> inlinedBlocks) {
+/// Handles alloca operations in the inlined blocks:
+/// - Moves all alloca operations with a constant size in the former entry block
+/// of the callee into the entry block of the caller, so they become part of
+/// the function prologue/epilogue during code generation.
+/// - Inserts lifetime intrinsics that limit the scope of inlined static allocas
+/// to the inlined blocks.
+/// - Inserts StackSave and StackRestore operations if dynamic allocas were
+/// inlined.
+static void
+handleInlinedAllocas(Operation *call,
+ iterator_range<Region::iterator> inlinedBlocks) {
Block *calleeEntryBlock = &(*inlinedBlocks.begin());
Block *callerEntryBlock = &(*calleeEntryBlock->getParent()->begin());
if (calleeEntryBlock == callerEntryBlock)
@@ -49,21 +55,43 @@ static void moveConstantAllocasToEntryBlock(
return;
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.
+ bool hasDynamicAlloca = false;
+ // Conservatively only move static 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)))
+ if (!matchPattern(allocaOp.getArraySize(), m_Constant(&arraySize))) {
+ hasDynamicAlloca = true;
continue;
+ }
bool shouldInsertLifetime =
arraySize.getValue() != 0 && !hasLifetimeMarkers(allocaOp);
shouldInsertLifetimes |= shouldInsertLifetime;
allocasToMove.emplace_back(allocaOp, arraySize, shouldInsertLifetime);
}
- if (allocasToMove.empty())
+ // Check the remaining inlined blocks for dynamic allocas as well.
+ for (Block &block : llvm::drop_begin(inlinedBlocks)) {
+ if (hasDynamicAlloca)
+ break;
+ hasDynamicAlloca =
+ llvm::any_of(block.getOps<LLVM::AllocaOp>(), [](auto allocaOp) {
+ return !matchPattern(allocaOp.getArraySize(), m_Constant());
+ });
+ }
+ if (allocasToMove.empty() && !hasDynamicAlloca)
return;
- OpBuilder builder(callerEntryBlock, callerEntryBlock->begin());
+ OpBuilder builder(calleeEntryBlock, calleeEntryBlock->begin());
+ Value stackPtr;
+ if (hasDynamicAlloca) {
+ // This may result in multiple stacksave/stackrestore intrinsics in the same
+ // scope if some are already present in the body of the caller. This is not
+ // invalid IR, but LLVM cleans these up in InstCombineCalls.cpp, along with
+ // other cases where the stacksave/stackrestore is redundant.
+ stackPtr = builder.create<LLVM::StackSaveOp>(
+ call->getLoc(), LLVM::LLVMPointerType::get(call->getContext()));
+ }
+ builder.setInsertionPoint(callerEntryBlock, callerEntryBlock->begin());
for (auto &[allocaOp, arraySize, shouldInsertLifetime] : allocasToMove) {
auto newConstant = builder.create<LLVM::ConstantOp>(
allocaOp->getLoc(), allocaOp.getArraySize().getType(), arraySize);
@@ -78,19 +106,20 @@ static void moveConstantAllocasToEntryBlock(
allocaOp->moveAfter(newConstant);
allocaOp.getArraySizeMutable().assign(newConstant.getResult());
}
- if (!shouldInsertLifetimes)
+ if (!shouldInsertLifetimes && !hasDynamicAlloca)
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());
+ if (hasDynamicAlloca)
+ builder.create<LLVM::StackRestoreOp>(call->getLoc(), stackPtr);
for (auto &[allocaOp, arraySize, shouldInsertLifetime] : allocasToMove) {
- if (!shouldInsertLifetime)
- continue;
- builder.create<LLVM::LifetimeEndOp>(
- allocaOp.getLoc(), arraySize.getValue().getLimitedValue(),
- allocaOp.getResult());
+ if (shouldInsertLifetime)
+ builder.create<LLVM::LifetimeEndOp>(
+ allocaOp.getLoc(), arraySize.getValue().getLimitedValue(),
+ allocaOp.getResult());
}
}
}
@@ -308,6 +337,8 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
LLVM::MemcpyOp,
LLVM::MemmoveOp,
LLVM::MemsetOp,
+ LLVM::StackRestoreOp,
+ LLVM::StackSaveOp,
LLVM::StoreOp,
LLVM::UnreachableOp>(op))
return true;
@@ -369,12 +400,7 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
void processInlinedCallBlocks(
Operation *call,
iterator_range<Region::iterator> inlinedBlocks) const override {
- // Alloca operations with a constant size that were in the entry block of
- // the callee should be moved to the entry block of the caller, as this will
- // fold into prologue/epilogue code during code generation.
- // This is not implemented as a standalone pattern because we need to know
- // which newly inlined block was previously the entry block of the callee.
- moveConstantAllocasToEntryBlock(inlinedBlocks);
+ handleInlinedAllocas(call, inlinedBlocks);
}
// Keeping this (immutable) state on the interface allows us to look up
diff --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir
index 873ab1d40901d..a1f4838dfee98 100644
--- a/mlir/test/Dialect/LLVMIR/inlining.mlir
+++ b/mlir/test/Dialect/LLVMIR/inlining.mlir
@@ -6,6 +6,7 @@
func.func @inner_func_inlinable(%ptr : !llvm.ptr) -> i32 {
%0 = llvm.mlir.constant(42 : i32) : i32
+ %stack = llvm.intr.stacksave : !llvm.ptr
llvm.store %0, %ptr { alignment = 8 } : i32, !llvm.ptr
%1 = llvm.load %ptr { alignment = 8 } : !llvm.ptr -> i32
llvm.intr.dbg.value #variable = %0 : i32
@@ -19,12 +20,14 @@ func.func @inner_func_inlinable(%ptr : !llvm.ptr) -> i32 {
^bb1:
llvm.unreachable
^bb2:
+ llvm.intr.stackrestore %stack : !llvm.ptr
return %1 : i32
}
// CHECK-LABEL: func.func @test_inline(
// CHECK-SAME: %[[PTR:[a-zA-Z0-9_]+]]
// CHECK: %[[CST:.*]] = llvm.mlir.constant(42
+// CHECK: %[[STACK:.+]] = llvm.intr.stacksave
// CHECK: llvm.store %[[CST]], %[[PTR]]
// CHECK: %[[RES:.+]] = llvm.load %[[PTR]]
// CHECK: llvm.intr.dbg.value #{{.+}} = %[[CST]]
@@ -33,6 +36,7 @@ func.func @inner_func_inlinable(%ptr : !llvm.ptr) -> i32 {
// CHECK: "llvm.intr.memmove"(%[[PTR]], %[[PTR]]
// CHECK: "llvm.intr.memcpy"(%[[PTR]], %[[PTR]]
// CHECK: llvm.unreachable
+// CHECK: llvm.intr.stackrestore %[[STACK]]
func.func @test_inline(%ptr : !llvm.ptr) -> i32 {
%0 = call @inner_func_inlinable(%ptr) : (!llvm.ptr) -> i32
return %0 : i32
@@ -253,20 +257,23 @@ llvm.func @test_inline(%cond : i1, %size : i32) -> f32 {
// CHECK: llvm.intr.lifetime.start
%0 = llvm.call @static_alloca() : () -> f32
// CHECK: llvm.intr.lifetime.end
- // CHECK: llvm.br
+ // CHECK: llvm.br ^[[BB3:[a-zA-Z0-9_]+]]
llvm.br ^bb3(%0: f32)
// CHECK: ^{{.+}}:
^bb2:
// Check that the dynamic alloca was inlined, but that it was not moved to the
// entry block.
+ // CHECK: %[[STACK:[a-zA-Z0-9_]+]] = llvm.intr.stacksave
// CHECK: llvm.add
- // CHECK-NEXT: llvm.alloca
+ // CHECK: llvm.alloca
+ // CHECK: llvm.intr.stackrestore %[[STACK]]
// CHECK-NOT: llvm.call @dynamic_alloca
%1 = llvm.call @dynamic_alloca(%size) : (i32) -> f32
- // CHECK: llvm.br
+ // CHECK: llvm.br ^[[BB3]]
llvm.br ^bb3(%1: f32)
- // CHECK: ^{{.+}}:
+ // CHECK: ^[[BB3]]
^bb3(%arg : f32):
+ // CHECK-NEXT: return
llvm.return %arg : f32
}
More information about the Mlir-commits
mailing list