blob: f809e3822caaf3310031cef94d895b8808e36be7 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
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<size_t, size_t> 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_
|
