1 /* Copyright 2018 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 #ifndef TENSORFLOW_CORE_COMMON_RUNTIME_SCOPED_ALLOCATOR_H_ 16 #define TENSORFLOW_CORE_COMMON_RUNTIME_SCOPED_ALLOCATOR_H_ 17 18 #include "tensorflow/core/framework/allocator.h" 19 #include "tensorflow/core/framework/tensor.h" 20 #include "tensorflow/core/lib/core/refcount.h" 21 #include "tensorflow/core/platform/mutex.h" 22 23 namespace tensorflow { 24 class ScopedAllocatorContainer; 25 class ScopedAllocatorInstance; 26 27 // Manages a single backing tensor and a collection of aliases. 28 class ScopedAllocator { 29 public: 30 static const int32 kInvalidId = 0; 31 static const size_t kMaxAlignment = 64; 32 33 // A subrange of the TensorBuffer associated with this object that 34 // will be the backing memory for one aliased tensor. 35 struct Field { 36 int32 scope_id; 37 size_t offset; 38 size_t bytes; 39 }; 40 // Field index that refers to backing tensor, not any aliased field. 41 static const int32 kBackingIndex = -1; 42 43 // backing_tensor is expected to be newly allocated by a ScopedAllocatorOp 44 // instance. It must be large enough to back all of the specified 45 // (offset, byte) ranges of the fields. 46 ScopedAllocator(const Tensor& backing_tensor, int32 scope_id, 47 const string& name, const gtl::ArraySlice<Field>& fields, 48 int32 expected_call_count, 49 ScopedAllocatorContainer* container); 50 51 // Automatically deletes when last use expires, or when 52 // ScopedAllocatorContainer decides to delete. 53 ~ScopedAllocator() LOCKS_EXCLUDED(mu_); 54 55 // For debugging: returns true iff p is a pointer that could have 56 // been returned by AllocateRaw. 57 bool VerifyPointer(const void* p); 58 bool VerifyTensor(const Tensor* t); 59 tensor()60 const Tensor& tensor() const { return backing_tensor_; } 61 name()62 const string& name() const { return name_; } 63 64 private: 65 friend class ScopedAllocatorInstance; 66 // Only ScopedAllocatorInstances can call AllocateRaw and DeallocateRaw on a 67 // ScopedAllocator 68 void* AllocateRaw(int32 field_index, size_t num_bytes) LOCKS_EXCLUDED(mu_); 69 void DeallocateRaw(void* p) LOCKS_EXCLUDED(mu_); 70 Tensor backing_tensor_; 71 TensorBuffer* tbuf_; 72 int32 id_; 73 string name_; 74 ScopedAllocatorContainer* container_; 75 std::vector<Field> fields_; 76 mutex mu_; 77 int32 expected_call_count_ GUARDED_BY(mu_); 78 int32 live_alloc_count_ GUARDED_BY(mu_); 79 }; 80 81 // An Allocator that will return a pointer into the backing buffer of 82 // a previously allocated tensor, allowing creation of an alias 83 // tensor. There is a one-to-one mapping between the fields of a 84 // ScopedAllocator and ScopedAllocatorInstances. There is also a one-to-one 85 // mapping between scope_ids and ScopedAllocatorInstances. It should be 86 // discarded immediately after a single use. 87 class ScopedAllocatorInstance : public Allocator { 88 public: 89 explicit ScopedAllocatorInstance(ScopedAllocator* sa, int32 field_index); 90 91 private: ~ScopedAllocatorInstance()92 ~ScopedAllocatorInstance() { VLOG(1) << "~ScopedAllocatorInstance " << this; } 93 94 public: 95 // When a ScopedAllocatorContainer "Drops" a scope_id, it calls DropFromTable 96 // on the underlying ScopedAllocatorInstance. If this instance has already 97 // deallocated the tensor slice, we can safely delete this. 98 void DropFromTable() LOCKS_EXCLUDED(mu_); 99 void* AllocateRaw(size_t alignment, size_t num_bytes) 100 LOCKS_EXCLUDED(mu_) override; AllocateRaw(size_t alignment,size_t num_bytes,const AllocationAttributes & allocator_attr)101 void* AllocateRaw(size_t alignment, size_t num_bytes, 102 const AllocationAttributes& allocator_attr) override { 103 return AllocateRaw(alignment, num_bytes); 104 } 105 void DeallocateRaw(void* p) LOCKS_EXCLUDED(mu_) override; TracksAllocationSizes()106 bool TracksAllocationSizes() override { return false; } ShouldAllocateEmptyTensors()107 bool ShouldAllocateEmptyTensors() override { return false; } RequestedSize(const void * ptr)108 size_t RequestedSize(const void* ptr) override { return 0; } AllocatedSize(const void * ptr)109 size_t AllocatedSize(const void* ptr) override { return 0; } AllocationId(const void * ptr)110 int64 AllocationId(const void* ptr) override { return 0; } AllocatedSizeSlow(const void * ptr)111 size_t AllocatedSizeSlow(const void* ptr) override { return 0; } 112 string Name() override; 113 114 private: 115 mutex mu_; 116 ScopedAllocator* scoped_allocator_; 117 int32 field_index_; 118 bool allocated_ GUARDED_BY(mu_); 119 bool deallocated_ GUARDED_BY(mu_); 120 bool in_table_ GUARDED_BY(mu_); 121 }; 122 123 } // namespace tensorflow 124 #endif // TENSORFLOW_CORE_COMMON_RUNTIME_SCOPED_ALLOCATOR_H_ 125