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 #ifndef TENSORFLOW_FRAMEWORK_RESOURCE_HANDLE_H_ 17 #define TENSORFLOW_FRAMEWORK_RESOURCE_HANDLE_H_ 18 19 #include <string> 20 21 #include "tensorflow/core/framework/resource_base.h" 22 #include "tensorflow/core/framework/tensor_shape.h" 23 #include "tensorflow/core/framework/type_index.h" 24 #include "tensorflow/core/framework/types.pb.h" 25 #include "tensorflow/core/platform/casts.h" 26 #include "tensorflow/core/platform/intrusive_ptr.h" 27 #include "tensorflow/core/platform/statusor.h" 28 #include "tensorflow/core/platform/tensor_coding.h" 29 #include "tensorflow/core/platform/types.h" 30 #include "tensorflow/core/util/managed_stack_trace.h" 31 32 namespace tensorflow { 33 34 class ResourceHandleProto; 35 36 // Class representing a handle to a tensorflow resource. Handles are 37 // not valid across executions, but can be serialized back and forth from within 38 // a single run (except for those created from MakeRefCountingHandle i.e. whose 39 // resource_ field is not empty). 40 // 41 // This is the native C++ class equivalent of ResourceHandleProto. They are 42 // separate so that kernels do not need to depend on protos. 43 class ResourceHandle { 44 public: 45 ResourceHandle(); 46 ResourceHandle(const ResourceHandleProto& proto); 47 ~ResourceHandle(); 48 49 // Unique name for the device containing the resource. device()50 const std::string& device() const { return device_; } 51 set_device(const std::string & device)52 void set_device(const std::string& device) { device_ = device; } 53 54 // Container in which this resource is placed. container()55 const std::string& container() const { return container_; } set_container(const std::string & container)56 void set_container(const std::string& container) { container_ = container; } 57 58 // Unique name of this resource. name()59 const std::string& name() const { return name_; } set_name(const std::string & name)60 void set_name(const std::string& name) { name_ = name; } 61 62 // Hash code for the type of the resource. Is only valid in the same device 63 // and in the same execution. hash_code()64 uint64 hash_code() const { return hash_code_; } set_hash_code(uint64 hash_code)65 void set_hash_code(uint64 hash_code) { hash_code_ = hash_code; } 66 67 // For debug-only, the name of the type pointed to by this handle, if 68 // available. maybe_type_name()69 const std::string& maybe_type_name() const { return maybe_type_name_; } set_maybe_type_name(const std::string & value)70 void set_maybe_type_name(const std::string& value) { 71 maybe_type_name_ = value; 72 } 73 74 // Data types and shapes for the underlying resource. dtypes_and_shapes()75 std::vector<DtypeAndPartialTensorShape> dtypes_and_shapes() const { 76 return dtypes_and_shapes_; 77 } set_dtypes_and_shapes(const std::vector<DtypeAndPartialTensorShape> & dtypes_and_shapes)78 void set_dtypes_and_shapes( 79 const std::vector<DtypeAndPartialTensorShape>& dtypes_and_shapes) { 80 dtypes_and_shapes_ = dtypes_and_shapes; 81 } 82 set_definition_stack_trace(const absl::optional<ManagedStackTrace> & definition_stack_trace)83 void set_definition_stack_trace( 84 const absl::optional<ManagedStackTrace>& definition_stack_trace) { 85 definition_stack_trace_ = definition_stack_trace; 86 } 87 definition_stack_trace()88 const absl::optional<ManagedStackTrace>& definition_stack_trace() const { 89 return definition_stack_trace_; 90 } 91 92 // Conversion to and from ResourceHandleProto 93 void AsProto(ResourceHandleProto* proto) const; 94 void FromProto(const ResourceHandleProto& proto); 95 96 // Serialization via ResourceHandleProto 97 std::string SerializeAsString() const; 98 bool ParseFromString(const std::string& s); 99 100 std::string DebugString() const; 101 SummarizeValue()102 std::string SummarizeValue() const { return "Resource Tensor"; } 103 104 // GUID for anonymous resources. Resources with this shared_name will have 105 // their shared_name replaced with a GUID at creation time 106 static constexpr const char* ANONYMOUS_NAME = 107 "cd2c89b7-88b7-44c8-ad83-06c2a9158347"; 108 109 // Creates a `ResourceHandle` that holds a pointer to a resource and takes 110 // ownership of it. Normally a `ResourceHandle` only contains the name (and 111 // some other metadata) of the resource. When created via this function, 112 // the handle will own the resource, in the sense that it will destroy the 113 // resource automatically when the resource is no longer needed. It does this 114 // via automatic ref-counting on the resource: when the handle is copied, it 115 // will call `Ref` on the resource (remember that all resources inherit from 116 // `ResourceBase` which inherits from `RefCounted`), and when the handle is 117 // destroyed, it will call `Unref` on the resource. When the last handle goes 118 // out of scope, the resource's ref-count will go down to zero and the 119 // resource will be destroyed. When calling this function, the `resource` 120 // argument should have a ref-count of one (which is the case when the 121 // resource is newly created). 122 // 123 // For those familiar with `ResourceMgr`, when you create a handle by the 124 // `MakeResourceHandle` function in resource_mgr.h, the handle doesn't hold a 125 // pointer to the resource, and the resource is owned by the resource manager 126 // and needs to be manually deleted by calling `ResourceMgr::Delete`. In 127 // contrast, a handle created by this function doesn't interact with 128 // `ResourceMgr` at all, and the resource shouldn't be registered in any 129 // resource manager (otherwise the resource manager will prevent the resource 130 // from being automatically deleted). 131 template <typename T> MakeRefCountingHandle(T * resource,const string & device_name,const std::vector<DtypeAndPartialTensorShape> & dtypes_and_shapes,const absl::optional<ManagedStackTrace> & definition_stack_trace)132 static ResourceHandle MakeRefCountingHandle( 133 T* resource, const string& device_name, 134 const std::vector<DtypeAndPartialTensorShape>& dtypes_and_shapes, 135 const absl::optional<ManagedStackTrace>& definition_stack_trace) { 136 return MakeRefCountingHandle(resource, device_name, TypeIndex::Make<T>(), 137 dtypes_and_shapes, definition_stack_trace); 138 } 139 140 static ResourceHandle MakeRefCountingHandle( 141 ResourceBase* resource, const string& device_name, 142 const TypeIndex& type_index, 143 const std::vector<DtypeAndPartialTensorShape>& dtypes_and_shapes, 144 const absl::optional<ManagedStackTrace>& definition_stack_trace); 145 146 // Pointer to the resource. resource()147 const core::IntrusivePtr<ResourceBase>& resource() const { return resource_; } 148 149 // Gets the resource pointer in `handle` as `T*`, or an error if the actual 150 // resource type is not `T`. 151 template <typename T> GetResource()152 StatusOr<T*> GetResource() const { 153 TF_RETURN_IF_ERROR(ValidateType<T>()); 154 return down_cast<T*>(resource_.get()); 155 } 156 157 // Validates that the resource type in `handle` is `T`. 158 template <typename T> ValidateType()159 Status ValidateType() const { 160 return ValidateType(TypeIndex::Make<T>()); 161 } 162 163 Status ValidateType(const TypeIndex& type_index) const; 164 165 // Generates unique IDs (e.g. for names of anonymous variables) 166 static int64 GenerateUniqueId(); 167 168 private: 169 std::string device_; 170 std::string container_; 171 std::string name_; 172 uint64 hash_code_ = 0; 173 std::string maybe_type_name_; 174 std::vector<DtypeAndPartialTensorShape> dtypes_and_shapes_; 175 absl::optional<ManagedStackTrace> definition_stack_trace_; 176 177 private: 178 // A smart pointer to the actual resource. When this field is not empty, the 179 // handle is in a "ref-counting" mode, owning the resource; otherwise it's in 180 // a "weak-ref" mode, only containing the name of the resource (conceptually a 181 // weak reference). 182 core::IntrusivePtr<ResourceBase> resource_; 183 static std::atomic<int64> current_id_; 184 }; 185 186 // For backwards compatibility for when this was a proto 187 std::string ProtoDebugString(const ResourceHandle& handle); 188 189 // Encodes a list of ResourceHandle protos in the given StringListEncoder. 190 void EncodeResourceHandleList(const ResourceHandle* p, int64_t n, 191 std::unique_ptr<port::StringListEncoder> e); 192 193 // Decodes a list of ResourceHandle protos from the given StringListDecoder. 194 bool DecodeResourceHandleList(std::unique_ptr<port::StringListDecoder> d, 195 ResourceHandle* ps, int64_t n); 196 197 } // namespace tensorflow 198 199 #endif // TENSORFLOW_FRAMEWORK_RESOURCE_HANDLE_H_ 200