diff options
Diffstat (limited to 'tensorflow/contrib/gdr/gdr_memory_manager.h')
-rw-r--r-- | tensorflow/contrib/gdr/gdr_memory_manager.h | 63 |
1 files changed, 63 insertions, 0 deletions
diff --git a/tensorflow/contrib/gdr/gdr_memory_manager.h b/tensorflow/contrib/gdr/gdr_memory_manager.h new file mode 100644 index 0000000000..7e9fe01e97 --- /dev/null +++ b/tensorflow/contrib/gdr/gdr_memory_manager.h @@ -0,0 +1,63 @@ +/* Copyright 2017 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 GDR_MEMORY_MANAGER_H_ +#define GDR_MEMORY_MANAGER_H_ + +#include "tensorflow/core/lib/core/status.h" + +namespace google { +namespace protobuf { +class Any; +} +} + +namespace tensorflow { + +class Device; +class DeviceContext; +class Tensor; + +// Abstract interface that handles out-of-band tensor transport. +// +// The transport options are encoded into a protocol buffer and transmitted via +// some other communication channels like RPC. +// See RecvTensorRequest in tensorflow/core/protobuf/worker.proto +class RemoteMemoryManager { + public: + virtual ~RemoteMemoryManager() {} + virtual Status Init() = 0; + virtual void Run() = 0; + virtual void Stop() = 0; + + // Encodes the tensor information to an arbitrary protocol buffer + // The protocol buffer needs to be transmitted via some other channel + virtual Status TransportOptionsFromTensor( + ::google::protobuf::Any* mutable_transport_options, const Tensor& tensor, + Device* device, DeviceContext* device_context, bool on_host) = 0; + + // Retrieve the tensor from the encoded protocol buffer + // Note that the tensor has to be allocated, but not initialized + virtual Status TensorFromTransportOptions( + Tensor* tensor, const ::google::protobuf::Any& transport_options, + Device* device, DeviceContext* device_context, bool on_host) = 0; +}; + +RemoteMemoryManager* CreateRemoteMemoryManager(const string& host, + const string& port); + +} // namespace tensorflow + +#endif // GDR_MEMORY_MANAGER_H_ |