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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
|
/* 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.
==============================================================================*/
// This header declares classes for the infeed manager and the infeed
// buffer that are used by the GPU runtime to transfer buffers into an
// executing GPU computation, e.g., to feed data into a while loop.
#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_INFEED_MANAGER_H_
#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_INFEED_MANAGER_H_
#include <deque>
#include <vector>
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/core/lib/gtl/flatset.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/stream_executor_no_cuda.h"
namespace xla {
namespace gpu {
// TODO(b/30467474) Once GPU infeed implementation settles, consider
// folding back the cpu and gpu infeed implementations into a generic
// one if possible.
//
// Current limitations:
// * Does not handle multiple devices/replicas.
//
// * Buffer space on GPU is allocated on every infeed enqueue request,
// and it does not handle the case when it runs out of
// memory. Potential solution is to pre-allocate a fixed amount of
// memory and block when that memory is full.
// Defines an infeed buffer that is passed to the runtime by
// the client. The client manages the memory of the buffer.
class InfeedBuffer {
public:
InfeedBuffer(se::StreamExecutor* executor, int64 length)
: executor_(executor), length_(length) {
device_memory_ = executor_->AllocateArray<uint8>(length);
CHECK(!device_memory_.is_null());
}
~InfeedBuffer() { executor_->Deallocate(&device_memory_); }
int64 length() const { return length_; }
// Callback to signal that this buffer is consumed. This helps the
// client to manage memory for the infeed buffers.
void Done() { delete this; }
se::DeviceMemoryBase* device_memory() { return &device_memory_; }
private:
se::StreamExecutor* executor_; // Not owned.
const int64 length_;
se::DeviceMemoryBase device_memory_;
};
// Client-side class used to enqueue infeed buffers.
class InfeedManager {
public:
InfeedManager();
// Calls the completion callback for any enqueued buffers that have
// not been dequeued by the runtime, and empties the infeed
// queue. Reset may not be called while a runtime computation is
// processing a dequeued buffer. The only safe way to ensure this
// condition is to call Reset when no computation is taking place.
void Reset();
// Adds a set of buffers to the infeed queue atomically. buffer->Done
// will be called when the buffer will no longer be accessed by the
// InfeedManager, either as a result of a call to Reset or because the
// runtime has dequeued and used the buffer.
void EnqueueBuffers(const std::vector<InfeedBuffer*>& buffers);
// Blocks until the infeed queue is non-empty, then returns the
// buffer at the head of the queue. Adds the current buffer to the
// to-be released set.
InfeedBuffer* BlockingDequeueBuffer();
// Releases a set of buffers from the to-be released set.
void ReleaseBuffers(const std::vector<InfeedBuffer*>& buffers);
// Returns a cached stream associated with an executor. Allocates a
// new stream on the first invocation. On subsequent invocations, if
// the cached executor is not the same as the requested executor,
// returns null.
se::Stream* GetStream(se::StreamExecutor* executor);
// Registers a callback that will be called when 'enqueued_buffer_' becomes
// empty.
void RegisterOnEmptyCallback(std::function<void()> callback);
private:
// TODO(b/30467474): Revisit if this mutex becomes a point of
// contention.
tensorflow::mutex mu_;
// Condition variable that is signaled every time a buffer is
// enqueued to an empty queue.
tensorflow::condition_variable cv_;
// InfeedBuffer* queue contents are not owned, but buffer->Done must
// be called when the buffer is no longer needed by the runtime.
std::deque<InfeedBuffer*> enqueued_buffer_;
// Buffers that are dequeued and currently being processed by the
// runtime. Not owned.
tensorflow::gtl::FlatSet<const InfeedBuffer*> dequeued_buffer_;
// Cached host to device stream for queuing infeed data.
std::unique_ptr<se::Stream> host_to_device_stream_;
// Executor that the host_to_device_stream belongs to. Not owned.
se::StreamExecutor* host_to_device_executor_;
// List of callbacks which will be called when 'enqueued_buffer_' becomes
// empty.
std::vector<std::function<void()>> on_empty_callbacks_;
};
// Singleton creator-or-accessor: Returns the GPU infeed manager.
InfeedManager* GetOrCreateInfeedManager();
} // namespace gpu
} // namespace xla
#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_INFEED_MANAGER_H_
|
