• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
16 // This approach to arenas overcomes many of the limitations described
17 // in the "Specialized allocators" section of
18 //     http://www.pdos.lcs.mit.edu/~dm/c++-new.html
19 //
20 // A somewhat similar approach to Gladiator, but for heap-detection, was
21 // suggested by Ron van der Wal and Scott Meyers at
22 //     http://www.aristeia.com/BookErrata/M27Comments_frames.html
23 
24 #include "tensorflow/core/lib/core/arena.h"
25 
26 #include <assert.h>
27 
28 #include <algorithm>
29 #include <vector>
30 
31 #include "tensorflow/core/lib/math/math_util.h"
32 #include "tensorflow/core/platform/logging.h"
33 #include "tensorflow/core/platform/macros.h"
34 #include "tensorflow/core/platform/mem.h"
35 
36 namespace tensorflow {
37 namespace core {
38 
39 // ----------------------------------------------------------------------
40 // Arena::Arena()
41 // Arena::~Arena()
42 //    Destroying the arena automatically calls Reset()
43 // ----------------------------------------------------------------------
44 
Arena(const size_t block_size)45 Arena::Arena(const size_t block_size)
46     : remaining_(0),
47       block_size_(block_size),
48       freestart_(nullptr),  // set for real in Reset()
49       blocks_alloced_(1),
50       overflow_blocks_(nullptr) {
51   assert(block_size > kDefaultAlignment);
52 
53   first_blocks_[0].mem =
54       reinterpret_cast<char*>(port::AlignedMalloc(block_size_, sizeof(void*)));
55 
56   first_blocks_[0].size = block_size_;
57 
58   Reset();
59 }
60 
~Arena()61 Arena::~Arena() {
62   FreeBlocks();
63   assert(overflow_blocks_ == nullptr);  // FreeBlocks() should do that
64   // The first X blocks stay allocated always by default.  Delete them now.
65   for (size_t i = 0; i < blocks_alloced_; ++i) {
66     port::AlignedFree(first_blocks_[i].mem);
67   }
68 }
69 
70 // Returns true iff it advances freestart_ to the first position
71 // satisfying alignment without exhausting the current block.
SatisfyAlignment(size_t alignment)72 bool Arena::SatisfyAlignment(size_t alignment) {
73   const size_t overage = reinterpret_cast<size_t>(freestart_) & (alignment - 1);
74   if (overage > 0) {
75     const size_t waste = alignment - overage;
76     if (waste >= remaining_) {
77       return false;
78     }
79     freestart_ += waste;
80     remaining_ -= waste;
81   }
82   DCHECK_EQ(size_t{0}, reinterpret_cast<size_t>(freestart_) & (alignment - 1));
83   return true;
84 }
85 
86 // ----------------------------------------------------------------------
87 // Arena::Reset()
88 //    Clears all the memory an arena is using.
89 // ----------------------------------------------------------------------
90 
Reset()91 void Arena::Reset() {
92   FreeBlocks();
93   freestart_ = first_blocks_[0].mem;
94   remaining_ = first_blocks_[0].size;
95 
96   // There is no guarantee the first block is properly aligned, so
97   // enforce that now.
98   CHECK(SatisfyAlignment(kDefaultAlignment));
99 
100   freestart_when_empty_ = freestart_;
101 }
102 
103 // ----------------------------------------------------------------------
104 // Arena::MakeNewBlock()
105 //    Our sbrk() equivalent.  We always make blocks of the same size
106 //    (though GetMemory() can also make a new block for really big
107 //    data.
108 // ----------------------------------------------------------------------
109 
MakeNewBlock(const uint32 alignment)110 void Arena::MakeNewBlock(const uint32 alignment) {
111   AllocatedBlock* block = AllocNewBlock(block_size_, alignment);
112   freestart_ = block->mem;
113   remaining_ = block->size;
114   CHECK(SatisfyAlignment(alignment));
115 }
116 
LeastCommonMultiple(uint32 a,uint32 b)117 static uint32 LeastCommonMultiple(uint32 a, uint32 b) {
118   if (a > b) {
119     return (a / MathUtil::GCD<uint32>(a, b)) * b;
120   } else if (a < b) {
121     return (b / MathUtil::GCD<uint32>(b, a)) * a;
122   } else {
123     return a;
124   }
125 }
126 
127 // -------------------------------------------------------------
128 // Arena::AllocNewBlock()
129 //    Adds and returns an AllocatedBlock.
130 //    The returned AllocatedBlock* is valid until the next call
131 //    to AllocNewBlock or Reset.  (i.e. anything that might
132 //    affect overflow_blocks_).
133 // -------------------------------------------------------------
134 
AllocNewBlock(const size_t block_size,const uint32 alignment)135 Arena::AllocatedBlock* Arena::AllocNewBlock(const size_t block_size,
136                                             const uint32 alignment) {
137   AllocatedBlock* block;
138   // Find the next block.
139   if (blocks_alloced_ < TF_ARRAYSIZE(first_blocks_)) {
140     // Use one of the pre-allocated blocks
141     block = &first_blocks_[blocks_alloced_++];
142   } else {  // oops, out of space, move to the vector
143     if (overflow_blocks_ == nullptr)
144       overflow_blocks_ = new std::vector<AllocatedBlock>;
145     // Adds another block to the vector.
146     overflow_blocks_->resize(overflow_blocks_->size() + 1);
147     // block points to the last block of the vector.
148     block = &overflow_blocks_->back();
149   }
150 
151   // NOTE(tucker): this utility is made slightly more complex by
152   // not disallowing the case where alignment > block_size.
153   // Can we, without breaking existing code?
154 
155   // Must be a multiple of kDefaultAlignment, unless requested
156   // alignment is 1, in which case we don't care at all.
157   uint32 adjusted_alignment =
158       (alignment > 1 ? LeastCommonMultiple(alignment, kDefaultAlignment) : 1);
159   // Required minimum alignment for port::AlignedMalloc().
160   adjusted_alignment =
161       std::max(adjusted_alignment, static_cast<uint32>(sizeof(void*)));
162 
163   CHECK_LE(adjusted_alignment, static_cast<uint32>(1 << 20))
164       << "Alignment on boundaries greater than 1MB not supported.";
165 
166   // If block_size > alignment we force block_size to be a multiple
167   // of alignment; if block_size < alignment we make no adjustment.
168   size_t adjusted_block_size = block_size;
169   if (adjusted_block_size > adjusted_alignment) {
170     const uint32 excess = adjusted_block_size % adjusted_alignment;
171     adjusted_block_size += (excess > 0 ? adjusted_alignment - excess : 0);
172   }
173   block->mem = reinterpret_cast<char*>(
174       port::AlignedMalloc(adjusted_block_size, adjusted_alignment));
175   block->size = adjusted_block_size;
176   CHECK(nullptr != block->mem) << "block_size=" << block_size
177                                << " adjusted_block_size=" << adjusted_block_size
178                                << " alignment=" << alignment
179                                << " adjusted_alignment=" << adjusted_alignment;
180 
181   return block;
182 }
183 
184 // ----------------------------------------------------------------------
185 // Arena::GetMemoryFallback()
186 //    We take memory out of our pool, aligned on the byte boundary
187 //    requested.  If we don't have space in our current pool, we
188 //    allocate a new block (wasting the remaining space in the
189 //    current block) and give you that.  If your memory needs are
190 //    too big for a single block, we make a special your-memory-only
191 //    allocation -- this is equivalent to not using the arena at all.
192 // ----------------------------------------------------------------------
193 
GetMemoryFallback(const size_t size,const int alignment)194 void* Arena::GetMemoryFallback(const size_t size, const int alignment) {
195   if (0 == size) {
196     return nullptr;  // stl/stl_alloc.h says this is okay
197   }
198 
199   // alignment must be a positive power of 2.
200   CHECK(alignment > 0 && 0 == (alignment & (alignment - 1)));
201 
202   // If the object is more than a quarter of the block size, allocate
203   // it separately to avoid wasting too much space in leftover bytes.
204   if (block_size_ == 0 || size > block_size_ / 4) {
205     return AllocNewBlock(size, alignment)->mem;
206   }
207 
208   // Enforce alignment on freestart_ then check for adequate space,
209   // which may require starting a new block.
210   if (!SatisfyAlignment(alignment) || size > remaining_) {
211     MakeNewBlock(alignment);
212   }
213   CHECK_LE(size, remaining_);
214 
215   remaining_ -= size;
216   void* result = freestart_;
217   freestart_ += size;
218 
219   return result;
220 }
221 
222 // ----------------------------------------------------------------------
223 // Arena::ReturnMemoryFallback()
224 // Arena::FreeBlocks()
225 //    Unlike GetMemory(), which does actual work, ReturnMemory() is a
226 //    no-op: we don't "free" memory until Reset() is called.  We do
227 //    update some stats, though.  Note we do no checking that the
228 //    pointer you pass in was actually allocated by us, or that it
229 //    was allocated for the size you say, so be careful here!
230 //       FreeBlocks() does the work for Reset(), actually freeing all
231 //    memory allocated in one fell swoop.
232 // ----------------------------------------------------------------------
233 
FreeBlocks()234 void Arena::FreeBlocks() {
235   for (size_t i = 1; i < blocks_alloced_; ++i) {  // keep first block alloced
236     port::AlignedFree(first_blocks_[i].mem);
237     first_blocks_[i].mem = nullptr;
238     first_blocks_[i].size = 0;
239   }
240   blocks_alloced_ = 1;
241   if (overflow_blocks_ != nullptr) {
242     std::vector<AllocatedBlock>::iterator it;
243     for (it = overflow_blocks_->begin(); it != overflow_blocks_->end(); ++it) {
244       port::AlignedFree(it->mem);
245     }
246     delete overflow_blocks_;  // These should be used very rarely
247     overflow_blocks_ = nullptr;
248   }
249 }
250 
251 }  // namespace core
252 }  // namespace tensorflow
253