[flang] Fix explicit space finalization

[peixin]

Summary

When assigning one scalar result to array section in forall construct
and the scalar result is generated with some temporary memory, the free
operation cannot be done after lowering the rhs expression since the
free operation will be outside the fir.do_loop region and the argument
of the free operation is illegal. Move the explicit space finalization
to lowerArrayExpression(rhs) when the rhs has scalar result.

[schweitz]

flang/lib/Lower/ConvertExpr.cpp

3371

Testing the rank all over the place doesn't seem correct. Expressions can contain subexpressions with or without rank and build temporaries accordingly.

Finalization should really happen unconditionally. (In the right place, of course.)

[peixin]

flang/lib/Lower/ConvertExpr.cpp

3371

Expressions can contain subexpressions with or without rank and build temporaries accordingly.

  integer :: arr(6), i
  forall (i=1:6:2)
    arr(i:i+1) = (/(func(j),j=i,i+1)/)
  end forall
 print *, arr
contains
  pure function func(a) result(res)
    integer, intent(in) :: a
    integer :: res
    res = a
  end function
end

Are you pointing to this?

1. I don't understand the case of subexpressions without rank.
2. To my understanding, the subexpressions do not matter. The key point is the expression of rhs. If its final result is scalar, one do_loop assignment will be performed from rhs to lhs and the allocmem is always inside the do_loop. Please correct me if I am wrong.

Finalization should really happen unconditionally. (In the right place, of course.)

If so, the solution would be better than this patch. If you already how to do it, can you take up this? This kind of solution is beyond my current knowledge of FIR lowering. Sorry about this.

[schweitz]

flang/lib/Lower/ConvertExpr.cpp

3371

Sure. Let me take a closer look here.

[peixin]

This is fixed by https://reviews.llvm.org/rG6822ed035f188aa55f9ac283d2eb20dcc7a9ca07.

flang/lib/Lower/ConvertExpr.cpp

void lowerArrayAssignment(const TL &lhs, const TR &rhs) {
     /// continuation is implicitly returned in `ccStoreToDest` and the ArrayLoad
     /// in `destination`.
     PushSemantics(ConstituentSemantics::ProjectedCopyInCopyOut);
     ccStoreToDest = genarr(lhs);
     determineShapeOfDest(lhs);
     semant = ConstituentSemantics::RefTransparent;
     ExtValue exv = lowerArrayExpression(rhs);
     if (explicitSpaceIsActive()) {
+      if constexpr (std::is_same_v<TR, Fortran::lower::SomeExpr>) {
+        if (rhs.Rank() > 0)
+          explicitSpace->finalizeContext();
+      } else if constexpr (std::is_same_v<TR, fir::ExtendedValue>) {
+        if (rhs.rank() > 0)
           explicitSpace->finalizeContext();
+      }
       builder.create<fir::ResultOp>(loc, fir::getBase(exv));
     } else {
       builder.create<fir::ArrayMergeStoreOp>(
           loc, destination, fir::getBase(exv), destination.getMemref(),
           destination.getSlice(), destination.getTypeparams());
     }
   }
 
void lowerAllocatableArrayAssignment(const Fortran::lower::SomeExpr &lhs,
     ccStoreToDest = genarr(realloc.newValue);
     // If the rhs is scalar, get shape from the allocatable ArrayLoad.
     if (destShape.empty())
       destShape = getShape(destination);
     // Finish lowering the loop nest.
     assert(destination && "destination must have been set");
     ExtValue exv = lowerArrayExpression(rhsCC, resultTy);
     if (explicitSpaceIsActive()) {
+      if (rhs.Rank() > 0)
         explicitSpace->finalizeContext();
       builder.create<fir::ResultOp>(loc, fir::getBase(exv));
     } else {
       builder.create<fir::ArrayMergeStoreOp>(
           loc, destination, fir::getBase(exv), destination.getMemref(),
           destination.getSlice(), destination.getTypeparams());
     }
     fir::factory::finalizeRealloc(builder, loc, mutableBox, lbounds,
                                   takeLboundsIfRealloc, realloc);
ccStoreToDest = [=](IterSpace iters) -> ExtValue {
       genScalarUserDefinedAssignmentCall(builder, loc, userAssignment, lhs,
                                          iters.elementExv());
       return modifiedArray;
     };
     determineShapeOfDest(lhs);
     semant = ConstituentSemantics::RefTransparent;
     auto exv = lowerArrayExpression(rhs);
     if (explicitSpaceIsActive()) {
+      if (rhs.Rank() > 0)
         explicitSpace->finalizeContext();
       builder.create<fir::ResultOp>(loc, fir::getBase(exv));
     } else {
       builder.create<fir::ArrayMergeStoreOp>(
           loc, destination, fir::getBase(exv), destination.getMemref(),
           destination.getSlice(), destination.getTypeparams());
     }
   }
 
ExtValue lowerArrayExpression(CC f, mlir::Type resultTy) {
     auto [iterSpace, insPt] = genIterSpace(resultTy);
     auto exv = f(iterSpace);
     iterSpace.setElement(std::move(exv));
     auto lambda = ccStoreToDest.hasValue()
                       ? ccStoreToDest.getValue()
                       : defaultStoreToDestination(/*substring=*/nullptr);
     mlir::Value updVal = fir::getBase(lambda(iterSpace));
     finalizeElementCtx();
+    if (explicitSpaceIsActive() && !resultTy.isa<fir::SequenceType>())
+      explicitSpace->finalizeContext();
     builder.create<fir::ResultOp>(loc, updVal);
     builder.restoreInsertionPoint(insPt);
     return abstractArrayExtValue(iterSpace.outerResult());
   }
 
   /// Compute the shape of a slice.
   llvm::SmallVector<mlir::Value> computeSliceShape(mlir::Value slice) {
     llvm::SmallVector<mlir::Value> slicedShape;

flang/test/Lower/forall/forall-3.f90

+! Test forall lowering for explicit space finalization
+
+! RUN: bbc -emit-fir %s -o - | FileCheck %s
+
+! CHECK-LABEL: func @_QPforall_scalar_to_array_section
+! CHECK:   %{{.*}} = fir.do_loop
+! CHECK:     %[[VAL_0:.*]] = fir.do_loop
+! CHECK:       %[[VAL_1:.*]] = fir.allocmem !fir.array<1xi32>
+! CHECK:       %[[VAL_2:.*]] = fir.if %{{.*}} -> (!fir.heap<!fir.array<1xi32>>) {
+! CHECK:         %[[VAL_3:.*]] = fir.convert %[[VAL_1]] : (!fir.heap<!fir.array<1xi32>>) -> !fir.ref<i8>
+! CHECK:         %[[VAL_4:.*]] = fir.call @realloc(%[[VAL_3]], %{{.*}}) : (!fir.ref<i8>, i64) -> !fir.ref<i8>
+! CHECK:         %[[VAL_5:.*]] = fir.convert %[[VAL_4]] : (!fir.ref<i8>) -> !fir.heap<!fir.array<1xi32>>
+! CHECK:         fir.result %[[VAL_5]] : !fir.heap<!fir.array<1xi32>>
+! CHECK:       } else {
+! CHECK:         fir.result %[[VAL_1:.*]] : !fir.heap<!fir.array<1xi32>>
+! CHECK:       }
+! CHECK:       %[[VAL_6:.*]] = fir.allocmem !fir.array<1xi32>
+! CHECK:       fir.freemem %[[VAL_6]] : !fir.heap<!fir.array<1xi32>>
+! CHECK:       fir.freemem %[[VAL_2]] : !fir.heap<!fir.array<1xi32>>
+! CHECK:       fir.result %{{.*}} : !fir.array<6xi32>
+! CHECK:     }
+! CHECK:     fir.result %[[VAL_0]] : !fir.array<6xi32>
+! CHECK:   }
+subroutine forall_scalar_to_array_section()
+  integer :: arr(6), i
+  forall (i=1:6)
+    arr(i:i+1) = func((/i/))
+  end forall
+
+contains
+  pure function func(a) result(res)
+    integer, intent(in) :: a(:)
+    integer :: res
+    res = a(1)
+  end function
+end subroutine
+
+! CHECK-LABEL: func @_QPforall_array_to_array_section
+! CHECK:  %{{.*}} = fir.do_loop
+! CHECK:    %[[VAL_0:.*]] = fir.allocmem !fir.array<1xi32>
+! CHECK:    %[[VAL_1:.*]] = fir.if %{{.*}} -> (!fir.heap<!fir.array<1xi32>>) {
+! CHECK:      %[[VAL_2:.*]] = fir.convert %[[VAL_0]] : (!fir.heap<!fir.array<1xi32>>) -> !fir.ref<i8>
+! CHECK:      %[[VAL_3:.*]] = fir.call @realloc(%[[VAL_2]], %{{.*}}) : (!fir.ref<i8>, i64) -> !fir.ref<i8>
+! CHECK:      %[[VAL_4:.*]] = fir.convert %[[VAL_3]] : (!fir.ref<i8>) -> !fir.heap<!fir.array<1xi32>>
+! CHECK:      fir.result %[[VAL_4]] : !fir.heap<!fir.array<1xi32>>
+! CHECK:    } else {
+! CHECK:      fir.result %[[VAL_0]] : !fir.heap<!fir.array<1xi32>>
+! CHECK:    }
+! CHECK:    %[[VAL_5:.*]] = fir.allocmem !fir.array<1xi32>
+! CHECK:    fir.freemem %[[VAL_5]] : !fir.heap<!fir.array<1xi32>>
+! CHECK:    fir.freemem %[[VAL_1]] : !fir.heap<!fir.array<1xi32>>
+! CHECK:    fir.result %{{.*}} : !fir.array<6xi32>
+! CHECK:  }
+subroutine forall_array_to_array_section()
+  integer :: arr(6), i
+  forall (i=1:6)
+    arr(i:i+1) = func((/i/))
+  end forall
+
+contains
+  pure function func(a) result(res)
+    integer, intent(in) :: a(:)
+    integer :: res(2)
+    res = a(1)
+  end function
+end subroutine