[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

Author: definelicht

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

@@ -2848,31 +2848,78 @@ 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());
   Block *callerEntryBlock = &(*calleeEntryBlock->getParent()->begin());
   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; });
   }
 

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
+}