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
31namespace tvm {
32namespace runtime {
33
34/*!
35 * \brief Abstract channel interface used to create RPCEndpoint.
36 */
37class 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 */
62class 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