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 <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 | |
31 | namespace tensorflow { |
32 | |
33 | namespace { |
34 | std::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 | |
52 | ResourceHandle::ResourceHandle() {} |
53 | |
54 | ResourceHandle::ResourceHandle(const ResourceHandleProto& proto) { |
55 | TF_CHECK_OK(FromProto(proto)); |
56 | } |
57 | |
58 | Status 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 | |
66 | ResourceHandle::~ResourceHandle() {} |
67 | |
68 | void 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 | |
81 | Status 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 | |
102 | string ResourceHandle::SerializeAsString() const { |
103 | ResourceHandleProto proto; |
104 | AsProto(&proto); |
105 | return proto.SerializeAsString(); |
106 | } |
107 | |
108 | bool ResourceHandle::ParseFromString(const string& s) { |
109 | ResourceHandleProto proto; |
110 | return proto.ParseFromString(s) && FromProto(proto).ok(); |
111 | } |
112 | |
113 | string 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 | } |
121 | string 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 | |
129 | ResourceHandle 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 | |
149 | Status 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 | |
162 | std::atomic<int64_t> ResourceHandle::current_id_; |
163 | |
164 | int64_t ResourceHandle::GenerateUniqueId() { return current_id_.fetch_add(1); } |
165 | |
166 | string ProtoDebugString(const ResourceHandle& handle) { |
167 | return handle.DebugString(); |
168 | } |
169 | |
170 | void 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 | |
180 | bool 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 | |