[MLIR][LLVM] Support inlining of LLVM::AllocaOp.

[definelicht]

Summary

If inlining into an LLVM function, move allocas of constant size from
the callee's entry block to the new entry block, as this will fold into
the prologue/epilogue code during code generation.

We still allow inlining allocas even if we cannot do this
post-processing (i.e., when we are not inlining into an LLVM::FuncOp),
since this is an optimization (and thus not necessary for correctness).

Depends on D141682

[definelicht]

Purged outdated comment.

[definelicht]

Update signature and tests.

Updated tests to not return alloca'ed pointers, since this is nonsensical code.
Changed function signature to take the inlined blocks directly.

[gysit]

LGTM!

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

2861–2862

nit: What about changing to

"Conservatively move only 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."

?

2953

ultra nit: alloca -> Alloca

[definelicht]

Update inline comment.

[definelicht]

Address ultra nit on capitalization.

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

struct LLVMOpAsmDialectInterface : public OpAsmDialectInterface {
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // DialectInlinerInterface
 //===----------------------------------------------------------------------===//
 
+/// Move all `LLVM::AllocaOp` operations with a constant size in `sourceBlock`
+/// to the end of `entryBlock`. If necessary, the operation defining the
+/// constant will be moved as well.
+static void moveConstantAllocasToEntryBlock(Block *sourceBlock,
+                                            Block *entryBlock) {
+  assert(entryBlock->isEntryBlock() &&
+         "Second argument must be an entry block.");
+  assert(sourceBlock->getParent() == entryBlock->getParent() &&
+         "Blocks must belong to the same region.");
+  if (sourceBlock == entryBlock)
+    // Nothing to do.
+    return;
+  SmallVector<std::pair<LLVM::AllocaOp, IntegerAttr>> allocasToMove;
+  // Don't walk nested regions of the block, since we can't be sure of their
+  // semantics.
+  for (auto allocaOp : sourceBlock->getOps<LLVM::AllocaOp>()) {
+    IntegerAttr arraySize;
+    if (matchPattern(allocaOp.getArraySize(), m_Constant(&arraySize)))
+      allocasToMove.emplace_back(allocaOp, arraySize);
+  }
+  OpBuilder builder(entryBlock, entryBlock->begin());
+  for (auto &[allocaOp, arraySize] : allocasToMove) {
+    auto newConstant = builder.create<LLVM::ConstantOp>(
+        allocaOp->getLoc(), allocaOp.getArraySize().getType(), arraySize);
+    allocaOp->moveAfter(newConstant);
+    allocaOp.getArraySizeMutable().assign(newConstant.getResult());
+  }
+}
+
 namespace {
 struct LLVMInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
   bool isLegalToInline(Operation *call, Operation *callable,
                        bool wouldBeCloned) const final {
     if (!wouldBeCloned)
       return false;
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>([](auto) { return true; })
+        .Case<LLVM::CallOp, LLVM::AllocaOp>([](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);
void handleTerminator(Operation *op,
     for (const auto &[dst, src] :
          llvm::zip(valuesToRepl, returnOp.getOperands()))
       dst.replaceAllUsesWith(src);
   }
 
 private:
   /// Returns true if all attributes of `callOp` are handled during inlining.
   [[nodiscard]] static bool isLegalToInlineCallAttributes(LLVM::CallOp callOp) {
     return all_of(callOp.getAttributeNames(), [&](StringRef attrName) {
       return llvm::StringSwitch<bool>(attrName)
           // TODO: Propagate and update branch weights.
           .Case("branch_weights", !callOp.getBranchWeights())
           .Case("callee", true)
           .Case("fastmathFlags", true)
           .Default(false);
     });
   }
return all_of(funcOp.getAttributeNames(), [&](StringRef attrName) {
           // personality is set.
           .Case("personality", !funcOp.getPersonality())
           // TODO: Handle result attributes.
           .Case("res_attrs", !funcOp.getResAttrs())
           .Case("sym_name", true)
           .Default(false);
     });
   }
+
+  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.
+    Block &calleeEntryBlock = *inlinedBlocks.begin();
+    Block &callerEntryBlock = *calleeEntryBlock.getParent()->begin();
+    moveConstantAllocasToEntryBlock(&calleeEntryBlock, &callerEntryBlock);
+  }
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // LLVMDialect initialization, type parsing, and registration.
 //===----------------------------------------------------------------------===//
 
 void LLVMDialect::initialize() {

mlir/test/Dialect/LLVMIR/inlining.mlir

 // CHECK-NEXT: return %[[RES]] : i32
 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() -> !llvm.ptr<f64> {
-  %0 = llvm.mlir.constant(0 : i32) : i32
-  %1 = llvm.alloca %0 x f64 : (i32) -> !llvm.ptr<f64>
-  return %1 : !llvm.ptr<f64>
+func.func @inner_func_not_inlinable() -> i32 {
+  %0 = llvm.inline_asm has_side_effects "foo", "bar" : () -> i32
+  return %0 : i32
 }
 
-// CHECK-LABEL: func.func @test_not_inline() -> !llvm.ptr<f64> {
-// CHECK-NEXT: %[[RES:.*]] = call @inner_func_not_inlinable() : () -> !llvm.ptr<f64>
-// CHECK-NEXT: return %[[RES]] : !llvm.ptr<f64>
-func.func @test_not_inline() -> !llvm.ptr<f64> {
-  %0 = call @inner_func_not_inlinable() : () -> !llvm.ptr<f64>
-  return %0 : !llvm.ptr<f64>
+// CHECK-LABEL: func.func @test_not_inline() -> i32 {
+// CHECK-NEXT: %[[RES:.*]] = call @inner_func_not_inlinable() : () -> i32
+// CHECK-NEXT: return %[[RES]] : i32
+func.func @test_not_inline() -> i32 {
+  %0 = call @inner_func_not_inlinable() : () -> i32
+  return %0 : i32
 }
 
 // -----
 
 llvm.metadata @metadata {
   llvm.access_group @group
   llvm.return
 }

 
 // CHECK-LABEL: llvm.func @caller
 // CHECK-NEXT: llvm.call @callee
 // CHECK-NEXT: llvm.return
 llvm.func @caller() {
   llvm.call @callee() { branch_weights = dense<42> : vector<1xi32> } : () -> ()
   llvm.return
 }
+
+// -----
+
+llvm.func @static_alloca() -> !llvm.ptr {
+  %0 = llvm.mlir.constant(4 : i32) : i32
+  %1 = llvm.alloca %0 x f32 : (i32) -> !llvm.ptr
+  llvm.return %1 : !llvm.ptr
+}
+
+llvm.func @dynamic_alloca(%size : i32) -> !llvm.ptr {
+  %0 = llvm.add %size, %size : i32
+  %1 = llvm.alloca %0 x f32 : (i32) -> !llvm.ptr
+  llvm.return %1 : !llvm.ptr
+}
+
+// CHECK-LABEL: llvm.func @test_inline
+// Check that the static alloca was moved to the entry block after inlining
+// with its size defined by a constant.
+// CHECK-NOT: ^{{.+}}:
+// CHECK-NEXT: llvm.mlir.constant
+// CHECK-NEXT: llvm.alloca
+// CHECK: llvm.cond_br
+// CHECK-NOT: llvm.call @static_alloca
+// Check that the dynamic alloca was inlined, but that it was not moved to the
+// entry block.
+// CHECK: llvm.add
+// CHECK-NEXT: llvm.alloca
+// CHECK-NOT: llvm.call @dynamic_alloca
+llvm.func @test_inline(%cond : i1, %size : i32) -> f32 {
+  llvm.cond_br %cond, ^bb1, ^bb2
+^bb1:
+  %0 = llvm.call @static_alloca() : () -> !llvm.ptr
+  llvm.br ^bb3(%0: !llvm.ptr)
+^bb2:
+  %1 = llvm.call @dynamic_alloca(%size) : (i32) -> !llvm.ptr
+  llvm.br ^bb3(%1: !llvm.ptr)
+^bb3(%ptr : !llvm.ptr):
+  %2 = llvm.load %ptr : !llvm.ptr -> f32
+  llvm.return %2 : f32
+}
+
+// -----
+
+llvm.func @static_alloca_not_in_entry(%cond : i1) -> f32 {
+  llvm.cond_br %cond, ^bb1, ^bb2
+^bb1:
+  %0 = llvm.mlir.constant(4 : i32) : i32
+  %1 = llvm.alloca %0 x f32 : (i32) -> !llvm.ptr
+  llvm.br ^bb3(%1: !llvm.ptr)
+^bb2:
+  %2 = llvm.mlir.constant(8 : i32) : i32
+  %3 = llvm.alloca %2 x f32 : (i32) -> !llvm.ptr
+  llvm.br ^bb3(%3: !llvm.ptr)
+^bb3(%ptr : !llvm.ptr):
+  %4 = llvm.load %ptr : !llvm.ptr -> f32
+  llvm.return %4 : f32
+}
+
+// CHECK-LABEL: llvm.func @test_inline
+// Make sure the alloca was not moved to the entry block.
+// CHECK-NOT: llvm.alloca
+// CHECK: llvm.cond_br
+// CHECK: llvm.alloca
+llvm.func @test_inline(%cond : i1) -> f32 {
+  %0 = llvm.call @static_alloca_not_in_entry(%cond) : (i1) -> f32
+  llvm.return %0 : f32
+}