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
|
/* Copyright 2016 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_FRAMEWORK_RESOURCE_OP_KERNEL_H_
#define TENSORFLOW_FRAMEWORK_RESOURCE_OP_KERNEL_H_
#include <string>
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
// ResourceOpKernel<T> is a virtual base class for resource op implementing
// interface type T. The inherited op looks up the resource name (determined by
// ContainerInfo), and creates a new resource if necessary.
//
// Requirements:
// - Op must be marked as stateful.
// - Op must have `container` and `shared_name` attributes. Empty `container`
// means using the default container. Empty `shared_name` means private
// resource.
// - Subclass must override CreateResource().
// - Subclass is encouraged to override VerifyResource().
template <typename T>
class ResourceOpKernel : public OpKernel {
public:
explicit ResourceOpKernel(OpKernelConstruction* context) : OpKernel(context) {
has_resource_type_ = (context->output_type(0) == DT_RESOURCE);
if (!has_resource_type_) {
// The resource variant of the op may be placed on non-CPU devices, but
// this allocation is always on the host. Fortunately we don't need it in
// the resource case.
OP_REQUIRES_OK(context,
context->allocate_persistent(DT_STRING, TensorShape({2}),
&handle_, nullptr));
}
}
// The resource is deleted from the resource manager only when it is private
// to kernel. Ideally the resource should be deleted when it is no longer held
// by anyone, but it would break backward compatibility.
~ResourceOpKernel() override {
if (resource_ != nullptr) {
resource_->Unref();
if (cinfo_.resource_is_private_to_kernel()) {
if (!cinfo_.resource_manager()
->template Delete<T>(cinfo_.container(), cinfo_.name())
.ok()) {
// Do nothing; the resource can have been deleted by session resets.
}
}
}
}
void Compute(OpKernelContext* context) override LOCKS_EXCLUDED(mu_) {
mutex_lock l(mu_);
if (resource_ == nullptr) {
ResourceMgr* mgr = context->resource_manager();
OP_REQUIRES_OK(context, cinfo_.Init(mgr, def()));
T* resource;
OP_REQUIRES_OK(
context,
mgr->LookupOrCreate<T>(cinfo_.container(), cinfo_.name(), &resource,
[this](T** ret) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
Status s = CreateResource(ret);
if (!s.ok() && *ret != nullptr) {
CHECK((*ret)->Unref());
}
return s;
}));
Status s = VerifyResource(resource);
if (TF_PREDICT_FALSE(!s.ok())) {
resource->Unref();
context->SetStatus(s);
return;
}
if (!has_resource_type_) {
auto h = handle_.AccessTensor(context)->template flat<string>();
h(0) = cinfo_.container();
h(1) = cinfo_.name();
}
resource_ = resource;
}
if (has_resource_type_) {
OP_REQUIRES_OK(context, MakeResourceHandleToOutput(
context, 0, cinfo_.container(), cinfo_.name(),
MakeTypeIndex<T>()));
} else {
context->set_output_ref(0, &mu_, handle_.AccessTensor(context));
}
}
protected:
// Variables accessible from subclasses.
mutex mu_;
ContainerInfo cinfo_ GUARDED_BY(mu_);
T* resource_ GUARDED_BY(mu_) = nullptr;
private:
// Must return a T descendant allocated with new that ResourceOpKernel will
// take ownership of.
virtual Status CreateResource(T** resource) EXCLUSIVE_LOCKS_REQUIRED(mu_) = 0;
// During the first Compute(), resource is either created or looked up using
// shared_name. In the latter case, the resource found should be verified if
// it is compatible with this op's configuration. The verification may fail in
// cases such as two graphs asking queues of the same shared name to have
// inconsistent capacities.
virtual Status VerifyResource(T* resource) { return Status::OK(); }
PersistentTensor handle_ GUARDED_BY(mu_);
// Is the output of the operator of type DT_RESOURCE?
bool has_resource_type_;
};
} // namespace tensorflow
#endif // TENSORFLOW_FRAMEWORK_RESOURCE_OP_KERNEL_H_
|
