[flang] Fix lowering of array paths in elemental calls

[luporl]

Summary

Elemental procedures may need their array arguments to be passed by
address. This is done by setting ArrayExprLowering::semant to a
value that corresponds to this semantics. Later, member functions
such as applyPathToArrayLoad() read this variable to generate FIR
instructions that match the needed behavior. The problem is that
the semant variable also affects how array paths are lowered. Thus,
if an index of the path is an array element, this will cause its
address to be used instead of its value, which usually results in a
segmentation fault at runtime.

Example: b(i:i) = elem_func(a(v(i):v(i)))

To fix this, ArrayExprLowering::nextPathSemant was added. When it's
set, the next array path is handled with the semantics specified by
it, while the elemental argument retains its original semantics.

Fixes https://github.com/llvm/llvm-project/issues/62981

[jeanPerier]

Looks good to me, thanks for the fix!
Note that this test is passing with HLFIR (flang-new -flang-experimental-hlfir) that is intended to simplify array expression lowering and fix issues there. I would not prioritize fix for bugs that are fixed with this option since we will transition to HLFIR by default as soon as we reach parity (we still have a few points that are behind like vector subscript support and optional argument in intrinsics).

[luporl]

In D153454#4443523, @jeanPerier wrote:

Looks good to me, thanks for the fix!
Note that this test is passing with HLFIR (flang-new -flang-experimental-hlfir) that is intended to simplify array expression lowering and fix issues there. I would not prioritize fix for bugs that are fixed with this option since we will transition to HLFIR by default as soon as we reach parity (we still have a few points that are behind like vector subscript support and optional argument in intrinsics).

Thanks for the review and for the information about HLFIR! I'll start to check and prioritize the issues that can also be reproduced with flang-new -flang-experimental-hlfir.

flang/lib/Lower/ConvertExpr.cpp

if ((semant == ConstituentSemantics::RefOpaque ||
             return createInMemoryScalarCopy(builder, loc, cc(iters));
           };
         }
         mlir::Type storageType =
             fir::unwrapSequenceType(converter.genType(toEvExpr(x)));
         return [=](IterSpace iters) -> ExtValue {
           return placeScalarValueInMemory(builder, loc, cc(iters), storageType);
         };
+      } else if (isArray(x)) {
+        // An array reference is needed, but the indices used in its path must
+        // still be retrieved by value.
+        assert(!nextPathSemant && "Next path semantics already set!");
+        nextPathSemant = ConstituentSemantics::RefTransparent;
+        CC cc = genarr(x);
+        assert(!nextPathSemant && "Next path semantics wasn't used!");
+        return cc;
       }
     }
     return genarr(x);
   }
 
   // A reference to a Fortran elemental intrinsic or intrinsic module procedure.
   CC genElementalIntrinsicProcRef(
       const Fortran::evaluate::ProcedureRef &procRef,
auto addComponentList = [&](mlir::Type ty, mlir::ValueRange vals) {
       } else {
         for (auto v : vals)
           components.suffixComponents.push_back(v);
       }
     };
     mlir::IndexType idxTy = builder.getIndexType();
     mlir::Value one = builder.createIntegerConstant(loc, idxTy, 1);
     bool atBase = true;
-    auto saveSemant = semant;
-    if (isProjectedCopyInCopyOut())
-      semant = ConstituentSemantics::RefTransparent;
+    PushSemantics(isProjectedCopyInCopyOut()
+                      ? ConstituentSemantics::RefTransparent
+                      : nextPathSemantics());
     unsigned index = 0;
     for (const auto &v : llvm::reverse(revPath)) {
       std::visit(
           Fortran::common::visitors{
               [&](const ImplicitSubscripts &) {
                 prefix = false;
                 ty = fir::unwrapSequenceType(ty);
               },
for (const auto &v : llvm::reverse(revPath)) {
                 } else {
                   TODO(loc, "other component type");
                 }
               }},
           v);
       atBase = false;
       ++index;
     }
-    semant = saveSemant;
     ty = fir::unwrapSequenceType(ty);
     components.applied = true;
     return ty;
   }
 
   llvm::SmallVector<mlir::Value> genSubstringBounds(ComponentPath &components) {
     llvm::SmallVector<mlir::Value> result;
     if (components.substring)
inline bool isReferentiallyOpaque() {
     return semant == ConstituentSemantics::RefOpaque;
   }
 
   /// Array appears in a context where it is passed as a VALUE argument.
   inline bool isValueAttribute() {
     return semant == ConstituentSemantics::ByValueArg;
   }
 
+  /// Semantics to use when lowering the next array path.
+  /// If no value was set, the path uses the same semantics as the array.
+  inline ConstituentSemantics nextPathSemantics() {
+    if (nextPathSemant) {
+      ConstituentSemantics sema = nextPathSemant.value();
+      nextPathSemant.reset();
+      return sema;
+    }
+
+    return semant;
+  }
+
   /// Can the loops over the expression be unordered?
   inline bool isUnordered() const { return unordered; }
 
   void setUnordered(bool b) { unordered = b; }
 
   inline bool isPointerAssignment() const { return lbounds.has_value(); }
 
   inline bool isBoundsSpec() const {

   llvm::SmallVector<mlir::Value> destShape;
   /// List of arrays in the expression that have been loaded.
   llvm::SmallVector<ArrayOperand> arrayOperands;
   /// If there is a user-defined iteration space, explicitShape will hold the
   /// information from the front end.
   Fortran::lower::ExplicitIterSpace *explicitSpace = nullptr;
   Fortran::lower::ImplicitIterSpace *implicitSpace = nullptr;
   ConstituentSemantics semant = ConstituentSemantics::RefTransparent;
+  std::optional<ConstituentSemantics> nextPathSemant;
   /// `lbounds`, `ubounds` are used in POINTER value assignments, which may only
   /// occur in an explicit iteration space.
   std::optional<llvm::SmallVector<mlir::Value>> lbounds;
   std::optional<llvm::SmallVector<mlir::Value>> ubounds;
   // Can the array expression be evaluated in any order?
   // Will be set to false if any of the expression parts prevent this.
   bool unordered = true;
   // ProcedureRef currently being lowered. Used to retrieve the iteration shape

flang/test/Lower/array-elemental-calls-3.f90

+! RUN: bbc -o - -emit-fir %s | FileCheck %s
+
+! Test lowering of elemental calls with array arguments that use array
+! elements as indices.
+! As reported in issue #62981, wrong code was being generated in this case.
+
+module test_ops
+  implicit none
+  interface
+    integer elemental function elem_func_i(i)
+      integer, intent(in) :: i
+    end function
+    real elemental function elem_func_r(r)
+      real, intent(in) :: r
+    end function
+  end interface
+
+  integer :: a(3), b(3), v(3), i, j, k, l
+  real :: x(2), y(2), u
+
+contains
+! CHECK-LABEL: func @_QMtest_opsPcheck_array_elems_as_indices() {
+subroutine check_array_elems_as_indices()
+! CHECK: %[[A_ADDR:.*]] = fir.address_of(@_QMtest_opsEa) : !fir.ref<!fir.array<3xi32>>
+! CHECK: %[[V_ADDR:.*]] = fir.address_of(@_QMtest_opsEv) : !fir.ref<!fir.array<3xi32>>
+! CHECK: %[[V:.*]] = fir.array_load %[[V_ADDR]](%{{.*}}) : (!fir.ref<!fir.array<3xi32>>, !fir.shape<1>) -> !fir.array<3xi32>
+! CHECK: %[[A:.*]] = fir.array_load %[[A_ADDR]](%{{.*}}) : (!fir.ref<!fir.array<3xi32>>, !fir.shape<1>) -> !fir.array<3xi32>
+! CHECK: fir.do_loop
+  forall (i=1:3)
+! CHECK: %{{.*}} = fir.array_fetch %[[V]], %{{.*}} : (!fir.array<3xi32>, index) -> i32
+! CHECK: fir.do_loop
+! CHECK: %[[ELEM:.*]] = fir.array_access %[[A]], %{{.*}} : (!fir.array<3xi32>, index) -> !fir.ref<i32>
+! CHECK: %{{.*}} = fir.call @_QPelem_func_i(%[[ELEM]]){{.*}} : (!fir.ref<i32>) -> i32
+    b(i:i) = elem_func_i(a(v(i):v(i)))
+  end forall
+end subroutine
+
+! CHECK-LABEL: func @_QMtest_opsPcheck_not_assert() {
+subroutine check_not_assert()
+  ! Implicit path.
+  b = 10 + elem_func_i(a)
+
+  ! Expression as argument, instead of variable.
+  forall (i=1:3)
+    b(i:i) = elem_func_i(a(i:i) + a(i:i))
+  end forall
+
+  ! Nested elemental function calls.
+  y = elem_func_r(cos(x))
+  y = elem_func_r(cos(x) + u)
+
+  ! Array constructors as elemental function arguments.
+  y = atan2( (/ (real(i, 4), i = 1, 2) /),
+             real( (/ (i, i = j, k, l) /), 4) )
+end subroutine
+
+end module