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
59
60
61
62
63
64
65
66
|
#pragma once
#ifdef USE_TENSORPIPE
#include <torch/csrc/distributed/rpc/utils.h>
namespace tensorpipe {
class Message;
} // namespace tensorpipe
namespace torch {
namespace distributed {
namespace rpc {
// A struct that holds pointers that keep alive all the memory that will be
// accessed by TensorPipe during a write operation.
struct TensorpipeWriteBuffers {
// Allocate on heap so pointers stay valid as we move the holder.
std::unique_ptr<MessageType> type;
std::unique_ptr<int64_t> id;
std::vector<char> payload;
std::vector<char> pickle;
// This contains the original tensors and the clones of the sparse tensors.
std::vector<torch::Tensor> tensors;
// This contains the copies of the data of the tensors that didn't own their
// memory, e.g., the ones created from torch::from_blob() with no deleter.
std::vector<std::vector<char>> copiedTensors;
};
// A struct that holds pointers that keep alive all the memory that will be
// accessed by TensorPipe during a read operation.
struct TensorpipeReadBuffers {
// Allocate on heap so pointers stay valid as we move the holder.
std::unique_ptr<MessageType> type;
std::unique_ptr<int64_t> id;
std::vector<char> payload;
std::vector<char> pickle;
std::vector<c10::DataPtr> tensors;
};
// Convert an RPC message into a TensorPipe message, plus a holder to all the
// data that must be kept alive while the write is performed asynchronously.
TORCH_API std::tuple<tensorpipe::Message, TensorpipeWriteBuffers>
tensorpipeSerialize(
Message&& rpcMessage,
std::vector<c10::DeviceIndex> devices = {});
// Allocate the buffers that will hold the incoming data. They will be managed
// by the returned holder, which must be kept alive until the asynchronous read
// has finished. Pointers to these buffers will be stored in-place in the
// TensorPipe message.
TORCH_API TensorpipeReadBuffers
tensorpipeAllocate(tensorpipe::Message& tpMessage);
// Convert a TensorPipe message back into an RPC message. This requires the data
// to be available and can thus only be performed once the asynchronous read has
// completed. The holder can be destroyed once this function returns.
TORCH_API Message tensorpipeDeserialize(
tensorpipe::Message&& tpMessage,
TensorpipeReadBuffers&& holder);
} // namespace rpc
} // namespace distributed
} // namespace torch
#endif // USE_TENSORPIPE
|
