[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/CMakeLists.txt

add_mlir_dialect_library(MLIRLLVMDialect
   LINK_COMPONENTS
   AsmParser
   BinaryFormat
   BitReader
   BitWriter
   Core
 
   LINK_LIBS PUBLIC
+  MLIRAnalysis
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIRDataLayoutInterfaces
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRSideEffectInterfaces
   MLIRSupport
   )

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

 //===- LLVMDialect.cpp - LLVM IR Ops and Dialect registration -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the types and operation details for the LLVM IR dialect in
 // MLIR, and the LLVM IR dialect.  It also registers the dialect.
 //
 //===----------------------------------------------------------------------===//
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "TypeDetail.h"
+#include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Matchers.h"
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.
+static bool hasLifetimeMarkers(LLVM::AllocaOp allocaOp) {
+  SetVector<Operation *> forwardSlice;
+  bool foundMarker = false;
+  getForwardSlice(allocaOp.getOperation(), &forwardSlice, [&](Operation *op) {
+    foundMarker |= isa<LLVM::LifetimeStartOp, LLVM::LifetimeEndOp>(op);
+    // We have to include `LLVM::AllocaOp` so `getForwardSlice` does not exclude
+    // the initial operation itself.
+    return !foundMarker && isa<LLVM::AllocaOp, LLVM::BitcastOp>(op);
+  });
+  return foundMarker;
+}
+
 /// 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);
   }
   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
+}