/* Copyright 2015 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.
|
==============================================================================*/
|
|
#include "tensorflow/core/kernels/queue_base.h"
|
|
#include <vector>
|
#include "tensorflow/core/framework/node_def.pb.h"
|
#include "tensorflow/core/framework/tensor_shape.h"
|
#include "tensorflow/core/lib/core/errors.h"
|
#include "tensorflow/core/platform/mutex.h"
|
#include "tensorflow/core/platform/types.h"
|
#include "tensorflow/core/util/batch_util.h"
|
|
namespace tensorflow {
|
|
namespace {
|
|
template <DataType DT>
|
Status HandleSliceToElement(const Tensor& parent, Tensor* element,
|
int64 index) {
|
typedef typename EnumToDataType<DT>::Type T;
|
DCHECK_NE(parent.dim_size(0), 0);
|
DCHECK_GE(index, 0);
|
if (element->NumElements() != (parent.NumElements() / parent.dim_size(0))) {
|
TensorShape chip_shape = parent.shape();
|
chip_shape.RemoveDim(0);
|
return errors::Internal(
|
"HandleSliceToElement Cannot copy slice: number of elements does not "
|
"match. Shapes are: [element]: ",
|
element->shape().DebugString(),
|
", [parent slice]: ", chip_shape.DebugString());
|
}
|
auto parent_as_matrix = parent.flat_outer_dims<T>();
|
element->flat<T>() = parent_as_matrix.chip(index, 0);
|
return Status::OK();
|
}
|
|
} // namespace
|
|
QueueBase::QueueBase(int32 capacity, const DataTypeVector& component_dtypes,
|
const std::vector<TensorShape>& component_shapes,
|
const string& name)
|
: capacity_(capacity),
|
component_dtypes_(component_dtypes),
|
component_shapes_(component_shapes),
|
name_(name),
|
closed_(false) {}
|
|
QueueBase::~QueueBase() {}
|
|
Status QueueBase::ValidateTupleCommon(const Tuple& tuple) const {
|
if (tuple.size() != static_cast<size_t>(num_components())) {
|
return errors::InvalidArgument(
|
"Wrong number of components in tuple. Expected ", num_components(),
|
", got ", tuple.size());
|
}
|
for (size_t i = 0; i < tuple.size(); ++i) {
|
if (tuple[i].dtype() != component_dtypes_[i]) {
|
return errors::InvalidArgument(
|
"Type mismatch in tuple component ", i, ". Expected ",
|
DataTypeString(component_dtypes_[i]), ", got ",
|
DataTypeString(tuple[i].dtype()));
|
}
|
}
|
return Status::OK();
|
}
|
|
// static
|
string QueueBase::ShapeListString(const gtl::ArraySlice<TensorShape>& shapes) {
|
string result = "[";
|
bool first = true;
|
for (const TensorShape& shape : shapes) {
|
strings::StrAppend(&result, (first ? "" : ", "), shape.DebugString());
|
first = false;
|
}
|
strings::StrAppend(&result, "]");
|
return result;
|
}
|
|
Status QueueBase::MatchesNodeDefOp(const NodeDef& node_def,
|
const string& op) const {
|
if (node_def.op() != op) {
|
return errors::InvalidArgument("Shared queue '", name_, "' has type '", op,
|
"' that does not match type of Node '",
|
node_def.name(), "': ", node_def.op());
|
}
|
return Status::OK();
|
}
|
|
Status QueueBase::MatchesNodeDefCapacity(const NodeDef& node_def,
|
int32 capacity) const {
|
int32 requested_capacity = -1;
|
TF_RETURN_IF_ERROR(GetNodeAttr(node_def, "capacity", &requested_capacity));
|
if (requested_capacity < 0) requested_capacity = kUnbounded;
|
if (requested_capacity != capacity) {
|
return errors::InvalidArgument("Shared queue '", name_, "' has capacity ",
|
capacity, " but requested capacity was ",
|
requested_capacity);
|
}
|
return Status::OK();
|
}
|
|
Status QueueBase::MatchesNodeDefTypes(const NodeDef& node_def) const {
|
DataTypeVector requested_dtypes;
|
TF_RETURN_IF_ERROR(
|
GetNodeAttr(node_def, "component_types", &requested_dtypes));
|
if (requested_dtypes != component_dtypes_) {
|
return errors::InvalidArgument("Shared queue '", name_,
|
"' has component types ",
|
DataTypeSliceString(component_dtypes_),
|
" but requested component types were ",
|
DataTypeSliceString(requested_dtypes));
|
}
|
return Status::OK();
|
}
|
|
Status QueueBase::MatchesNodeDefShapes(const NodeDef& node_def) const {
|
std::vector<TensorShape> requested_shapes;
|
TF_RETURN_IF_ERROR(GetNodeAttr(node_def, "shapes", &requested_shapes));
|
if (requested_shapes != component_shapes_) {
|
return errors::InvalidArgument("Shared queue '", name_,
|
"' has component shapes ",
|
ShapeListString(component_shapes_),
|
" but requested component shapes were ",
|
ShapeListString(requested_shapes));
|
}
|
return Status::OK();
|
}
|
|
// TODO(mrry): If these checks become a bottleneck, find a way to
|
// reduce the number of times that they are called.
|
Status QueueBase::ValidateTuple(const Tuple& tuple) {
|
TF_RETURN_IF_ERROR(ValidateTupleCommon(tuple));
|
if (specified_shapes()) {
|
for (size_t i = 0; i < tuple.size(); ++i) {
|
if (!component_shapes_[i].IsSameSize(tuple[i].shape())) {
|
return errors::InvalidArgument(
|
"Shape mismatch in tuple component ", i, ". Expected ",
|
component_shapes_[i].DebugString(), ", got ",
|
tuple[i].shape().DebugString());
|
}
|
}
|
}
|
return Status::OK();
|
}
|
|
// TODO(mrry): If these checks become a bottleneck, find a way to
|
// reduce the number of times that they are called.
|
Status QueueBase::ValidateManyTuple(const Tuple& tuple) {
|
TF_RETURN_IF_ERROR(ValidateTupleCommon(tuple));
|
const int64 batch_size = tuple[0].dim_size(0);
|
if (specified_shapes()) {
|
for (size_t i = 0; i < tuple.size(); ++i) {
|
// Expected shape is [batch_size] + component_shapes_[i]
|
const TensorShape expected_shape = ManyOutShape(i, batch_size);
|
if (!expected_shape.IsSameSize(tuple[i].shape())) {
|
return errors::InvalidArgument("Shape mismatch in tuple component ", i,
|
". Expected ",
|
expected_shape.DebugString(), ", got ",
|
tuple[i].shape().DebugString());
|
}
|
}
|
} else {
|
for (size_t i = 1; i < tuple.size(); ++i) {
|
if (tuple[i].dim_size(0) != batch_size) {
|
return errors::InvalidArgument(
|
"All input tensors must have the same size in the 0th ",
|
"dimension. Component ", i, " has ", tuple[i].dim_size(0),
|
", and should have ", batch_size);
|
}
|
}
|
}
|
return Status::OK();
|
}
|
|
void QueueBase::Cancel(Action action, CancellationManager* cancellation_manager,
|
CancellationToken token) {
|
DoneCallback callback = nullptr;
|
{
|
mutex_lock lock(mu_);
|
std::deque<Attempt>* attempts =
|
action == kEnqueue ? &enqueue_attempts_ : &dequeue_attempts_;
|
|
for (Attempt& attempt : *attempts) {
|
if (attempt.cancellation_manager == cancellation_manager &&
|
attempt.cancellation_token == token) {
|
if (!attempt.is_cancelled) {
|
attempt.is_cancelled = true;
|
if (action == kEnqueue) {
|
attempt.context->SetStatus(
|
errors::Cancelled("Enqueue operation was cancelled"));
|
} else {
|
attempt.context->SetStatus(
|
errors::Cancelled("Dequeue operation was cancelled"));
|
}
|
std::swap(callback, attempt.done_callback);
|
}
|
break;
|
}
|
}
|
}
|
if (callback) {
|
callback();
|
FlushUnlocked();
|
}
|
}
|
|
void QueueBase::CloseAndCancel() {
|
std::vector<DoneCallback> callbacks;
|
{
|
mutex_lock lock(mu_);
|
closed_ = true;
|
for (Attempt& attempt : enqueue_attempts_) {
|
if (!attempt.is_cancelled) {
|
attempt.is_cancelled = true;
|
attempt.context->SetStatus(
|
errors::Cancelled("Enqueue operation was cancelled"));
|
callbacks.emplace_back(std::move(attempt.done_callback));
|
}
|
}
|
}
|
for (const DoneCallback& callback : callbacks) {
|
callback();
|
}
|
FlushUnlocked();
|
}
|
|
void QueueBase::Close(OpKernelContext* ctx, bool cancel_pending_enqueues,
|
DoneCallback callback) {
|
if (cancel_pending_enqueues) {
|
CloseAndCancel();
|
callback();
|
} else {
|
{
|
mutex_lock lock(mu_);
|
enqueue_attempts_.emplace_back(
|
0, callback, ctx, nullptr, CancellationManager::kInvalidToken,
|
[this](Attempt* attempt) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
|
if (closed_) {
|
attempt->context->SetStatus(
|
errors::Cancelled("Queue '", name_, "' is already closed."));
|
} else {
|
closed_ = true;
|
}
|
return kComplete;
|
});
|
}
|
FlushUnlocked();
|
}
|
}
|
|
bool QueueBase::TryAttemptLocked(Action action,
|
std::vector<CleanUp>* clean_up) {
|
std::deque<Attempt>* attempts =
|
action == kEnqueue ? &enqueue_attempts_ : &dequeue_attempts_;
|
|
bool progress = false;
|
bool done = false;
|
while (!done && !attempts->empty()) {
|
if (attempts->front().is_cancelled) {
|
if (action == kEnqueue) {
|
if (closed_) {
|
VLOG(1) << "Skipping cancelled enqueue attempt";
|
} else {
|
LOG(WARNING)
|
<< name_
|
<< ": Skipping cancelled enqueue attempt with queue not closed";
|
}
|
} else {
|
if (closed_) {
|
VLOG(1) << "Skipping cancelled dequeue attempt";
|
} else {
|
LOG(WARNING)
|
<< name_
|
<< ": Skipping cancelled dequeue attempt with queue not closed";
|
}
|
}
|
attempts->pop_front();
|
} else {
|
Attempt* cur_attempt = &attempts->front();
|
switch (cur_attempt->run_callback(cur_attempt)) {
|
case kNoProgress:
|
done = true;
|
break;
|
case kProgress:
|
done = true;
|
progress = true;
|
break;
|
case kComplete:
|
progress = true;
|
clean_up->emplace_back(std::move(cur_attempt->done_callback),
|
cur_attempt->cancellation_token,
|
cur_attempt->context->cancellation_manager());
|
attempts->pop_front();
|
break;
|
}
|
}
|
}
|
return progress;
|
}
|
|
void QueueBase::FlushUnlocked() {
|
std::vector<CleanUp> clean_up;
|
Ref();
|
{
|
mutex_lock lock(mu_);
|
bool changed;
|
do {
|
changed = TryAttemptLocked(kEnqueue, &clean_up);
|
changed = TryAttemptLocked(kDequeue, &clean_up) || changed;
|
} while (changed);
|
}
|
Unref();
|
for (const auto& to_clean : clean_up) {
|
if (to_clean.to_deregister != CancellationManager::kInvalidToken) {
|
// NOTE(mrry): We can safely ignore the return value of
|
// DeregisterCallback because the mutex mu_ ensures that the
|
// cleanup action only executes once.
|
to_clean.cm->DeregisterCallback(to_clean.to_deregister);
|
}
|
to_clean.finished();
|
}
|
}
|
|
Status QueueBase::CopySliceToElement(const Tensor& parent, Tensor* element,
|
int64 index) {
|
return batch_util::CopySliceToElement(parent, element, index);
|
}
|
|
/* static */
|
Status QueueBase::CopyElementToSlice(const Tensor& element, Tensor* parent,
|
int64 index) {
|
return batch_util::CopyElementToSlice(element, parent, index);
|
}
|
|
} // namespace tensorflow
|
