From f41959ccb2d9d4c722fe8fc3351401d53bcf4900 Mon Sep 17 00:00:00 2001 From: Manjunath Kudlur Date: Fri, 6 Nov 2015 16:27:58 -0800 Subject: TensorFlow: Initial commit of TensorFlow library. TensorFlow is an open source software library for numerical computation using data flow graphs. Base CL: 107276108 --- tensorflow/core/framework/tracking_allocator.h | 80 ++++++++++++++++++++++++++ 1 file changed, 80 insertions(+) create mode 100644 tensorflow/core/framework/tracking_allocator.h (limited to 'tensorflow/core/framework/tracking_allocator.h') diff --git a/tensorflow/core/framework/tracking_allocator.h b/tensorflow/core/framework/tracking_allocator.h new file mode 100644 index 0000000000..f809e3822c --- /dev/null +++ b/tensorflow/core/framework/tracking_allocator.h @@ -0,0 +1,80 @@ +#ifndef TENSORFLOW_FRAMEWORK_TRACKING_ALLOCATOR_H_ +#define TENSORFLOW_FRAMEWORK_TRACKING_ALLOCATOR_H_ + +#include "tensorflow/core/framework/allocator.h" +#include "tensorflow/core/lib/core/refcount.h" +#include "tensorflow/core/platform/port.h" +#include "tensorflow/core/platform/thread_annotations.h" + +namespace tensorflow { + +// TrackingAllocator is a wrapper for an Allocator. It keeps a running +// count of the number of bytes allocated through the wrapper. It is +// used by the Executor to "charge" allocations to particular Op +// executions. Each Op gets a separate TrackingAllocator wrapper +// around the underlying allocator. +// +// The implementation assumes the invariant that all calls to +// AllocateRaw by an Op (or work items spawned by the Op) will occur +// before the Op's Compute method returns. Thus the high watermark is +// established once Compute returns. +// +// DeallocateRaw can be called long after the Op has finished, +// e.g. when an output tensor is deallocated, and the wrapper cannot +// be deleted until the last of these calls has occurred. The +// TrackingAllocator keeps track of outstanding calls using a +// reference count, and deletes itself once the last call has been +// received and the high watermark has been retrieved. +class TrackingAllocator : public Allocator { + public: + explicit TrackingAllocator(Allocator* allocator); + string Name() override { return allocator_->Name(); } + void* AllocateRaw(size_t alignment, size_t num_bytes) override; + void DeallocateRaw(void* ptr) override; + bool TracksAllocationSizes() override; + size_t RequestedSize(void* ptr) override; + size_t AllocatedSize(void* ptr) override; + + // If the underlying allocator tracks allocation sizes, this returns + // a pair where the first value is the total number of bytes + // allocated through this wrapper, and the second value is the high + // watermark of bytes allocated through this wrapper. If the + // underlying allocator does not track allocation sizes the first + // value is the total number of bytes requested through this wrapper + // and the second is 0. + // + // After GetSizesAndUnref is called, the only further calls allowed + // on this wrapper are calls to DeallocateRaw with pointers that + // were allocated by this wrapper and have not yet been + // deallocated. After this call completes and all allocated pointers + // have been deallocated the wrapper will delete itself. + std::pair GetSizesAndUnRef(); + + private: + ~TrackingAllocator() override {} + bool UnRef() EXCLUSIVE_LOCKS_REQUIRED(mu_); + + Allocator* allocator_; // not owned. + mutex mu_; + // the number of calls to AllocateRaw that have not yet been matched + // by a corresponding call to DeAllocateRaw, plus 1 if the Executor + // has not yet read out the high watermark. + int ref_ GUARDED_BY(mu_); + // the current number of outstanding bytes that have been allocated + // by this wrapper, or 0 if the underlying allocator does not track + // allocation sizes. + size_t allocated_ GUARDED_BY(mu_); + // the maximum number of outstanding bytes that have been allocated + // by this wrapper, or 0 if the underlying allocator does not track + // allocation sizes. + size_t high_watermark_ GUARDED_BY(mu_); + // the total number of bytes that have been allocated by this + // wrapper if the underlying allocator tracks allocation sizes, + // otherwise the total number of bytes that have been requested by + // this allocator. + size_t total_bytes_ GUARDED_BY(mu_); +}; + +} // end namespace tensorflow + +#endif // TENSORFLOW_FRAMEWORK_TRACKING_ALLOCATOR_H_ -- cgit v1.2.3