Lower user defined assignments inside WHERE and FORALL

jeanPerier · jeanPerier · commit 9b5ca9014b58 · 2021-09-30T13:52:50.000+02:00
-  Update the ScalarArrayExprLowering to add a userAssign entry point
   that uses the CustomCopyinCopyOut semantics.
-  Thread explicit and implicit context in createSubroutineCall.
-  When the explicit context is active (inside a FORALL), and when the
   assignment LHS is a scalar, use the ScalarArrayExprLowering newly
   added userAssign entry point.
-  Thread implicit and explicit in lowerElementalUserAssignment to cover
   the array case inside WHERE and FORALL.
-  Ensure analyzeExplicitSpace (FORALL analysis) sees and caches the same
   expression address pointers as the one used in lowering. This
   requires looking at the ProcRef operands of evaluate:Assign, and not
   its LHS/RHS. This also requires not wrapping the ProcRef in an Expr
   when threading it, because it would make a copy.
-  Apply CustomCopyinCopyOut in ArrayExprLowering::applyPathToArrayLoad
   to create fir.array_modify for array assignments inside FORALL.

The case where a non elemental user defined array assignment is used
inside FORALL is left TODO, because it is not clear to me how one can
operate on sub-array parts inside the loops.
diff --git a/flang/include/flang/Lower/ConvertExpr.h b/flang/include/flang/Lower/ConvertExpr.h
@@ -205,7 +205,9 @@ createSomeArrayBox(AbstractConverter &converter,
 /// returns, the returned value indicates which label the code should jump to.
 /// The returned value is null otherwise.
 mlir::Value createSubroutineCall(AbstractConverter &converter,
-                                 const evaluate::Expr<evaluate::SomeType> &call,
+                                 const evaluate::ProcedureRef &call,
+                                 ExplicitIterSpace &explicitIterSpace,
+                                 ImplicitIterSpace &implicitIterSpace,
                                  SymMap &symMap, StatementContext &stmtCtx,
                                  bool isUserDefAssignment);
 
diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
@@ -661,9 +661,9 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     setCurrentPosition(stmt.v.source);
     assert(stmt.typedCall && "Call was not analyzed");
     // Call statement lowering shares code with function call lowering.
-    Fortran::semantics::SomeExpr expr{*stmt.typedCall};
     auto res = Fortran::lower::createSubroutineCall(
-        *this, expr, localSymbols, stmtCtx, /*isUserDefAssignment=*/false);
+        *this, *stmt.typedCall, explicitIterSpace, implicitIterSpace,
+        localSymbols, stmtCtx, /*isUserDefAssignment=*/false);
     if (!res)
       return; // "Normal" subroutine call.
     // Call with alternate return specifiers.
@@ -1941,18 +1941,12 @@ class FirConverter : public Fortran::lower::AbstractConverter {
             // [2] User defined assignment. If the context is a scalar
             // expression then call the procedure.
             [&](const Fortran::evaluate::ProcedureRef &procRef) {
-              if (implicitIterationSpace())
-                TODO(loc, "user defined assignment within WHERE");
-
-              Fortran::semantics::SomeExpr expr{procRef};
               auto &ctx = explicitIterationSpace()
                               ? explicitIterSpace.stmtContext()
                               : stmtCtx;
               Fortran::lower::createSubroutineCall(
-                  *this, expr, localSymbols, ctx, /*isUserDefAssignment=*/true);
-              if (explicitIterationSpace())
-                builder->create<fir::ResultOp>(
-                    loc, explicitIterSpace.getInnerArgs());
+                  *this, procRef, explicitIterSpace, implicitIterSpace,
+                  localSymbols, ctx, /*isUserDefAssignment=*/true);
             },
 
             // [3] Pointer assignment with possibly empty bounds-spec. R1035: a
@@ -2595,8 +2589,24 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     explicitIterSpace.exprBase(&e, LHS);
   }
   void analyzeExplicitSpace(const Fortran::evaluate::Assignment *assign) {
-    analyzeExplicitSpace</*LHS=*/true>(assign->lhs);
-    analyzeExplicitSpace(assign->rhs);
+    auto analyzeAssign = [&](const Fortran::lower::SomeExpr &lhs,
+                             const Fortran::lower::SomeExpr &rhs) {
+      analyzeExplicitSpace</*LHS=*/true>(lhs);
+      analyzeExplicitSpace(rhs);
+    };
+    std::visit(
+        Fortran::common::visitors{
+            [&](const Fortran::evaluate::ProcedureRef &procRef) {
+              // Ensure the procRef expressions are the one being visited.
+              assert(procRef.arguments().size() == 2);
+              const auto *lhs = procRef.arguments()[0].value().UnwrapExpr();
+              const auto *rhs = procRef.arguments()[1].value().UnwrapExpr();
+              assert(lhs && rhs &&
+                     "user defined assignment arguments must be expressions");
+              analyzeAssign(*lhs, *rhs);
+            },
+            [&](const auto &) { analyzeAssign(assign->lhs, assign->rhs); }},
+        assign->u);
     explicitIterSpace.endAssign();
   }
   void analyzeExplicitSpace(const Fortran::parser::ForallAssignmentStmt &stmt) {
diff --git a/flang/lib/Lower/ConvertExpr.cpp b/flang/lib/Lower/ConvertExpr.cpp
@@ -271,25 +271,6 @@ static fir::ExtendedValue arrayLoadExtValue(fir::FirOpBuilder &builder,
             "properties are explicit, assumed, deferred, or ?");
 }
 
-/// Convert the result of a fir.array_modify to an ExtendedValue given the
-/// related fir.array_load.
-static fir::ExtendedValue arrayModifyToExv(fir::FirOpBuilder &builder,
-                                           mlir::Location loc,
-                                           fir::ArrayLoadOp load,
-                                           mlir::Value elementAddr) {
-  auto eleTy = fir::unwrapPassByRefType(elementAddr.getType());
-  if (fir::isa_char(eleTy)) {
-    auto len = fir::factory::CharacterExprHelper{builder, loc}.getLength(
-        load.memref());
-    if (!len) {
-      assert(load.typeparams().size() == 1 && "length must be in array_load");
-      len = load.typeparams()[0];
-    }
-    return fir::CharBoxValue{elementAddr, len};
-  }
-  return elementAddr;
-}
-
 /// Is this a call to an elemental procedure with at least one array argument ?
 static bool
 isElementalProcWithArrayArgs(const Fortran::evaluate::ProcedureRef &procRef) {
@@ -2488,6 +2469,25 @@ static void genScalarUserDefinedAssignmentCall(fir::FirOpBuilder &builder,
   builder.create<fir::CallOp>(loc, func, mlir::ValueRange{lhsArg, rhsArg});
 }
 
+/// Convert the result of a fir.array_modify to an ExtendedValue given the
+/// related fir.array_load.
+static fir::ExtendedValue arrayModifyToExv(fir::FirOpBuilder &builder,
+                                           mlir::Location loc,
+                                           fir::ArrayLoadOp load,
+                                           mlir::Value elementAddr) {
+  auto eleTy = fir::unwrapPassByRefType(elementAddr.getType());
+  if (fir::isa_char(eleTy)) {
+    auto len = fir::factory::CharacterExprHelper{builder, loc}.getLength(
+        load.memref());
+    if (!len) {
+      assert(load.typeparams().size() == 1 && "length must be in array_load");
+      len = load.typeparams()[0];
+    }
+    return fir::CharBoxValue{elementAddr, len};
+  }
+  return elementAddr;
+}
+
 //===----------------------------------------------------------------------===//
 //
 // Lowering of scalar expressions in an explicit iteration space context.
@@ -2540,7 +2540,8 @@ class ScalarArrayExprLowering {
     // 3) Finalize the inner context.
     expSpace.finalizeContext();
     // 4) Thread the array value updated forward. Note: the lhs might be
-    // ill-formed, in which case there is no array to thread.
+    // ill-formed (performing scalar assignment in an array context),
+    // in which case there is no array to thread.
     if (auto updateOp = mlir::dyn_cast<fir::ArrayUpdateOp>(
             fir::getBase(lexv).getDefiningOp())) {
       auto oldInnerArg = updateOp.sequence();
@@ -2557,6 +2558,45 @@ class ScalarArrayExprLowering {
     return lexv;
   }
 
+  ExtValue userAssign(mlir::FuncOp userAssignment,
+                      const Fortran::lower::SomeExpr &lhs,
+                      const Fortran::lower::SomeExpr &rhs) {
+    auto loc = getLoc();
+    semant = ConstituentSemantics::RefTransparent;
+    // 1) Lower the rhs expression with array_fetch op(s).
+    auto rexv = lower(rhs);
+    // 2) Lower the lhs expression to an array_modify.
+    semant = ConstituentSemantics::CustomCopyInCopyOut;
+    auto lexv = lower(lhs);
+    bool isIllFormedLHS = false;
+    // 3) Insert the call
+    if (auto modifyOp = mlir::dyn_cast<fir::ArrayModifyOp>(
+            fir::getBase(lexv).getDefiningOp())) {
+      auto oldInnerArg = modifyOp.sequence();
+      auto offset = expSpace.argPosition(oldInnerArg);
+      expSpace.setInnerArg(offset, fir::getBase(lexv));
+      auto exv =
+          arrayModifyToExv(builder, loc, expSpace.getLhsLoad(0).getValue(),
+                           modifyOp.getResult(0));
+      genScalarUserDefinedAssignmentCall(builder, loc, userAssignment, exv,
+                                         rexv);
+    } else {
+      // LHS is ill formed, it is a scalar with no references to FORALL
+      // subscripts, so there is actually no array assignment here. The user
+      // code is probably bad, but still insert user assignment call since it
+      // was not rejected by semantics (a warning was emitted).
+      isIllFormedLHS = true;
+      genScalarUserDefinedAssignmentCall(builder, getLoc(), userAssignment,
+                                         lexv, rexv);
+    }
+    // 4) Finalize the inner context.
+    expSpace.finalizeContext();
+    // 5). Thread the array value updated forward.
+    if (!isIllFormedLHS)
+      builder.create<fir::ResultOp>(getLoc(), fir::getBase(lexv));
+    return lexv;
+  }
+
 private:
   bool pathIsEmpty() { return reversePath.empty(); }
 
@@ -2605,9 +2645,9 @@ class ScalarArrayExprLowering {
   ExtValue applyPathToArrayLoad(fir::ArrayLoadOp load) {
     auto loc = getLoc();
     ExtValue result;
+    auto path = lowerPath(load.getType());
     if (semant == ConstituentSemantics::ProjectedCopyInCopyOut) {
       auto innerArg = expSpace.findArgumentOfLoad(load);
-      auto path = lowerPath(load.getType());
       auto eleTy = fir::applyPathToType(innerArg.getType(), path);
       auto toTy = adjustedArrayElementType(eleTy);
       auto castedElement = builder.createConvert(loc, toTy, elementalValue);
@@ -2618,8 +2658,22 @@ class ScalarArrayExprLowering {
       update->setAttr(fir::factory::attrFortranArrayOffsets(),
                       builder.getUnitAttr());
       result = arrayLoadExtValue(builder, loc, load, {}, update);
+    } else if (semant == ConstituentSemantics::CustomCopyInCopyOut) {
+      // Create an array_modify to get the LHS element address and indicate
+      // the assignment, and create the call to the user defined assignment.
+      auto innerArg = expSpace.findArgumentOfLoad(load);
+      auto eleTy = fir::applyPathToType(innerArg.getType(), path);
+      auto refEleTy =
+          fir::isa_ref_type(eleTy) ? eleTy : builder.getRefType(eleTy);
+      auto arrModify = builder.create<fir::ArrayModifyOp>(
+          loc, mlir::TypeRange{refEleTy, innerArg.getType()}, innerArg, path,
+          load.typeparams());
+      // Flag the offsets as "Fortran" as they are not zero-origin.
+      arrModify->setAttr(fir::factory::attrFortranArrayOffsets(),
+                         builder.getUnitAttr());
+      result =
+          arrayLoadExtValue(builder, loc, load, {}, arrModify.getResult(1));
     } else {
-      auto path = lowerPath(load.getType());
       auto eleTy = fir::applyPathToType(load.getType(), path);
       assert(eleTy && "path did not apply to type");
       auto resTy = adjustedArrayElementType(eleTy);
@@ -3715,13 +3769,16 @@ class ArrayExprLowering {
     builder.restoreInsertionPoint(insPt);
   }
 
-  static void lowerElementalUserAssignment(
-      Fortran::lower::AbstractConverter &converter,
-      Fortran::lower::SymMap &symMap, Fortran::lower::StatementContext &stmtCtx,
-      const Fortran::evaluate::Expr<Fortran::evaluate::SomeType> &call) {
+  static void
+  lowerElementalUserAssignment(Fortran::lower::AbstractConverter &converter,
+                               Fortran::lower::SymMap &symMap,
+                               Fortran::lower::StatementContext &stmtCtx,
+                               Fortran::lower::ExplicitIterSpace &explicitSpace,
+                               Fortran::lower::ImplicitIterSpace &implicitSpace,
+                               const Fortran::evaluate::ProcedureRef &procRef) {
     ArrayExprLowering ael(converter, stmtCtx, symMap,
-                          ConstituentSemantics::CustomCopyInCopyOut);
-    auto procRef = std::get<Fortran::evaluate::ProcedureRef>(call.u);
+                          ConstituentSemantics::CustomCopyInCopyOut,
+                          &explicitSpace, &implicitSpace);
     assert(procRef.arguments().size() == 2);
     const auto *lhs = procRef.arguments()[0].value().UnwrapExpr();
     const auto *rhs = procRef.arguments()[1].value().UnwrapExpr();
@@ -3752,9 +3809,12 @@ class ArrayExprLowering {
     determineShapeOfDest(lhs);
     semant = ConstituentSemantics::RefTransparent;
     auto exv = lowerArrayExpression(rhs);
-    builder.create<fir::ArrayMergeStoreOp>(
-        loc, destination, fir::getBase(exv), destination.memref(),
-        destination.slice(), destination.typeparams());
+    if (explicitSpaceIsActive())
+      builder.create<fir::ResultOp>(loc, fir::getBase(exv));
+    else
+      builder.create<fir::ArrayMergeStoreOp>(
+          loc, destination, fir::getBase(exv), destination.memref(),
+          destination.slice(), destination.typeparams());
   }
 
   /// Compute the shape of a slice.
@@ -5914,6 +5974,25 @@ class ArrayExprLowering {
         return arrayLoadExtValue(builder, loc, load, {}, update);
       };
     }
+    if (semant == ConstituentSemantics::CustomCopyInCopyOut) {
+      // Create an array_modify to get the LHS element address and indicate
+      // the assignment, and create the call to the user defined assignment.
+      destination = load;
+      auto innerArg = explicitSpace->findArgumentOfLoad(load);
+      return [=](IterSpace iters) mutable {
+        auto [path, eleTy] = lowerPath(loc, revPath, load.getType(), iters);
+        auto refEleTy =
+            fir::isa_ref_type(eleTy) ? eleTy : builder.getRefType(eleTy);
+        auto arrModify = builder.create<fir::ArrayModifyOp>(
+            loc, mlir::TypeRange{refEleTy, innerArg.getType()}, innerArg, path,
+            load.typeparams());
+        // Flag the offsets as "Fortran" as they are not zero-origin.
+        arrModify->setAttr(fir::factory::attrFortranArrayOffsets(),
+                           builder.getUnitAttr());
+        return arrayLoadExtValue(builder, loc, load, {},
+                                 arrModify.getResult(1));
+      };
+    }
     return [=](IterSpace iters) mutable {
       auto [path, eleTy] = lowerPath(loc, revPath, load.getType(), iters);
       auto resTy = adjustedArrayElementType(eleTy);
@@ -6232,26 +6311,61 @@ fir::MutableBoxValue Fortran::lower::createMutableBox(
 }
 
 mlir::Value Fortran::lower::createSubroutineCall(
-    AbstractConverter &converter,
-    const evaluate::Expr<evaluate::SomeType> &call, SymMap &symMap,
-    StatementContext &stmtCtx, bool isUserDefAssignment) {
+    AbstractConverter &converter, const evaluate::ProcedureRef &call,
+    ExplicitIterSpace &explicitIterSpace, ImplicitIterSpace &implicitIterSpace,
+    SymMap &symMap, StatementContext &stmtCtx, bool isUserDefAssignment) {
   auto loc = converter.getCurrentLocation();
+
+  if (isUserDefAssignment) {
+    assert(call.arguments().size() == 2);
+    const auto *lhs = call.arguments()[0].value().UnwrapExpr();
+    const auto *rhs = call.arguments()[1].value().UnwrapExpr();
+    assert(lhs && rhs &&
+           "user defined assignment arguments must be expressions");
+    if (call.IsElemental() && lhs->Rank() > 0) {
+      // Elemental user defined assignment has special requirements to deal with
+      // LHS/RHS overlaps. See 10.2.1.5 p2.
+      ArrayExprLowering::lowerElementalUserAssignment(
+          converter, symMap, stmtCtx, explicitIterSpace, implicitIterSpace,
+          call);
+    } else if (explicitIterSpace.isActive() && lhs->Rank() == 0) {
+      // Scalar defined assignment (elemental or not) in a FORALL context.
+      auto func = Fortran::lower::CallerInterface(call, converter).getFuncOp();
+      ScalarArrayExprLowering sael(converter, symMap, explicitIterSpace,
+                                   stmtCtx);
+      sael.userAssign(func, *lhs, *rhs);
+    } else if (explicitIterSpace.isActive()) {
+      // TODO: need to array fetch/modify sub-arrays ?
+      TODO(loc, "non elemental user defined array assignment inside FORALL");
+    } else {
+      if (!implicitIterSpace.empty())
+        fir::emitFatalError(
+            loc,
+            "C1032: user defined assignment inside WHERE must be elemental");
+      // Non elemental user defined assignment outside of FORALL and WHERE.
+      // FIXME: The non elemental user defined assignment case with array
+      // arguments must be take into account potential overlap. So far the front
+      // end does not add parentheses around the RHS argument in the call as it
+      // should according to 15.4.3.4.3 p2.
+      Fortran::semantics::SomeExpr expr{call};
+      Fortran::lower::createSomeExtendedExpression(loc, converter, expr, symMap,
+                                                   stmtCtx);
+    }
+    return {};
+  }
+
+  assert(implicitIterSpace.empty() && !explicitIterSpace.isActive() &&
+         "subroutine calls are not allowed inside WHERE and FORALL");
+
   if (isElementalProcWithArrayArgs(call)) {
-    // Elemental user defined assignment has special requirements to deal with
-    // LHS/RHS overlaps. See 10.2.1.5 p2.
-    if (isUserDefAssignment)
-      ArrayExprLowering::lowerElementalUserAssignment(converter, symMap,
-                                                      stmtCtx, call);
-    else
-      ArrayExprLowering::lowerArrayElementalSubroutine(converter, symMap,
-                                                       stmtCtx, call);
-    return mlir::Value{};
+    Fortran::semantics::SomeExpr expr{call};
+    ArrayExprLowering::lowerArrayElementalSubroutine(converter, symMap, stmtCtx,
+                                                     expr);
+    return {};
   }
-  // FIXME: The non elemental user defined assignment case with array arguments
-  // must be take into account potential overlap. So far the front end does not
-  // add parentheses around the RHS argument in the call as it should according
-  // to 15.4.3.4.3 p2.
-  auto res = Fortran::lower::createSomeExtendedExpression(loc, converter, call,
+  // Simple subroutine call, with potential alternate return.
+  Fortran::semantics::SomeExpr expr{call};
+  auto res = Fortran::lower::createSomeExtendedExpression(loc, converter, expr,
                                                           symMap, stmtCtx);
   return fir::getBase(res);
 }
diff --git a/flang/test/Lower/array-user-def-assignments.f90 b/flang/test/Lower/array-user-def-assignments.f90
