[flang][do concurrent] Extned getAllocaBlock() and emit yields correctly (#146853)

Handles some loose ends in `do concurrent` reduction declarations. This
PR extends `getAllocaBlock` to handle declare ops, and also emit
`fir.yield` in all regions.

Author: ergawy

flang/lib/Lower/Support/ReductionProcessor.cpp

@@ -529,7 +529,9 @@ static void createReductionAllocAndInitRegions(
         converter, loc, type, initValue, initBlock,
         reductionDecl.getInitializerAllocArg(),
         reductionDecl.getInitializerMoldArg(), reductionDecl.getCleanupRegion(),
-        DeclOperationKind::Reduction);
+        DeclOperationKind::Reduction, /*sym=*/nullptr,
+        /*cannotHaveLowerBounds=*/false,
+        /*isDoConcurrent*/ std::is_same_v<OpType, fir::DeclareReductionOp>);
   }
 
   if (fir::isa_trivial(ty)) {

flang/lib/Optimizer/Builder/FIRBuilder.cpp

@@ -286,6 +286,9 @@ mlir::Block *fir::FirOpBuilder::getAllocaBlock() {
   if (auto firLocalOp = getRegion().getParentOfType<fir::LocalitySpecifierOp>())
     return &getRegion().front();
 
+  if (auto firLocalOp = getRegion().getParentOfType<fir::DeclareReductionOp>())
+    return &getRegion().front();
+
   return getEntryBlock();
 }
 

flang/test/HLFIR/fir-reduction-alloca-block.fir

@@ -0,0 +1,31 @@
+// 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.declare_reduction @add_reduction_byref_box_heap_UxUxf32 : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>> alloc {
+  %0 = fir.alloca !fir.box<!fir.heap<!fir.array<?x?xf32>>>
+  fir.yield(%0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>)
+} init {
+^bb0(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>, %arg1: !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>):
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = fir.load %arg1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>
+  hlfir.assign %cst to %0 : f32, !fir.box<!fir.heap<!fir.array<?x?xf32>>>
+  fir.yield(%arg1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>)
+} combiner {
+^bb0(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>, %arg1: !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>):
+  fir.yield(%arg0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>)
+}
+
+// CHECK-LABEL:   fir.declare_reduction @add_reduction_byref_box_heap_UxUxf32 : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>> alloc {
+// CHECK:           %[[VAL_0:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?x?xf32>>>
+// CHECK:           fir.yield(%[[VAL_0]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>)
+
+// CHECK-LABEL:   } init {
+// CHECK:         ^bb0(%[[VAL_0:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>, %[[VAL_1:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>):
+// CHECK:           %[[VAL_2:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?x?xf32>>>
+// CHECK:           fir.yield(%[[VAL_1]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>)
+
+// CHECK-LABEL:   } combiner {
+// CHECK:         ^bb0(%[[VAL_0:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>, %[[VAL_1:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>):
+// CHECK:           fir.yield(%[[VAL_0]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>)
+// CHECK:         }

flang/test/Lower/do_concurrent_reduce_allocatable.f90

@@ -0,0 +1,22 @@
+! RUN: %flang_fc1 -emit-hlfir -o - %s | FileCheck %s
+
+subroutine do_concurrent_allocatable
+  integer :: i
+  real, allocatable, dimension(:,:) :: x
+
+  do concurrent (i = 1:10) reduce(+: x)
+  end do
+end subroutine
+
+! CHECK: fir.declare_reduction @[[RED_OP:.*]] : ![[RED_TYPE:.*]] alloc {
+! CHECK:   %[[ALLOC:.*]] = fir.alloca
+! CHECK:   fir.yield(%[[ALLOC]] : ![[RED_TYPE]])
+! CHECK: } init {
+! CHECK: ^bb0(%{{.*}}: ![[RED_TYPE]], %[[RED_ARG:.*]]: ![[RED_TYPE]]):
+! CHECK:   fir.yield(%[[RED_ARG]] : !{{.*}})
+! CHECK: } combiner {
+! CHECK: ^bb0(%[[COMB_RES:.*]]: ![[RED_TYPE]], %{{.*}}: ![[RED_TYPE]]):
+! CHECK:   fir.yield(%[[COMB_RES]] : !{{.*}})
+! CHECK: } cleanup {
+! CHECK:   fir.yield
+! CHECK: }