[flang][runtime] Add dynamically allocated temporary storage

These functions will be used as part of the HLFIR lowering for
forall/where. The contents of the API were requested by @jeanPerier.

The API is designed around that use case, assuming that the caller knows
through some side channel what size to allocate for boxes returned from
the pop() function.

Differential Revision: https://reviews.llvm.org/D150050

flang/include/flang/Runtime/temporary-stack.h

@@ -0,0 +1,67 @@
+//===-- include/flang/Runtime/temporary-stack.h -----------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+// Runtime functions for storing a dynamically resizable number of temporaries.
+// For use in HLFIR lowering.
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_RUNTIME_TEMPORARY_STACK_H_
+#define FORTRAN_RUNTIME_TEMPORARY_STACK_H_
+
+#include "flang/Runtime/entry-names.h"
+#include <stdint.h>
+
+namespace Fortran::runtime {
+class Descriptor;
+extern "C" {
+
+// Stores both the descriptor and a copy of the value in a dynamically resizable
+// data structure identified by opaquePtr. All value stacks must be destroyed
+// at the end of their lifetime and not used afterwards.
+// Popped descriptors point to the copy of the value, not the original address
+// of the value. This copy is dynamically allocated, it is up to the caller to
+// free the value pointed to by the box. The copy operation is a simple memcpy.
+// The sourceFile and line number used when creating the stack are shared for
+// all operations.
+// Opaque pointers returned from these are incompatible with those returned by
+// the flavours for storing descriptors.
+[[nodiscard]] void *RTNAME(CreateValueStack)(
+    const char *sourceFile = nullptr, int line = 0);
+void RTNAME(PushValue)(void *opaquePtr, const Descriptor &value);
+// Note: retValue should be large enough to hold the right number of dimensions,
+// and the optional descriptor addendum
+void RTNAME(PopValue)(void *opaquePtr, Descriptor &retValue);
+// Return the i'th element into retValue (which must be the right size). An
+// exact copy of this descriptor remains in this storage so this one should not
+// be deallocated
+void RTNAME(ValueAt)(void *opaquePtr, uint64_t i, Descriptor &retValue);
+void RTNAME(DestroyValueStack)(void *opaquePtr);
+
+// Stores descriptors value in a dynamically resizable data structure identified
+// by opaquePtr. All descriptor stacks must be destroyed at the end of their
+// lifetime and not used afterwards.
+// Popped descriptors are identical to those which were pushed.
+// The sourceFile and line number used when creating the stack are shared for
+// all operations.
+// Opaque pointers returned from these are incompatible with those returned by
+// the flavours for storing both descriptors and values.
+[[nodiscard]] void *RTNAME(CreateDescriptorStack)(
+    const char *sourceFile = nullptr, int line = 0);
+void RTNAME(PushDescriptor)(void *opaquePtr, const Descriptor &value);
+// Note: retValue should be large enough to hold the right number of dimensions,
+// and the optional descriptor addendum
+void RTNAME(PopDescriptor)(void *opaquePtr, Descriptor &retValue);
+// Return the i'th element into retValue (which must be the right size). An
+// exact copy of this descriptor remains in this storage so this one should not
+// be deallocated
+void RTNAME(DescriptorAt)(void *opaquePtr, uint64_t i, Descriptor &retValue);
+void RTNAME(DestroyDescriptorStack)(void *opaquePtr);
+
+} // extern "C"
+} // namespace Fortran::runtime
+
+#endif // FORTRAN_RUNTIME_TEMPORARY_STACK_H_

flang/runtime/CMakeLists.txt

@@ -144,6 +144,7 @@ add_flang_library(FortranRuntime
   stop.cpp
   sum.cpp
   support.cpp
+  temporary-stack.cpp
   terminator.cpp
   time-intrinsic.cpp
   tools.cpp

flang/runtime/temporary-stack.cpp

@@ -0,0 +1,222 @@
+//===-- runtime/temporary-stack.cpp ---------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+// Implements std::vector like storage for a dynamically resizable number of
+// temporaries. For use in HLFIR lowering.
+
+#include "flang/Runtime/temporary-stack.h"
+#include "terminator.h"
+#include "flang/ISO_Fortran_binding.h"
+#include "flang/Runtime/assign.h"
+#include "flang/Runtime/descriptor.h"
+#include "flang/Runtime/memory.h"
+
+namespace {
+
+using namespace Fortran::runtime;
+
+// the number of elements to allocate when first creating the vector
+constexpr size_t INITIAL_ALLOC = 8;
+
+/// To store C style data. Does not run constructors/destructors.
+/// Not using std::vector to avoid linking the runtime library to stdc++
+template <bool COPY_VALUES> class DescriptorStorage final {
+  using size_type = uint64_t; // see checkedMultiply()
+
+  size_type capacity_{0};
+  size_type size_{0};
+  Descriptor **data_{nullptr};
+  Terminator terminator_;
+
+  // return true on overflow
+  static bool checkedMultiply(size_type x, size_type y, size_type &res);
+
+  void resize(size_type newCapacity);
+
+  Descriptor *cloneDescriptor(const Descriptor &source);
+
+public:
+  DescriptorStorage(const char *sourceFile, int line);
+  ~DescriptorStorage();
+
+  // `new` but using the runtime allocation API
+  static inline DescriptorStorage *allocate(const char *sourceFile, int line) {
+    Terminator term{sourceFile, line};
+    void *ptr = AllocateMemoryOrCrash(term, sizeof(DescriptorStorage));
+    return new (ptr) DescriptorStorage{sourceFile, line};
+  }
+
+  // `delete` but using the runtime allocation API
+  static inline void destroy(DescriptorStorage *instance) {
+    instance->~DescriptorStorage();
+    FreeMemory(instance);
+  }
+
+  // clones a descriptor into this storage
+  void push(const Descriptor &source);
+
+  // out must be big enough to hold a descriptor of the right rank and addendum
+  void pop(Descriptor &out);
+
+  // out must be big enough to hold a descriptor of the right rank and addendum
+  void at(size_type i, Descriptor &out);
+};
+
+using ValueStack = DescriptorStorage</*COPY_VALUES=*/true>;
+using DescriptorStack = DescriptorStorage</*COPY_VALUES=*/false>;
+} // namespace
+
+template <bool COPY_VALUES>
+bool DescriptorStorage<COPY_VALUES>::checkedMultiply(
+    size_type x, size_type y, size_type &res) {
+  // TODO: c++20 [[unlikely]]
+  if (x > UINT64_MAX / y) {
+    return true;
+  }
+  res = x * y;
+  return false;
+}
+
+template <bool COPY_VALUES>
+void DescriptorStorage<COPY_VALUES>::resize(size_type newCapacity) {
+  if (newCapacity <= capacity_) {
+    return;
+  }
+  size_type bytes;
+  if (checkedMultiply(newCapacity, sizeof(Descriptor *), bytes)) {
+    terminator_.Crash("temporary-stack: out of memory");
+  }
+  Descriptor **newData =
+      static_cast<Descriptor **>(AllocateMemoryOrCrash(terminator_, bytes));
+  memcpy(newData, data_, capacity_ * sizeof(Descriptor *));
+  FreeMemory(data_);
+  data_ = newData;
+  capacity_ = newCapacity;
+}
+
+template <bool COPY_VALUES>
+Descriptor *DescriptorStorage<COPY_VALUES>::cloneDescriptor(
+    const Descriptor &source) {
+  const std::size_t bytes = source.SizeInBytes();
+  void *memory = AllocateMemoryOrCrash(terminator_, bytes);
+  Descriptor *desc = new (memory) Descriptor{source};
+  return desc;
+}
+
+template <bool COPY_VALUES>
+DescriptorStorage<COPY_VALUES>::DescriptorStorage(
+    const char *sourceFile, int line)
+    : terminator_{sourceFile, line} {
+  resize(INITIAL_ALLOC);
+}
+
+template <bool COPY_VALUES>
+DescriptorStorage<COPY_VALUES>::~DescriptorStorage() {
+  for (size_type i = 0; i < size_; ++i) {
+    Descriptor *element = data_[i];
+    if constexpr (COPY_VALUES) {
+      element->Destroy(false, true);
+    }
+    FreeMemory(element);
+  }
+  FreeMemory(data_);
+}
+
+template <bool COPY_VALUES>
+void DescriptorStorage<COPY_VALUES>::push(const Descriptor &source) {
+  if (size_ == capacity_) {
+    size_type newSize;
+    if (checkedMultiply(capacity_, 2, newSize)) {
+      terminator_.Crash("temporary-stack: out of address space");
+    }
+    resize(newSize);
+  }
+  data_[size_] = cloneDescriptor(source);
+  Descriptor &box = *data_[size_];
+  size_ += 1;
+
+  if constexpr (COPY_VALUES) {
+    // copy the data pointed to by the box
+    box.set_base_addr(nullptr);
+    box.Allocate();
+    RTNAME(AssignTemporary)
+    (box, source, terminator_.sourceFileName(), terminator_.sourceLine());
+  }
+}
+
+template <bool COPY_VALUES>
+void DescriptorStorage<COPY_VALUES>::pop(Descriptor &out) {
+  if (size_ == 0) {
+    terminator_.Crash("temporary-stack: pop empty storage");
+  }
+  size_ -= 1;
+  Descriptor *ptr = data_[size_];
+  out = *ptr; // Descriptor::operator= handles the different sizes
+  FreeMemory(ptr);
+}
+
+template <bool COPY_VALUES>
+void DescriptorStorage<COPY_VALUES>::at(size_type i, Descriptor &out) {
+  if (i >= size_) {
+    terminator_.Crash("temporary-stack: out of bounds access");
+  }
+  Descriptor *ptr = data_[i];
+  out = *ptr; // Descriptor::operator= handles the different sizes
+}
+
+inline static ValueStack *getValueStorage(void *opaquePtr) {
+  return static_cast<ValueStack *>(opaquePtr);
+}
+inline static DescriptorStack *getDescriptorStorage(void *opaquePtr) {
+  return static_cast<DescriptorStack *>(opaquePtr);
+}
+
+namespace Fortran::runtime {
+extern "C" {
+void *RTNAME(CreateValueStack)(const char *sourceFile, int line) {
+  return ValueStack::allocate(sourceFile, line);
+}
+
+void RTNAME(PushValue)(void *opaquePtr, const Descriptor &value) {
+  getValueStorage(opaquePtr)->push(value);
+}
+
+void RTNAME(PopValue)(void *opaquePtr, Descriptor &value) {
+  getValueStorage(opaquePtr)->pop(value);
+}
+
+void RTNAME(ValueAt)(void *opaquePtr, uint64_t i, Descriptor &value) {
+  getValueStorage(opaquePtr)->at(i, value);
+}
+
+void RTNAME(DestroyValueStack)(void *opaquePtr) {
+  ValueStack::destroy(getValueStorage(opaquePtr));
+}
+
+void *RTNAME(CreateDescriptorStack)(const char *sourceFile, int line) {
+  return DescriptorStack::allocate(sourceFile, line);
+}
+
+void RTNAME(PushDescriptor)(void *opaquePtr, const Descriptor &value) {
+  getDescriptorStorage(opaquePtr)->push(value);
+}
+
+void RTNAME(PopDescriptor)(void *opaquePtr, Descriptor &value) {
+  getDescriptorStorage(opaquePtr)->pop(value);
+}
+
+void RTNAME(DescriptorAt)(void *opaquePtr, uint64_t i, Descriptor &value) {
+  getValueStorage(opaquePtr)->at(i, value);
+}
+
+void RTNAME(DestroyDescriptorStack)(void *opaquePtr) {
+  DescriptorStack::destroy(getDescriptorStorage(opaquePtr));
+}
+
+} // extern "C"
+} // namespace Fortran::runtime

flang/unittests/Runtime/CMakeLists.txt

@@ -25,6 +25,7 @@ add_flang_unittest(FlangRuntimeTests
   RuntimeCrashTest.cpp
   Stop.cpp
   Time.cpp
+  TemporaryStack.cpp
   Transformational.cpp
 )