1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/core/framework/resource_handle.h"
17
18 #include "tensorflow/core/framework/resource_handle.pb.h"
19 #include "tensorflow/core/lib/core/errors.h"
20 #include "tensorflow/core/lib/strings/strcat.h"
21
22 namespace tensorflow {
23
24 // Must be declared here for pre-C++17 compatibility.
25 /* static */ constexpr const char* ResourceHandle::ANONYMOUS_NAME;
26
ResourceHandle()27 ResourceHandle::ResourceHandle() {}
28
ResourceHandle(const ResourceHandleProto & proto)29 ResourceHandle::ResourceHandle(const ResourceHandleProto& proto) {
30 FromProto(proto);
31 }
32
~ResourceHandle()33 ResourceHandle::~ResourceHandle() {}
34
AsProto(ResourceHandleProto * proto) const35 void ResourceHandle::AsProto(ResourceHandleProto* proto) const {
36 proto->set_device(device());
37 proto->set_container(container());
38 proto->set_name(name());
39 proto->set_hash_code(hash_code());
40 proto->set_maybe_type_name(maybe_type_name());
41 for (const auto& dtype_and_shape_pair : dtypes_and_shapes_) {
42 auto dtype_and_shape = proto->add_dtypes_and_shapes();
43 dtype_and_shape->set_dtype(dtype_and_shape_pair.dtype);
44 dtype_and_shape_pair.shape.AsProto(dtype_and_shape->mutable_shape());
45 }
46 }
47
FromProto(const ResourceHandleProto & proto)48 void ResourceHandle::FromProto(const ResourceHandleProto& proto) {
49 set_device(proto.device());
50 set_container(proto.container());
51 set_name(proto.name());
52 set_hash_code(proto.hash_code());
53 set_maybe_type_name(proto.maybe_type_name());
54 std::vector<DtypeAndPartialTensorShape> dtypes_and_shapes;
55 for (const auto& dtype_and_shape : proto.dtypes_and_shapes()) {
56 DataType dtype = dtype_and_shape.dtype();
57 PartialTensorShape shape(dtype_and_shape.shape());
58 dtypes_and_shapes.push_back(DtypeAndPartialTensorShape{dtype, shape});
59 }
60 dtypes_and_shapes_ = std::move(dtypes_and_shapes);
61 }
62
SerializeAsString() const63 string ResourceHandle::SerializeAsString() const {
64 ResourceHandleProto proto;
65 AsProto(&proto);
66 return proto.SerializeAsString();
67 }
68
ParseFromString(const string & s)69 bool ResourceHandle::ParseFromString(const string& s) {
70 ResourceHandleProto proto;
71 const bool status = proto.ParseFromString(s);
72 if (status) FromProto(proto);
73 return status;
74 }
75
DebugString() const76 string ResourceHandle::DebugString() const {
77 return strings::StrCat("device: ", device(), " container: ", container(),
78 " name: ", name(), " hash_code: ", hash_code(),
79 " maybe_type_name: ", maybe_type_name());
80 }
81
MakeRefCountingHandle(ResourceBase * resource,const string & device_name,const TypeIndex & type_index,const std::vector<DtypeAndPartialTensorShape> & dtypes_and_shapes,const absl::optional<ManagedStackTrace> & definition_stack_trace)82 ResourceHandle ResourceHandle::MakeRefCountingHandle(
83 ResourceBase* resource, const string& device_name,
84 const TypeIndex& type_index,
85 const std::vector<DtypeAndPartialTensorShape>& dtypes_and_shapes,
86 const absl::optional<ManagedStackTrace>& definition_stack_trace) {
87 ResourceHandle result;
88 result.resource_.reset(resource, /*add_ref=*/false);
89 result.set_device(device_name);
90 // "container" is only a ResourceMgr-only concept
91 result.set_container("");
92 result.set_definition_stack_trace(definition_stack_trace);
93 result.set_name(strings::StrCat("_AnonymousResource", GenerateUniqueId()));
94 result.set_hash_code(type_index.hash_code());
95 result.set_maybe_type_name(type_index.name());
96 result.set_dtypes_and_shapes(dtypes_and_shapes);
97 return result;
98 }
99
ValidateType(const TypeIndex & type_index) const100 Status ResourceHandle::ValidateType(const TypeIndex& type_index) const {
101 if (type_index.hash_code() != hash_code()) {
102 return errors::InvalidArgument(
103 "Trying to access a handle's resource using the wrong type. ",
104 "The handle points to a resource (name '", name(), "') of type '",
105 maybe_type_name(), "' (hash code ", hash_code(),
106 ") but you are trying to access the resource as type '",
107 type_index.name(), "' (hash code ", type_index.hash_code(), ")");
108 }
109 return Status::OK();
110 }
111
112 std::atomic<int64> ResourceHandle::current_id_;
113
GenerateUniqueId()114 int64 ResourceHandle::GenerateUniqueId() { return current_id_.fetch_add(1); }
115
ProtoDebugString(const ResourceHandle & handle)116 string ProtoDebugString(const ResourceHandle& handle) {
117 return handle.DebugString();
118 }
119
EncodeResourceHandleList(const ResourceHandle * p,int64_t n,std::unique_ptr<port::StringListEncoder> e)120 void EncodeResourceHandleList(const ResourceHandle* p, int64_t n,
121 std::unique_ptr<port::StringListEncoder> e) {
122 ResourceHandleProto proto;
123 for (int i = 0; i < n; ++i) {
124 p[i].AsProto(&proto);
125 e->Append(proto);
126 }
127 e->Finalize();
128 }
129
DecodeResourceHandleList(std::unique_ptr<port::StringListDecoder> d,ResourceHandle * ps,int64_t n)130 bool DecodeResourceHandleList(std::unique_ptr<port::StringListDecoder> d,
131 ResourceHandle* ps, int64_t n) {
132 std::vector<uint32> sizes(n);
133 if (!d->ReadSizes(&sizes)) return false;
134
135 ResourceHandleProto proto;
136 for (int i = 0; i < n; ++i) {
137 if (!proto.ParseFromArray(d->Data(sizes[i]), sizes[i])) {
138 return false;
139 }
140 ps[i].FromProto(proto);
141 }
142 return true;
143 }
144
145 } // namespace tensorflow
146