[flang] implement simple pointer assignments inside FORALL (#129522)

The semantic of pointer assignments inside FORALL requires evaluating
the targets (RHS) and pointer variables (LHS) of all iterations before
evaluating the assignments.

In practice, if the compiler can prove that the RHS and LHS evaluations
are not impacted by the assignments, the evaluation of the FORALL
assignment statement can be done in a single loop. However, if the
compiler cannot prove this, it needs to "save" the addresses of the
targets and/or the pointer descriptors of each iterations before doing
the assignments.

This patch implements the most common cases where there is no lower bound
spec, no bounds remapping, the LHS is not polymorphic, and the RHS is
not NULL.

The HLFIR operation used to represent assignments inside FORALL can be
used for pointer assignments to (the only difference being that the LHS
is a descriptor address).

The analysis for intrinsic assignment can be reused, with the
distinction that the RHS data is not read during the assignment.

The logic that is used to save LHS in intrinsic assignments inside
FORALL is extracted to be used for the RHS of pointer assignments when
needed (saving a descriptor value).
Pointer assignment LHS are just descriptor addresses and are saved as
int_ptr values.

Author: jeanPerier

flang/include/flang/Optimizer/Builder/HLFIRTools.h

@@ -249,8 +249,11 @@ mlir::Value genVariableBoxChar(mlir::Location loc, fir::FirOpBuilder &builder,
                                hlfir::Entity var);
 
 /// Get or create a fir.box or fir.class from a variable.
+/// A fir.box with different attributes that \p var can be created
+/// using \p forceBoxType.
 hlfir::Entity genVariableBox(mlir::Location loc, fir::FirOpBuilder &builder,
-                             hlfir::Entity var);
+                             hlfir::Entity var,
+                             fir::BaseBoxType forceBoxType = {});
 
 /// If the entity is a variable, load its value (dereference pointers and
 /// allocatables if needed). Do nothing if the entity is already a value, and

flang/include/flang/Optimizer/Builder/TemporaryStorage.h

@@ -179,6 +179,8 @@ class AnyValueStack {
 /// type. Fetching variable N will return a variable with the same address,
 /// dynamic type, bounds, and type parameters as the Nth variable that was
 /// pushed. It is implemented using runtime.
+/// Note that this is not meant to save POINTER or ALLOCATABLE descriptor
+/// addresses, use AnyDescriptorAddressStack instead.
 class AnyVariableStack {
 public:
   AnyVariableStack(mlir::Location loc, fir::FirOpBuilder &builder,
@@ -203,6 +205,21 @@ class AnyVariableStack {
   mlir::Value retValueBox;
 };
 
+/// Data structure to stack descriptor addresses. It stores the descriptor
+/// addresses as int_ptr values under the hood.
+class AnyDescriptorAddressStack : public AnyValueStack {
+public:
+  AnyDescriptorAddressStack(mlir::Location loc, fir::FirOpBuilder &builder,
+                            mlir::Type descriptorAddressType);
+
+  void pushValue(mlir::Location loc, fir::FirOpBuilder &builder,
+                 mlir::Value value);
+  mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder);
+
+private:
+  mlir::Type descriptorAddressType;
+};
+
 class TemporaryStorage;
 
 /// Data structure to stack vector subscripted entity shape and
@@ -264,7 +281,8 @@ class TemporaryStorage {
 
 private:
   std::variant<HomogeneousScalarStack, SimpleCopy, SSARegister, AnyValueStack,
-               AnyVariableStack, AnyVectorSubscriptStack>
+               AnyVariableStack, AnyVectorSubscriptStack,
+               AnyDescriptorAddressStack>
       impl;
 };
 } // namespace fir::factory

flang/include/flang/Optimizer/Dialect/FIRType.h

@@ -53,6 +53,9 @@ class BaseBoxType : public mlir::Type {
   /// Is this the box for an assumed rank?
   bool isAssumedRank() const;
 
+  /// Is this a box for a pointer?
+  bool isPointer() const;
+
   /// Return the same type, except for the shape, that is taken the shape
   /// of shapeMold.
   BaseBoxType getBoxTypeWithNewShape(mlir::Type shapeMold) const;

flang/include/flang/Optimizer/HLFIR/HLFIROps.td

@@ -1377,7 +1377,7 @@ def hlfir_RegionAssignOp : hlfir_Op<"region_assign", [hlfir_OrderedAssignmentTre
         regions.push_back(&getUserDefinedAssignment());
     }
     mlir::Region* getSubTreeRegion() { return nullptr; }
-
+    bool isPointerAssignment();
   }];
 
   let hasCustomAssemblyFormat = 1;

flang/lib/Lower/Bridge.cpp

@@ -4355,6 +4355,62 @@ class FirConverter : public Fortran::lower::AbstractConverter {
                                         stmtCtx);
   }
 
+  void genForallPointerAssignment(
+      mlir::Location loc, const Fortran::evaluate::Assignment &assign,
+      const Fortran::evaluate::Assignment::BoundsSpec &lbExprs) {
+    if (Fortran::evaluate::IsProcedureDesignator(assign.rhs))
+      TODO(loc, "procedure pointer assignment inside FORALL");
+    std::optional<Fortran::evaluate::DynamicType> lhsType =
+        assign.lhs.GetType();
+    // Polymorphic pointer assignment is delegated to the runtime, and
+    // PointerAssociateLowerBounds needs the lower bounds as arguments, so they
+    // must be preserved.
+    if (lhsType && lhsType->IsPolymorphic())
+      TODO(loc, "polymorphic pointer assignment in FORALL");
+    // Nullification is special, there is no RHS that can be prepared,
+    // need to encode it in HLFIR.
+    if (Fortran::evaluate::UnwrapExpr<Fortran::evaluate::NullPointer>(
+            assign.rhs))
+      TODO(loc, "NULL pointer assignment in FORALL");
+    // Lower bounds could be "applied" when preparing RHS, but in order
+    // to deal with the polymorphic case and to reuse existing pointer
+    // assignment helpers in HLFIR codegen, it is better to keep them
+    // separate.
+    if (!lbExprs.empty())
+      TODO(loc, "Pointer assignment with new lower bounds inside FORALL");
+    // Otherwise, this is a "dumb" pointer assignment that can be represented
+    // with hlfir.region_assign with descriptor address/value and later
+    // implemented with a store.
+    auto regionAssignOp = builder->create<hlfir::RegionAssignOp>(loc);
+
+    // Lower LHS in its own region.
+    builder->createBlock(&regionAssignOp.getLhsRegion());
+    Fortran::lower::StatementContext lhsContext;
+    hlfir::Entity lhs = Fortran::lower::convertExprToHLFIR(
+        loc, *this, assign.lhs, localSymbols, lhsContext);
+
+    auto lhsYieldOp = builder->create<hlfir::YieldOp>(loc, lhs);
+    Fortran::lower::genCleanUpInRegionIfAny(
+        loc, *builder, lhsYieldOp.getCleanup(), lhsContext);
+
+    // Lower RHS in its own region.
+    builder->createBlock(&regionAssignOp.getRhsRegion());
+    Fortran::lower::StatementContext rhsContext;
+    hlfir::Entity rhs = Fortran::lower::convertExprToHLFIR(
+        loc, *this, assign.rhs, localSymbols, rhsContext);
+    // Create pointer descriptor value from the RHS.
+    if (rhs.isMutableBox())
+      rhs = hlfir::Entity{builder->create<fir::LoadOp>(loc, rhs)};
+    auto lhsBoxType =
+        llvm::cast<fir::BaseBoxType>(fir::unwrapRefType(lhs.getType()));
+    mlir::Value newBox = hlfir::genVariableBox(loc, *builder, rhs, lhsBoxType);
+    auto rhsYieldOp = builder->create<hlfir::YieldOp>(loc, newBox);
+    Fortran::lower::genCleanUpInRegionIfAny(
+        loc, *builder, rhsYieldOp.getCleanup(), rhsContext);
+
+    builder->setInsertionPointAfter(regionAssignOp);
+  }
+
   // Create the 2 x newRank array with the bounds to be passed to the runtime as
   // a descriptor.
   mlir::Value createBoundArray(llvm::ArrayRef<mlir::Value> lbounds,
@@ -4793,13 +4849,16 @@ class FirConverter : public Fortran::lower::AbstractConverter {
               },
               [&](const Fortran::evaluate::Assignment::BoundsSpec &lbExprs) {
                 if (isInsideHlfirForallOrWhere())
-                  TODO(loc, "pointer assignment inside FORALL");
-                genPointerAssignment(loc, assign, lbExprs);
+                  genForallPointerAssignment(loc, assign, lbExprs);
+                else
+                  genPointerAssignment(loc, assign, lbExprs);
               },
               [&](const Fortran::evaluate::Assignment::BoundsRemapping
                       &boundExprs) {
                 if (isInsideHlfirForallOrWhere())
-                  TODO(loc, "pointer assignment inside FORALL");
+                  TODO(
+                      loc,
+                      "pointer assignment with bounds remapping inside FORALL");
                 genPointerAssignment(loc, assign, boundExprs);
               },
           },

flang/lib/Optimizer/Builder/HLFIRTools.cpp

@@ -349,26 +349,54 @@ mlir::Value hlfir::genVariableBoxChar(mlir::Location loc,
                                           lengths[0]);
 }
 
+static hlfir::Entity changeBoxAttributes(mlir::Location loc,
+                                         fir::FirOpBuilder &builder,
+                                         hlfir::Entity var,
+                                         fir::BaseBoxType forceBoxType) {
+  assert(llvm::isa<fir::BaseBoxType>(var.getType()) && "expect box type");
+  // Propagate lower bounds.
+  mlir::Value shift;
+  llvm::SmallVector<mlir::Value> lbounds =
+      getNonDefaultLowerBounds(loc, builder, var);
+  if (!lbounds.empty())
+    shift = builder.genShift(loc, lbounds);
+  auto rebox = builder.create<fir::ReboxOp>(loc, forceBoxType, var, shift,
+                                            /*slice=*/nullptr);
+  return hlfir::Entity{rebox};
+}
+
 hlfir::Entity hlfir::genVariableBox(mlir::Location loc,
                                     fir::FirOpBuilder &builder,
-                                    hlfir::Entity var) {
+                                    hlfir::Entity var,
+                                    fir::BaseBoxType forceBoxType) {
   assert(var.isVariable() && "must be a variable");
   var = hlfir::derefPointersAndAllocatables(loc, builder, var);
-  if (mlir::isa<fir::BaseBoxType>(var.getType()))
-    return var;
+  if (mlir::isa<fir::BaseBoxType>(var.getType())) {
+    if (!forceBoxType || forceBoxType == var.getType())
+      return var;
+    return changeBoxAttributes(loc, builder, var, forceBoxType);
+  }
   // Note: if the var is not a fir.box/fir.class at that point, it has default
   // lower bounds and is not polymorphic.
   mlir::Value shape =
       var.isArray() ? hlfir::genShape(loc, builder, var) : mlir::Value{};
   llvm::SmallVector<mlir::Value> typeParams;
-  auto maybeCharType =
-      mlir::dyn_cast<fir::CharacterType>(var.getFortranElementType());
+  mlir::Type elementType =
+      forceBoxType ? fir::getFortranElementType(forceBoxType.getEleTy())
+                   : var.getFortranElementType();
+  auto maybeCharType = mlir::dyn_cast<fir::CharacterType>(elementType);
   if (!maybeCharType || maybeCharType.hasDynamicLen())
     hlfir::genLengthParameters(loc, builder, var, typeParams);
   mlir::Value addr = var.getBase();
   if (mlir::isa<fir::BoxCharType>(var.getType()))
     addr = genVariableRawAddress(loc, builder, var);
   mlir::Type boxType = fir::BoxType::get(var.getElementOrSequenceType());
+  if (forceBoxType) {
+    boxType = forceBoxType;
+    mlir::Type baseType =
+        fir::ReferenceType::get(fir::unwrapRefType(forceBoxType.getEleTy()));
+    addr = builder.createConvert(loc, baseType, addr);
+  }
   auto embox =
       builder.create<fir::EmboxOp>(loc, boxType, addr, shape,
                                    /*slice=*/mlir::Value{}, typeParams);

flang/lib/Optimizer/Builder/TemporaryStorage.cpp

@@ -355,3 +355,27 @@ void fir::factory::AnyVectorSubscriptStack::destroy(
   static_cast<AnyVariableStack *>(this)->destroy(loc, builder);
   shapeTemp->destroy(loc, builder);
 }
+
+//===----------------------------------------------------------------------===//
+// fir::factory::AnyDescriptorAddressStack implementation.
+//===----------------------------------------------------------------------===//
+
+fir::factory::AnyDescriptorAddressStack::AnyDescriptorAddressStack(
+    mlir::Location loc, fir::FirOpBuilder &builder,
+    mlir::Type descriptorAddressType)
+    : AnyValueStack(loc, builder, builder.getIntPtrType()),
+      descriptorAddressType{descriptorAddressType} {}
+
+void fir::factory::AnyDescriptorAddressStack::pushValue(
+    mlir::Location loc, fir::FirOpBuilder &builder, mlir::Value variable) {
+  mlir::Value cast =
+      builder.createConvert(loc, builder.getIntPtrType(), variable);
+  static_cast<AnyValueStack *>(this)->pushValue(loc, builder, cast);
+}
+
+mlir::Value
+fir::factory::AnyDescriptorAddressStack::fetch(mlir::Location loc,
+                                               fir::FirOpBuilder &builder) {
+  mlir::Value addr = static_cast<AnyValueStack *>(this)->fetch(loc, builder);
+  return builder.createConvert(loc, descriptorAddressType, addr);
+}

flang/lib/Optimizer/Dialect/FIRType.cpp

@@ -1359,6 +1359,10 @@ bool fir::BaseBoxType::isAssumedRank() const {
   return false;
 }
 
+bool fir::BaseBoxType::isPointer() const {
+  return llvm::isa<fir::PointerType>(getEleTy());
+}
+
 //===----------------------------------------------------------------------===//
 // FIROpsDialect
 //===----------------------------------------------------------------------===//

flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp

@@ -1891,6 +1891,20 @@ llvm::LogicalResult hlfir::RegionAssignOp::verify() {
   return mlir::success();
 }
 
+bool hlfir::RegionAssignOp::isPointerAssignment() {
+  if (!getUserDefinedAssignment().empty())
+    return false;
+  hlfir::YieldOp yieldOp =
+      mlir::dyn_cast_or_null<hlfir::YieldOp>(getTerminator(getLhsRegion()));
+  if (!yieldOp)
+    return false;
+  mlir::Type lhsType = yieldOp.getEntity().getType();
+  if (!hlfir::isBoxAddressType(lhsType))
+    return false;
+  auto baseBoxType = llvm::cast<fir::BaseBoxType>(fir::unwrapRefType(lhsType));
+  return baseBoxType.isPointer();
+}
+
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//