[Mlir-commits] [mlir] [mlir][llvm] Improve alloca handling during inlining (PR #75961)

[Tobias Gysi]

https://github.com/gysit updated https://github.com/llvm/llvm-project/pull/75961

>From d73d6002ebe23fcb66f7805170bdcbec83531b26 Mon Sep 17 00:00:00 2001
From: Tobias Gysi <tobias.gysi at nextsilicon.com>
Date: Tue, 19 Dec 2023 19:03:12 +0000
Subject: [PATCH] [mlir][llvm] Better alloca handling when inlining

This revision changes the alloca handling in the LLVM inliner.
It ensures that alloca operations, even those nested within a
region operation, can be relocated to the entry block.

While the LLVM dialect does not have any region operations,
the inlining interface may be used on IR that mixes different
dialects.
---
 mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp | 27 ++++++++++++++++++---
 mlir/test/Dialect/LLVMIR/inlining.mlir      | 27 +++++++++++++++++++++
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
index 6e9019f932aa8c..5a89d81b4ee3ac 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
@@ -50,11 +50,29 @@ static bool hasLifetimeMarkers(LLVM::AllocaOp allocaOp) {
 static void
 handleInlinedAllocas(Operation *call,
                      iterator_range<Region::iterator> inlinedBlocks) {
+  // The traversal's objective is to locate the entry block of the closest call
+  // site ancestor that has the automatic allocation scope trait. In pure LLVM
+  // dialect programs, this is the LLVM function containing the call site.
+  // However, mixed dialect programs may have other allocation scopes.
+  assert(call->getParentWithTrait<OpTrait::AutomaticAllocationScope>() &&
+         "expected call site to be in an automatic allocation scope");
+  Block *allocaScopeEntryBlock = nullptr;
+  Operation *currentOp = call;
+  while (Operation *parentOp = currentOp->getParentOp()) {
+    if (parentOp->hasTrait<OpTrait::AutomaticAllocationScope>()) {
+      allocaScopeEntryBlock = &currentOp->getParentRegion()->front();
+      break;
+    }
+    currentOp = parentOp;
+  }
+
+  // Avoid relocating the alloca operations if the call has been inlined into
+  // the entry block of an alloca scope, which is typically the encompassing
+  // LLVM function, or if the relevant alloca scope cannot be identified.
   Block *calleeEntryBlock = &(*inlinedBlocks.begin());
-  Block *callerEntryBlock = &(*calleeEntryBlock->getParent()->begin());
-  if (calleeEntryBlock == callerEntryBlock)
-    // Nothing to do.
+  if (!allocaScopeEntryBlock || allocaScopeEntryBlock == calleeEntryBlock)
     return;
+
   SmallVector<std::tuple<LLVM::AllocaOp, IntegerAttr, bool>> allocasToMove;
   bool shouldInsertLifetimes = false;
   bool hasDynamicAlloca = false;
@@ -93,7 +111,8 @@ handleInlinedAllocas(Operation *call,
     stackPtr = builder.create<LLVM::StackSaveOp>(
         call->getLoc(), LLVM::LLVMPointerType::get(call->getContext()));
   }
-  builder.setInsertionPoint(callerEntryBlock, callerEntryBlock->begin());
+  builder.setInsertionPoint(allocaScopeEntryBlock,
+                            allocaScopeEntryBlock->begin());
   for (auto &[allocaOp, arraySize, shouldInsertLifetime] : allocasToMove) {
     auto newConstant = builder.create<LLVM::ConstantOp>(
         allocaOp->getLoc(), allocaOp.getArraySize().getType(), arraySize);
diff --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir
index b684be1f9626b1..e52142dd60949c 100644
--- a/mlir/test/Dialect/LLVMIR/inlining.mlir
+++ b/mlir/test/Dialect/LLVMIR/inlining.mlir
@@ -324,6 +324,33 @@ llvm.func @test_inline(%cond0 : i1, %cond1 : i1, %funcArg : f32) -> f32 {
 
 // -----
 
+llvm.func @static_alloca(%cond: i1) -> f32 {
+  %0 = llvm.mlir.constant(4 : i32) : i32
+  %1 = llvm.alloca %0 x f32 : (i32) -> !llvm.ptr
+  %2 = llvm.load %1 : !llvm.ptr -> f32
+  llvm.return %2 : f32
+}
+
+// CHECK-LABEL: llvm.func @test_inline
+llvm.func @test_inline(%cond0 : i1) {
+  // Ensure that the alloca has been relocated to the entry block of the
+  // nearest ancestor that serves as an allocation scope, in this case,
+  // the function itself.
+  // CHECK: %[[ALLOCA:.+]] = llvm.alloca
+  // CHECK: "test.one_region_op"() ({
+  "test.one_region_op"() ({
+    %0 = llvm.call @static_alloca(%cond0) : (i1) -> f32
+    // CHECK-NEXT: llvm.intr.lifetime.start 4, %[[ALLOCA]]
+    // CHECK-NEXT: %[[RES:.+]] = llvm.load %[[ALLOCA]]
+    // CHECK-NEXT: llvm.intr.lifetime.end 4, %[[ALLOCA]]
+    // CHECK-NEXT: test.region_yield %[[RES]]
+    test.region_yield %0 : f32
+  }) : () -> ()
+  llvm.return
+}
+
+// -----
+
 llvm.func @alloca_with_lifetime(%cond: i1) -> f32 {
   %0 = llvm.mlir.constant(4 : i32) : i32
   %1 = llvm.alloca %0 x f32 : (i32) -> !llvm.ptr