1/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#include "tensorflow/core/framework/resource_handle.h"
17
18#include <string>
19#include <utility>
20#include <vector>
21
22#include "absl/strings/str_format.h"
23#include "tensorflow/core/framework/resource_handle.pb.h"
24#include "tensorflow/core/framework/tensor_shape.h"
25#include "tensorflow/core/lib/core/errors.h"
26#include "tensorflow/core/lib/strings/strcat.h"
27#include "tensorflow/core/platform/demangle.h"
28#include "tensorflow/core/platform/errors.h"
29#include "tensorflow/core/platform/macros.h"
30
31namespace tensorflow {
32
33namespace {
34std::string DtypeAndShapesToString(
35 const std::vector<DtypeAndPartialTensorShape>& dtype_and_shapes) {
36 std::vector<std::string> dtype_and_shape_strings;
37 dtype_and_shape_strings.reserve(dtype_and_shapes.size());
38 for (const DtypeAndPartialTensorShape& dtype_and_shape : dtype_and_shapes) {
39 // Note that it is a bit unfortunate to return int/enum as dtype, given we
40 // can't directly use DataTypeString due to circular dependency.
41 dtype_and_shape_strings.push_back(
42 absl::StrFormat("DType enum: %d, Shape: %s", dtype_and_shape.dtype,
43 dtype_and_shape.shape.DebugString()));
44 }
45 return absl::StrFormat("[ %s ]", absl::StrJoin(dtype_and_shape_strings, ","));
46}
47} // namespace
48
49// Must be declared here for pre-C++17 compatibility.
50/* static */ constexpr const char* ResourceHandle::ANONYMOUS_NAME;
51
52ResourceHandle::ResourceHandle() {}
53
54ResourceHandle::ResourceHandle(const ResourceHandleProto& proto) {
55 TF_CHECK_OK(FromProto(proto));
56}
57
58Status ResourceHandle::BuildResourceHandle(const ResourceHandleProto& proto,
59 ResourceHandle* out) {
60 if (out == nullptr)
61 return errors::Internal(
62 "BuildResourceHandle() was called with nullptr for the output");
63 return out->FromProto(proto);
64}
65
66ResourceHandle::~ResourceHandle() {}
67
68void ResourceHandle::AsProto(ResourceHandleProto* proto) const {
69 proto->set_device(device());
70 proto->set_container(container());
71 proto->set_name(name());
72 proto->set_hash_code(hash_code());
73 proto->set_maybe_type_name(maybe_type_name());
74 for (const auto& dtype_and_shape_pair : dtypes_and_shapes_) {
75 auto dtype_and_shape = proto->add_dtypes_and_shapes();
76 dtype_and_shape->set_dtype(dtype_and_shape_pair.dtype);
77 dtype_and_shape_pair.shape.AsProto(dtype_and_shape->mutable_shape());
78 }
79}
80
81Status ResourceHandle::FromProto(const ResourceHandleProto& proto) {
82 set_device(proto.device());
83 set_container(proto.container());
84 set_name(proto.name());
85 set_hash_code(proto.hash_code());
86 set_maybe_type_name(proto.maybe_type_name());
87 std::vector<DtypeAndPartialTensorShape> dtypes_and_shapes;
88 for (const auto& dtype_and_shape : proto.dtypes_and_shapes()) {
89 DataType dtype = dtype_and_shape.dtype();
90 PartialTensorShape shape;
91 Status s = PartialTensorShape::BuildPartialTensorShape(
92 dtype_and_shape.shape(), &shape);
93 if (!s.ok()) {
94 return s;
95 }
96 dtypes_and_shapes.push_back(DtypeAndPartialTensorShape{dtype, shape});
97 }
98 dtypes_and_shapes_ = std::move(dtypes_and_shapes);
99 return OkStatus();
100}
101
102string ResourceHandle::SerializeAsString() const {
103 ResourceHandleProto proto;
104 AsProto(&proto);
105 return proto.SerializeAsString();
106}
107
108bool ResourceHandle::ParseFromString(const string& s) {
109 ResourceHandleProto proto;
110 return proto.ParseFromString(s) && FromProto(proto).ok();
111}
112
113string ResourceHandle::DebugString() const {
114 return absl::StrFormat(
115 "device: %s container: %s name: %s hash_code: 0x%X maybe_type_name %s, "
116 "dtype and shapes : %s",
117 device(), container(), name(), hash_code(),
118 port::Demangle(maybe_type_name()),
119 DtypeAndShapesToString(dtypes_and_shapes()));
120}
121string ResourceHandle::SummarizeValue() const {
122 return absl::StrFormat(
123 "ResourceHandle(name=\"%s\", device=\"%s\", container=\"%s\", "
124 "type=\"%s\", dtype and shapes : \"%s\")",
125 name(), device(), container(), port::Demangle(maybe_type_name()),
126 DtypeAndShapesToString(dtypes_and_shapes()));
127}
128
129ResourceHandle ResourceHandle::MakeRefCountingHandle(
130 ResourceBase* resource, const string& device_name,
131 const TypeIndex& type_index,
132 const std::vector<DtypeAndPartialTensorShape>& dtypes_and_shapes,
133 const absl::optional<ManagedStackTrace>& definition_stack_trace) {
134 ResourceHandle result;
135 result.resource_.reset(resource, /*add_ref=*/false);
136 result.set_device(device_name);
137 // All resources owned by anonymous handles are put into the same container,
138 // and they get process-unique handle names.
139 result.set_container("Anonymous");
140 result.set_definition_stack_trace(definition_stack_trace);
141 result.set_name(
142 absl::StrFormat("Resource-%d-at-%p", GenerateUniqueId(), resource));
143 result.set_hash_code(type_index.hash_code());
144 result.set_maybe_type_name(type_index.name());
145 result.set_dtypes_and_shapes(dtypes_and_shapes);
146 return result;
147}
148
149Status ResourceHandle::ValidateType(const TypeIndex& type_index) const {
150 if (type_index.hash_code() != hash_code()) {
151 return errors::InvalidArgument(
152 "Trying to access a handle's resource using the wrong type. ",
153 "The handle points to a resource (name '", name(), "') of type '",
154 port::Demangle(maybe_type_name()), "' (hash code ", hash_code(),
155 ") but you are trying to access the resource as type '",
156 port::Demangle(type_index.name()), "' (hash code ",
157 type_index.hash_code(), ")");
158 }
159 return OkStatus();
160}
161
162std::atomic<int64_t> ResourceHandle::current_id_;
163
164int64_t ResourceHandle::GenerateUniqueId() { return current_id_.fetch_add(1); }
165
166string ProtoDebugString(const ResourceHandle& handle) {
167 return handle.DebugString();
168}
169
170void EncodeResourceHandleList(const ResourceHandle* p, int64_t n,
171 std::unique_ptr<port::StringListEncoder> e) {
172 ResourceHandleProto proto;
173 for (int i = 0; i < n; ++i) {
174 p[i].AsProto(&proto);
175 e->Append(proto);
176 }
177 e->Finalize();
178}
179
180bool DecodeResourceHandleList(std::unique_ptr<port::StringListDecoder> d,
181 ResourceHandle* ps, int64_t n) {
182 std::vector<uint32> sizes(n);
183 if (!d->ReadSizes(&sizes)) return false;
184
185 ResourceHandleProto proto;
186 for (int i = 0; i < n; ++i) {
187 if (!proto.ParseFromArray(d->Data(sizes[i]), sizes[i])) {
188 return false;
189 }
190 if (!ps[i].FromProto(proto).ok()) {
191 return false;
192 }
193 }
194 return true;
195}
196
197} // namespace tensorflow
198