[flang] Inline array size call when dim is compile time constant

[razvanlupusoru]

Summary

Instead of calling _FortranASizeDim, we can instead load extent
directly from descriptor. Add this support for cases where dim
is a known constant at compile time.

[clementval]

Looks good. Can you just add a lit test for that in flang/test/Lower/Intrinsics/size.f90?

[razvanlupusoru]

Added lit test. Previously this test was calling _FortranASizeDim but now loads from descriptor directly.

[clementval]

LGTM

flang/lib/Optimizer/Builder/IntrinsicCall.cpp

IntrinsicLibrary::genSize(mlir::Type resultType,
   // an array of SIZE with DIM in most cases, but it may not be
   // possible in some cases like when in SIZE(function_call()).
   if (isStaticallyAbsent(args, 1))
     return builder.createConvert(loc, resultType,
                                  fir::runtime::genSize(builder, loc, array));
 
   // Get the DIM argument.
   mlir::Value dim = fir::getBase(args[1]);
+  if (std::optional<std::int64_t> cstDim = fir::getIntIfConstant(dim)) {
+    // If it is a compile time constant, skip the runtime call.
+    return builder.createConvert(loc, resultType,
+                                 fir::factory::readExtent(builder, loc,
+                                                          fir::BoxValue{array},
+                                                          cstDim.value() - 1));
+  }
   if (!fir::isa_ref_type(dim.getType()))
     return builder.createConvert(
         loc, resultType, fir::runtime::genSizeDim(builder, loc, array, dim));
 
   mlir::Value isDynamicallyAbsent = builder.genIsNullAddr(loc, dim);
   return builder
       .genIfOp(loc, {resultType}, isDynamicallyAbsent,
                /*withElseRegion=*/true)

flang/test/Lower/Intrinsics/ubound.f90

 ! CHECK:         %[[VAL_12:.*]] = fir.convert %[[VAL_3]]
 ! CHECK:         %[[VAL_13:.*]] = fir.convert %[[VAL_10]]
 ! CHECK:         %[[VAL_14:.*]] = fir.call @_FortranALboundDim(%[[VAL_11]], %[[VAL_12]], %[[VAL_13]], %{{.*}}) {{.*}}: (!fir.box<none>, i32, !fir.ref<i8>, i32) -> i64
 ! CHECK:         %[[VAL_15:.*]] = arith.subi %[[VAL_14]], %{{.*}} : i64
 ! CHECK:         %[[VAL_16:.*]] = arith.addi %[[VAL_15]], %[[VAL_8]] : i64
 ! CHECK:         fir.store %[[VAL_16]] to %{{.*}} : !fir.ref<i64>
   res = ubound(a, dim, 8)
 end subroutine
+
+
+! CHECK-LABEL: func @_QPubound_test_const_dim(
+subroutine ubound_test_const_dim(array)
+  real :: array(11:)
+  integer :: res
+! Should not call _FortranASizeDim when dim is compile time constant. But instead load from descriptor directly.
+! CHECK:         %[[C0:.*]] = arith.constant 0 : index
+! CHECK:         %[[DIMS:.*]]:3 = fir.box_dims %arg0, %[[C0]] : (!fir.box<!fir.array<?xf32>>, index) -> (index, index, index)
+! CHECK:         %{{.*}} = fir.convert %[[DIMS]]#1 : (index) -> i32
+  res = ubound(array, 1)
+end subroutine