[flang] Lower storage_size intrinsic for polymorphic entities

[clementval]

Summary

Lower STOARGE_SIZE intrinsic when the argument is polymoprhic
or unlimited polymorphic. STOARGE_SIZE for monomorphic entity is folded
by the frontend.

[vzakhari]

nit: typo in description: STOARGE_SIZE

flang/lib/Lower/IntrinsicCall.cpp

4831

Is it legal to dereference args[1] if we allow args.size() to be 1 at line 4801?

[clementval]

flang/lib/Lower/IntrinsicCall.cpp

4831

Good catch it should be isStaticallyAbsent(args, 1);

[clementval]

update hanlding of optional kind

[vzakhari]

Thanks!

flang/include/flang/Optimizer/Builder/MutableBox.h

mlir::Value getMutableIRBox(fir::FirOpBuilder &builder, mlir::Location loc,
                             const fir::MutableBoxValue &box);
 
 /// Generate allocation or association status test and returns the resulting
 /// i1. This is testing this for a valid/non-null base address value.
 mlir::Value genIsAllocatedOrAssociatedTest(fir::FirOpBuilder &builder,
                                            mlir::Location loc,
                                            const fir::MutableBoxValue &box);
 
+/// Generate allocation or association status test and returns the resulting
+/// i1. This is testing this for a valid/non-null base address value.
+mlir::Value genIsNotAllocatedOrAssociatedTest(fir::FirOpBuilder &builder,
+                                              mlir::Location loc,
+                                              const fir::MutableBoxValue &box);
+
 } // namespace fir::factory
 
 #endif // FORTRAN_OPTIMIZER_BUILDER_MUTABLEBOX_H

flang/lib/Lower/IntrinsicCall.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "flang/Lower/IntrinsicCall.h"
 #include "flang/Common/static-multimap-view.h"
 #include "flang/Lower/Mangler.h"
 #include "flang/Lower/Runtime.h"
 #include "flang/Lower/StatementContext.h"
+#include "flang/Lower/Support/Utils.h"
 #include "flang/Lower/SymbolMap.h"
 #include "flang/Optimizer/Builder/BoxValue.h"
 #include "flang/Optimizer/Builder/Character.h"
 #include "flang/Optimizer/Builder/Complex.h"
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Builder/MutableBox.h"
 #include "flang/Optimizer/Builder/Runtime/Allocatable.h"
 #include "flang/Optimizer/Builder/Runtime/Character.h"
struct IntrinsicLibrary {
   template <typename Shift>
   mlir::Value genShift(mlir::Type resultType, llvm::ArrayRef<mlir::Value>);
   mlir::Value genShiftA(mlir::Type resultType, llvm::ArrayRef<mlir::Value>);
   mlir::Value genSign(mlir::Type, llvm::ArrayRef<mlir::Value>);
   fir::ExtendedValue genSize(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genSpacing(mlir::Type resultType,
                          llvm::ArrayRef<mlir::Value> args);
   fir::ExtendedValue genSpread(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
+  fir::ExtendedValue genStorageSize(mlir::Type,
+                                    llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genSum(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   void genSystemClock(llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genTrailz(mlir::Type, llvm::ArrayRef<mlir::Value>);
   fir::ExtendedValue genTransfer(mlir::Type,
                                  llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genTranspose(mlir::Type,
                                   llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genTrim(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
&I::genSize,
        {"dim", asAddr, handleDynamicOptional},
        {"kind", asValue}}},
      /*isElemental=*/false},
     {"spacing", &I::genSpacing},
     {"spread",
      &I::genSpread,
      {{{"source", asBox}, {"dim", asValue}, {"ncopies", asValue}}},
      /*isElemental=*/false},
+    {"storage_size",
+     &I::genStorageSize,
+     {{{"a", asInquired}, {"kind", asValue}}},
+     /*isElemental=*/false},
     {"sum",
      &I::genSum,
      {{{"array", asBox},
        {"dim", asValue},
        {"mask", asBox, handleDynamicOptional}}},
      /*isElemental=*/false},
     {"system_clock",
      &I::genSystemClock,
IntrinsicLibrary::genSpread(mlir::Type resultType,
   mlir::Value resultIrBox =
       fir::factory::getMutableIRBox(builder, loc, resultMutableBox);
 
   fir::runtime::genSpread(builder, loc, resultIrBox, source, dim, ncopies);
 
   return readAndAddCleanUp(resultMutableBox, resultType, "SPREAD");
 }
 
+// STORAGE_SIZE
+fir::ExtendedValue
+IntrinsicLibrary::genStorageSize(mlir::Type resultType,
+                                 llvm::ArrayRef<fir::ExtendedValue> args) {
+  assert(args.size() == 2 || args.size() == 1);
+  mlir::Value box = fir::getBase(args[0]);
+  mlir::Type boxTy = box.getType();
+  mlir::Type kindTy = builder.getDefaultIntegerType();
+  bool needRuntimeCheck = false;
+  std::string errorMsg;
+
+  if (fir::isUnlimitedPolymorphicType(boxTy) &&
+      (fir::isAllocatableType(boxTy) || fir::isPointerType(boxTy))) {
+    needRuntimeCheck = true;
+    errorMsg =
+        fir::isPointerType(boxTy)
+            ? "unlimited polymorphic disassociated POINTER in STORAGE_SIZE"
+            : "unlimited polymorphic unallocated ALLOCATABLE in STORAGE_SIZE";
+  } else if (fir::isPolymorphicType(boxTy) && fir::isPointerType(boxTy)) {
+    needRuntimeCheck = true;
+    errorMsg = "polymorphic disassociated POINTER in STORAGE_SIZE";
+  }
+  const fir::MutableBoxValue *mutBox = args[0].getBoxOf<fir::MutableBoxValue>();
+  if (needRuntimeCheck && mutBox) {
+    mlir::Value isNotAllocOrAssoc =
+        fir::factory::genIsNotAllocatedOrAssociatedTest(builder, loc, *mutBox);
+    builder.genIfThen(loc, isNotAllocOrAssoc)
+        .genThen([&]() {
+          fir::runtime::genReportFatalUserError(builder, loc, errorMsg);
+        })
+        .end();
+  }
+
+  // Handle optional kind argument
+  bool absentKind = isStaticallyAbsent(args, 1);
+  if (!absentKind) {
+    mlir::Operation *defKind = fir::getBase(args[1]).getDefiningOp();
+    assert(mlir::isa<mlir::arith::ConstantOp>(*defKind) &&
+           "kind not a constant");
+    auto constOp = mlir::dyn_cast<mlir::arith::ConstantOp>(*defKind);
+    kindTy = builder.getIntegerType(
+        builder.getKindMap().getIntegerBitsize(fir::toInt(constOp)));
+  }
+
+  if (box.getType().isa<fir::ReferenceType>())
+    box = builder.create<fir::LoadOp>(loc, box);
+  mlir::Value eleSize = builder.create<fir::BoxEleSizeOp>(loc, kindTy, box);
+  mlir::Value c8 = builder.createIntegerConstant(loc, kindTy, 8);
+  return builder.create<mlir::arith::MulIOp>(loc, eleSize, c8);
+}
+
 // SUM
 fir::ExtendedValue
 IntrinsicLibrary::genSum(mlir::Type resultType,
                          llvm::ArrayRef<fir::ExtendedValue> args) {
   return genReduction(fir::runtime::genSum, fir::runtime::genSumDim, "SUM",
                       resultType, args);
 }
 

flang/lib/Optimizer/Builder/MutableBox.cpp

 mlir::Value
 fir::factory::genIsAllocatedOrAssociatedTest(fir::FirOpBuilder &builder,
                                              mlir::Location loc,
                                              const fir::MutableBoxValue &box) {
   auto addr = MutablePropertyReader(builder, loc, box).readBaseAddress();
   return builder.genIsNotNullAddr(loc, addr);
 }
 
+mlir::Value fir::factory::genIsNotAllocatedOrAssociatedTest(
+    fir::FirOpBuilder &builder, mlir::Location loc,
+    const fir::MutableBoxValue &box) {
+  auto addr = MutablePropertyReader(builder, loc, box).readBaseAddress();
+  return builder.genIsNullAddr(loc, addr);
+}
+
 /// Generate finalizer call and inlined free. This does not check that the
 /// address was allocated.
 static void genFinalizeAndFree(fir::FirOpBuilder &builder, mlir::Location loc,
                                mlir::Value addr) {
   // TODO: call finalizer if any.
 
   // A heap (ALLOCATABLE) object may have been converted to a ptr (POINTER),
   // so make sure the heap type is restored before deallocation.

flang/test/Lower/Intrinsics/storage_size.f90

+! RUN: bbc -emit-fir -polymorphic-type %s -o - | FileCheck %s
+
+module storage_size_test
+  type :: p1
+    integer :: a
+  end type
+
+  type, extends(p1) :: p2
+    integer :: b
+  end type
+
+contains
+
+  integer function unlimited_polymorphic_pointer(p) result(size)
+    class(*), pointer :: p
+    size = storage_size(p)
+  end function
+
+! CHECK-LABEL: func.func @_QMstorage_size_testPunlimited_polymorphic_pointer(
+! CHECK-SAME: %[[P:.*]]: !fir.ref<!fir.class<!fir.ptr<none>>> {fir.bindc_name = "p"}) -> i32 {
+! CHECK: %[[SIZE:.*]] = fir.alloca i32 {bindc_name = "size", uniq_name = "_QMstorage_size_testFunlimited_polymorphic_pointerEsize"}
+! CHECK: %[[LOAD_P:.*]] = fir.load %[[P]] : !fir.ref<!fir.class<!fir.ptr<none>>>
+! CHECK: %[[P_ADDR:.*]] = fir.box_addr %[[LOAD_P]] : (!fir.class<!fir.ptr<none>>) -> !fir.ptr<none>
+! CHECK: %[[P_ADDR_I64:.*]] = fir.convert %[[P_ADDR]] : (!fir.ptr<none>) -> i64
+! CHECK: %[[C0:.*]] = arith.constant 0 : i64
+! CHECK: %[[IS_NULL_ADDR:.*]] = arith.cmpi eq, %[[P_ADDR_I64]], %[[C0]] : i64
+! CHECK: fir.if %[[IS_NULL_ADDR]] {
+! CHECK:   %{{.*}} = fir.call @_FortranAReportFatalUserError(%{{.*}}, %{{.*}}, %{{.*}}) {{.*}} : (!fir.ref<i8>, !fir.ref<i8>, i32) -> none
+! CHECK: }
+! CHECK: %[[LOAD_P:.*]] = fir.load %[[P]] : !fir.ref<!fir.class<!fir.ptr<none>>>
+! CHECK: %[[ELE_SIZE:.*]] = fir.box_elesize %[[LOAD_P]] : (!fir.class<!fir.ptr<none>>) -> i32
+! CHECK: %[[C8:.*]] = arith.constant 8 : i32
+! CHECK: %[[BITS:.*]] = arith.muli %[[ELE_SIZE]], %[[C8]] : i32
+! CHECK: fir.store %[[BITS]] to %[[SIZE]] : !fir.ref<i32>
+! CHECK: %[[RES:.*]] = fir.load %[[SIZE]] : !fir.ref<i32>
+! CHECK: return %[[RES]] : i32
+
+  integer function unlimited_polymorphic_allocatable(p) result(size)
+    class(*), allocatable :: p
+    size = storage_size(p)
+  end function
+
+! CHECK-LABEL: func.func @_QMstorage_size_testPunlimited_polymorphic_allocatable(
+! CHECK-SAME: %[[P:.*]]: !fir.ref<!fir.class<!fir.heap<none>>> {fir.bindc_name = "p"}) -> i32 {
+! CHECK: %[[SIZE:.*]] = fir.alloca i32 {bindc_name = "size", uniq_name = "_QMstorage_size_testFunlimited_polymorphic_allocatableEsize"}
+! CHECK: %[[LOAD_P:.*]] = fir.load %[[P]] : !fir.ref<!fir.class<!fir.heap<none>>>
+! CHECK: %[[P_ADDR:.*]] = fir.box_addr %[[LOAD_P]] : (!fir.class<!fir.heap<none>>) -> !fir.heap<none>
+! CHECK: %[[P_ADDR_I64:.*]] = fir.convert %[[P_ADDR]] : (!fir.heap<none>) -> i64
+! CHECK: %[[C0:.*]] = arith.constant 0 : i64
+! CHECK: %[[IS_NULL_ADDR:.*]] = arith.cmpi eq, %[[P_ADDR_I64]], %[[C0]] : i64
+! CHECK: fir.if %[[IS_NULL_ADDR]] {
+! CHECK:   %{{.*}} = fir.call @_FortranAReportFatalUserError(%{{.*}}, %{{.*}}, %{{.*}}) {{.*}} : (!fir.ref<i8>, !fir.ref<i8>, i32) -> none
+! CHECK: }
+! CHECK: %[[LOAD_P:.*]] = fir.load %[[P]] : !fir.ref<!fir.class<!fir.heap<none>>>
+! CHECK: %[[ELE_SIZE:.*]] = fir.box_elesize %[[LOAD_P]] : (!fir.class<!fir.heap<none>>) -> i32
+! CHECK: %[[C8:.*]] = arith.constant 8 : i32
+! CHECK: %[[BITS:.*]] = arith.muli %[[ELE_SIZE]], %[[C8]] : i32
+! CHECK: fir.store %[[BITS]] to %[[SIZE]] : !fir.ref<i32>
+! CHECK: %[[RES:.*]] = fir.load %[[SIZE]] : !fir.ref<i32>
+! CHECK: return %[[RES]] : i32
+
+  integer function polymorphic_pointer(p) result(size)
+    class(p1), pointer :: p
+    size = storage_size(p)
+  end function
+
+! CHECK-LABEL: func.func @_QMstorage_size_testPpolymorphic_pointer(
+! CHECK-SAME: %[[P:.*]]: !fir.ref<!fir.class<!fir.ptr<!fir.type<_QMstorage_size_testTp1{a:i32}>>>> {fir.bindc_name = "p"}) -> i32 {
+! CHECK: %[[SIZE:.*]] = fir.alloca i32 {bindc_name = "size", uniq_name = "_QMstorage_size_testFpolymorphic_pointerEsize"}
+! CHECK: %[[LOAD_P:.*]] = fir.load %[[P]] : !fir.ref<!fir.class<!fir.ptr<!fir.type<_QMstorage_size_testTp1{a:i32}>>>>
+! CHECK: %[[P_ADDR:.*]] = fir.box_addr %[[LOAD_P]] : (!fir.class<!fir.ptr<!fir.type<_QMstorage_size_testTp1{a:i32}>>>) -> !fir.ptr<!fir.type<_QMstorage_size_testTp1{a:i32}>>
+! CHECK: %[[P_ADDR_I64:.*]] = fir.convert %[[P_ADDR]] : (!fir.ptr<!fir.type<_QMstorage_size_testTp1{a:i32}>>) -> i64
+! CHECK: %[[C0:.*]] = arith.constant 0 : i64
+! CHECK: %[[IS_NULL_ADDR:.*]] = arith.cmpi eq, %[[P_ADDR_I64]], %[[C0]] : i64
+! CHECK: fir.if %[[IS_NULL_ADDR]] {
+! CHECK:   %{{.*}} = fir.call @_FortranAReportFatalUserError(%{{.*}}, %{{.*}}, %{{.*}}) {{.*}} : (!fir.ref<i8>, !fir.ref<i8>, i32) -> none
+! CHECK: }
+! CHECK: %[[LOAD_P:.*]] = fir.load %[[P]] : !fir.ref<!fir.class<!fir.ptr<!fir.type<_QMstorage_size_testTp1{a:i32}>>>>
+! CHECK: %[[ELE_SIZE:.*]] = fir.box_elesize %[[LOAD_P]] : (!fir.class<!fir.ptr<!fir.type<_QMstorage_size_testTp1{a:i32}>>>) -> i32
+! CHECK: %[[C8:.*]] = arith.constant 8 : i32
+! CHECK: %[[BITS:.*]] = arith.muli %[[ELE_SIZE]], %[[C8]] : i32
+! CHECK: fir.store %[[BITS]] to %[[SIZE]] : !fir.ref<i32>
+! CHECK: %[[RES:.*]] = fir.load %[[SIZE]] : !fir.ref<i32>
+! CHECK: return %[[RES]] : i32
+
+  integer function polymorphic(p) result(size)
+    class(p1) :: p
+    size = storage_size(p)
+  end function
+
+! CHECK-LABEL: func.func @_QMstorage_size_testPpolymorphic(
+! CHECK-SAME: %[[P:.*]]: !fir.class<!fir.type<_QMstorage_size_testTp1{a:i32}>> {fir.bindc_name = "p"}) -> i32 {
+! CHECK: %[[SIZE:.*]] = fir.alloca i32 {bindc_name = "size", uniq_name = "_QMstorage_size_testFpolymorphicEsize"}
+! CHECK: %[[ELE_SIZE:.*]] = fir.box_elesize %[[P]] : (!fir.class<!fir.type<_QMstorage_size_testTp1{a:i32}>>) -> i32
+! CHECK: %[[C8:.*]] = arith.constant 8 : i32
+! CHECK: %[[BITS:.*]] = arith.muli %[[ELE_SIZE]], %[[C8]] : i32
+! CHECK: fir.store %[[BITS]] to %[[SIZE]] : !fir.ref<i32>
+! CHECK: %[[RES:.*]] = fir.load %[[SIZE]] : !fir.ref<i32>
+! CHECK: return %[[RES]] : i32
+
+  integer(8) function polymorphic_rank(p) result(size)
+    class(p1) :: p
+    size = storage_size(p, 8)
+  end function
+
+! CHECK-LABEL: func.func @_QMstorage_size_testPpolymorphic_rank(
+! CHECK-SAME: %[[P:.*]]: !fir.class<!fir.type<_QMstorage_size_testTp1{a:i32}>> {fir.bindc_name = "p"}) -> i64 {
+! CHECK: %[[SIZE:.*]] = fir.alloca i64 {bindc_name = "size", uniq_name = "_QMstorage_size_testFpolymorphic_rankEsize"}
+! CHECK: %[[ELE_SIZE:.*]] = fir.box_elesize %[[P]] : (!fir.class<!fir.type<_QMstorage_size_testTp1{a:i32}>>) -> i64
+! CHECK: %[[C8:.*]] = arith.constant 8 : i64
+! CHECK: %[[BITS:.*]] = arith.muli %[[ELE_SIZE]], %[[C8]] : i64
+! CHECK: fir.store %[[BITS]] to %[[SIZE]] : !fir.ref<i64>
+! CHECK: %[[RES:.*]] = fir.load %[[SIZE]] : !fir.ref<i64>
+! CHECK: return %[[RES]] : i64
+
+end module