[Support][Allocator] Add C++17 memory_resource adaptor for LLVM allocators

[jfurtek]

Summary

C++17 added support for polymorphic memory resources to the standard library:

https://isocpp.org/files/papers/N3916.pdf

Prior to polymorphic memory resources, the allocator was part of the compile-
time type for C++ containers. (As an example, std::vector instances with
different allocators were distinct C++ types.) With polymorphic memory
resources, the memory_resource abstract base class provides an interface that
different allocators can implement, allowing containers with different
allocators to have the same type. Runtime allocator dispatch occurs via a
virtual function call.

This diff adds a memory_resource adaptor for LLVM allocators (and in
particular, the BumpPtrAllocator). This is an opt-in modification: no existing
uses of the ubiquitous BumpPtrAllocator are affected.

With the new adaptor, LLVM code that uses the BumpPtrAllocator can make use of
the standard C++ library containers with PMR support (i.e. std::pmr::vector<>,
std::pmr::list<>, std::pmr::string, ...), with container allocations using
the LLVM BumpPtrAllocator.

Future MLIR diffs may use this feature. Those are being kept separate from this
diff to facilitate review.

[tschuett]

Have you checked that the minimal supported compilers/versions support pmr? I have the latest Xcode with libc++ version 15 and there is no pmr support.

[jfurtek]

Have you checked that the minimal supported compilers/versions support pmr? I have the latest Xcode with libc++ version 15 and there is no pmr support.

I had not checked, but you are right: there is a red square here:

https://en.cppreference.com/w/cpp/compiler_support/17#C.2B.2B17_library_features

It looks there might be support in Xcode 15 Beta, whereas LLVM 17 must work with 10.0. So I'm only off by 5 or 6 versions. ;-)

https://developer.apple.com/documentation/xcode-release-notes/xcode-15-release-notes
https://developer.apple.com/xcode/cpp/

Closing.

llvm/include/llvm/Support/Allocator.h

 #include "llvm/Support/AllocatorBase.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/MathExtras.h"
 #include <algorithm>
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
 #include <iterator>
+#include <memory_resource>
 #include <optional>
 #include <utility>
 
 namespace llvm {
 
 namespace detail {
 
 // We call out to an external function to actually print the message as the
void DestroyAll() {
 
     Allocator.Reset();
   }
 
   /// Allocate space for an array of objects without constructing them.
   T *Allocate(size_t num = 1) { return Allocator.Allocate<T>(num); }
 };
 
+// This class inherits from the C++17 polymorphic memory resource abstract base
+// class to adapt LLVM allocators for use with C++ standard libraries that have
+// pmr support (e.g. std::pmr::vector<>, std::pmr::list<>, ...). This allows
+// containers with different PMR allocators to interact, at the cost of a
+// virtual function call.
+template <class TAlloc>
+class AllocMemoryResource : public std::pmr::memory_resource {
+public:
+  AllocMemoryResource(TAlloc *alloc) noexcept : Alloc(alloc) {}
+
+protected:
+  void *do_allocate(size_t Bytes, size_t AlignTo) override {
+    return Alloc->Allocate(Bytes, AlignTo);
+  };
+  void do_deallocate(void *p, size_t Bytes, size_t AlignTo) override {
+    Alloc->Deallocate(p, Bytes, AlignTo);
+  };
+  bool
+  do_is_equal(const std::pmr::memory_resource &other) const noexcept override {
+    return (static_cast<const std::pmr::memory_resource *>(this) == &other);
+  }
+
+private:
+  TAlloc *Alloc;
+};
+
+// This class inherits from a template allocator class, as well as the
+// AllocMemoryResource class, allowing LLVM allocators to opt-in to supporting
+// both the LLVM Allocator (Allocate()/Deallocate()) interface as well as the
+// C++17 polymorphic memory resource interface.
+template <class TAlloc>
+class AllocatorWithMemoryResource : public TAlloc,
+                                    public AllocMemoryResource<TAlloc> {
+public:
+  AllocatorWithMemoryResource()
+      : AllocMemoryResource<TAlloc>(static_cast<TAlloc *>(this)) {}
+};
+
 } // end namespace llvm
 
 template <typename AllocatorT, size_t SlabSize, size_t SizeThreshold,
           size_t GrowthDelay>
 void *
 operator new(size_t Size,
              llvm::BumpPtrAllocatorImpl<AllocatorT, SlabSize, SizeThreshold,
                                         GrowthDelay> &Allocator) {

llvm/unittests/Support/AllocatorTest.cpp

 //===- llvm/unittest/Support/AllocatorTest.cpp - BumpPtrAllocator tests ---===//
 //
 // 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 "llvm/Support/Allocator.h"
 #include "gtest/gtest.h"
 #include <cstdlib>
+#include <numeric>
 
 using namespace llvm;
 
 namespace {
 
 TEST(AllocatorTest, Basics) {
   BumpPtrAllocator Alloc;
   int *a = (int*)Alloc.Allocate(sizeof(int), alignof(int));
TEST(AllocatorTest, TestBigAlignment) {
   // Now the big chunk with a big alignment.
   (void)Alloc.Allocate(3000, 2048);
 
   // We test that the last slab size is not the default 4096 byte slab, but
   // rather a custom sized slab that is larger.
   EXPECT_GT(MockSlabAllocator::GetLastSlabSize(), 4096u);
 }
 
+// Test an adaptor to allow using LLVM allocators with C++ standard library
+// containers.
+TEST(AllocatorTest, MemoryResource) {
+  // Declare a BumpPtrAllocator instance with an associated PMR memory resource.
+  // (This is similar to using a C++17 monotonic_buffer_resource, but instead
+  // uses the LLVM BumpPtrAllocator.)
+  using Allocator = AllocatorWithMemoryResource<BumpPtrAllocator>;
+  Allocator Alloc;
+  std::pmr::vector<size_t> vec(&Alloc);
+  // Populate the vector with increasing values
+  vec.resize(64);
+  std::iota(vec.begin(), vec.end(), 0);
+  for (size_t i = 0; i < vec.size(); ++i) {
+    // Verify the vector values
+    EXPECT_EQ(vec.at(i), i);
+    // Verify that the vector elements are part of the provide allocator
+    EXPECT_TRUE(Alloc.identifyObject(&(vec[i])).has_value());
+  }
+}
+
 }  // anonymous namespace