[flang] implement simple pointer assignments inside FORALL

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

@@ -1358,6 +1358,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
 //===----------------------------------------------------------------------===//