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

[Tobias Gysi]

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

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.

>From bc056c4a9de81e008849d9bf42e8f1467fc5b244 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] Improve alloca handling during 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 |  7 ++++--
 mlir/test/Dialect/LLVMIR/inlining.mlir      | 26 +++++++++++++++++++++
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
index 6e9019f932aa8c..dd2b06a8981684 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
@@ -51,10 +51,13 @@ static void
 handleInlinedAllocas(Operation *call,
                      iterator_range<Region::iterator> inlinedBlocks) {
   Block *calleeEntryBlock = &(*inlinedBlocks.begin());
-  Block *callerEntryBlock = &(*calleeEntryBlock->getParent()->begin());
-  if (calleeEntryBlock == callerEntryBlock)
+  Block *callerEntryBlock = calleeEntryBlock;
+  if (auto funcOp = call->getParentOfType<LLVM::LLVMFuncOp>())
+    callerEntryBlock = &funcOp.getFunctionBody().front();
+  if (calleeEntryBlock == callerEntryBlock) {
     // Nothing to do.
     return;
+  }
   SmallVector<std::tuple<LLVM::AllocaOp, IntegerAttr, bool>> allocasToMove;
   bool shouldInsertLifetimes = false;
   bool hasDynamicAlloca = false;
diff --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir
index b684be1f9626b1..87e0ffb0a7e9d0 100644
--- a/mlir/test/Dialect/LLVMIR/inlining.mlir
+++ b/mlir/test/Dialect/LLVMIR/inlining.mlir
@@ -324,6 +324,32 @@ 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) {
+  // Make sure the alloca has been moved to the entry block even if the
+  // callsite is nested in a region operation.
+  // 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