[MLIR][LLVM] Inline LLVM::StackSaveOp and LLVM::StackRestoreOp.

[definelicht]

Summary

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.

[definelicht]

Update naming and documentation to reflect dynamic alloca support.

[Dinistro]

LGTM, just added a minor NIT comment.

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

81

NIT: I would prefer an llvm::any_of here.

[gysit]

LGTM!

It would be good to double check before landing if LLVM somehow handles existing stack save and restore calls in the callable.

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

43

nit: I would start a new sentence after here (during code generation. The handler inserts ...)

57–58

nit: part of the entry block and not dynamic?

74

nit: I would drop this check and move the hasDynamicAlloca check from the end to the beginning of the outer loop to save one level of indention.

[definelicht]

Address review comments and add additional inline documentation.

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

while (!stack.empty()) {
     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, and insert
-/// lifetime intrinsics that limit their scope to the inlined blocks.
-static void moveConstantAllocasToEntryBlock(
+/// 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)
     // Nothing to do.
     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);
     // 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)
+  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;
+      if (shouldInsertLifetime)
         builder.create<LLVM::LifetimeEndOp>(
             allocaOp.getLoc(), arraySize.getValue().getLimitedValue(),
             allocaOp.getResult());
     }
   }
 }
 
 /// Tries to find and return the alignment of the pointer `value` by looking for
 /// an alignment attribute on the defining allocation op or function argument.
 /// If no such attribute is found, returns 1 (i.e., assume that no alignment is
 /// guaranteed).
if (isa<LLVM::AllocaOp,
             LLVM::DbgDeclareOp,
             LLVM::DbgValueOp,
             LLVM::LifetimeEndOp,
             LLVM::LifetimeStartOp,
             LLVM::LoadOp,
             LLVM::MemcpyOp,
             LLVM::MemmoveOp,
             LLVM::MemsetOp,
+            LLVM::StackRestoreOp,
+            LLVM::StackSaveOp,
             LLVM::StoreOp,
             LLVM::UnreachableOp>(op))
       return true;
     // clang-format on
     LLVM_DEBUG(llvm::dbgs()
                << "Cannot inline: unhandled side effecting operation \""
                << op->getName() << "\"\n");
     return false;
if (std::optional<NamedAttribute> attr =
                                  requestedAlignment);
     }
     return argument;
   }
 
   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
   // StringAttrs instead of looking up strings, since StringAttrs are bound to
   // the current context and thus cannot be initialized as static fields.
   const DenseSet<StringAttr> disallowedFunctionAttrs;
 };
 
 } // end anonymous namespace
 
 void LLVM::detail::addLLVMInlinerInterface(LLVM::LLVMDialect *dialect) {
   dialect->addInterfaces<LLVMInlinerInterface>();
 }

mlir/test/Dialect/LLVMIR/inlining.mlir

 // RUN: mlir-opt %s -inline -split-input-file | FileCheck %s
 
 #file = #llvm.di_file<"foo.mlir" in "/foo/">
 #variable = #llvm.di_local_variable<scope = #file>
 #variableAddr = #llvm.di_local_variable<scope = #file>
 
 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
   llvm.intr.dbg.declare #variableAddr = %ptr : !llvm.ptr
   %byte = llvm.mlir.constant(43 : i8) : i8
   %volatile = llvm.mlir.constant(1 : i1) : i1
   "llvm.intr.memset"(%ptr, %byte, %0, %volatile) : (!llvm.ptr, i8, i32, i1) -> ()
   "llvm.intr.memmove"(%ptr, %ptr, %0, %volatile) : (!llvm.ptr, !llvm.ptr, i32, i1) -> ()
   "llvm.intr.memcpy"(%ptr, %ptr, %0, %volatile) : (!llvm.ptr, !llvm.ptr, i32, i1) -> ()
   llvm.cond_br %volatile, ^bb1, ^bb2
 ^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]]
 // CHECK: llvm.intr.dbg.declare #{{.+}} = %[[PTR]]
 // CHECK: "llvm.intr.memset"(%[[PTR]]
 // 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
 }
 
 // -----
 
 func.func @inner_func_not_inlinable() -> i32 {
llvm.func @test_inline(%cond : i1, %size : i32) -> f32 {
   // 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
+  // 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
 }
 
 // -----
 
 llvm.func @static_alloca_not_in_entry(%cond : i1) -> f32 {
   llvm.cond_br %cond, ^bb1, ^bb2
 ^bb1: