[flang] Add ability to have special allocator for descriptor data (#100690)

This patch enhances the descriptor with the ability to have specialized
allocator. The allocators are registered in a dedicated registry and the
index of the desired allocator is stored in the descriptor. The default
allocator, std::malloc, is registered at index 0.

In order to have this allocator index in the descriptor, the f18Addendum
field is repurposed to be able to hold the presence flag for the
addendum (lsb) and the allocator index.

Since this is a change in the semantic and name of the 7th field of the
descriptor, the CFI_VERSION is bumped to the date of the initial change.

This patch only adds the ability to have this features as part of the
descriptor but does not add specific allocator yet. CUDA fortran will be
the first user of this feature to allocate descriptor data in the
different type of device memory base on the CUDA attribute.

---------

Co-authored-by: Slava Zakharin <[email protected]>

Author: clementval

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));