1 | /* |
2 | * Licensed to the Apache Software Foundation (ASF) under one |
3 | * or more contributor license agreements. See the NOTICE file |
4 | * distributed with this work for additional information |
5 | * regarding copyright ownership. The ASF licenses this file |
6 | * to you under the Apache License, Version 2.0 (the |
7 | * "License"); you may not use this file except in compliance |
8 | * with the License. You may obtain a copy of the License at |
9 | * |
10 | * http://www.apache.org/licenses/LICENSE-2.0 |
11 | * |
12 | * Unless required by applicable law or agreed to in writing, |
13 | * software distributed under the License is distributed on an |
14 | * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
15 | * KIND, either express or implied. See the License for the |
16 | * specific language governing permissions and limitations |
17 | * under the License. |
18 | */ |
19 | |
20 | /*! |
21 | * \file rpc_channel.h |
22 | * \brief Communication endpoints to connect local and remote RPC sessions. |
23 | */ |
24 | #ifndef TVM_RUNTIME_RPC_RPC_CHANNEL_H_ |
25 | #define TVM_RUNTIME_RPC_RPC_CHANNEL_H_ |
26 | |
27 | #include <tvm/runtime/packed_func.h> |
28 | |
29 | #include <utility> |
30 | |
31 | namespace tvm { |
32 | namespace runtime { |
33 | |
34 | /*! |
35 | * \brief Abstract channel interface used to create RPCEndpoint. |
36 | */ |
37 | class RPCChannel { |
38 | public: |
39 | /*! \brief virtual destructor */ |
40 | virtual ~RPCChannel() {} |
41 | /*! |
42 | * \brief Send data over to the channel. |
43 | * \param data The data pointer. |
44 | * \param size The size fo the data. |
45 | * \return The actual bytes sent. |
46 | */ |
47 | virtual size_t Send(const void* data, size_t size) = 0; |
48 | /*! |
49 | * \brief Recv data from channel. |
50 | * |
51 | * \param data The data pointer. |
52 | * \param size The size fo the data. |
53 | * \return The actual bytes received. |
54 | */ |
55 | virtual size_t Recv(void* data, size_t size) = 0; |
56 | }; |
57 | |
58 | /*! |
59 | * \brief RPC channel which callback |
60 | * frontend (Python/Java/etc.)'s send & recv function |
61 | */ |
62 | class CallbackChannel final : public RPCChannel { |
63 | public: |
64 | /*! |
65 | * \brief Constructor. |
66 | * |
67 | * \param fsend The send function, takes in a TVMByteArray and returns the |
68 | * number of bytes sent in that array. Returns -1 if error happens. |
69 | * \param frecv The recv function, takes an expected maximum size, and return |
70 | * a byte array with the actual amount of data received. |
71 | */ |
72 | explicit CallbackChannel(PackedFunc fsend, PackedFunc frecv) |
73 | : fsend_(std::move(fsend)), frecv_(std::move(frecv)) {} |
74 | |
75 | ~CallbackChannel() {} |
76 | /*! |
77 | * \brief Send data over to the channel. |
78 | * \param data The data pointer. |
79 | * \param size The size fo the data. |
80 | * \return The actual bytes sent. |
81 | */ |
82 | size_t Send(const void* data, size_t size) final; |
83 | /*! |
84 | * \brief Recv data from channel. |
85 | * |
86 | * \param data The data pointer. |
87 | * \param size The size fo the data. |
88 | * \return The actual bytes received. |
89 | */ |
90 | size_t Recv(void* data, size_t size) final; |
91 | |
92 | private: |
93 | PackedFunc fsend_; |
94 | PackedFunc frecv_; |
95 | }; |
96 | |
97 | } // namespace runtime |
98 | } // namespace tvm |
99 | #endif // TVM_RUNTIME_RPC_RPC_CHANNEL_H_ |
100 | |