[flang] Pointer assignment with remapping involcing polymorphic entities

[clementval]

Summary

Lower pointer assignment with remapping involving polymorphic entities
to runtime call to PointerAssociateRemapping.
For the time being all pointer assignment involcing polymorphic entities are
done with the runtime call. When lhs is not unlimited polymorphic
we might be able to do it inlined as well.

[jeanPerier]

Small nits, LGTM otherwise.

flang/lib/Lower/Bridge.cpp

2802

You can remove the "may" from the comment if they are indeed getting a special treatment here.

2808

You may want to leave a TODO for the unlikely case where this happens in a FORALL (explicitIterationSpace()) and overlaps between the RHSs and LHSs needs to be taken into account.

[PeteSteinfeld]

All builds and tests correctly and looks good.

[clementval]

Add TODO for FORALL

[clementval]

Rebase

flang/include/flang/Lower/Runtime.h

 void genUnlockStatement(AbstractConverter &, const parser::UnlockStmt &);
 void genPauseStatement(AbstractConverter &, const parser::PauseStmt &);
 
 mlir::Value genAssociated(fir::FirOpBuilder &, mlir::Location,
                           mlir::Value pointer, mlir::Value target);
 
 void genPointerAssociate(fir::FirOpBuilder &, mlir::Location,
                          mlir::Value pointer, mlir::Value target);
+void genPointerAssociateRemapping(fir::FirOpBuilder &, mlir::Location,
+                                  mlir::Value pointer, mlir::Value target,
+                                  mlir::Value bounds);
 
 mlir::Value genCpuTime(fir::FirOpBuilder &, mlir::Location);
 void genDateAndTime(fir::FirOpBuilder &, mlir::Location,
                     llvm::Optional<fir::CharBoxValue> date,
                     llvm::Optional<fir::CharBoxValue> time,
                     llvm::Optional<fir::CharBoxValue> zone, mlir::Value values);
 
 void genRandomInit(fir::FirOpBuilder &, mlir::Location, mlir::Value repeatable,

flang/lib/Lower/Bridge.cpp

std::visit(
               Fortran::lower::associateMutableBox(*this, loc, lhs, assign.rhs,
                                                   lbounds, stmtCtx);
             },
 
             // [4] Pointer assignment with bounds-remapping. R1036: a
             // bounds-remapping is a pair, lower bound and upper bound.
             [&](const Fortran::evaluate::Assignment::BoundsRemapping
                     &boundExprs) {
-              std::optional<Fortran::evaluate::DynamicType> lhsType =
-                  assign.lhs.GetType();
-              std::optional<Fortran::evaluate::DynamicType> rhsType =
-                  assign.rhs.GetType();
-              // Polymorphic lhs/rhs may need more care. See F2018 10.2.2.3.
-              if ((lhsType && lhsType->IsPolymorphic()) ||
-                  (rhsType && rhsType->IsPolymorphic()))
-                TODO(loc, "pointer assignment involving polymorphic entity");
-
               llvm::SmallVector<mlir::Value> lbounds;
               llvm::SmallVector<mlir::Value> ubounds;
               for (const std::pair<Fortran::evaluate::ExtentExpr,
                                    Fortran::evaluate::ExtentExpr> &pair :
                    boundExprs) {
                 const Fortran::evaluate::ExtentExpr &lbExpr = pair.first;
                 const Fortran::evaluate::ExtentExpr &ubExpr = pair.second;
                 lbounds.push_back(
                     fir::getBase(genExprValue(toEvExpr(lbExpr), stmtCtx)));
                 ubounds.push_back(
                     fir::getBase(genExprValue(toEvExpr(ubExpr), stmtCtx)));
               }
+
+              std::optional<Fortran::evaluate::DynamicType> lhsType =
+                  assign.lhs.GetType();
+              std::optional<Fortran::evaluate::DynamicType> rhsType =
+                  assign.rhs.GetType();
+              // Polymorphic lhs/rhs need more care. See F2018 10.2.2.3.
+              if ((lhsType && lhsType->IsPolymorphic()) ||
+                  (rhsType && rhsType->IsPolymorphic())) {
+                if (explicitIterationSpace())
+                  TODO(loc, "polymorphic pointer assignment in FORALL");
+                
+                mlir::Value lhs = genExprMutableBox(loc, assign.lhs).getAddr();
+                mlir::Value rhs =
+                    fir::getBase(genExprBox(loc, assign.rhs, stmtCtx));
+
+                // Create the 2xnewRank array with the bounds to be passed to
+                // the runtime as a descriptor.
+                assert(lbounds.size() && ubounds.size());
+                fir::SequenceType::Shape shape(2, lbounds.size());
+                mlir::Type indexTy = builder->getIndexType();
+                mlir::Type boundArrayTy =
+                    fir::SequenceType::get(shape, builder->getI64Type());
+                mlir::Value boundArray =
+                    builder->create<fir::AllocaOp>(loc, boundArrayTy);
+                mlir::Value array =
+                    builder->create<fir::UndefOp>(loc, boundArrayTy);
+                llvm::SmallVector<mlir::Value> exts;
+                mlir::Value c2 =
+                    builder->createIntegerConstant(loc, indexTy, 2);
+                for (unsigned i = 0; i < lbounds.size(); ++i) {
+                  array = builder->create<fir::InsertValueOp>(
+                      loc, boundArrayTy, array, lbounds[i],
+                      builder->getArrayAttr(
+                          {builder->getIntegerAttr(builder->getIndexType(),
+                                                   static_cast<int>(i)),
+                           builder->getIntegerAttr(builder->getIndexType(),
+                                                   static_cast<int>(0))}));
+                  array = builder->create<fir::InsertValueOp>(
+                      loc, boundArrayTy, array, ubounds[i],
+                      builder->getArrayAttr(
+                          {builder->getIntegerAttr(builder->getIndexType(),
+                                                   static_cast<int>(i)),
+                           builder->getIntegerAttr(builder->getIndexType(),
+                                                   static_cast<int>(1))}));
+                  exts.push_back(c2);
+                }
+                builder->create<fir::StoreOp>(loc, array, boundArray);
+                mlir::Type boxTy = fir::BoxType::get(boundArrayTy);
+                mlir::Value shapeOp = builder->genShape(loc, exts);
+                mlir::Value boundsDesc = builder->create<fir::EmboxOp>(
+                    loc, boxTy, boundArray, shapeOp);
+                Fortran::lower::genPointerAssociateRemapping(*builder, loc, lhs,
+                                                             rhs, boundsDesc);
+                return;
+              }
               if (explicitIterationSpace()) {
                 // Pointer assignment in FORALL context. Copy the rhs box value
                 // into the lhs box variable.
                 genArrayAssignment(assign, stmtCtx, lbounds, ubounds);
                 return;
               }
               fir::MutableBoxValue lhs = genExprMutableBox(loc, assign.lhs);
               if (Fortran::evaluate::UnwrapExpr<Fortran::evaluate::NullPointer>(

flang/lib/Lower/Runtime.cpp

void Fortran::lower::genPointerAssociate(fir::FirOpBuilder &builder,
                                          mlir::Value target) {
   mlir::func::FuncOp func =
       fir::runtime::getRuntimeFunc<mkRTKey(PointerAssociate)>(loc, builder);
   llvm::SmallVector<mlir::Value> args = fir::runtime::createArguments(
       builder, loc, func.getFunctionType(), pointer, target);
   builder.create<fir::CallOp>(loc, func, args).getResult(0);
 }
 
+void Fortran::lower::genPointerAssociateRemapping(fir::FirOpBuilder &builder,
+                                                  mlir::Location loc,
+                                                  mlir::Value pointer,
+                                                  mlir::Value target,
+                                                  mlir::Value bounds) {
+  mlir::func::FuncOp func =
+      fir::runtime::getRuntimeFunc<mkRTKey(PointerAssociateRemapping)>(loc,
+                                                                       builder);
+  auto fTy = func.getFunctionType();
+  auto sourceFile = fir::factory::locationToFilename(builder, loc);
+  auto sourceLine =
+      fir::factory::locationToLineNo(builder, loc, fTy.getInput(4));
+  llvm::SmallVector<mlir::Value> args = fir::runtime::createArguments(
+      builder, loc, func.getFunctionType(), pointer, target, bounds, sourceFile,
+      sourceLine);
+  builder.create<fir::CallOp>(loc, func, args).getResult(0);
+}
+
 mlir::Value Fortran::lower::genCpuTime(fir::FirOpBuilder &builder,
                                        mlir::Location loc) {
   mlir::func::FuncOp func =
       fir::runtime::getRuntimeFunc<mkRTKey(CpuTime)>(loc, builder);
   return builder.create<fir::CallOp>(loc, func, std::nullopt).getResult(0);
 }
 
 void Fortran::lower::genDateAndTime(fir::FirOpBuilder &builder,

flang/test/Lower/polymorphic.f90

 ! CHECK: %[[T:.*]] = fir.alloca !fir.array<10x20x!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>> {bindc_name = "t", fir.target, uniq_name = "_QMpolymorphic_testFassign_polymorphic_allocatableEt"}
 ! CHECK: %[[SHAPE:.*]] = fir.shape %[[C10]], %[[C20]] : (index, index) -> !fir.shape<2>
 ! CHECK: %[[BOXED_T:.*]] = fir.embox %[[T]](%[[SHAPE]]) : (!fir.ref<!fir.array<10x20x!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>, !fir.shape<2>) -> !fir.box<!fir.array<10x20x!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>
 ! CHECK: %[[CONV_C:.*]] = fir.convert %[[C]] : (!fir.ref<!fir.class<!fir.heap<!fir.array<?x?x!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>>>) -> !fir.ref<!fir.box<none>>
 ! CHECK: %[[CONV_BOXED_T:.*]] = fir.convert %[[BOXED_T]] : (!fir.box<!fir.array<10x20x!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>) -> !fir.box<none>
 ! CHECK: %{{.*}} = fir.call @_FortranAAssign(%[[CONV_C]], %[[CONV_BOXED_T]], %{{.*}}, %{{.*}}) fastmath<contract> : (!fir.ref<!fir.box<none>>, !fir.box<none>, !fir.ref<i8>, i32) -> none
 ! CHECK: return
 
+  subroutine pointer_assign_remap()
+    class(p1), pointer :: a(:)
+    class(p1), pointer :: p(:,:)
+    allocate(a(100))
+    p(1:10,1:10) => a
+  end subroutine
+
+! CHECK-LABEL: func.func @_QMpolymorphic_testPpointer_assign_remap() {
+! CHECK: %[[A:.*]] = fir.alloca !fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>> {bindc_name = "a", uniq_name = "_QMpolymorphic_testFpointer_assign_remapEa"}
+! CHECK: %[[P:.*]] = fir.alloca !fir.class<!fir.ptr<!fir.array<?x?x!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>> {bindc_name = "p", uniq_name = "_QMpolymorphic_testFpointer_assign_remapEp"}
+! CHECK: %[[C1_0:.*]] = arith.constant 1 : i64
+! CHECK: %[[C10_0:.*]] = arith.constant 10 : i64
+! CHECK: %[[C1_1:.*]] = arith.constant 1 : i64
+! CHECK: %[[C10_1:.*]] = arith.constant 10 : i64
+! CHECK: %[[LOAD_A:.*]] = fir.load %[[A]] : !fir.ref<!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>>>
+! CHECK: %[[REBOX_A:.*]] = fir.rebox %[[LOAD_A]](%{{.*}}) : (!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>>, !fir.shift<1>) -> !fir.class<!fir.array<?x!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>
+! CHECK: %[[BOUND_ARRAY:.*]] = fir.alloca !fir.array<2x2xi64>
+! CHECK: %[[ARRAY:.*]] = fir.undefined !fir.array<2x2xi64>
+! CHECK: %[[C2:.*]] = arith.constant 2 : index
+! CHECK: %[[ARRAY0:.*]] = fir.insert_value %[[ARRAY]], %[[C1_0]], [0 : index, 0 : index] : (!fir.array<2x2xi64>, i64) -> !fir.array<2x2xi64>
+! CHECK: %[[ARRAY1:.*]] = fir.insert_value %[[ARRAY0]], %[[C10_0]], [0 : index, 1 : index] : (!fir.array<2x2xi64>, i64) -> !fir.array<2x2xi64>
+! CHECK: %[[ARRAY2:.*]] = fir.insert_value %[[ARRAY1]], %[[C1_1]], [1 : index, 0 : index] : (!fir.array<2x2xi64>, i64) -> !fir.array<2x2xi64>
+! CHECK: %[[ARRAY3:.*]] = fir.insert_value %[[ARRAY2]], %[[C10_1]], [1 : index, 1 : index] : (!fir.array<2x2xi64>, i64) -> !fir.array<2x2xi64>
+! CHECK: fir.store %[[ARRAY3]] to %[[BOUND_ARRAY]] : !fir.ref<!fir.array<2x2xi64>>
+! CHECK: %[[BOUND_ARRAY_SHAPE:.*]] = fir.shape %[[C2]], %[[C2]] : (index, index) -> !fir.shape<2>
+! CHECK: %[[BOXED_BOUND_ARRAY:.*]] = fir.embox %[[BOUND_ARRAY]](%[[BOUND_ARRAY_SHAPE]]) : (!fir.ref<!fir.array<2x2xi64>>, !fir.shape<2>) -> !fir.box<!fir.array<2x2xi64>>
+! CHECK: %[[ARG0:.*]] = fir.convert %[[P]] : (!fir.ref<!fir.class<!fir.ptr<!fir.array<?x?x!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>>>) -> !fir.ref<!fir.box<none>>
+! CHECK: %[[ARG1:.*]] = fir.convert %[[REBOX_A]] : (!fir.class<!fir.array<?x!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>) -> !fir.box<none>
+! CHECK: %[[ARG2:.*]] = fir.convert %[[BOXED_BOUND_ARRAY]] : (!fir.box<!fir.array<2x2xi64>>) -> !fir.box<none>
+! CHECK:  %{{.*}} = fir.call @_FortranAPointerAssociateRemapping(%[[ARG0]], %[[ARG1]], %[[ARG2]], %{{.*}}, %{{.*}}) fastmath<contract> : (!fir.ref<!fir.box<none>>, !fir.box<none>, !fir.box<none>, !fir.ref<i8>, i32) -> none
+
 end module