[flang] Allow scalar boxed record type in intrinsic elemental lowering

[clementval]

Summary

Relax a bit the condition added in D144417 and allow scalar polymorphic entities
and boxed scalar record type.

[jeanPerier]

Looks good

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

 
 /// Return true iff `ty` is !fir.box<none>.
 bool isBoxNone(mlir::Type ty);
 
 /// Return true iff `ty` is the type of a boxed record type.
 /// e.g. !fir.box<!fir.type<derived>>
 bool isBoxedRecordType(mlir::Type ty);
 
+/// Return true iff `ty` is a scalar boxed record type.
+/// e.g. !fir.box<!fir.type<derived>>
+///      !fir.box<!fir.heap<!fir.type<derived>>>
+///      !fir.class<!fir.type<derived>>
+bool isScalarBoxedRecordType(mlir::Type ty);
+
 /// Return the nested RecordType if one if found. Return ty otherwise.
 mlir::Type getDerivedType(mlir::Type ty);
 
 /// Return true iff `ty` is the type of an polymorphic entity or
 /// value.
 bool isPolymorphicType(mlir::Type ty);
 
 /// Return true iff `ty` is the type of an unlimited polymorphic entity or

flang/lib/Optimizer/Builder/IntrinsicCall.cpp

 template <>
 fir::ExtendedValue
 IntrinsicLibrary::genElementalCall<IntrinsicLibrary::ExtendedGenerator>(
     ExtendedGenerator generator, llvm::StringRef name, mlir::Type resultType,
     llvm::ArrayRef<fir::ExtendedValue> args, bool outline) {
   for (const fir::ExtendedValue &arg : args) {
     auto *box = arg.getBoxOf<fir::BoxValue>();
     if (!arg.getUnboxed() && !arg.getCharBox() &&
-        !(box && fir::isPolymorphicType(fir::getBase(*box).getType())))
+        !(box && fir::isScalarBoxedRecordType(fir::getBase(*box).getType())))
       fir::emitFatalError(loc, "nonscalar intrinsic argument");
   }
   if (outline)
     return outlineInExtendedWrapper(generator, name, resultType, args);
   return std::invoke(generator, *this, resultType, args);
 }
 
 template <>

flang/lib/Optimizer/Dialect/FIRType.cpp

if (auto boxTy = ty.dyn_cast<fir::BoxType>()) {
     if (boxTy.getEleTy().isa<fir::RecordType>())
       return true;
     mlir::Type innerType = boxTy.unwrapInnerType();
     return innerType && innerType.isa<fir::RecordType>();
   }
   return false;
 }
 
+bool isScalarBoxedRecordType(mlir::Type ty) {
+  if (auto refTy = fir::dyn_cast_ptrEleTy(ty))
+    ty = refTy;
+  if (auto boxTy = ty.dyn_cast<fir::BaseBoxType>()) {
+    if (boxTy.getEleTy().isa<fir::RecordType>())
+      return true;
+    if (auto heapTy = boxTy.getEleTy().dyn_cast<fir::HeapType>())
+      return heapTy.getEleTy().isa<fir::RecordType>();
+    if (auto ptrTy = boxTy.getEleTy().dyn_cast<fir::PointerType>())
+      return ptrTy.getEleTy().isa<fir::RecordType>();
+  }
+  return false;
+}
+
 static bool isAssumedType(mlir::Type ty) {
   if (auto boxTy = ty.dyn_cast<fir::BoxType>()) {
     if (boxTy.getEleTy().isa<mlir::NoneType>())
       return true;
     if (auto seqTy = boxTy.getEleTy().dyn_cast<fir::SequenceType>())
       return seqTy.getEleTy().isa<mlir::NoneType>();
   }
   return false;

flang/test/Lower/polymorphic-temp.f90

 ! CHECK: %[[LOAD_A1:.*]] = fir.load %[[COORD_A]] : !fir.ref<i32>
 ! CHECK: %[[FIELD_A:.*]] = fir.field_index a, !fir.type<_QMpoly_tmpTp1{a:i32}>
 ! CHECK: %[[COORD_A:.*]] = fir.coordinate_of %[[ARG1]], %[[FIELD_A]] : (!fir.class<!fir.type<_QMpoly_tmpTp1{a:i32}>>, !fir.field) -> !fir.ref<i32>
 ! CHECK: %[[LOAD_A2:.*]] = fir.load %[[COORD_A]] : !fir.ref<i32>
 ! CHECK: %[[CMPI:.*]] = arith.cmpi sgt, %[[LOAD_A1]], %[[LOAD_A2]] : i32
 ! CHECK: %[[SELECT:.*]] = arith.select %[[CMPI]], %[[ARG0]], %[[ARG1]] : !fir.class<!fir.type<_QMpoly_tmpTp1{a:i32}>>
 ! CHECK: fir.call @_QMpoly_tmpPcheck_scalar(%[[SELECT]]) {{.*}} : (!fir.class<!fir.type<_QMpoly_tmpTp1{a:i32}>>) -> ()
 
+  subroutine test_merge_intrinsic2(a, b, i)
+    class(p1), allocatable, intent(in) :: a
+    type(p1), allocatable :: b
+    integer, intent(in) :: i
+
+    call check_scalar(merge(a, b, i==1))
+  end subroutine
+
+
+! CHECK-LABEL: func.func @_QMpoly_tmpPtest_merge_intrinsic2(
+! CHECK-SAME: %[[A:.*]]: !fir.ref<!fir.class<!fir.heap<!fir.type<_QMpoly_tmpTp1{a:i32}>>>> {fir.bindc_name = "a"}, %[[B:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.type<_QMpoly_tmpTp1{a:i32}>>>> {fir.bindc_name = "b"}, %[[I:.*]]: !fir.ref<i32> {fir.bindc_name = "i"}) {
+! CHECK: %[[LOAD_A:.*]] = fir.load %[[A]] : !fir.ref<!fir.class<!fir.heap<!fir.type<_QMpoly_tmpTp1{a:i32}>>>>
+! CHECK: %[[LOAD_B:.*]] = fir.load %[[B]] : !fir.ref<!fir.box<!fir.heap<!fir.type<_QMpoly_tmpTp1{a:i32}>>>>
+! CHECK: %[[LOAD_I:.*]] = fir.load %[[I]] : !fir.ref<i32>
+! CHECK: %[[C1:.*]] = arith.constant 1 : i32
+! CHECK: %[[CMPI:.*]] = arith.cmpi eq, %[[LOAD_I]], %[[C1]] : i32
+! CHECK: %[[B_CONV:.*]] = fir.convert %[[LOAD_B]] : (!fir.box<!fir.heap<!fir.type<_QMpoly_tmpTp1{a:i32}>>>) -> !fir.class<!fir.heap<!fir.type<_QMpoly_tmpTp1{a:i32}>>>
+! CHECK: %{{.*}} = arith.select %[[CMPI]], %[[LOAD_A]], %[[B_CONV]] : !fir.class<!fir.heap<!fir.type<_QMpoly_tmpTp1{a:i32}>>>
+
 end module

flang/unittests/Optimizer/FIRTypesTest.cpp

TEST_F(FIRTypesTest, isBoxedRecordType) {
   // TYPE(T), DIMENSION(:)
   ty = fir::BoxType::get(seqRecTy);
   EXPECT_TRUE(fir::isBoxedRecordType(ty));
 
   EXPECT_FALSE(fir::isBoxedRecordType(fir::BoxType::get(
       fir::ReferenceType::get(mlir::IntegerType::get(&context, 32)))));
 }
 
+// Test fir::isScalarBoxedRecordType from flang/Optimizer/Dialect/FIRType.h.
+TEST_F(FIRTypesTest, isScalarBoxedRecordType) {
+  mlir::Type recTy = fir::RecordType::get(&context, "dt");
+  mlir::Type seqRecTy =
+      fir::SequenceType::get({fir::SequenceType::getUnknownExtent()}, recTy);
+  mlir::Type ty = fir::BoxType::get(recTy);
+  EXPECT_TRUE(fir::isScalarBoxedRecordType(ty));
+  EXPECT_TRUE(fir::isScalarBoxedRecordType(fir::ReferenceType::get(ty)));
+
+  // CLASS(T), ALLOCATABLE
+  ty = fir::ClassType::get(fir::HeapType::get(recTy));
+  EXPECT_TRUE(fir::isScalarBoxedRecordType(ty));
+
+  // TYPE(T), ALLOCATABLE
+  ty = fir::BoxType::get(fir::HeapType::get(recTy));
+  EXPECT_TRUE(fir::isScalarBoxedRecordType(ty));
+
+  // TYPE(T), POINTER
+  ty = fir::BoxType::get(fir::PointerType::get(recTy));
+  EXPECT_TRUE(fir::isScalarBoxedRecordType(ty));
+
+  // CLASS(T), POINTER
+  ty = fir::ClassType::get(fir::PointerType::get(recTy));
+  EXPECT_TRUE(fir::isScalarBoxedRecordType(ty));
+
+  // TYPE(T), DIMENSION(10)
+  ty = fir::BoxType::get(fir::SequenceType::get({10}, recTy));
+  EXPECT_FALSE(fir::isScalarBoxedRecordType(ty));
+
+  // TYPE(T), DIMENSION(:)
+  ty = fir::BoxType::get(seqRecTy);
+  EXPECT_FALSE(fir::isScalarBoxedRecordType(ty));
+
+  EXPECT_FALSE(fir::isScalarBoxedRecordType(fir::BoxType::get(
+      fir::ReferenceType::get(mlir::IntegerType::get(&context, 32)))));
+}
+
 TEST_F(FIRTypesTest, updateTypeForUnlimitedPolymorphic) {
   // RecordType are not changed.
 
   // !fir.tyep<T> -> !fir.type<T>
   mlir::Type recTy = fir::RecordType::get(&context, "dt");
   EXPECT_EQ(recTy, fir::updateTypeForUnlimitedPolymorphic(recTy));
 
   // !fir.array<2x!fir.type<T>> -> !fir.array<2x!fir.type<T>>