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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
|
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_ALLOCATION_TRACKER_H_
#define TENSORFLOW_COMPILER_XLA_SERVICE_ALLOCATION_TRACKER_H_
#include <map>
#include <memory>
#include <set>
#include <string>
#include <vector>
#include "tensorflow/compiler/xla/service/backend.h"
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/platform/types.h"
namespace xla {
// Tracks allocations for the XLA service; allocations can be registered
// with shape/device/tag and resolved from a handle for later use.
class AllocationTracker {
public:
// The allocator is used for deallocating memory when allocations are
// deregistered. All registered allocations must have the same platform as the
// allocator.
AllocationTracker(Backend* backend) : backend_(backend), next_handle_(1) {}
// Registers a shaped buffer of device memory, and returns a corresponding
// handle that can be used for talking to XLA clients.
StatusOr<GlobalDataHandle> Register(
std::unique_ptr<ShapedBuffer> shaped_buffer, const string& tag);
// Unregister the allocation for the given data handle.
Status Unregister(const GlobalDataHandle& data);
// Returns a vector of global data handles that point to the tuple elements.
StatusOr<std::vector<GlobalDataHandle>> DeconstructTuple(
const GlobalDataHandle& Data);
// Resolve a handle from an XLA client to a shaped buffer, or provide an error
// status to say whether it was not found (or found, but found deallocated).
StatusOr<const ShapedBuffer*> Resolve(const GlobalDataHandle& data);
private:
// Data structure encapsulating single memory allocation on the device.
struct Allocation {
// The pointer to this allocation.
perftools::gputools::DeviceMemoryBase device_memory;
// The device that the memory is allocated on.
int device_ordinal;
// This is the number of times this memory allocation is refered to by
// registered data handles.
int ref_count;
};
// Internal helper which resolves the given GlobalDataHandle to a
// ShapedBuffer.
StatusOr<ShapedBuffer*> ResolveInternal(const GlobalDataHandle& data)
EXCLUSIVE_LOCKS_REQUIRED(mutex_);
// Internal helper which registers a shaped buffer.
StatusOr<GlobalDataHandle> RegisterInternal(
std::unique_ptr<ShapedBuffer> shaped_buffer, const string& tag)
EXCLUSIVE_LOCKS_REQUIRED(mutex_);
// Adds the given device address to the allocation tracker, or if it already
// exists, then increment it's reference count.
void AddAllocationOrIncrementRefCount(
perftools::gputools::DeviceMemoryBase device_memory, int device_ordinal)
EXCLUSIVE_LOCKS_REQUIRED(mutex_);
// Decrements the reference count of the given device memory. Then, if it is
// zero, deallocate the memory.
Status DecrementRefCount(perftools::gputools::DeviceMemoryBase device_memory,
int device_ordinal) EXCLUSIVE_LOCKS_REQUIRED(mutex_);
// A map from device memory opaque value to allocation. One such map is
// maintained per device ordinal.
using AllocationMap = tensorflow::gtl::FlatMap<const void*, Allocation>;
tensorflow::mutex mutex_;
// Backend to use with this tracker. The backend supplies the memory allocator
// to use when deallocating memory.
Backend* backend_;
// The next handle to assign to an allocation, guarded by the same mutex as
// the mapping as they'll be mutated at the same time.
int64 next_handle_ GUARDED_BY(mutex_);
// A map from device ordinal to AllocationMap.
tensorflow::gtl::FlatMap<int, AllocationMap> opaque_to_allocation_map_
GUARDED_BY(mutex_);
// A map from data handle to ShapedBuffer.
tensorflow::gtl::FlatMap<int64, std::unique_ptr<ShapedBuffer>>
handle_to_shaped_buffer_ GUARDED_BY(mutex_);
TF_DISALLOW_COPY_AND_ASSIGN(AllocationTracker);
};
} // namespace xla
#endif // TENSORFLOW_COMPILER_XLA_SERVICE_ALLOCATION_TRACKER_H_
|
