[flang-commits] [flang] [flang][runtime] add SHAPE runtime interface (PR #94702)

[via flang-commits]

https://github.com/jeanPerier created https://github.com/llvm/llvm-project/pull/94702

Add SHAPE runtime API (will be used for assumed-rank, lowering is generating other cases inline).

I tried to make it in a way were there is no dynamic allocation in the runtime/deallocation expected to be inserted by inline code for arrays that we know are small (lowering will just always stack allocate a rank 15 array to avoid dynamic stack allocation or heap allocation).

Please advise if you would rather stick with the approach that was taken in UBOUND (that I plan to update if you are OK with the new approach).

>From b57b90cc7dfe26fa2adc75b5151dc26ca99860f0 Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Thu, 6 Jun 2024 16:57:16 -0700
Subject: [PATCH] [flang][runtime] add SHAPE runtime interface

---
 flang/include/flang/Runtime/inquiry.h |  9 ++++++--
 flang/runtime/inquiry.cpp             | 19 ++++++++++++++++
 flang/unittests/Runtime/Inquiry.cpp   | 31 +++++++++++++++++++++++++++
 3 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/flang/include/flang/Runtime/inquiry.h b/flang/include/flang/Runtime/inquiry.h
index 3fe670b0fae31..7161d1e41c4bb 100644
--- a/flang/include/flang/Runtime/inquiry.h
+++ b/flang/include/flang/Runtime/inquiry.h
@@ -23,13 +23,18 @@ extern "C" {
 
 std::int64_t RTDECL(LboundDim)(const Descriptor &array, int dim,
     const char *sourceFile = nullptr, int line = 0);
-void RTDECL(Ubound)(Descriptor &result, const Descriptor &array, int kind,
-    const char *sourceFile = nullptr, int line = 0);
+
+void RTDECL(Shape)(void *result, const Descriptor &array, int kind);
+
 std::int64_t RTDECL(Size)(
     const Descriptor &array, const char *sourceFile = nullptr, int line = 0);
+
 std::int64_t RTDECL(SizeDim)(const Descriptor &array, int dim,
     const char *sourceFile = nullptr, int line = 0);
 
+void RTDECL(Ubound)(Descriptor &result, const Descriptor &array, int kind,
+    const char *sourceFile = nullptr, int line = 0);
+
 } // extern "C"
 } // namespace Fortran::runtime
 #endif // FORTRAN_RUNTIME_INQUIRY_H_
diff --git a/flang/runtime/inquiry.cpp b/flang/runtime/inquiry.cpp
index 2b59a1cfab1a9..ea114174de7fd 100644
--- a/flang/runtime/inquiry.cpp
+++ b/flang/runtime/inquiry.cpp
@@ -18,6 +18,15 @@
 
 namespace Fortran::runtime {
 
+template <int KIND> struct RawStoreIntegerAt {
+  RT_API_ATTRS void operator()(
+      void *contiguousIntegerArray, std::size_t at, std::int64_t value) const {
+    reinterpret_cast<Fortran::runtime::CppTypeFor<
+        Fortran::common::TypeCategory::Integer, KIND> *>(
+        contiguousIntegerArray)[at] = value;
+  }
+};
+
 extern "C" {
 std::int64_t RTDEF(LboundDim)(
     const Descriptor &array, int dim, const char *sourceFile, int line) {
@@ -76,5 +85,15 @@ 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__};
+  INTERNAL_CHECK(array.rank() <= common::maxRank);
+  for (SubscriptValue i{0}; i < array.rank(); ++i) {
+    const Dimension &dimension{array.GetDimension(i)};
+    Fortran::runtime::ApplyIntegerKind<RawStoreIntegerAt, void>(
+        kind, terminator, result, i, dimension.Extent());
+  }
+}
+
 } // extern "C"
 } // namespace Fortran::runtime
diff --git a/flang/unittests/Runtime/Inquiry.cpp b/flang/unittests/Runtime/Inquiry.cpp
index 5b97bb239f49b..665a930ee4ff9 100644
--- a/flang/unittests/Runtime/Inquiry.cpp
+++ b/flang/unittests/Runtime/Inquiry.cpp
@@ -76,3 +76,34 @@ TEST(Inquiry, Size) {
   EXPECT_EQ(RTNAME(SizeDim)(*array, 2, __FILE__, __LINE__), std::int64_t{3});
   EXPECT_EQ(RTNAME(Size)(*array, __FILE__, __LINE__), std::int64_t{6});
 }
+
+TEST(Inquiry, Shape) {
+  // ARRAY  1 3 5
+  //        2 4 6
+  auto array{MakeArray<TypeCategory::Integer, 4>(
+      std::vector<int>{2, 3}, std::vector<std::int32_t>{1, 2, 3, 4, 5, 6})};
+
+  // SHAPE(ARRAY, KIND=1)
+  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);
+  EXPECT_EQ(*int8Result->ZeroBasedIndexedElement<std::int8_t>(0), 2);
+  EXPECT_EQ(*int8Result->ZeroBasedIndexedElement<std::int8_t>(1), 3);
+
+  // SHAPE(ARRAY, KIND=4)
+  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);
+  EXPECT_EQ(*int32Result->ZeroBasedIndexedElement<std::int32_t>(0), 2);
+  EXPECT_EQ(*int32Result->ZeroBasedIndexedElement<std::int32_t>(1), 3);
+
+  // SHAPE(ARRAY, KIND=8)
+  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);
+  EXPECT_EQ(*int64Result->ZeroBasedIndexedElement<std::int64_t>(0), 2);
+  EXPECT_EQ(*int64Result->ZeroBasedIndexedElement<std::int64_t>(1), 3);
+}