OmniSciDB  f17484ade4
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
ArenaAllocator.h
Go to the documentation of this file.
1 /*
2  * Copyright 2022 HEAVY.AI, Inc.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #pragma once
18 
19 #include "DataMgr/DataMgr.h"
20 #include "Shared/checked_alloc.h"
21 
22 template <class T>
23 class SysAllocator {
24  public:
25  using Self = SysAllocator;
26  using value_type = T;
27 
28  constexpr SysAllocator() = default;
29 
30  constexpr SysAllocator(SysAllocator const&) = default;
31 
32  template <class U>
33  constexpr SysAllocator(const SysAllocator<U>&) noexcept {}
34 
35  [[nodiscard]] T* allocate(size_t count) {
36  return reinterpret_cast<T*>(checked_malloc(count));
37  }
38 
39  void deallocate(T* p, size_t /* count */) { free(p); }
40 
41  friend bool operator==(Self const&, Self const&) noexcept { return true; }
42  friend bool operator!=(Self const&, Self const&) noexcept { return false; }
43 };
44 
45 class Arena {
46  public:
47  enum class MemoryType { DRAM, PMEM };
48  virtual ~Arena() {}
49  virtual void* allocate(size_t size) = 0;
50  virtual void* allocateAndZero(const size_t size) = 0;
51  virtual size_t bytesUsed() const = 0;
52  virtual MemoryType getMemoryType() const = 0;
53 };
54 
55 #ifdef HAVE_FOLLY
56 
57 #include <folly/Memory.h>
58 #include <folly/memory/Arena.h>
59 
60 using AllocatorType = char;
61 
62 constexpr size_t kArenaBlockOverhead =
63  folly::Arena<::SysAllocator<AllocatorType>>::kBlockOverhead;
64 
69 class DramArena : public Arena, public folly::Arena<::SysAllocator<AllocatorType>> {
70  public:
71  explicit DramArena(size_t min_block_size = static_cast<size_t>(1ULL << 32) +
72  kBlockOverhead,
73  size_t size_limit = kNoSizeLimit,
74  size_t max_align = kDefaultMaxAlign)
75  : folly::Arena<::SysAllocator<AllocatorType>>({},
76  min_block_size,
77  size_limit,
78  max_align) {}
79  ~DramArena() override {}
80 
81  void* allocate(size_t size) override {
82  return folly::Arena<::SysAllocator<AllocatorType>>::allocate(size);
83  }
84 
85  void* allocateAndZero(const size_t size) override {
86  auto ret = allocate(size);
87  std::memset(ret, 0, size);
88  return ret;
89  }
90 
91  size_t bytesUsed() const override {
92  return folly::Arena<::SysAllocator<AllocatorType>>::bytesUsed();
93  }
94 
95  MemoryType getMemoryType() const override { return MemoryType::DRAM; }
96 };
97 
98 template <>
99 struct folly::ArenaAllocatorTraits<::SysAllocator<AllocatorType>> {
100  static size_t goodSize(const ::SysAllocator<AllocatorType>& /* alloc */, size_t size) {
101  return folly::goodMallocSize(size);
102  }
103 };
104 
105 #else
106 
107 constexpr size_t kArenaBlockOverhead = 0;
108 
114 class DramArena : public Arena {
115  public:
116  explicit DramArena(size_t min_block_size = 1ULL << 32, size_t size_limit = 0)
117  : size_limit_(size_limit), size_(0) {}
118 
119  ~DramArena() override {
120  for (auto [ptr, size] : allocations_) {
121  allocator_.deallocate(ptr, 0);
122  size_ -= size;
123  }
124  }
125 
126  void* allocate(size_t num_bytes) override {
127  if (size_limit_ != 0 && size_ + num_bytes > size_limit_) {
128  throw OutOfHostMemory(num_bytes);
129  }
130  auto ret = allocator_.allocate(num_bytes);
131  size_ += num_bytes;
132  allocations_.push_back({ret, num_bytes});
133  return ret;
134  }
135 
136  void* allocateAndZero(const size_t size) override {
137  auto ret = allocate(size);
138  std::memset(ret, 0, size);
139  return ret;
140  }
141 
142  size_t bytesUsed() const override { return size_; }
143 
144  MemoryType getMemoryType() const override { return MemoryType::DRAM; }
145 
146  private:
147  size_t size_limit_;
148  size_t size_;
150  std::vector<std::pair<void*, size_t>> allocations_;
151 };
152 
153 #endif
constexpr size_t kArenaBlockOverhead
void * allocate(size_t num_bytes) override
virtual MemoryType getMemoryType() const =0
friend bool operator!=(Self const &, Self const &) noexcept
constexpr SysAllocator()=default
T * allocate(size_t count)
virtual size_t bytesUsed() const =0
void deallocate(T *p, size_t)
size_t size_
virtual void * allocate(size_t size)=0
size_t size_limit_
size_t bytesUsed() const override
void * checked_malloc(const size_t size)
Definition: checked_alloc.h:45
SysAllocator< void > allocator_
constexpr SysAllocator(const SysAllocator< U > &) noexcept
DramArena(size_t min_block_size=1ULL<< 32, size_t size_limit=0)
virtual ~Arena()
void * allocateAndZero(const size_t size) override
~DramArena() override
friend bool operator==(Self const &, Self const &) noexcept
MemoryType getMemoryType() const override
std::vector< std::pair< void *, size_t > > allocations_
virtual void * allocateAndZero(const size_t size)=0