[flang] Add ability to have special allocator for descriptor data

flang/include/flang/ISO_Fortran_binding.h

@@ -30,7 +30,7 @@
 #endif
 
 /* 18.5.4 */
-#define CFI_VERSION 20180515
+#define CFI_VERSION 20240719
 
 #define CFI_MAX_RANK 15
 typedef unsigned char CFI_rank_t;
@@ -146,7 +146,9 @@ extern "C++" template <typename T> struct FlexibleArray : T {
   CFI_rank_t rank; /* [0 .. CFI_MAX_RANK] */ \
   CFI_type_t type; \
   CFI_attribute_t attribute; \
-  unsigned char f18Addendum;
+  /* This encodes both the presence of the f18Addendum and the index of the \
+   * allocator used to managed memory of the data hold by the descriptor. */ \
+  unsigned char extra;
 
 typedef struct CFI_cdesc_t {
   _CFI_CDESC_T_HEADER_MEMBERS

flang/include/flang/Optimizer/CodeGen/TBAABuilder.h

@@ -55,7 +55,7 @@ namespace fir {
 //                  <@type_desc_3, 20>, // rank
 //                  <@type_desc_3, 21>, // type
 //                  <@type_desc_3, 22>, // attribute
-//                  <@type_desc_3, 23>} // f18Addendum
+//                  <@type_desc_3, 23>} // extra
 //   }
 //   llvm.tbaa_type_desc @type_desc_5 {
 //       id = "CFI_cdesc_t_dim1",
@@ -65,7 +65,7 @@ namespace fir {
 //                  <@type_desc_3, 20>, // rank
 //                  <@type_desc_3, 21>, // type
 //                  <@type_desc_3, 22>, // attribute
-//                  <@type_desc_3, 23>, // f18Addendum
+//                  <@type_desc_3, 23>, // extra
 //                  <@type_desc_3, 24>, // dim[0].lower_bound
 //                  <@type_desc_3, 32>, // dim[0].extent
 //                  <@type_desc_3, 40>} // dim[0].sm
@@ -78,7 +78,7 @@ namespace fir {
 //                  <@type_desc_3, 20>, // rank
 //                  <@type_desc_3, 21>, // type
 //                  <@type_desc_3, 22>, // attribute
-//                  <@type_desc_3, 23>, // f18Addendum
+//                  <@type_desc_3, 23>, // extra
 //                  <@type_desc_3, 24>, // dim[0].lower_bound
 //                  <@type_desc_3, 32>, // dim[0].extent
 //                  <@type_desc_3, 40>, // dim[0].sm

flang/include/flang/Optimizer/CodeGen/TypeConverter.h

@@ -29,7 +29,7 @@ static constexpr unsigned kVersionPosInBox = 2;
 static constexpr unsigned kRankPosInBox = 3;
 static constexpr unsigned kTypePosInBox = 4;
 static constexpr unsigned kAttributePosInBox = 5;
-static constexpr unsigned kF18AddendumPosInBox = 6;
+static constexpr unsigned kExtraPosInBox = 6;
 static constexpr unsigned kDimsPosInBox = 7;
 static constexpr unsigned kOptTypePtrPosInBox = 8;
 static constexpr unsigned kOptRowTypePosInBox = 9;

flang/include/flang/Runtime/descriptor.h

@@ -92,8 +92,8 @@ class Dimension {
 // The storage for this object follows the last used dim[] entry in a
 // Descriptor (CFI_cdesc_t) generic descriptor.  Space matters here, since
 // descriptors serve as POINTER and ALLOCATABLE components of derived type
-// instances.  The presence of this structure is implied by the flag
-// CFI_cdesc_t.f18Addendum, and the number of elements in the len_[]
+// instances.  The presence of this structure is encoded in the
+// CFI_cdesc_t.extra field, and the number of elements in the len_[]
 // array is determined by derivedType_->LenParameters().
 class DescriptorAddendum {
 public:
@@ -339,14 +339,14 @@ class Descriptor {
       const SubscriptValue *, const int *permutation = nullptr) const;
 
   RT_API_ATTRS DescriptorAddendum *Addendum() {
-    if (raw_.f18Addendum != 0) {
+    if (HasAddendum()) {
       return reinterpret_cast<DescriptorAddendum *>(&GetDimension(rank()));
     } else {
       return nullptr;
     }
   }
   RT_API_ATTRS const DescriptorAddendum *Addendum() const {
-    if (raw_.f18Addendum != 0) {
+    if (HasAddendum()) {
       return reinterpret_cast<const DescriptorAddendum *>(
           &GetDimension(rank()));
     } else {
@@ -420,6 +420,27 @@ class Descriptor {
 
   void Dump(FILE * = stdout) const;
 
+// Value of the addendum presence flag.
+#define _CFI_ADDENDUM_FLAG 1
+// Number of bits needed to be shifted when manipulating the allocator index.
+#define _CFI_ALLOCATOR_IDX_SHIFT 1
+// Allocator index mask.
+#define _CFI_ALLOCATOR_IDX_MASK 0b00001110
+
+  RT_API_ATTRS inline bool HasAddendum() const {
+    return raw_.extra & _CFI_ADDENDUM_FLAG;
+  }
+  RT_API_ATTRS inline void SetHasAddendum() {
+    raw_.extra |= _CFI_ADDENDUM_FLAG;
+  }
+  RT_API_ATTRS inline int GetAllocIdx() const {
+    return (raw_.extra & _CFI_ALLOCATOR_IDX_MASK) >> _CFI_ALLOCATOR_IDX_SHIFT;
+  }
+  RT_API_ATTRS inline void SetAllocIdx(int pos) {
+    raw_.extra &= ~_CFI_ALLOCATOR_IDX_MASK; // Clear the allocator index bits.
+    raw_.extra |= (pos << _CFI_ALLOCATOR_IDX_SHIFT);
+  }
+
 private:
   ISO::CFI_cdesc_t raw_;
 };

flang/lib/Optimizer/CodeGen/CodeGen.cpp

@@ -1241,9 +1241,10 @@ struct EmboxCommonConversion : public fir::FIROpConversion<OP> {
     descriptor =
         insertField(rewriter, loc, descriptor, {kAttributePosInBox},
                     this->genI32Constant(loc, rewriter, getCFIAttr(boxTy)));
+
     const bool hasAddendum = fir::boxHasAddendum(boxTy);
     descriptor =
-        insertField(rewriter, loc, descriptor, {kF18AddendumPosInBox},
+        insertField(rewriter, loc, descriptor, {kExtraPosInBox},
                     this->genI32Constant(loc, rewriter, hasAddendum ? 1 : 0));
 
     if (hasAddendum) {

flang/lib/Optimizer/CodeGen/TypeConverter.cpp

@@ -177,7 +177,7 @@ bool LLVMTypeConverter::requiresExtendedDesc(mlir::Type boxElementType) const {
 // the addendum defined in descriptor.h.
 mlir::Type LLVMTypeConverter::convertBoxTypeAsStruct(BaseBoxType box,
                                                      int rank) const {
-  // (base_addr*, elem_len, version, rank, type, attribute, f18Addendum, [dim]
+  // (base_addr*, elem_len, version, rank, type, attribute, extra, [dim]
   llvm::SmallVector<mlir::Type> dataDescFields;
   mlir::Type ele = box.getEleTy();
   // remove fir.heap/fir.ref/fir.ptr
@@ -206,9 +206,9 @@ mlir::Type LLVMTypeConverter::convertBoxTypeAsStruct(BaseBoxType box,
   // attribute
   dataDescFields.push_back(
       getDescFieldTypeModel<kAttributePosInBox>()(&getContext()));
-  // f18Addendum
+  // extra
   dataDescFields.push_back(
-      getDescFieldTypeModel<kF18AddendumPosInBox>()(&getContext()));
+      getDescFieldTypeModel<kExtraPosInBox>()(&getContext()));
   // [dims]
   if (rank == unknownRank()) {
     if (auto seqTy = mlir::dyn_cast<SequenceType>(ele))

flang/runtime/CMakeLists.txt

@@ -107,6 +107,7 @@ add_subdirectory(Float128Math)
 
 set(sources
   ISO_Fortran_binding.cpp
+  allocator-registry.cpp
   allocatable.cpp
   array-constructor.cpp
   assign.cpp
@@ -178,6 +179,7 @@ include(AddFlangOffloadRuntime)
 set(supported_files
   ISO_Fortran_binding.cpp
   allocatable.cpp
+  allocator-registry.cpp
   array-constructor.cpp
   assign.cpp
   buffer.cpp

flang/runtime/ISO_Fortran_util.h

@@ -86,7 +86,7 @@ static inline RT_API_ATTRS void EstablishDescriptor(CFI_cdesc_t *descriptor,
   descriptor->rank = rank;
   descriptor->type = type;
   descriptor->attribute = attribute;
-  descriptor->f18Addendum = 0;
+  descriptor->extra = 0;
   std::size_t byteSize{elem_len};
   constexpr std::size_t lower_bound{0};
   if (base_addr) {

flang/runtime/allocator-registry.cpp

@@ -0,0 +1,42 @@
+//===-- runtime/allocator-registry.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
+//
+//===----------------------------------------------------------------------===//
+
+#include "allocator-registry.h"
+#include "terminator.h"
+
+namespace Fortran::runtime {
+
+#ifndef FLANG_RUNTIME_NO_GLOBAL_VAR_DEFS
+RT_OFFLOAD_VAR_GROUP_BEGIN
+RT_VAR_ATTRS AllocatorRegistry allocatorRegistry;
+RT_OFFLOAD_VAR_GROUP_END
+#endif // FLANG_RUNTIME_NO_GLOBAL_VAR_DEFS
+
+RT_OFFLOAD_API_GROUP_BEGIN
+RT_API_ATTRS void AllocatorRegistry::Register(int pos, Allocator_t allocator) {
+  // pos 0 is reserved for the default allocator and is registered in the
+  // struct ctor.
+  INTERNAL_CHECK(pos > 0 && pos < MAX_ALLOCATOR);
+  allocators[pos] = allocator;
+}
+
+RT_API_ATTRS AllocFct AllocatorRegistry::GetAllocator(int pos) {
+  INTERNAL_CHECK(pos >= 0 && pos < MAX_ALLOCATOR);
+  AllocFct f{allocators[pos].alloc};
+  INTERNAL_CHECK(f != nullptr);
+  return f;
+}
+
+RT_API_ATTRS FreeFct AllocatorRegistry::GetDeallocator(int pos) {
+  INTERNAL_CHECK(pos >= 0 && pos < MAX_ALLOCATOR);
+  FreeFct f{allocators[pos].free};
+  INTERNAL_CHECK(f != nullptr);
+  return f;
+}
+RT_OFFLOAD_API_GROUP_END
+} // namespace Fortran::runtime

flang/runtime/allocator-registry.h

@@ -0,0 +1,55 @@
+//===-- runtime/allocator-registry.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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_RUNTIME_ALLOCATOR_H_
+#define FORTRAN_RUNTIME_ALLOCATOR_H_
+
+#include "flang/Common/api-attrs.h"
+#include <cstdlib>
+#include <vector>
+
+#define MAX_ALLOCATOR 5
+
+namespace Fortran::runtime {
+
+using AllocFct = void *(*)(std::size_t);
+using FreeFct = void (*)(void *);
+
+typedef struct Allocator_t {
+  AllocFct alloc{nullptr};
+  FreeFct free{nullptr};
+} Allocator_t;
+
+#ifdef RT_DEVICE_COMPILATION
+static RT_API_ATTRS void *MallocWrapper(std::size_t size) {
+  return std::malloc(size);
+}
+static RT_API_ATTRS void FreeWrapper(void *p) { return std::free(p); }
+#endif
+
+struct AllocatorRegistry {
+#ifdef RT_DEVICE_COMPILATION
+  RT_API_ATTRS constexpr AllocatorRegistry()
+      : allocators{{&MallocWrapper, &FreeWrapper}} {}
+#else
+  constexpr AllocatorRegistry() { allocators[0] = {&std::malloc, &std::free}; };
+#endif
+  RT_API_ATTRS void Register(int, Allocator_t);
+  RT_API_ATTRS AllocFct GetAllocator(int pos);
+  RT_API_ATTRS FreeFct GetDeallocator(int pos);
+
+  Allocator_t allocators[MAX_ALLOCATOR];
+};
+
+RT_OFFLOAD_VAR_GROUP_BEGIN
+extern RT_VAR_ATTRS AllocatorRegistry allocatorRegistry;
+RT_OFFLOAD_VAR_GROUP_END
+
+} // namespace Fortran::runtime
+
+#endif // FORTRAN_RUNTIME_ALLOCATOR_H_

flang/runtime/descriptor.cpp

@@ -8,6 +8,7 @@
 
 #include "flang/Runtime/descriptor.h"
 #include "ISO_Fortran_util.h"
+#include "allocator-registry.h"
 #include "derived.h"
 #include "memory.h"
 #include "stat.h"
@@ -50,7 +51,9 @@ RT_API_ATTRS void Descriptor::Establish(TypeCode t, std::size_t elementBytes,
       GetDimension(j).SetByteStride(0);
     }
   }
-  raw_.f18Addendum = addendum;
+  if (addendum) {
+    SetHasAddendum();
+  }
   DescriptorAddendum *a{Addendum()};
   RUNTIME_CHECK(terminator, addendum == (a != nullptr));
   if (a) {
@@ -162,7 +165,9 @@ RT_API_ATTRS int Descriptor::Allocate() {
   // Zero size allocation is possible in Fortran and the resulting
   // descriptor must be allocated/associated. Since std::malloc(0)
   // result is implementation defined, always allocate at least one byte.
-  void *p{byteSize ? std::malloc(byteSize) : std::malloc(1)};
+
+  AllocFct alloc{allocatorRegistry.GetAllocator(GetAllocIdx())};
+  void *p{alloc(byteSize ? byteSize : 1)};
   if (!p) {
     return CFI_ERROR_MEM_ALLOCATION;
   }
@@ -204,7 +209,8 @@ RT_API_ATTRS int Descriptor::Deallocate() {
   if (!descriptor.base_addr) {
     return CFI_ERROR_BASE_ADDR_NULL;
   } else {
-    std::free(descriptor.base_addr);
+    FreeFct free{allocatorRegistry.GetDeallocator(GetAllocIdx())};
+    free(descriptor.base_addr);
     descriptor.base_addr = nullptr;
     return CFI_SUCCESS;
   }
@@ -290,7 +296,9 @@ void Descriptor::Dump(FILE *f) const {
   std::fprintf(f, "  rank      %d\n", static_cast<int>(raw_.rank));
   std::fprintf(f, "  type      %d\n", static_cast<int>(raw_.type));
   std::fprintf(f, "  attribute %d\n", static_cast<int>(raw_.attribute));
-  std::fprintf(f, "  addendum  %d\n", static_cast<int>(raw_.f18Addendum));
+  std::fprintf(f, "  extra     %d\n", static_cast<int>(raw_.extra));
+  std::fprintf(f, "    addendum  %d\n", static_cast<int>(HasAddendum()));
+  std::fprintf(f, "    alloc_idx %d\n", static_cast<int>(GetAllocIdx()));
   for (int j{0}; j < raw_.rank; ++j) {
     std::fprintf(f, "  dim[%d] lower_bound %jd\n", j,
         static_cast<std::intmax_t>(raw_.dim[j].lower_bound));