[flang-commits] [flang] [flang][fir] Provide allocation block for `fir.local` when required (PR #144521)

[Kareem Ergawy via flang-commits]

https://github.com/ergawy updated https://github.com/llvm/llvm-project/pull/144521

>From dceb16bdee9eb224879dac0f21747e1da61cf9ea Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Tue, 17 Jun 2025 07:24:43 -0500
Subject: [PATCH 1/2] [flang][fir] Provide allocation block for `fir.local`
 when required

Extends `fir::FirOpBuilder::getAllocaBlock()` to support `fir.local`.
This allows us to retrieve an allocation block when needed for
`fir.local`.
---
 flang/lib/Optimizer/Builder/FIRBuilder.cpp  |  5 +++
 flang/test/HLFIR/fir-local-alloca-block.fir | 34 +++++++++++++++++++++
 2 files changed, 39 insertions(+)
 create mode 100644 flang/test/HLFIR/fir-local-alloca-block.fir

diff --git a/flang/lib/Optimizer/Builder/FIRBuilder.cpp b/flang/lib/Optimizer/Builder/FIRBuilder.cpp
index 584f3c8ee310e..d1f43c2177ed8 100644
--- a/flang/lib/Optimizer/Builder/FIRBuilder.cpp
+++ b/flang/lib/Optimizer/Builder/FIRBuilder.cpp
@@ -283,6 +283,11 @@ mlir::Block *fir::FirOpBuilder::getAllocaBlock() {
   if (auto doConcurentOp = getRegion().getParentOfType<fir::DoConcurrentOp>())
     return doConcurentOp.getBody();
 
+  if (auto firLocalOp =
+          getRegion().getParentOfType<fir::LocalitySpecifierOp>()) {
+    return &getRegion().front();
+  }
+
   return getEntryBlock();
 }
 
diff --git a/flang/test/HLFIR/fir-local-alloca-block.fir b/flang/test/HLFIR/fir-local-alloca-block.fir
new file mode 100644
index 0000000000000..9d76e86fec3d9
--- /dev/null
+++ b/flang/test/HLFIR/fir-local-alloca-block.fir
@@ -0,0 +1,34 @@
+// Tests that `fir.local` ops are able to provide an alloca block when required.
+
+// RUN: fir-opt %s -convert-hlfir-to-fir | FileCheck %s
+
+fir.local {type = local_init} @localizer : !fir.box<!fir.array<1xi32>> copy {
+^bb0(%arg0: !fir.ref<!fir.box<!fir.array<1xi32>>>, %arg1: !fir.ref<!fir.box<!fir.array<1xi32>>>):
+  %0 = fir.load %arg0 : !fir.ref<!fir.box<!fir.array<1xi32>>>
+  hlfir.assign %0 to %arg1 : !fir.box<!fir.array<1xi32>>, !fir.ref<!fir.box<!fir.array<1xi32>>>
+  fir.yield(%arg1 : !fir.ref<!fir.box<!fir.array<1xi32>>>)
+}
+
+func.func @foo() {
+  %c1 = arith.constant 1 : index
+  %0 = fir.alloca !fir.box<!fir.array<1xi32>>
+  fir.do_concurrent {
+    fir.do_concurrent.loop (%arg0) = (%c1) to (%c1) step (%c1) local(@localizer %0 -> %arg1 : !fir.ref<!fir.box<!fir.array<1xi32>>>) {
+    }
+  }
+  return
+}
+
+// CHECK:  fir.local {type = local_init} @localizer : ![[TYPE:fir.box<!fir.array<1xi32>>]] copy {
+// CHECK:    ^bb0(%[[VAL_0:.*]]: !fir.ref<![[TYPE]]>, %[[VAL_1:.*]]: !fir.ref<![[TYPE]]>):
+// CHECK:      %[[VAL_2:.*]] = fir.alloca ![[TYPE]]
+// CHECK:      %[[VAL_3:.*]] = fir.load %[[VAL_0]] : !fir.ref<![[TYPE]]>
+// CHECK:      %[[VAL_4:.*]] = arith.constant 0 : index
+// CHECK:      %[[VAL_5:.*]]:3 = fir.box_dims %[[VAL_3]], %[[VAL_4]] : (![[TYPE]], index) -> (index, index, index)
+// CHECK:      %[[VAL_6:.*]] = fir.load %[[VAL_1]] : !fir.ref<![[TYPE]]>
+// CHECK:      fir.store %[[VAL_6]] to %[[VAL_2]] : !fir.ref<![[TYPE]]>
+// CHECK:      %[[VAL_10:.*]] = fir.convert %[[VAL_2]] : (!fir.ref<![[TYPE]]>) -> !fir.ref<!fir.box<none>>
+// CHECK:      %[[VAL_11:.*]] = fir.convert %[[VAL_3]] : (![[TYPE]]) -> !fir.box<none>
+// CHECK:      fir.call @_FortranAAssign(%[[VAL_10]], %[[VAL_11]], %{{.*}}, %{{.*}})
+// CHECK:      fir.yield(%[[VAL_1]] : !fir.ref<![[TYPE]]>)
+// CHECK:  }

>From abb6d264937b52fb76209996ce32d67b1e5dbd36 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Wed, 18 Jun 2025 02:34:07 -0500
Subject: [PATCH 2/2] remove braces

---
 flang/lib/Optimizer/Builder/FIRBuilder.cpp | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/flang/lib/Optimizer/Builder/FIRBuilder.cpp b/flang/lib/Optimizer/Builder/FIRBuilder.cpp
index d1f43c2177ed8..6ac87067f6511 100644
--- a/flang/lib/Optimizer/Builder/FIRBuilder.cpp
+++ b/flang/lib/Optimizer/Builder/FIRBuilder.cpp
@@ -283,10 +283,8 @@ mlir::Block *fir::FirOpBuilder::getAllocaBlock() {
   if (auto doConcurentOp = getRegion().getParentOfType<fir::DoConcurrentOp>())
     return doConcurentOp.getBody();
 
-  if (auto firLocalOp =
-          getRegion().getParentOfType<fir::LocalitySpecifierOp>()) {
+  if (auto firLocalOp = getRegion().getParentOfType<fir::LocalitySpecifierOp>())
     return &getRegion().front();
-  }
 
   return getEntryBlock();
 }