diff options
| -rw-r--r-- | tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc | 15 | ||||
| -rw-r--r-- | tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.h | 26 |
2 files changed, 27 insertions, 14 deletions
diff --git a/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc b/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc index cd82361030..2ec149ee8e 100644 --- a/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc +++ b/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc @@ -159,12 +159,14 @@ void* GPUBFCAllocator::AllocateRawInternal(size_t unused_alignment, // Start searching from the first bin for the smallest chunk that fits // rounded_bytes. Bin* b = it->second; - for (GPUBFCAllocator::Chunk* chunk : b->free_chunks) { + for (auto citer = b->free_chunks.begin(); citer != b->free_chunks.end(); + ++citer) { + GPUBFCAllocator::Chunk* chunk = (*citer); DCHECK(!chunk->in_use()); if (chunk->size >= rounded_bytes) { // We found an existing chunk that fits us that wasn't in use, so remove // it from the free bin structure prior to using. - RemoveFreeChunkFromBin(chunk); + RemoveFreeChunkIterFromBin(&b->free_chunks, citer); // If we can break the size of the chunk into two reasonably // large pieces, do so. @@ -299,6 +301,15 @@ void GPUBFCAllocator::InsertFreeChunkIntoBin(GPUBFCAllocator::Chunk* c) { new_bin->free_chunks.insert(c); } +void GPUBFCAllocator::RemoveFreeChunkIterFromBin( + GPUBFCAllocator::Bin::FreeChunkSet* free_chunks, + const GPUBFCAllocator::Bin::FreeChunkSet::iterator& citer) { + GPUBFCAllocator::Chunk* c = *citer; + CHECK(!c->in_use() && c->bin); + free_chunks->erase(citer); + c->bin = nullptr; +} + void GPUBFCAllocator::RemoveFreeChunkFromBin(GPUBFCAllocator::Chunk* c) { CHECK(!c->in_use() && c->bin); int count = c->bin->free_chunks.erase(c); diff --git a/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.h b/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.h index 5551c8bdb7..6205abbc87 100644 --- a/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.h +++ b/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.h @@ -125,17 +125,6 @@ class GPUBFCAllocator : public VisitableAllocator { return dbg; } }; - - Chunk* AllocateNewChunk(size_t num_bytes); - void SplitChunk(Chunk* c, size_t num_bytes) EXCLUSIVE_LOCKS_REQUIRED(lock_); - void Merge(Chunk* c1, Chunk* c2) EXCLUSIVE_LOCKS_REQUIRED(lock_); - void FreeAndMaybeCoalesce(Chunk* c) EXCLUSIVE_LOCKS_REQUIRED(lock_); - void InsertFreeChunkIntoBin(Chunk* c) EXCLUSIVE_LOCKS_REQUIRED(lock_); - void RemoveFreeChunkFromBin(Chunk* c); - void DeleteChunk(Chunk* c) EXCLUSIVE_LOCKS_REQUIRED(lock_); - - void DumpMemoryLog(size_t num_bytes) EXCLUSIVE_LOCKS_REQUIRED(lock_); - // A Bin is a collection of similar-sized free chunks. struct Bin { // All chunks in this bin have >= bin_size memory. @@ -151,13 +140,26 @@ class GPUBFCAllocator : public VisitableAllocator { } }; + typedef std::set<Chunk*, ChunkComparator> FreeChunkSet; // List of free chunks within the bin, sorted by chunk size. // Chunk * not owned. - std::set<Chunk*, ChunkComparator> free_chunks; + FreeChunkSet free_chunks; explicit Bin(size_t bs) : bin_size(bs) {} }; + Chunk* AllocateNewChunk(size_t num_bytes); + void SplitChunk(Chunk* c, size_t num_bytes) EXCLUSIVE_LOCKS_REQUIRED(lock_); + void Merge(Chunk* c1, Chunk* c2) EXCLUSIVE_LOCKS_REQUIRED(lock_); + void FreeAndMaybeCoalesce(Chunk* c) EXCLUSIVE_LOCKS_REQUIRED(lock_); + void InsertFreeChunkIntoBin(Chunk* c) EXCLUSIVE_LOCKS_REQUIRED(lock_); + void RemoveFreeChunkIterFromBin(Bin::FreeChunkSet* free_chunks, + const Bin::FreeChunkSet::iterator& c); + void RemoveFreeChunkFromBin(Chunk* c); + void DeleteChunk(Chunk* c) EXCLUSIVE_LOCKS_REQUIRED(lock_); + + void DumpMemoryLog(size_t num_bytes) EXCLUSIVE_LOCKS_REQUIRED(lock_); + GPUAllocatorRetry retry_helper_; // Structures immutable after construction |