[flang] lower SHAPE with assumed-rank arguments #94812

jeanPerier · 2024-06-07T22:28:51Z

Allocate result statically on the stack (using max rank) and use the runtime to fill it in correctly.
Runtime commit change is reviewed here: #94781

llvmbot · 2024-06-07T22:29:24Z

@llvm/pr-subscribers-flang-runtime

@llvm/pr-subscribers-flang-fir-hlfir

Author: None (jeanPerier)

Changes

Allocate result statically on the stack (using max rank) and use the runtime to fill it in correctly.
Runtime commit change is reviewed here: #94781

Full diff: https://github.com/llvm/llvm-project/pull/94812.diff

7 Files Affected:

(modified) flang/include/flang/Optimizer/Builder/Runtime/Inquiry.h (+6)
(modified) flang/include/flang/Runtime/inquiry.h (+2-1)
(modified) flang/lib/Optimizer/Builder/IntrinsicCall.cpp (+32-2)
(modified) flang/lib/Optimizer/Builder/Runtime/Inquiry.cpp (+14)
(modified) flang/runtime/inquiry.cpp (+3-2)
(added) flang/test/Lower/HLFIR/assumed-rank-inquiries-3.f90 (+56)
(modified) flang/unittests/Runtime/Inquiry.cpp (+6-3)

diff --git a/flang/include/flang/Optimizer/Builder/Runtime/Inquiry.h b/flang/include/flang/Optimizer/Builder/Runtime/Inquiry.h
index 132592a0197f8..5f14d7781004b 100644
--- a/flang/include/flang/Optimizer/Builder/Runtime/Inquiry.h
+++ b/flang/include/flang/Optimizer/Builder/Runtime/Inquiry.h
@@ -32,6 +32,12 @@ mlir::Value genLboundDim(fir::FirOpBuilder &builder, mlir::Location loc,
 void genUbound(fir::FirOpBuilder &builder, mlir::Location loc,
                mlir::Value resultBox, mlir::Value array, mlir::Value kind);
 
+/// Generate call to `Shape` runtime routine.
+/// First argument is a raw pointer to the result array storage that
+/// must be allocated by the caller.
+void genShape(fir::FirOpBuilder &builder, mlir::Location loc,
+              mlir::Value resultAddr, mlir::Value arrayt, mlir::Value kind);
+
 /// Generate call to `Size` runtime routine. This routine is a specialized
 /// version when the DIM argument is not specified by the user.
 mlir::Value genSize(fir::FirOpBuilder &builder, mlir::Location loc,
diff --git a/flang/include/flang/Runtime/inquiry.h b/flang/include/flang/Runtime/inquiry.h
index 7161d1e41c4bb..0a3cd51236fa3 100644
--- a/flang/include/flang/Runtime/inquiry.h
+++ b/flang/include/flang/Runtime/inquiry.h
@@ -24,7 +24,8 @@ extern "C" {
 std::int64_t RTDECL(LboundDim)(const Descriptor &array, int dim,
     const char *sourceFile = nullptr, int line = 0);
 
-void RTDECL(Shape)(void *result, const Descriptor &array, int kind);
+void RTDECL(Shape)(void *result, const Descriptor &array, int kind,
+    const char *sourceFile = nullptr, int line = 0);
 
 std::int64_t RTDECL(Size)(
     const Descriptor &array, const char *sourceFile = nullptr, int line = 0);
diff --git a/flang/lib/Optimizer/Builder/IntrinsicCall.cpp b/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
index 861b26de06370..b3e1ee3da3a77 100644
--- a/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
+++ b/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
@@ -5992,15 +5992,45 @@ mlir::Value IntrinsicLibrary::genSetExponent(mlir::Type resultType,
                                    fir::getBase(args[1])));
 }
 
+/// Generate runtime call to inquire about all the bounds/extents of an
+/// assumed-rank array.
+template <typename Func>
+static fir::ExtendedValue genAssumedRankBoundInquiry(
+    fir::FirOpBuilder &builder, mlir::Location loc, mlir::Type resultType,
+    llvm::ArrayRef<fir::ExtendedValue> args, int kindPos, Func genRtCall) {
+  const fir::ExtendedValue &array = args[0];
+  // Allocate an array with the maximum rank, that is big enough to hold the
+  // result but still "small" (15 elements). Static size alloca make stack
+  // analysis/manipulation easier.
+  mlir::Type resultElementType = fir::unwrapSequenceType(resultType);
+  mlir::Type allocSeqType =
+      fir::SequenceType::get({Fortran::common::maxRank}, resultElementType);
+  mlir::Value resultStorage = builder.createTemporary(loc, allocSeqType);
+  mlir::Value arrayBox = builder.createBox(loc, array);
+  mlir::Value kind = isStaticallyAbsent(args, kindPos)
+                         ? builder.createIntegerConstant(
+                               loc, builder.getI32Type(),
+                               builder.getKindMap().defaultIntegerKind())
+                         : fir::getBase(args[kindPos]);
+  genRtCall(builder, loc, resultStorage, arrayBox, kind);
+  mlir::Type baseType =
+      fir::ReferenceType::get(builder.getVarLenSeqTy(resultElementType));
+  mlir::Value resultBase = builder.createConvert(loc, baseType, resultStorage);
+  mlir::Value rank =
+      builder.create<fir::BoxRankOp>(loc, builder.getIndexType(), arrayBox);
+  return fir::ArrayBoxValue{resultBase, {rank}};
+}
+
 // SHAPE
 fir::ExtendedValue
 IntrinsicLibrary::genShape(mlir::Type resultType,
                            llvm::ArrayRef<fir::ExtendedValue> args) {
   assert(args.size() >= 1);
   const fir::ExtendedValue &array = args[0];
+  if (array.hasAssumedRank())
+    return genAssumedRankBoundInquiry(builder, loc, resultType, args,
+                                      /*kindPos=*/1, fir::runtime::genShape);
   int rank = array.rank();
-  if (rank == 0)
-    TODO(loc, "shape intrinsic lowering with assumed-rank source");
   mlir::Type indexType = builder.getIndexType();
   mlir::Type extentType = fir::unwrapSequenceType(resultType);
   mlir::Type seqType = fir::SequenceType::get(
diff --git a/flang/lib/Optimizer/Builder/Runtime/Inquiry.cpp b/flang/lib/Optimizer/Builder/Runtime/Inquiry.cpp
index 16f63bea4617a..34c4020b5907c 100644
--- a/flang/lib/Optimizer/Builder/Runtime/Inquiry.cpp
+++ b/flang/lib/Optimizer/Builder/Runtime/Inquiry.cpp
@@ -87,3 +87,17 @@ mlir::Value fir::runtime::genIsContiguous(fir::FirOpBuilder &builder,
   auto args = fir::runtime::createArguments(builder, loc, fTy, array);
   return builder.create<fir::CallOp>(loc, isContiguousFunc, args).getResult(0);
 }
+
+void fir::runtime::genShape(fir::FirOpBuilder &builder, mlir::Location loc,
+                            mlir::Value resultAddr, mlir::Value array,
+                            mlir::Value kind) {
+  mlir::func::FuncOp func =
+      fir::runtime::getRuntimeFunc<mkRTKey(Shape)>(loc, builder);
+  auto fTy = func.getFunctionType();
+  auto sourceFile = fir::factory::locationToFilename(builder, loc);
+  auto sourceLine =
+      fir::factory::locationToLineNo(builder, loc, fTy.getInput(4));
+  auto args = fir::runtime::createArguments(
+      builder, loc, fTy, resultAddr, array, kind, sourceFile, sourceLine);
+  builder.create<fir::CallOp>(loc, func, args).getResult(0);
+}
diff --git a/flang/runtime/inquiry.cpp b/flang/runtime/inquiry.cpp
index ea114174de7fd..443e6291e5e23 100644
--- a/flang/runtime/inquiry.cpp
+++ b/flang/runtime/inquiry.cpp
@@ -85,8 +85,9 @@ std::int64_t RTDEF(SizeDim)(
   return static_cast<std::int64_t>(dimension.Extent());
 }
 
-void RTDEF(Shape)(void *result, const Descriptor &array, int kind) {
-  Terminator terminator{__FILE__, __LINE__};
+void RTDEF(Shape)(void *result, const Descriptor &array, int kind,
+    const char *sourceFile, int line) {
+  Terminator terminator{sourceFile, line};
   INTERNAL_CHECK(array.rank() <= common::maxRank);
   for (SubscriptValue i{0}; i < array.rank(); ++i) {
     const Dimension &dimension{array.GetDimension(i)};
diff --git a/flang/test/Lower/HLFIR/assumed-rank-inquiries-3.f90 b/flang/test/Lower/HLFIR/assumed-rank-inquiries-3.f90
new file mode 100644
index 0000000000000..bbeff5ff05191
--- /dev/null
+++ b/flang/test/Lower/HLFIR/assumed-rank-inquiries-3.f90
@@ -0,0 +1,56 @@
+! Test shape lowering for assumed-rank
+! RUN: bbc -emit-hlfir -o - %s -allow-assumed-rank | FileCheck %s
+
+subroutine test_shape(x)
+  real :: x(..)
+  call takes_integer_array(shape(x))
+end subroutine
+! CHECK-LABEL:   func.func @_QPtest_shape(
+! CHECK:           %[[VAL_1:.*]] = fir.alloca !fir.array<15xi32>
+! CHECK:           %[[VAL_4:.*]] = arith.constant 4 : i32
+! CHECK:           %[[VAL_7:.*]] = fir.convert %[[VAL_1]] : (!fir.ref<!fir.array<15xi32>>) -> !fir.llvm_ptr<i8>
+! CHECK:           %[[VAL_8:.*]] = fir.convert %[[VAL_3:.*]] : (!fir.box<!fir.array<*:f32>>) -> !fir.box<none>
+! CHECK:           %[[VAL_10:.*]] = fir.call @_FortranAShape(%[[VAL_7]], %[[VAL_8]], %[[VAL_4]], %{{.*}}, %{{.*}})
+! CHECK:           %[[VAL_11:.*]] = fir.convert %[[VAL_1]] : (!fir.ref<!fir.array<15xi32>>) -> !fir.ref<!fir.array<?xi32>>
+! CHECK:           %[[VAL_12:.*]] = fir.box_rank %[[VAL_3]] : (!fir.box<!fir.array<*:f32>>) -> index
+! CHECK:           %[[VAL_13:.*]] = fir.shape %[[VAL_12]] : (index) -> !fir.shape<1>
+! CHECK:           %[[VAL_14:.*]]:2 = hlfir.declare %[[VAL_11]](%[[VAL_13]]) {uniq_name = ".tmp.intrinsic_result"} : (!fir.ref<!fir.array<?xi32>>, !fir.shape<1>) -> (!fir.box<!fir.array<?xi32>>, !fir.ref<!fir.array<?xi32>>)
+! CHECK:           %[[VAL_15:.*]] = arith.constant false
+! CHECK:           %[[VAL_16:.*]] = hlfir.as_expr %[[VAL_14]]#0 move %[[VAL_15]] : (!fir.box<!fir.array<?xi32>>, i1) -> !hlfir.expr<?xi32>
+! CHECK:           %[[VAL_17:.*]]:3 = hlfir.associate %[[VAL_16]](%[[VAL_13]]) {adapt.valuebyref} : (!hlfir.expr<?xi32>, !fir.shape<1>) -> (!fir.box<!fir.array<?xi32>>, !fir.ref<!fir.array<?xi32>>, i1)
+! CHECK:           fir.call @_QPtakes_integer_array(%[[VAL_17]]#1) fastmath<contract> : (!fir.ref<!fir.array<?xi32>>) -> ()
+! CHECK:           hlfir.end_associate %[[VAL_17]]#1, %[[VAL_17]]#2 : !fir.ref<!fir.array<?xi32>>, i1
+! CHECK:           hlfir.destroy %[[VAL_16]] : !hlfir.expr<?xi32>
+! CHECK:           return
+! CHECK:         }
+
+subroutine test_shape_kind(x)
+  real :: x(..)
+  call takes_integer8_array(shape(x, kind=8))
+end subroutine
+! CHECK-LABEL:   func.func @_QPtest_shape_kind(
+! CHECK:           %[[VAL_1:.*]] = fir.alloca !fir.array<15xi64>
+! CHECK:           %[[VAL_4:.*]] = arith.constant 8 : i32
+! CHECK:           %[[VAL_7:.*]] = fir.convert %[[VAL_1]] : (!fir.ref<!fir.array<15xi64>>) -> !fir.llvm_ptr<i8>
+! CHECK:           %[[VAL_8:.*]] = fir.convert %[[VAL_3:.*]] : (!fir.box<!fir.array<*:f32>>) -> !fir.box<none>
+! CHECK:           %[[VAL_10:.*]] = fir.call @_FortranAShape(%[[VAL_7]], %[[VAL_8]], %[[VAL_4]], %{{.*}}, %{{.*}})
+! CHECK:           %[[VAL_11:.*]] = fir.convert %[[VAL_1]] : (!fir.ref<!fir.array<15xi64>>) -> !fir.ref<!fir.array<?xi64>>
+! CHECK:           %[[VAL_12:.*]] = fir.box_rank %[[VAL_3]] : (!fir.box<!fir.array<*:f32>>) -> index
+! CHECK:           %[[VAL_13:.*]] = fir.shape %[[VAL_12]] : (index) -> !fir.shape<1>
+! CHECK:           %[[VAL_14:.*]]:2 = hlfir.declare %[[VAL_11]](%[[VAL_13]]) {uniq_name = ".tmp.intrinsic_result"} : (!fir.ref<!fir.array<?xi64>>, !fir.shape<1>) -> (!fir.box<!fir.array<?xi64>>, !fir.ref<!fir.array<?xi64>>)
+
+subroutine test_shape_2(x)
+  real, pointer :: x(..)
+  call takes_integer_array(shape(x))
+end subroutine
+! CHECK-LABEL:   func.func @_QPtest_shape_2(
+! CHECK:           %[[VAL_1:.*]] = fir.alloca !fir.array<15xi32>
+! CHECK:           %[[VAL_4:.*]] = fir.load %[[VAL_3:.*]] : !fir.ref<!fir.box<!fir.ptr<!fir.array<*:f32>>>>
+! CHECK:           %[[VAL_5:.*]] = arith.constant 4 : i32
+! CHECK:           %[[VAL_8:.*]] = fir.convert %[[VAL_1]] : (!fir.ref<!fir.array<15xi32>>) -> !fir.llvm_ptr<i8>
+! CHECK:           %[[VAL_9:.*]] = fir.convert %[[VAL_4]] : (!fir.box<!fir.ptr<!fir.array<*:f32>>>) -> !fir.box<none>
+! CHECK:           %[[VAL_11:.*]] = fir.call @_FortranAShape(%[[VAL_8]], %[[VAL_9]], %[[VAL_5]], %{{.*}}, %{{.*}})
+! CHECK:           %[[VAL_12:.*]] = fir.convert %[[VAL_1]] : (!fir.ref<!fir.array<15xi32>>) -> !fir.ref<!fir.array<?xi32>>
+! CHECK:           %[[VAL_13:.*]] = fir.box_rank %[[VAL_4]] : (!fir.box<!fir.ptr<!fir.array<*:f32>>>) -> index
+! CHECK:           %[[VAL_14:.*]] = fir.shape %[[VAL_13]] : (index) -> !fir.shape<1>
+! CHECK:           %[[VAL_15:.*]]:2 = hlfir.declare %[[VAL_12]](%[[VAL_14]]) {uniq_name = ".tmp.intrinsic_result"} : (!fir.ref<!fir.array<?xi32>>, !fir.shape<1>) -> (!fir.box<!fir.array<?xi32>>, !fir.ref<!fir.array<?xi32>>)
diff --git a/flang/unittests/Runtime/Inquiry.cpp b/flang/unittests/Runtime/Inquiry.cpp
index 665a930ee4ff9..220ebb765f8c7 100644
--- a/flang/unittests/Runtime/Inquiry.cpp
+++ b/flang/unittests/Runtime/Inquiry.cpp
@@ -87,7 +87,8 @@ TEST(Inquiry, Shape) {
   auto int8Result{
       MakeArray<TypeCategory::Integer, 1>(std::vector<int>{array->rank()},
           std::vector<std::int8_t>(array->rank(), 0))};
-  RTNAME(Shape)(int8Result->raw().base_addr, *array, /*KIND=*/1);
+  RTNAME(Shape)
+  (int8Result->raw().base_addr, *array, /*KIND=*/1, __FILE__, __LINE__);
   EXPECT_EQ(*int8Result->ZeroBasedIndexedElement<std::int8_t>(0), 2);
   EXPECT_EQ(*int8Result->ZeroBasedIndexedElement<std::int8_t>(1), 3);
 
@@ -95,7 +96,8 @@ TEST(Inquiry, Shape) {
   auto int32Result{
       MakeArray<TypeCategory::Integer, 4>(std::vector<int>{array->rank()},
           std::vector<std::int32_t>(array->rank(), 0))};
-  RTNAME(Shape)(int32Result->raw().base_addr, *array, /*KIND=*/4);
+  RTNAME(Shape)
+  (int32Result->raw().base_addr, *array, /*KIND=*/4, __FILE__, __LINE__);
   EXPECT_EQ(*int32Result->ZeroBasedIndexedElement<std::int32_t>(0), 2);
   EXPECT_EQ(*int32Result->ZeroBasedIndexedElement<std::int32_t>(1), 3);
 
@@ -103,7 +105,8 @@ TEST(Inquiry, Shape) {
   auto int64Result{
       MakeArray<TypeCategory::Integer, 8>(std::vector<int>{array->rank()},
           std::vector<std::int64_t>(array->rank(), 0))};
-  RTNAME(Shape)(int64Result->raw().base_addr, *array, /*KIND=*/8);
+  RTNAME(Shape)
+  (int64Result->raw().base_addr, *array, /*KIND=*/8, __FILE__, __LINE__);
   EXPECT_EQ(*int64Result->ZeroBasedIndexedElement<std::int64_t>(0), 2);
   EXPECT_EQ(*int64Result->ZeroBasedIndexedElement<std::int64_t>(1), 3);
 }

clementval

LGTM

jeanPerier added 2 commits June 7, 2024 15:04

[flang] add source to SHAPE API

6131bc7

[flang] lower SHAPE with assumed-rank arguments

7b347e7

jeanPerier requested a review from clementval June 7, 2024 22:28

llvmbot added flang:runtime flang Flang issues not falling into any other category flang:fir-hlfir labels Jun 7, 2024

clementval approved these changes Jun 8, 2024

View reviewed changes

jeanPerier merged commit 0257f9c into llvm:main Jun 10, 2024
11 checks passed

jeanPerier deleted the jp-assumed-rank-shape-3 branch June 10, 2024 08:22

HerrCai0907 mentioned this pull request Jun 13, 2024

tidy #95384

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[flang] lower SHAPE with assumed-rank arguments #94812

[flang] lower SHAPE with assumed-rank arguments #94812

jeanPerier commented Jun 7, 2024

llvmbot commented Jun 7, 2024 •

edited

Loading

clementval left a comment

[flang] lower SHAPE with assumed-rank arguments #94812

[flang] lower SHAPE with assumed-rank arguments #94812

Conversation

jeanPerier commented Jun 7, 2024

llvmbot commented Jun 7, 2024 • edited Loading

clementval left a comment

Choose a reason for hiding this comment

llvmbot commented Jun 7, 2024 •

edited

Loading