From db59654e236780ea713eb1015334dc4ee555f576 Mon Sep 17 00:00:00 2001 From: Johannes de Fine Licht Date: Mon, 30 Jan 2023 14:33:02 +0100 Subject: [PATCH] [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 --- mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp | 60 ++++++++++++++++++++--- mlir/test/Dialect/LLVMIR/inlining.mlir | 78 ++++++++++++++++++++++++++++++ 2 files changed, 132 insertions(+), 6 deletions(-) diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp index 017a59b..1940bcd 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 stack(allocaOp->getUsers().begin(), + allocaOp->getUsers().end()); + while (!stack.empty()) { + Operation *op = stack.pop_back_val(); + if (isa(op)) + return true; + if (isa(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 inlinedBlocks) { Block *calleeEntryBlock = &(*inlinedBlocks.begin()); @@ -2857,22 +2874,52 @@ static void moveConstantAllocasToEntryBlock( if (calleeEntryBlock == callerEntryBlock) // Nothing to do. return; - SmallVector> allocasToMove; + SmallVector> 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()) { 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( 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( + 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()) + continue; + builder.setInsertionPoint(block.getTerminator()); + for (auto &[allocaOp, arraySize, shouldInsertLifetime] : allocasToMove) { + if (!shouldInsertLifetime) + continue; + builder.create( + allocaOp.getLoc(), arraySize.getValue().getLimitedValue(), + allocaOp.getResult()); + } + } } namespace { @@ -2912,7 +2959,8 @@ struct LLVMInlinerInterface : public DialectInlinerInterface { return false; return true; }) - .Case([](auto) { return true; }) + .Case([](auto) { return true; }) .Default([](auto) { return false; }); } diff --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir index 5e5c93b..6563243 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 +} -- 2.7.4