[lowering] Use fir.array_modify to lower elemental defined assignments

- Add a CustomCopyInCopyOut ConstituentSemantics to create fir.array_modify
in the expression lowering.
- Add a helper function to generate user assignment function call
(packaging the arguments) based on an LHS and RHS ExtendedValue for
the elements.
- Use the address returned by the fir.array_modify modify to build the
  LHS ExtendedValue and used the added helper function to generate the
  call.

- Add some better support to get the length form a memref mlir::Value
  (do not only try to get a constant length, but also read it from the
  fir.box if the memref is a fir.box). Move this code to
  CharacterExprHelper.

This is so far only enabled outside of Forall/Where contexts, another
place creating fir.array_update in these contexts will need to be
updated to handle the CustomCopyInCopyOut.

Author: jeanPerier

flang/include/flang/Optimizer/Builder/Character.h

@@ -153,6 +153,12 @@ class CharacterExprHelper {
   /// Returns integer value held in a character singleton.
   mlir::Value extractCodeFromSingleton(mlir::Value singleton);
 
+  /// Create a value for the length of a character based on its memory reference
+  /// that may be a boxchar, box or !fir.[ptr|ref|heap]<fir.char<kind, len>>. If
+  /// the memref is a simple address and the length is not constant in type, the
+  /// returned length will be empty.
+  mlir::Value getLength(mlir::Value memref);
+
   /// Compute length given a fir.box describing a character entity.
   /// It adjusts the length from the number of bytes per the descriptor
   /// to the number of characters per the Fortran KIND.

flang/lib/Lower/ConvertExpr.cpp

@@ -136,6 +136,10 @@ enum class ConstituentSemantics {
   // Similar to CopyInCopyOut but `a` may be a transient projection (rather than
   // a whole array).
   ProjectedCopyInCopyOut,
+  // Similar to ProjectedCopyInCopyOut, except the merge value is not assigned
+  // automatically by the framework. Instead, and address for `[xs]` is made
+  // accessible so that custom assignments to `[xs]` can be implemented.
+  CustomCopyInCopyOut,
   // Referentially opaque. Refers to the address of `x_i`.
   RefOpaque
 };
@@ -219,18 +223,19 @@ static fir::ExtendedValue arrayLoadExtValue(fir::FirOpBuilder &builder,
   // Recover the extended value from the load.
   assert(!load.slice() && "slice is not allowed");
   auto arrTy = load.getType();
-  auto idxTy = builder.getIndexType();
   if (!path.empty()) {
     auto ty = fir::applyPathToType(arrTy, path);
     if (!ty)
       fir::emitFatalError(loc, "path does not apply to type");
     if (!ty.isa<fir::SequenceType>()) {
-      if (auto charTy = ty.dyn_cast<fir::CharacterType>()) {
-        // ???: Is this in CharacterExprHelper?
-        auto len = charTy.hasConstantLen()
-                       ? builder.createIntegerConstant(
-                             loc, idxTy, ty.cast<fir::CharacterType>().getLen())
-                       : load.typeparams()[0];
+      if (fir::isa_char(ty)) {
+        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{newBase, len};
       }
       return newBase;
@@ -250,11 +255,13 @@ static fir::ExtendedValue arrayLoadExtValue(fir::FirOpBuilder &builder,
   }
   auto extents = fir::factory::getExtents(load.shape());
   auto lbounds = fir::factory::getOrigins(load.shape());
-  if (auto charTy = eleTy.dyn_cast<fir::CharacterType>()) {
-    auto len = charTy.hasConstantLen()
-                   ? builder.createIntegerConstant(
-                         loc, idxTy, eleTy.cast<fir::CharacterType>().getLen())
-                   : load.typeparams()[0];
+  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::CharArrayBoxValue{newBase, len, extents, lbounds};
   }
   if (load.typeparams().empty()) {
@@ -264,6 +271,25 @@ 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) {
@@ -1835,7 +1861,7 @@ class ScalarExprLowering {
       });
       /// Result lengths parameters should not be provided to box storage
       /// allocation and save_results, but they are still useful information to
-      /// keep in the ExtentdedValue if non-deferred.
+      /// keep in the ExtendedValue if non-deferred.
       if (!type.isa<fir::BoxType>())
         resultLengths = lengths;
       auto temp =
@@ -2418,6 +2444,50 @@ static mlir::Value adjustedArrayElement(mlir::Location loc,
   return builder.createConvert(loc, adjustedArrayElementType(eleTy), val);
 }
 
+/// Helper to generate calls to scalar user defined assignment procedures.
+static void genScalarUserDefinedAssignmentCall(fir::FirOpBuilder &builder,
+                                               mlir::Location loc,
+                                               mlir::FuncOp func,
+                                               const fir::ExtendedValue &lhs,
+                                               const fir::ExtendedValue &rhs) {
+  auto prepareUserDefinedArg =
+      [](fir::FirOpBuilder &builder, mlir::Location loc,
+         const fir::ExtendedValue &value, mlir::Type argType) -> mlir::Value {
+    if (argType.isa<fir::BoxCharType>()) {
+      const auto *charBox = value.getCharBox();
+      assert(charBox && "argument type mismatch in elemental user assignment");
+      return fir::factory::CharacterExprHelper{builder, loc}.createEmbox(
+          *charBox);
+    }
+    if (argType.isa<fir::BoxType>()) {
+      auto box = builder.createBox(loc, value);
+      return builder.createConvert(loc, argType, box);
+    }
+    // Simple pass by address.
+    auto argBaseType = fir::unwrapRefType(argType);
+    assert(!fir::hasDynamicSize(argBaseType));
+    auto from = fir::getBase(value);
+    if (argBaseType != fir::unwrapRefType(from.getType())) {
+      // With logicals, it is possible that from is i1 here.
+      if (fir::isa_ref_type(from.getType()))
+        from = builder.create<fir::LoadOp>(loc, from);
+      from = builder.createConvert(loc, argBaseType, from);
+    }
+    if (!fir::isa_ref_type(from.getType())) {
+      auto temp = builder.createTemporary(loc, argBaseType);
+      builder.create<fir::StoreOp>(loc, from, temp);
+      from = temp;
+    }
+    return builder.createConvert(loc, argType, from);
+  };
+  assert(func.getNumArguments() == 2);
+  auto lhsType = func.getType().getInput(0);
+  auto rhsType = func.getType().getInput(1);
+  auto lhsArg = prepareUserDefinedArg(builder, loc, lhs, lhsType);
+  auto rhsArg = prepareUserDefinedArg(builder, loc, rhs, rhsType);
+  builder.create<fir::CallOp>(loc, func, mlir::ValueRange{lhsArg, rhsArg});
+}
+
 //===----------------------------------------------------------------------===//
 //
 // Lowering of scalar expressions in an explicit iteration space context.
@@ -2477,6 +2547,12 @@ class ScalarArrayExprLowering {
       auto offset = expSpace.argPosition(oldInnerArg);
       expSpace.setInnerArg(offset, fir::getBase(lexv));
       builder.create<fir::ResultOp>(getLoc(), fir::getBase(lexv));
+    } else if (auto updateOp = mlir::dyn_cast<fir::ArrayModifyOp>(
+                   fir::getBase(lexv).getDefiningOp())) {
+      auto oldInnerArg = updateOp.sequence();
+      auto offset = expSpace.argPosition(oldInnerArg);
+      expSpace.setInnerArg(offset, fir::getBase(lexv));
+      builder.create<fir::ResultOp>(getLoc(), fir::getBase(lexv));
     }
     return lexv;
   }
@@ -3619,6 +3695,8 @@ class ArrayExprLowering {
     return abstractArrayExtValue(iterSpace.outerResult());
   }
 
+  /// Lower an elemental subroutine call with at least one array argument.
+  /// Not for user defined assignments.
   static void lowerArrayElementalSubroutine(
       Fortran::lower::AbstractConverter &converter,
       Fortran::lower::SymMap &symMap, Fortran::lower::StatementContext &stmtCtx,
@@ -3628,7 +3706,6 @@ class ArrayExprLowering {
     ael.lowerArrayElementalSubroutine(call);
   }
 
-  // ! Not for user defined assignment elemental subroutine.
   void lowerArrayElementalSubroutine(
       const Fortran::evaluate::Expr<Fortran::evaluate::SomeType> &call) {
     auto f = genarr(call);
@@ -3638,6 +3715,48 @@ 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) {
+    ArrayExprLowering ael(converter, stmtCtx, symMap,
+                          ConstituentSemantics::CustomCopyInCopyOut);
+    auto procRef = std::get<Fortran::evaluate::ProcedureRef>(call.u);
+    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");
+    auto func = Fortran::lower::CallerInterface(procRef, converter).getFuncOp();
+    ael.lowerElementalUserAssignment(func, *lhs, *rhs);
+  }
+
+  void lowerElementalUserAssignment(
+      mlir::FuncOp userAssignment,
+      const Fortran::evaluate::Expr<Fortran::evaluate::SomeType> &lhs,
+      const Fortran::evaluate::Expr<Fortran::evaluate::SomeType> &rhs) {
+    auto loc = getLoc();
+    PushSemantics(ConstituentSemantics::CustomCopyInCopyOut);
+    auto genArrayModify = genarr(lhs);
+    ccStoreToDest = [=](IterSpace iters) -> ExtValue {
+      auto modifiedArray = genArrayModify(iters);
+      auto arrayModify = mlir::dyn_cast_or_null<fir::ArrayModifyOp>(
+          fir::getBase(modifiedArray).getDefiningOp());
+      assert(arrayModify && "must be created by ArrayModifyOp");
+      auto lhs =
+          arrayModifyToExv(builder, loc, destination, arrayModify.getResult(0));
+      genScalarUserDefinedAssignmentCall(builder, loc, userAssignment, lhs,
+                                         iters.elementExv());
+      return modifiedArray;
+    };
+    determineShapeOfDest(lhs);
+    semant = ConstituentSemantics::RefTransparent;
+    auto exv = lowerArrayExpression(rhs);
+    builder.create<fir::ArrayMergeStoreOp>(
+        loc, destination, fir::getBase(exv), destination.memref(),
+        destination.slice(), destination.typeparams());
+  }
+
   /// Compute the shape of a slice.
   llvm::SmallVector<mlir::Value> computeSliceShape(mlir::Value slice) {
     llvm::SmallVector<mlir::Value> slicedShape;
@@ -4994,6 +5113,23 @@ class ArrayExprLowering {
         return abstractArrayExtValue(arrUpdate);
       };
     }
+    if (isCustomCopyInCopyOut()) {
+      // Create an array_modify to get the LHS element address and indicate
+      // the assignment, the actual assignment must be implemented in
+      // ccStoreToDest.
+      destination = arrLoad;
+      return [=](IterSpace iters) -> ExtValue {
+        auto innerArg = iters.innerArgument();
+        auto resTy = innerArg.getType();
+        auto eleTy = fir::applyPathToType(resTy, iters.iterVec());
+        auto refEleTy =
+            fir::isa_ref_type(eleTy) ? eleTy : builder.getRefType(eleTy);
+        auto arrModify = builder.create<fir::ArrayModifyOp>(
+            loc, mlir::TypeRange{refEleTy, resTy}, innerArg, iters.iterVec(),
+            destination.typeparams());
+        return abstractArrayExtValue(arrModify.getResult(1));
+      };
+    }
     if (isCopyInCopyOut()) {
       // Semantics are copy-in copy-out.
       // The continuation simply forwards the result of the `array_load` Op,
@@ -5914,6 +6050,10 @@ class ArrayExprLowering {
     return semant == ConstituentSemantics::ProjectedCopyInCopyOut;
   }
 
+  bool isCustomCopyInCopyOut() {
+    return semant == ConstituentSemantics::CustomCopyInCopyOut;
+  }
+
   /// Array appears in a context where it must be boxed.
   bool isBoxValue() { return semant == ConstituentSemantics::BoxValue; }
 
@@ -6100,9 +6240,11 @@ mlir::Value Fortran::lower::createSubroutineCall(
     // Elemental user defined assignment has special requirements to deal with
     // LHS/RHS overlaps. See 10.2.1.5 p2.
     if (isUserDefAssignment)
-      TODO(converter.getCurrentLocation(), "elemental user defined assignment");
-    ArrayExprLowering::lowerArrayElementalSubroutine(converter, symMap, stmtCtx,
-                                                     call);
+      ArrayExprLowering::lowerElementalUserAssignment(converter, symMap,
+                                                      stmtCtx, call);
+    else
+      ArrayExprLowering::lowerArrayElementalSubroutine(converter, symMap,
+                                                       stmtCtx, call);
     return mlir::Value{};
   }
   // FIXME: The non elemental user defined assignment case with array arguments

flang/lib/Optimizer/Builder/Character.cpp

@@ -705,3 +705,19 @@ fir::factory::CharacterExprHelper::readLengthFromBox(mlir::Value box) {
   }
   return size;
 }
+
+mlir::Value fir::factory::CharacterExprHelper::getLength(mlir::Value memref) {
+  auto memrefType = memref.getType();
+  auto charType = recoverCharacterType(memrefType);
+  assert(charType && "must be a character type");
+  if (charType.hasConstantLen())
+    return builder.createIntegerConstant(loc, builder.getCharacterLengthType(),
+                                         charType.getLen());
+  if (memrefType.isa<fir::BoxType>())
+    return readLengthFromBox(memref);
+  if (memrefType.isa<fir::BoxCharType>())
+    return createUnboxChar(memref).second;
+
+  // Length cannot be deduced from memref.
+  return {};
+}