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
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
|
/* Copyright 2015 Google Inc. 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 THIRD_PARTY_TENSORFLOW_CORE_KERNELS_QUEUE_BASE_H_
#define THIRD_PARTY_TENSORFLOW_CORE_KERNELS_QUEUE_BASE_H_
#include <deque>
#include <vector>
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/queue_interface.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
#include "tensorflow/core/platform/port.h"
#include "tensorflow/core/public/tensor.h"
#include "tensorflow/core/public/tensor_shape.h"
namespace tensorflow {
// Functionality common to asynchronous QueueInterface implementations.
class QueueBase : public QueueInterface {
public:
// As a possible value of 'capacity'.
static const int32 kUnbounded = INT_MAX;
// Args:
// component_dtypes: The types of each component in a queue-element tuple.
// component_shapes: The shapes of each component in a queue-element tuple,
// which must either be empty (if the shapes are not specified) or
// or have the same size as component_dtypes.
// name: A name to use for the queue.
QueueBase(int32 capacity, const DataTypeVector& component_dtypes,
const std::vector<TensorShape>& component_shapes,
const string& name);
// Implementations of QueueInterface methods --------------------------------
const DataTypeVector& component_dtypes() const override {
return component_dtypes_;
}
Status ValidateTuple(const Tuple& tuple) override;
Status ValidateManyTuple(const Tuple& tuple) override;
void Close(OpKernelContext* ctx, bool cancel_pending_enqueues,
DoneCallback callback) override;
// Other public methods -----------------------------------------------------
const std::vector<TensorShape>& component_shapes() const {
return component_shapes_;
}
int32 capacity() const { return capacity_; }
protected:
enum Action { kEnqueue, kDequeue };
enum RunResult { kNoProgress, kProgress, kComplete };
// Tries to enqueue/dequeue (or close) based on whatever is at the
// front of enqueue_attempts_/dequeue_attempts_. Appends to
// *finished the callback for any finished attempt (so it may be
// called once mu_ is released). Returns true if any progress was
// made.
struct CleanUp {
CleanUp(DoneCallback&& f, CancellationToken ct, CancellationManager* cm)
: finished(f), to_deregister(ct), cm(cm) {}
DoneCallback finished;
CancellationToken to_deregister;
CancellationManager* cm;
};
// Returns the number of components in a queue-element tuple.
int32 num_components() const { return component_dtypes_.size(); }
// True if shapes were specified. If so, inputs will be validated
// against them, etc.
bool specified_shapes() const { return component_shapes_.size() > 0; }
// Code common to Validate*Tuple().
Status ValidateTupleCommon(const Tuple& tuple) const;
TensorShape ManyOutShape(int i, int64 batch_size) {
TensorShape shape({batch_size});
shape.AppendShape(component_shapes_[i]);
return shape;
}
// Copies the index^th slice (in the first dimension) of parent into element.
static Status CopySliceToElement(const Tensor& parent, Tensor* element,
int index);
// Copies element into the index^th slice (in the first dimension) of parent.
static Status CopyElementToSlice(const Tensor& element, Tensor* parent,
int index);
void Cancel(Action action, CancellationToken token);
// Helper for cancelling all pending Enqueue(Many) operations when
// Close is called with cancel_pending_enqueues.
void CloseAndCancel();
bool TryAttemptLocked(Action action, std::vector<CleanUp>* clean_up)
EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Tries to make progress on the enqueues or dequeues at the front
// of the *_attempts_ queues.
void FlushUnlocked();
~QueueBase() override {}
// Helpers for implementing MatchesNodeDef().
static string ShapeListString(const gtl::ArraySlice<TensorShape>& shapes);
Status MatchesNodeDefOp(const NodeDef& node_def, const string& op) const;
Status MatchesNodeDefCapacity(const NodeDef& node_def, int32 capacity) const;
Status MatchesNodeDefTypes(const NodeDef& node_def) const;
Status MatchesNodeDefShapes(const NodeDef& node_def) const;
protected:
const int32 capacity_;
const DataTypeVector component_dtypes_;
const std::vector<TensorShape> component_shapes_;
const string name_;
mutex mu_;
bool closed_ GUARDED_BY(mu_);
struct Attempt;
typedef std::function<RunResult(Attempt*)> RunCallback;
struct Attempt {
int32 elements_requested;
DoneCallback done_callback; // must be run outside mu_
OpKernelContext* context;
CancellationToken cancellation_token;
RunCallback run_callback; // must be run while holding mu_
bool is_cancelled;
Tuple tuple;
Attempt(int32 elements_requested, DoneCallback done_callback,
OpKernelContext* context, CancellationToken cancellation_token,
RunCallback run_callback)
: elements_requested(elements_requested),
done_callback(done_callback),
context(context),
cancellation_token(cancellation_token),
run_callback(run_callback),
is_cancelled(false) {}
};
std::deque<Attempt> enqueue_attempts_ GUARDED_BY(mu_);
std::deque<Attempt> dequeue_attempts_ GUARDED_BY(mu_);
TF_DISALLOW_COPY_AND_ASSIGN(QueueBase);
};
} // namespace tensorflow
#endif // THIRD_PARTY_TENSORFLOW_CORE_KERNELS_QUEUE_BASE_H_
|
