Home | History | Annotate | Download | only in stream_executor 
      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 #include "tensorflow/stream_executor/temporary_memory_manager.h"
     17 
     18 #include "tensorflow/stream_executor/platform/logging.h"
     19 #include "tensorflow/stream_executor/lib/stringprintf.h"
     20 #include "tensorflow/stream_executor/lib/ptr_util.h"
     21 #include "tensorflow/stream_executor/stream.h"
     22 #include "tensorflow/stream_executor/stream_executor_pimpl.h"
     23 
     24 namespace perftools {
     25 namespace gputools {
     26 namespace internal {
     27 
     28 void TemporaryMemoryManager::ForceDeallocateAll() {
     29   mutex_lock lock(mutex_);
     30   VLOG(1) << "force-deallocating " << records_.size() << " remaining records";
     31   for (auto it = records_.begin(); it != records_.end(); ++it) {
     32     DeviceMemoryBase device_memory = it->first;
     33     stream_->parent()->Deallocate(&device_memory);
     34   }
     35 }
     36 
     37 void TemporaryMemoryManager::MarkFinalized(
     38     const DeviceMemoryBase& device_memory, uint64 generation, bool must_exist) {
     39   mutex_lock lock(mutex_);
     40   auto it = records_.find(device_memory);
     41   if (it == records_.end()) {
     42     if (must_exist) {
     43       LOG(FATAL) << "attempted to mark finalization for temporary "
     44                     "memory that does not exist";
     45     }
     46     return;
     47   }
     48   it->second.finalized = true;
     49 }
     50 
     51 void TemporaryMemoryManager::DeallocateFinalizedTemporaries() {
     52   mutex_lock lock(mutex_);
     53   int deallocated_count = 0;
     54   for (auto it = records_.begin(); it != records_.end();) {
     55     if (it->second.finalized) {
     56       DeviceMemoryBase device_memory = it->first;
     57       stream_->parent()->Deallocate(&device_memory);
     58       ++deallocated_count;
     59       it = records_.erase(it);
     60     } else {
     61       ++it;
     62     }
     63   }
     64   VLOG(1) << "deallocated " << deallocated_count << " finalized temporaries";
     65 }
     66 
     67 bool TemporaryMemoryManager::IsFinalized(const DeviceMemoryBase& device_memory,
     68                                          uint64 allocation_generation) const {
     69   mutex_lock lock(mutex_);
     70   auto it = records_.find(device_memory);
     71   if (it == records_.end()) {
     72     return true;  // If there's no record present it's vacuously finalized.
     73   }
     74 
     75   if (it->second.allocation_generation == allocation_generation) {
     76     return it->second.finalized;
     77   }
     78 
     79   // If the allocation generation did not match, it's vacuously true.
     80   return true;
     81 }
     82 
     83 bool TemporaryMemoryManager::HasAllocated(const DeviceMemoryBase& device_memory,
     84                                           uint64 generation) const {
     85   mutex_lock lock(mutex_);
     86   auto it = records_.find(device_memory);
     87   if (it == records_.end()) {
     88     return false;
     89   }
     90   return it->second.allocation_generation == generation;
     91 }
     92 
     93 port::StatusOr<std::unique_ptr<TemporaryDeviceMemoryBase>>
     94 TemporaryMemoryManager::AllocateArrayBase(uint64 element_count,
     95                                           uint64 element_size) {
     96   uint64 byte_size = element_count * element_size;
     97   DeviceMemoryBase device_memory =
     98       stream_->parent()->AllocateArray<uint8>(byte_size);
     99   if (device_memory == nullptr) {
    100     return port::Status(port::error::RESOURCE_EXHAUSTED,
    101                         port::StrCat("could not allocate temporary memory of ",
    102                                      byte_size, " bytes"));
    103   }
    104 
    105   uint64 generation;
    106 
    107   // Add the record before instantiating the device memory instance so we can
    108   // check the allocation invariant at TemporaryDeviceMemory construction time.
    109   {
    110     mutex_lock lock(mutex_);
    111     generation = ++generation_;
    112     DCHECK(records_.find(device_memory) == records_.end());
    113     records_[device_memory] = {generation,
    114                                /*finalized=*/false};
    115   }
    116 
    117   VLOG(1) << port::Printf(
    118       "stream %p allocated temporary device memory at %p (size %llu) in "
    119       "generation %llu",
    120       stream_, device_memory.opaque(), byte_size, generation);
    121   std::unique_ptr<TemporaryDeviceMemoryBase> result(
    122       new TemporaryDeviceMemoryBase(stream_, device_memory, generation));
    123   return std::move(result);
    124 }
    125 
    126 }  // namespace internal
    127 }  // namespace gputools
    128 }  // namespace perftools
    129