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 // Temporary memories are used to allocate scratch space required by an 17 // operation about to be enqueued onto a stream. 18 // 19 // std::unique_ptr<TemporaryDeviceMemory<float>> temporary_memory = 20 // stream.AllocateTemporaryArray<float>(1024).ConsumeValueOrDie(); 21 // // ... enqueue stuff onto the stream using the temporary memory ... 22 // // Note that the memory is accessible via 23 // // temporary_memory->device_memory() and similar. 24 // 25 // // Finalize the temporary memory. The underlying device memory may 26 // // be released any time after this program point, as another thread may 27 // // call Stream::BlockHostUntilDone, causing synchronization. This 28 // // finalization also happens automatically for the user if the unique_ptr 29 // // goes out of scope. 30 // temporary_memory.Finalize(); 31 // 32 // WARNING: do NOT hold onto the device memory associated with temporary_memory 33 // after finalization. If temporary_memory->device_memory() is used after the 34 // temporary memory is finalized, it will cause a DCHECK failure. 35 // 36 // Note that standard usage takes advantage of the type-safe wrapper, 37 // TemporaryDeviceMemory<T>, defined below. 38 // 39 // Also see tests for executable sample usage. 40 41 #ifndef TENSORFLOW_STREAM_EXECUTOR_TEMPORARY_DEVICE_MEMORY_H_ 42 #define TENSORFLOW_STREAM_EXECUTOR_TEMPORARY_DEVICE_MEMORY_H_ 43 44 #include "tensorflow/stream_executor/device_memory.h" 45 46 namespace perftools { 47 namespace gputools { 48 49 class Stream; 50 namespace internal { 51 class TemporaryMemoryManager; 52 } 53 54 // Untyped base class (analogous to a void*) for temporary device memory 55 // allocations associated with a stream. 56 class TemporaryDeviceMemoryBase { 57 public: 58 // Marks the temporary memory as finalized if it is not already marked as 59 // such. 60 ~TemporaryDeviceMemoryBase(); 61 62 // Precondition: !IsFinalized() 63 DeviceMemoryBase* mutable_device_memory(); 64 65 // Precondition: !IsFinalized() 66 const DeviceMemoryBase& device_memory() const; 67 68 // "Finalizes" this temporary memory, making it acceptable to release at the 69 // next stream synchronization point -- the device memory can be reclaimed at 70 // any time after the temporary memory is marked as finalized (e.g. if a 71 // separate thread is calls Stream::BlockHostUntilDone). This may only be 72 // called once -- see the precondition below. 73 // 74 // Precondition: !IsFinalized() 75 void Finalize(); 76 77 // Returns true iff the temporary memory is finalized (that is, the user is 78 // done referring to the temporary device memory, and thus it can be released 79 // at the next stream synchronization point). 80 bool IsFinalized() const; 81 82 // Returns true iff the temporary memory is still allocated. 83 // 84 // Note: this is a polling call, no guarantee is made that the temporary 85 // memory is still allocated after the call has completed. 86 bool IsAllocated() const; 87 88 private: 89 friend class internal::TemporaryMemoryManager; 90 friend class TemporaryDeviceMemoryTest; 91 92 // Note: construction DCHECKs that the memory is known-allocated in the 93 // stream's temporary-allocation-manager. 94 TemporaryDeviceMemoryBase(Stream* parent, DeviceMemoryBase device_memory, 95 uint64 allocation_generation); 96 97 // The device memory region that has allocated. 98 DeviceMemoryBase device_memory_; 99 100 // The generation counter value for the temporary memory record in the 101 // temporary memory manager. 102 uint64 allocation_generation_; 103 104 // The stream that this temporary memory was allocated for. 105 Stream* parent_; 106 }; 107 108 // Type-safe wrapper around the base type (which is analogous to a void*). 109 template <typename T> 110 class TemporaryDeviceMemory : public TemporaryDeviceMemoryBase { 111 public: 112 // Type-safe wrapper around TemporaryDeviceMemoryBase::mutable_device_memory. mutable_device_memory()113 DeviceMemory<T>* mutable_device_memory() { 114 StaticSlicingAssertionDummy(); 115 return reinterpret_cast<DeviceMemory<T>*>( 116 TemporaryDeviceMemoryBase::mutable_device_memory()); 117 } 118 119 // Type-safe wrapper around TemporaryDeviceMemoryBase::device_memory. device_memory()120 const DeviceMemory<T>& device_memory() const { 121 StaticSlicingAssertionDummy(); 122 return reinterpret_cast<const DeviceMemory<T>&>( 123 TemporaryDeviceMemoryBase::device_memory()); 124 } 125 126 private: StaticSlicingAssertionDummy()127 static void StaticSlicingAssertionDummy() { 128 static_assert( 129 sizeof(TemporaryDeviceMemory) == sizeof(TemporaryDeviceMemoryBase), 130 "derived class is simply a wrapper, no members may be added due to " 131 "slicing"); 132 } 133 }; 134 135 } // namespace gputools 136 } // namespace perftools 137 138 #endif // TENSORFLOW_STREAM_EXECUTOR_TEMPORARY_DEVICE_MEMORY_H_ 139