1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
|
/* 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 "google/protobuf/any.pb.h"
#include "tensorflow/core/lib/core/status.h"
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_
|
