[flang] add support for procedure pointer assignment inside FORALL

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

@@ -60,7 +60,7 @@ class Entity : public mlir::Value {
   bool isVariable() const { return !isValue(); }
   bool isMutableBox() const { return hlfir::isBoxAddressType(getType()); }
   bool isProcedurePointer() const {
-    return fir::isBoxProcAddressType(getType());
+    return hlfir::isFortranProcedurePointerType(getType());
   }
   bool isBoxAddressOrValue() const {
     return hlfir::isBoxAddressOrValueType(getType());

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

@@ -180,7 +180,7 @@ class AnyValueStack {
 /// 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.
+/// addresses, use AnyAddressStack instead.
 class AnyVariableStack {
 public:
   AnyVariableStack(mlir::Location loc, fir::FirOpBuilder &builder,
@@ -205,19 +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 {
+/// Data structure to stack simple addresses (C pointers). It can be used to
+/// store data base addresses, descriptor addresses, procedure addresses, and
+/// pointer procedure address. It stores the addresses as int_ptr values under
+/// the hood.
+class AnyAddressStack : public AnyValueStack {
 public:
-  AnyDescriptorAddressStack(mlir::Location loc, fir::FirOpBuilder &builder,
-                            mlir::Type descriptorAddressType);
+  AnyAddressStack(mlir::Location loc, fir::FirOpBuilder &builder,
+                  mlir::Type addressType);
 
   void pushValue(mlir::Location loc, fir::FirOpBuilder &builder,
                  mlir::Value value);
   mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder);
 
 private:
-  mlir::Type descriptorAddressType;
+  mlir::Type addressType;
 };
 
 class TemporaryStorage;
@@ -281,8 +283,7 @@ class TemporaryStorage {
 
 private:
   std::variant<HomogeneousScalarStack, SimpleCopy, SSARegister, AnyValueStack,
-               AnyVariableStack, AnyVectorSubscriptStack,
-               AnyDescriptorAddressStack>
+               AnyVariableStack, AnyVectorSubscriptStack, AnyAddressStack>
       impl;
 };
 } // namespace fir::factory

flang/include/flang/Optimizer/HLFIR/HLFIRDialect.h

@@ -82,6 +82,17 @@ inline bool isPolymorphicType(mlir::Type type) {
   return fir::isPolymorphicType(type);
 }
 
+/// Is this the FIR type of a Fortran procedure pointer?
+inline bool isFortranProcedurePointerType(mlir::Type type) {
+  return fir::isBoxProcAddressType(type);
+}
+
+inline bool isFortranPointerObjectType(mlir::Type type) {
+  auto boxTy =
+      llvm::dyn_cast_or_null<fir::BaseBoxType>(fir::dyn_cast_ptrEleTy(type));
+  return boxTy && boxTy.isPointer();
+}
+
 /// Is this an SSA value type for the value of a Fortran procedure
 /// designator ?
 inline bool isFortranProcedureValue(mlir::Type type) {

flang/include/flang/Optimizer/HLFIR/HLFIROpBase.td

@@ -91,10 +91,9 @@ def IsFortranVariablePred
 def AnyFortranVariable : Type<IsFortranVariablePred, "any HLFIR variable type">;
 
 
-def AnyFortranValue : TypeConstraint<Or<[AnyLogicalLike.predicate,
-    AnyIntegerLike.predicate, AnyRealLike.predicate,
-    AnyFirComplexLike.predicate,
-    hlfir_ExprType.predicate]>, "any Fortran value type">;
+def IsFortranValuePred : CPred<"::hlfir::isFortranValueType($_self)">;
+def AnyFortranValue
+        : TypeConstraint<IsFortranValuePred, "any Fortran value type">;
 
 
 def AnyFortranEntity : TypeConstraint<Or<[AnyFortranVariable.predicate,

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

@@ -1378,6 +1378,8 @@ def hlfir_RegionAssignOp : hlfir_Op<"region_assign", [hlfir_OrderedAssignmentTre
     }
     mlir::Region* getSubTreeRegion() { return nullptr; }
     bool isPointerAssignment();
+    bool isPointerObjectAssignment();
+    bool isProcedurePointerAssignment();
   }];
 
   let hasCustomAssemblyFormat = 1;

flang/lib/Lower/Bridge.cpp

@@ -4353,8 +4353,6 @@ class FirConverter : public Fortran::lower::AbstractConverter {
   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
@@ -4383,27 +4381,38 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     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;
+    mlir::Value rhs =
+        genForallPointerAssignmentRhs(loc, lhs, assign, rhsContext);
+    auto rhsYieldOp = builder->create<hlfir::YieldOp>(loc, rhs);
+    Fortran::lower::genCleanUpInRegionIfAny(
+        loc, *builder, rhsYieldOp.getCleanup(), rhsContext);
+
+    builder->setInsertionPointAfter(regionAssignOp);
+  }
+
+  mlir::Value
+  genForallPointerAssignmentRhs(mlir::Location loc, mlir::Value lhs,
+                                const Fortran::evaluate::Assignment &assign,
+                                Fortran::lower::StatementContext &rhsContext) {
+    if (Fortran::evaluate::IsProcedureDesignator(assign.rhs))
+      return fir::getBase(Fortran::lower::convertExprToAddress(
+          loc, *this, assign.rhs, localSymbols, rhsContext));
+    // Data target.
     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);
+    return hlfir::genVariableBox(loc, *builder, rhs, lhsBoxType);
   }
 
   // Create the 2 x newRank array with the bounds to be passed to the runtime as

flang/lib/Optimizer/Builder/TemporaryStorage.cpp

@@ -357,25 +357,33 @@ void fir::factory::AnyVectorSubscriptStack::destroy(
 }
 
 //===----------------------------------------------------------------------===//
-// fir::factory::AnyDescriptorAddressStack implementation.
+// fir::factory::AnyAddressStack implementation.
 //===----------------------------------------------------------------------===//
 
-fir::factory::AnyDescriptorAddressStack::AnyDescriptorAddressStack(
-    mlir::Location loc, fir::FirOpBuilder &builder,
-    mlir::Type descriptorAddressType)
+fir::factory::AnyAddressStack::AnyAddressStack(mlir::Location loc,
+                                               fir::FirOpBuilder &builder,
+                                               mlir::Type addressType)
     : 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);
+      addressType{addressType} {}
+
+void fir::factory::AnyAddressStack::pushValue(mlir::Location loc,
+                                              fir::FirOpBuilder &builder,
+                                              mlir::Value variable) {
+  mlir::Value cast = variable;
+  if (auto boxProcType = llvm::dyn_cast<fir::BoxProcType>(variable.getType())) {
+    cast =
+        builder.create<fir::BoxAddrOp>(loc, boxProcType.getEleTy(), variable);
+  }
+  cast = builder.createConvert(loc, builder.getIntPtrType(), cast);
   static_cast<AnyValueStack *>(this)->pushValue(loc, builder, cast);
 }
 
-mlir::Value
-fir::factory::AnyDescriptorAddressStack::fetch(mlir::Location loc,
-                                               fir::FirOpBuilder &builder) {
+mlir::Value fir::factory::AnyAddressStack::fetch(mlir::Location loc,
+                                                 fir::FirOpBuilder &builder) {
   mlir::Value addr = static_cast<AnyValueStack *>(this)->fetch(loc, builder);
-  return builder.createConvert(loc, descriptorAddressType, addr);
+  if (auto boxProcType = llvm::dyn_cast<fir::BoxProcType>(addressType)) {
+    mlir::Value cast = builder.createConvert(loc, boxProcType.getEleTy(), addr);
+    return builder.create<fir::EmboxProcOp>(loc, boxProcType, cast);
+  }
+  return builder.createConvert(loc, addressType, addr);
 }

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

@@ -1891,18 +1891,33 @@ llvm::LogicalResult hlfir::RegionAssignOp::verify() {
   return mlir::success();
 }
 
-bool hlfir::RegionAssignOp::isPointerAssignment() {
+static mlir::Type
+getNonVectorSubscriptedLhsType(hlfir::RegionAssignOp regionAssign) {
+  hlfir::YieldOp yieldOp = mlir::dyn_cast_or_null<hlfir::YieldOp>(
+      getTerminator(regionAssign.getLhsRegion()));
+  return yieldOp ? yieldOp.getEntity().getType() : mlir::Type{};
+}
+
+bool hlfir::RegionAssignOp::isPointerObjectAssignment() {
   if (!getUserDefinedAssignment().empty())
     return false;
-  hlfir::YieldOp yieldOp =
-      mlir::dyn_cast_or_null<hlfir::YieldOp>(getTerminator(getLhsRegion()));
-  if (!yieldOp)
+  mlir::Type lhsType = getNonVectorSubscriptedLhsType(*this);
+  return lhsType && hlfir::isFortranPointerObjectType(lhsType);
+}
+
+bool hlfir::RegionAssignOp::isProcedurePointerAssignment() {
+  if (!getUserDefinedAssignment().empty())
     return false;
-  mlir::Type lhsType = yieldOp.getEntity().getType();
-  if (!hlfir::isBoxAddressType(lhsType))
+  mlir::Type lhsType = getNonVectorSubscriptedLhsType(*this);
+  return lhsType && hlfir::isFortranProcedurePointerType(lhsType);
+}
+
+bool hlfir::RegionAssignOp::isPointerAssignment() {
+  if (!getUserDefinedAssignment().empty())
     return false;
-  auto baseBoxType = llvm::cast<fir::BaseBoxType>(fir::unwrapRefType(lhsType));
-  return baseBoxType.isPointer();
+  mlir::Type lhsType = getNonVectorSubscriptedLhsType(*this);
+  return lhsType && (hlfir::isFortranPointerObjectType(lhsType) ||
+                     hlfir::isFortranProcedurePointerType(lhsType));
 }
 
 //===----------------------------------------------------------------------===//

flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp

@@ -1277,11 +1277,13 @@ void OrderedAssignmentRewriter::saveNonVectorSubscriptedAddress(
         [&] { temp = insertSavedEntity(region, fir::factory::SSARegister{}); });
   else
     doBeforeLoopNest([&] {
-      if (var.isMutableBox())
-        temp =
-            insertSavedEntity(region, fir::factory::AnyDescriptorAddressStack{
-                                          loc, builder, var.getType()});
+      if (var.isMutableBox() || var.isProcedure() || var.isProcedurePointer())
+        // Store single C pointer to entity.
+        temp = insertSavedEntity(
+            region, fir::factory::AnyAddressStack{loc, builder, var.getType()});
       else
+        // Store the base address and dynamic shape/length/type information
+        // as descriptor.
         temp = insertSavedEntity(region, fir::factory::AnyVariableStack{
                                              loc, builder, var.getType()});
     });