| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you 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. |
| */ |
| |
| /*! |
| * \file rpc_reference.h |
| * \brief Common header defining the communication code used in the RPC protocol. |
| */ |
| #ifndef TVM_RUNTIME_MINRPC_RPC_REFERENCE_H_ |
| #define TVM_RUNTIME_MINRPC_RPC_REFERENCE_H_ |
| |
| #include <tvm/ffi/container/tensor.h> |
| |
| namespace tvm { |
| namespace ffi { |
| // Forward declare TVM Object to use `Object*` in RPC protocol. |
| class Object; |
| } // namespace ffi |
| |
| namespace runtime { |
| |
| /*! \brief The current RPC procotol version. */ |
| constexpr const char* kRPCProtocolVer = "0.8.0"; |
| |
| // When tvm.rpc.server.GetCRTMaxPacketSize global function is not registered. |
| const uint64_t kRPCMaxTransferSizeBytesDefault = UINT64_MAX; |
| |
| /*! \brief The RPC code */ |
| enum class RPCCode : int { |
| kNone, |
| kShutdown, |
| kInitServer, |
| kCallFunc, |
| kReturn, |
| kException, |
| kCopyFromRemote, |
| kCopyToRemote, |
| kCopyAck, |
| // The following are syscall code that can send over CallRemote |
| kSyscallCodeStart, |
| kGetGlobalFunc = kSyscallCodeStart, |
| kFreeHandle, |
| kDevSetDevice, |
| kDevGetAttr, |
| kDevAllocData, |
| kDevFreeData, |
| kDevStreamSync, |
| kCopyAmongRemote, |
| kDevAllocDataWithScope, |
| kDevCreateStream, |
| kDevFreeStream, |
| kDevSetStream, |
| kDevGetCurrentStream, |
| }; |
| |
| /*! |
| * \brief List of potential error status during rpc communication. |
| */ |
| enum class RPCServerStatus : int { |
| kSuccess = 0, |
| kInvalidTypeCodeObject, |
| kInvalidTypeCodeTensor, |
| kInvalidDLTensorFieldStride, |
| kInvalidDLTensorFieldByteOffset, |
| kUnknownTypeIndex, |
| kUnknownRPCCode, |
| kRPCCodeNotSupported, |
| kUnknownRPCSyscall, |
| kCheckError, |
| kReadError, |
| kWriteError, |
| kAllocError |
| }; |
| |
| inline const char* RPCCodeToString(RPCCode code) { |
| switch (code) { |
| case RPCCode::kShutdown: |
| return "kShutdown"; |
| case RPCCode::kInitServer: |
| return "kInitServer"; |
| case RPCCode::kCallFunc: |
| return "kCallFunc"; |
| case RPCCode::kReturn: |
| return "kReturn"; |
| case RPCCode::kException: |
| return "kException"; |
| case RPCCode::kCopyFromRemote: |
| return "kCopyFromRemote"; |
| case RPCCode::kCopyToRemote: |
| return "kCopyToRemote"; |
| case RPCCode::kCopyAck: |
| return "kCopyAck"; |
| // The following are syscall code that can send over CallRemote |
| case RPCCode::kGetGlobalFunc: |
| return "kGetGlobalFunc"; |
| case RPCCode::kFreeHandle: |
| return "kFreeHandle"; |
| case RPCCode::kDevSetDevice: |
| return "kDevSetDevice"; |
| case RPCCode::kDevGetAttr: |
| return "kDevGetAttr"; |
| case RPCCode::kDevAllocData: |
| return "kDevAllocData"; |
| case RPCCode::kDevFreeData: |
| return "kDevFreeData"; |
| case RPCCode::kDevCreateStream: |
| return "kDevCreateStream"; |
| case RPCCode::kDevFreeStream: |
| return "kDevFreeStream"; |
| case RPCCode::kDevStreamSync: |
| return "kDevStreamSync"; |
| case RPCCode::kDevSetStream: |
| return "kDevSetStream"; |
| case RPCCode::kCopyAmongRemote: |
| return "kCopyAmongRemote"; |
| case RPCCode::kDevAllocDataWithScope: |
| return "kDevAllocDataWithScope"; |
| default: |
| return ""; |
| } |
| } |
| |
| /*! |
| * \brief Convert RPC server status to string. |
| * \param status The status. |
| * \return The corresponding string. |
| */ |
| inline const char* RPCServerStatusToString(RPCServerStatus status) { |
| switch (status) { |
| case RPCServerStatus::kSuccess: |
| return "kSuccess"; |
| case RPCServerStatus::kInvalidTypeCodeObject: |
| return "kInvalidTypeCodeObject"; |
| case RPCServerStatus::kInvalidTypeCodeTensor: |
| return "kInvalidTypeCodeTensor"; |
| case RPCServerStatus::kInvalidDLTensorFieldStride: |
| return "kInvalidDLTensorFieldStride"; |
| case RPCServerStatus::kInvalidDLTensorFieldByteOffset: { |
| return "kInvalidDLTensorFieldByteOffset"; |
| } |
| case RPCServerStatus::kUnknownTypeIndex: |
| return "kUnknownTypeIndex"; |
| case RPCServerStatus::kUnknownRPCCode: |
| return "kUnknownRPCCode"; |
| case RPCServerStatus::kRPCCodeNotSupported: |
| return "RPCCodeNotSupported"; |
| case RPCServerStatus::kUnknownRPCSyscall: |
| return "kUnknownRPCSyscall"; |
| case RPCServerStatus::kCheckError: |
| return "kCheckError"; |
| case RPCServerStatus::kReadError: |
| return "kReadError"; |
| case RPCServerStatus::kWriteError: |
| return "kWriteError"; |
| case RPCServerStatus::kAllocError: |
| return "kAllocError"; |
| default: |
| return ""; |
| } |
| } |
| |
| /*! |
| * \brief Reference implementation of the communication protocol. |
| * |
| * \note The implementation is intentionally written via template |
| * so it can be used in a dependency free setting. |
| * |
| * \sa src/runtime/rpc/device/min_rpc_server.h |
| */ |
| struct RPCReference { |
| /*! |
| * \brief Auxiliary class to get the packed sequence. |
| * \tparam TChannel The channel to throw errror. |
| */ |
| template <typename TChannel> |
| struct PackedSeqNumBytesGetter { |
| public: |
| explicit PackedSeqNumBytesGetter(TChannel* channel) : channel_(channel) {} |
| |
| template <typename T> |
| void Write(const T& value) { |
| num_bytes_ += sizeof(T); |
| } |
| |
| template <typename T> |
| void WriteArray(const T* value, size_t num) { |
| num_bytes_ += sizeof(T) * num; |
| } |
| |
| void WriteFFIAny(const TVMFFIAny* obj) { num_bytes_ += channel_->GetFFIAnyProtocolBytes(obj); } |
| |
| void ThrowError(RPCServerStatus status) { channel_->ThrowError(status); } |
| |
| uint64_t num_bytes() const { return num_bytes_; } |
| |
| private: |
| TChannel* channel_; |
| uint64_t num_bytes_{0}; |
| }; |
| |
| /*! |
| * \return the length of the str. |
| * \param str the string. |
| * \return The length. |
| */ |
| static uint64_t StrLength(const char* str) { |
| uint64_t len = 0; |
| while (str[len] != '\0') ++len; |
| return len; |
| } |
| |
| /*! |
| * \brief Get the total nbytes to be sent in the packed sequence. |
| * |
| * \param arg_values The values to be sent over. |
| * \param type_codes The type codes to be sent over. |
| * \param num_args Number of argument. |
| * \param client_mode Whether it is a client to server call. |
| * \param channel The communication channel handler. |
| * \tparam TChannel The type of the communication channel. |
| * \return The total number of bytes. |
| */ |
| template <typename TChannel> |
| static uint64_t PackedSeqGetNumBytes(const TVMFFIAny* packed_args, int num_args, bool client_mode, |
| TChannel* channel) { |
| PackedSeqNumBytesGetter<TChannel> getter(channel); |
| SendPackedSeq(packed_args, num_args, client_mode, &getter); |
| return getter.num_bytes(); |
| } |
| |
| template <typename TChannelPtr> |
| static void SendDLTensor(TChannelPtr channel, DLTensor* arr) { |
| DLDevice dev; |
| uint64_t data; |
| // When we return Tensor, we directly return |
| // the space and the context |
| // The client will be further wrapping |
| dev = arr->device; |
| data = reinterpret_cast<uint64_t>(arr->data); |
| channel->Write(data); |
| channel->Write(dev); |
| channel->Write(arr->ndim); |
| channel->Write(arr->dtype); |
| channel->WriteArray(arr->shape, arr->ndim); |
| if (!ffi::IsContiguous(*arr)) { |
| channel->ThrowError(RPCServerStatus::kInvalidDLTensorFieldStride); |
| } |
| channel->Write(arr->byte_offset); |
| return; |
| } |
| |
| template <typename TChannelPtr> |
| static DLTensor* ReceiveDLTensor(TChannelPtr channel) { |
| uint64_t handle; |
| channel->Read(&handle); |
| DLTensor* arr = channel->template ArenaAlloc<DLTensor>(1); |
| DLTensor& tensor = *arr; |
| tensor.data = reinterpret_cast<void*>(handle); |
| channel->Read(&(tensor.device)); |
| channel->Read(&(tensor.ndim)); |
| channel->Read(&(tensor.dtype)); |
| tensor.shape = channel->template ArenaAlloc<int64_t>(tensor.ndim); |
| channel->ReadArray(tensor.shape, tensor.ndim); |
| tensor.strides = nullptr; |
| channel->Read(&(tensor.byte_offset)); |
| return arr; |
| } |
| |
| /*! |
| * \brief Send packed argument sequnce to the other peer. |
| * |
| * This function serves as the foundational communication primitive between peers. |
| * |
| * TVMValue sequence encoding protocol(according to the type): |
| * |
| * - int/float/uint/bytes/str: Serialize all content. |
| * - DLTensor: send meta-data, send data handle as opaque handle(via uint64_t) |
| * - OpaqueHandle: send as uint64_t |
| * - ModuleHandle, PackedFuncHandle: send as uint64_t, |
| * The support to Module/PackedFuncHandle are reserved for arguments |
| * in the CallFunc from a client to server only. |
| * Note that we cannot simply take these argument out(as the handle) |
| * refers to a value on the remote(instead of local). |
| * |
| * \param packed_args The values to be sent over. |
| * \param num_args Number of argument. |
| * \param client_mode Whether it is a client to server call. |
| * \param channel The communication channel handler. |
| * \tparam TChannel The type of the communication channel. |
| */ |
| template <typename TChannel> |
| static void SendPackedSeq(const TVMFFIAny* packed_args, int num_args, bool client_mode, |
| TChannel* channel) { |
| channel->Write(num_args); |
| |
| // Argument packing. |
| for (int i = 0; i < num_args; ++i) { |
| int32_t type_index = packed_args[i].type_index; |
| channel->template Write<int32_t>(type_index); |
| switch (type_index) { |
| case ffi::TypeIndex::kTVMFFINone: { |
| break; |
| } |
| case ffi::TypeIndex::kTVMFFIBool: |
| case ffi::TypeIndex::kTVMFFIInt: |
| case ffi::TypeIndex::kTVMFFIFloat: { |
| channel->template Write<int64_t>(packed_args[i].v_int64); |
| break; |
| } |
| case ffi::TypeIndex::kTVMFFIOpaquePtr: { |
| // always send handle in 64 bit. |
| uint64_t handle = reinterpret_cast<uint64_t>(packed_args[i].v_ptr); |
| channel->template Write<int64_t>(handle); |
| break; |
| } |
| case ffi::TypeIndex::kTVMFFIDataType: { |
| channel->Write(packed_args[i].v_dtype); |
| // padding |
| int32_t padding = 0; |
| channel->template Write<int32_t>(padding); |
| break; |
| } |
| case ffi::TypeIndex::kTVMFFIDevice: { |
| channel->Write(packed_args[i].v_device); |
| break; |
| } |
| |
| case ffi::TypeIndex::kTVMFFIFunction: |
| case ffi::TypeIndex::kTVMFFIModule: { |
| if (!client_mode) { |
| channel->ThrowError(RPCServerStatus::kInvalidTypeCodeObject); |
| } |
| // always send handle in 64 bit. |
| uint64_t handle = reinterpret_cast<uint64_t>(packed_args[i].v_obj); |
| channel->Write(handle); |
| break; |
| } |
| |
| case ffi::TypeIndex::kTVMFFITensor: { |
| channel->ThrowError(RPCServerStatus::kInvalidTypeCodeTensor); |
| break; |
| } |
| case ffi::TypeIndex::kTVMFFIDLTensorPtr: { |
| DLTensor* arr = static_cast<DLTensor*>(packed_args[i].v_ptr); |
| SendDLTensor(channel, arr); |
| break; |
| } |
| case ffi::TypeIndex::kTVMFFIRawStr: { |
| const char* s = packed_args[i].v_c_str; |
| uint64_t len = StrLength(s); |
| channel->Write(len); |
| channel->WriteArray(s, len); |
| break; |
| } |
| case ffi::TypeIndex::kTVMFFIByteArrayPtr: { |
| TVMFFIByteArray* bytes = static_cast<TVMFFIByteArray*>(packed_args[i].v_ptr); |
| uint64_t len = bytes->size; |
| channel->Write(len); |
| channel->WriteArray(bytes->data, len); |
| break; |
| } |
| default: { |
| channel->WriteFFIAny(&(packed_args[i])); |
| break; |
| } |
| } |
| } |
| } |
| |
| /*! |
| * \brief Receive packed seq from the channel. |
| * |
| * \param out_packed_args The values to be received. |
| * \param out_num_args Number of argument. |
| * \param channel The communication channel handler. |
| * \tparam TChannel The type of the communication channel. |
| * \note The temporary space are populated via an arena inside channel. |
| */ |
| template <typename TChannel> |
| static void RecvPackedSeq(TVMFFIAny** out_packed_args, int32_t* out_num_args, TChannel* channel) { |
| // receive number of args |
| int32_t num_args; |
| channel->Read(&num_args); |
| *out_num_args = num_args; |
| if (num_args == 0) { |
| *out_packed_args = nullptr; |
| return; |
| } |
| |
| TVMFFIAny* packed_args = channel->template ArenaAlloc<TVMFFIAny>(num_args); |
| *out_packed_args = packed_args; |
| |
| // receive arguments |
| for (int32_t i = 0; i < num_args; ++i) { |
| int32_t type_index; |
| channel->Read(&type_index); |
| packed_args[i].type_index = type_index; |
| packed_args[i].zero_padding = 0; |
| // clear to ensure compact for 32 bit platform |
| packed_args[i].v_int64 = 0; |
| switch (type_index) { |
| case ffi::TypeIndex::kTVMFFINone: { |
| break; |
| } |
| case ffi::TypeIndex::kTVMFFIBool: |
| case ffi::TypeIndex::kTVMFFIInt: |
| case ffi::TypeIndex::kTVMFFIFloat: { |
| channel->template Read<int64_t>(&(packed_args[i].v_int64)); |
| break; |
| } |
| case ffi::TypeIndex::kTVMFFIOpaquePtr: { |
| uint64_t handle; |
| channel->Read(&handle); |
| packed_args[i].v_ptr = reinterpret_cast<void*>(handle); |
| break; |
| } |
| case ffi::TypeIndex::kTVMFFIDataType: { |
| channel->Read(&(packed_args[i].v_dtype)); |
| int32_t padding = 0; |
| channel->template Read<int32_t>(&padding); |
| break; |
| } |
| case ffi::TypeIndex::kTVMFFIDevice: { |
| channel->Read(&(packed_args[i].v_device)); |
| break; |
| } |
| case ffi::TypeIndex::kTVMFFIFunction: |
| case ffi::TypeIndex::kTVMFFIModule: { |
| // always send handle in 64 bit. |
| uint64_t handle; |
| channel->Read(&handle); |
| packed_args[i].v_obj = reinterpret_cast<TVMFFIObject*>(handle); |
| break; |
| } |
| case ffi::TypeIndex::kTVMFFIRawStr: { |
| uint64_t len; |
| channel->Read(&len); |
| char* str = channel->template ArenaAlloc<char>(len + 1); |
| str[len] = '\0'; |
| channel->ReadArray(str, len); |
| packed_args[i].v_c_str = str; |
| break; |
| } |
| case ffi::TypeIndex::kTVMFFIByteArrayPtr: { |
| uint64_t len; |
| channel->Read(&len); |
| TVMFFIByteArray* arr = channel->template ArenaAlloc<TVMFFIByteArray>(1); |
| char* data = channel->template ArenaAlloc<char>(len); |
| arr->size = len; |
| arr->data = data; |
| channel->ReadArray(data, len); |
| packed_args[i].v_ptr = arr; |
| break; |
| } |
| case ffi::TypeIndex::kTVMFFIDLTensorPtr: { |
| packed_args[i].v_ptr = ReceiveDLTensor(channel); |
| break; |
| } |
| default: { |
| if (type_index >= ffi::TypeIndex::kTVMFFIStaticObjectBegin || |
| type_index == ffi::TypeIndex::kTVMFFISmallStr || |
| type_index == ffi::TypeIndex::kTVMFFISmallBytes) { |
| channel->ReadFFIAny(&(packed_args[i])); |
| } else { |
| channel->ThrowError(RPCServerStatus::kUnknownTypeIndex); |
| } |
| break; |
| } |
| } |
| } |
| } |
| |
| /*! |
| * \brief Return an exception packet. |
| * |
| * \param msg The error message. |
| * \param channel The communication channel handler. |
| * \tparam TChannel The type of the communication channel. |
| */ |
| template <typename TChannel> |
| static void ReturnException(const char* msg, TChannel* channel) { |
| RPCCode code = RPCCode::kException; |
| int32_t num_args = 1; |
| int32_t type_index = ffi::TypeIndex::kTVMFFIRawStr; |
| uint64_t len = StrLength(msg); |
| |
| uint64_t packet_nbytes = |
| sizeof(code) + sizeof(num_args) + sizeof(type_index) + sizeof(len) + len; |
| |
| channel->MessageStart(packet_nbytes); |
| channel->Write(packet_nbytes); |
| channel->Write(code); |
| channel->Write(num_args); |
| channel->Write(type_index); |
| channel->Write(len); |
| channel->WriteArray(msg, len); |
| channel->MessageDone(); |
| } |
| |
| /*! |
| * \brief Return a normal packed sequence packet. |
| * |
| * \param msg The error message. |
| * \param channel The communication channel handler. |
| * \tparam TChannel The type of the communication channel. |
| */ |
| template <typename TChannel> |
| static void ReturnPackedSeq(const TVMFFIAny* packed_args, int num_args, TChannel* channel) { |
| RPCCode code = RPCCode::kReturn; |
| |
| uint64_t packet_nbytes = |
| sizeof(code) + PackedSeqGetNumBytes(packed_args, num_args, false, channel); |
| |
| channel->MessageStart(packet_nbytes); |
| channel->Write(packet_nbytes); |
| channel->Write(code); |
| SendPackedSeq(packed_args, num_args, false, channel); |
| channel->MessageDone(); |
| } |
| |
| /*! |
| * \brief Return a null(void) packet. |
| * |
| * \param channel The communication channel handler. |
| * \tparam TChannel The type of the communication channel. |
| */ |
| template <typename TChannel> |
| static void ReturnVoid(TChannel* channel) { |
| int32_t num_args = 1; |
| int32_t type_index = ffi::TypeIndex::kTVMFFINone; |
| RPCCode code = RPCCode::kReturn; |
| |
| uint64_t packet_nbytes = sizeof(code) + sizeof(num_args) + sizeof(type_index); |
| |
| channel->MessageStart(packet_nbytes); |
| channel->Write(packet_nbytes); |
| channel->Write(code); |
| channel->Write(num_args); |
| channel->Write(type_index); |
| channel->MessageDone(); |
| } |
| }; |
| |
| } // namespace runtime |
| } // namespace tvm |
| |
| #endif // TVM_RUNTIME_MINRPC_RPC_REFERENCE_H_ |
