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/compiler/xla/stream_executor/temporary_memory_manager.h"
17
18 #include "absl/strings/str_cat.h"
19 #include "absl/strings/str_format.h"
20 #include "tensorflow/compiler/xla/stream_executor/platform/logging.h"
21 #include "tensorflow/compiler/xla/stream_executor/stream.h"
22 #include "tensorflow/compiler/xla/stream_executor/stream_executor_pimpl.h"
23
24 namespace stream_executor {
25 namespace internal {
26
ForceDeallocateAll()27 void TemporaryMemoryManager::ForceDeallocateAll() {
28 absl::MutexLock lock(&mutex_);
29 VLOG(1) << "force-deallocating " << records_.size() << " remaining records";
30 for (auto it = records_.begin(); it != records_.end(); ++it) {
31 DeviceMemoryBase device_memory = it->first;
32 stream_->parent()->Deallocate(&device_memory);
33 }
34 }
35
MarkFinalized(const DeviceMemoryBase & device_memory,uint64_t generation,bool must_exist)36 void TemporaryMemoryManager::MarkFinalized(
37 const DeviceMemoryBase& device_memory, uint64_t generation,
38 bool must_exist) {
39 absl::MutexLock 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
DeallocateFinalizedTemporaries()51 void TemporaryMemoryManager::DeallocateFinalizedTemporaries() {
52 absl::MutexLock 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
IsFinalized(const DeviceMemoryBase & device_memory,uint64_t allocation_generation) const67 bool TemporaryMemoryManager::IsFinalized(const DeviceMemoryBase& device_memory,
68 uint64_t allocation_generation) const {
69 absl::MutexLock 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
HasAllocated(const DeviceMemoryBase & device_memory,uint64_t generation) const83 bool TemporaryMemoryManager::HasAllocated(const DeviceMemoryBase& device_memory,
84 uint64_t generation) const {
85 absl::MutexLock 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>>
AllocateArrayBase(uint64_t element_count,uint64_t element_size)94 TemporaryMemoryManager::AllocateArrayBase(uint64_t element_count,
95 uint64_t element_size) {
96 uint64_t 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 absl::StrCat("could not allocate temporary memory of ",
102 byte_size, " bytes"));
103 }
104
105 uint64_t 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 absl::MutexLock 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) << absl::StreamFormat(
118 "stream %p allocated temporary device memory at %p (size %u) in "
119 "generation %u",
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 stream_executor
128