[MLIR][LLVM] Generate LLVM lifetime intrinsics while inlining.

[definelicht]

Summary

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

[definelicht]

Drop MLIRAnalysis dependency in favor of simple stack-based implementation.

[gysit]

LGTM!

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

struct LLVMOpAsmDialectInterface : public OpAsmDialectInterface {
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // 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 {
 struct LLVMInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
   bool isLegalToInline(Operation *call, Operation *callable,
                        bool wouldBeCloned) const final {
return llvm::TypeSwitch<Operation *, bool>(op)
           if (memOp.getAccessGroups())
             return false;
           if (memOp.getAliasScopes())
             return false;
           if (memOp.getNoaliasScopes())
             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; });
   }
 
   /// Handle the given inlined return by replacing it with a branch. This
   /// overload is called when the inlined region has more than one block.
   void handleTerminator(Operation *op, Block *newDest) const final {
     // Only return needs to be handled here.
     auto returnOp = dyn_cast<LLVM::ReturnOp>(op);

mlir/test/Dialect/LLVMIR/inlining.mlir

llvm.func @test_inline(%cond : i1, %size : i32) -> f32 {
   // CHECK-NOT: ^{{.+}}:
   // CHECK-NEXT: llvm.mlir.constant
   // CHECK-NEXT: llvm.alloca
   // CHECK: llvm.cond_br
   llvm.cond_br %cond, ^bb1, ^bb2
   // 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: ^{{.+}}:
 ^bb2:
   // Check that the dynamic alloca was inlined, but that it was not moved to the
   // entry block.
   // CHECK: llvm.add
   // CHECK-NEXT: llvm.alloca

 llvm.func @test_inline(%cond : i1) -> f32 {
   // Make sure the alloca was not moved to the entry block.
   // CHECK-NOT: llvm.alloca
   // CHECK: llvm.cond_br
   // CHECK: llvm.alloca
   %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
+}