1 //===- PerThreadBumpPtrAllocator.h ------------------------------*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #ifndef LLVM_SUPPORT_PERTHREADBUMPPTRALLOCATOR_H
10 #define LLVM_SUPPORT_PERTHREADBUMPPTRALLOCATOR_H
11 
12 #include "llvm/Support/Allocator.h"
13 #include "llvm/Support/Parallel.h"
14 
15 namespace llvm {
16 namespace parallel {
17 
18 /// PerThreadAllocator is used in conjunction with ThreadPoolExecutor to allow
19 /// per-thread allocations. It wraps a possibly thread-unsafe allocator,
20 /// e.g. BumpPtrAllocator. PerThreadAllocator must be used with only main thread
21 /// or threads created by ThreadPoolExecutor, as it utilizes getThreadIndex,
22 /// which is set by ThreadPoolExecutor. To work properly, ThreadPoolExecutor
23 /// should be initialized before PerThreadAllocator is created.
24 /// TODO: The same approach might be implemented for ThreadPool.
25 
26 template <typename AllocatorTy>
27 class PerThreadAllocator
28     : public AllocatorBase<PerThreadAllocator<AllocatorTy>> {
29 public:
30   PerThreadAllocator()
31       : NumOfAllocators(parallel::getThreadCount()),
32         Allocators(std::make_unique<AllocatorTy[]>(NumOfAllocators)) {}
33 
34   /// \defgroup Methods which could be called asynchronously:
35   ///
36   /// @{
37 
38   using AllocatorBase<PerThreadAllocator<AllocatorTy>>::Allocate;
39 
40   using AllocatorBase<PerThreadAllocator<AllocatorTy>>::Deallocate;
41 
42   /// Allocate \a Size bytes of \a Alignment aligned memory.
43   void *Allocate(size_t Size, size_t Alignment) {
44     assert(getThreadIndex() < NumOfAllocators);
45     return Allocators[getThreadIndex()].Allocate(Size, Alignment);
46   }
47 
48   /// Deallocate \a Ptr to \a Size bytes of memory allocated by this
49   /// allocator.
50   void Deallocate(const void *Ptr, size_t Size, size_t Alignment) {
51     assert(getThreadIndex() < NumOfAllocators);
52     return Allocators[getThreadIndex()].Deallocate(Ptr, Size, Alignment);
53   }
54 
55   /// Return allocator corresponding to the current thread.
56   AllocatorTy &getThreadLocalAllocator() {
57     assert(getThreadIndex() < NumOfAllocators);
58     return Allocators[getThreadIndex()];
59   }
60 
61   // Return number of used allocators.
62   size_t getNumberOfAllocators() const { return NumOfAllocators; }
63   /// @}
64 
65   /// \defgroup Methods which could not be called asynchronously:
66   ///
67   /// @{
68 
69   /// Reset state of allocators.
70   void Reset() {
71     for (size_t Idx = 0; Idx < getNumberOfAllocators(); Idx++)
72       Allocators[Idx].Reset();
73   }
74 
75   /// Return total memory size used by all allocators.
76   size_t getTotalMemory() const {
77     size_t TotalMemory = 0;
78 
79     for (size_t Idx = 0; Idx < getNumberOfAllocators(); Idx++)
80       TotalMemory += Allocators[Idx].getTotalMemory();
81 
82     return TotalMemory;
83   }
84 
85   /// Return allocated size by all allocators.
86   size_t getBytesAllocated() const {
87     size_t BytesAllocated = 0;
88 
89     for (size_t Idx = 0; Idx < getNumberOfAllocators(); Idx++)
90       BytesAllocated += Allocators[Idx].getBytesAllocated();
91 
92     return BytesAllocated;
93   }
94 
95   /// Set red zone for all allocators.
96   void setRedZoneSize(size_t NewSize) {
97     for (size_t Idx = 0; Idx < getNumberOfAllocators(); Idx++)
98       Allocators[Idx].setRedZoneSize(NewSize);
99   }
100 
101   /// Print statistic for each allocator.
102   void PrintStats() const {
103     for (size_t Idx = 0; Idx < getNumberOfAllocators(); Idx++) {
104       errs() << "\n Allocator " << Idx << "\n";
105       Allocators[Idx].PrintStats();
106     }
107   }
108   /// @}
109 
110 protected:
111   size_t NumOfAllocators;
112   std::unique_ptr<AllocatorTy[]> Allocators;
113 };
114 
115 using PerThreadBumpPtrAllocator = PerThreadAllocator<BumpPtrAllocator>;
116 
117 } // end namespace parallel
118 } // end namespace llvm
119 
120 #endif // LLVM_SUPPORT_PERTHREADBUMPPTRALLOCATOR_H
121