[Hexagon] Adjust RPC read buffer size from python

python/tvm/contrib/hexagon/session.py

@@ -57,10 +57,12 @@ def __init__(
         remote_kw: dict,
         session_name: str = "hexagon-rpc",
         remote_stack_size_bytes: int = 128 * 1024,
+        rpc_receive_buffer_size_bytes: int = 2 * 1024 * 1024,
     ):
         self._launcher = launcher
         self._session_name = session_name
         self._remote_stack_size_bytes = remote_stack_size_bytes
+        self._rpc_receive_buffer_size_bytes = rpc_receive_buffer_size_bytes
         self._remote_kw = remote_kw
         self._rpc = None
         self.device = None
@@ -81,6 +83,7 @@ def __enter__(self):
                     self._session_name,
                     self._remote_stack_size_bytes,
                     os.environ.get("HEXAGON_SIM_ARGS", ""),
+                    self._rpc_receive_buffer_size_bytes,
                 ],
             )
             self.device = self._rpc.hexagon(0)

src/runtime/hexagon/rpc/android/session.cc

@@ -45,13 +45,19 @@ namespace hexagon {
 
 class HexagonTransportChannel : public RPCChannel {
  public:
-  explicit HexagonTransportChannel(const std::string& uri, int remote_stack_size_bytes) {
+  explicit HexagonTransportChannel(const std::string& uri, int remote_stack_size_bytes,
+                                   uint32_t receive_buf_size_bytes) {
     if (_handle != AEE_EUNKNOWN) return;
 
     enable_unsigned_pd(true);
     set_remote_stack_size(remote_stack_size_bytes);
+
     AEEResult rc = hexagon_rpc_open(uri.c_str(), &_handle);
     ICHECK(rc == AEE_SUCCESS) << "hexagon_rpc_open failed. URI: " << uri.c_str();
+
+    rc = hexagon_rpc_init(_handle, receive_buf_size_bytes);
+    ICHECK(rc == AEE_SUCCESS) << "hexagon_rpc_set_receive_buf_size failed. receive_buf_size_bytes: "
+                              << receive_buf_size_bytes;
   }
 
   size_t Send(const void* data, size_t size) override {
@@ -105,10 +111,15 @@ class HexagonTransportChannel : public RPCChannel {
 
 TVM_REGISTER_GLOBAL("tvm.contrib.hexagon.create_hexagon_session")
     .set_body([](TVMArgs args, TVMRetValue* rv) {
+      ICHECK(args.size() >= 4) << args.size() << " is less than 4";
+
       std::string session_name = args[0];
       int remote_stack_size_bytes = args[1];
+      // For simulator, the third parameter is sim_args, ignore it.
+      int hexagon_rpc_receive_buf_size_bytes = args[3];
       HexagonTransportChannel* hexagon_channel =
-          new HexagonTransportChannel(hexagon_rpc_URI CDSP_DOMAIN, remote_stack_size_bytes);
+          new HexagonTransportChannel(hexagon_rpc_URI CDSP_DOMAIN, remote_stack_size_bytes,
+                                      static_cast<uint32_t>(hexagon_rpc_receive_buf_size_bytes));
       std::unique_ptr<RPCChannel> channel(hexagon_channel);
       auto ep = RPCEndpoint::Create(std::move(channel), session_name, "", NULL);
       auto sess = CreateClientSession(ep);

src/runtime/hexagon/rpc/hexagon/rpc_server.cc

@@ -40,8 +40,7 @@ extern "C" {
 #include "../../hexagon/hexagon_common.h"
 #include "hexagon_rpc.h"
 
-// TODO(mehrdadh): make this configurable.
-#define TVM_HEXAGON_RPC_BUFF_SIZE_BYTES 2 * 1024 * 1024
+#define TVM_HEXAGON_RPC_BUFF_DEFAULT_SIZE_BYTES 2 * 1024 * 1024
 
 // TODO(csulivan,adstraw,kparzysz-quic) This should be set on a TVM-wide basis.
 #if defined(__hexagon__)
@@ -196,10 +195,16 @@ class HexagonRPCServer {
 }  // namespace tvm
 
 namespace {
-tvm::runtime::hexagon::HexagonRPCServer* get_hexagon_rpc_server() {
-  static tvm::runtime::hexagon::HexagonRPCServer g_hexagon_rpc_server(
-      new uint8_t[TVM_HEXAGON_RPC_BUFF_SIZE_BYTES], TVM_HEXAGON_RPC_BUFF_SIZE_BYTES);
-  return &g_hexagon_rpc_server;
+static tvm::runtime::hexagon::HexagonRPCServer* g_hexagon_rpc_server;
+tvm::runtime::hexagon::HexagonRPCServer* get_hexagon_rpc_server(
+    uint32_t rpc_receive_buff_size_bytes = TVM_HEXAGON_RPC_BUFF_DEFAULT_SIZE_BYTES) {
+  if (g_hexagon_rpc_server) {
+    return g_hexagon_rpc_server;
+  }
+  static tvm::runtime::hexagon::HexagonRPCServer hexagon_rpc_server(
+      new uint8_t[rpc_receive_buff_size_bytes], rpc_receive_buff_size_bytes);
+  g_hexagon_rpc_server = &hexagon_rpc_server;
+  return g_hexagon_rpc_server;
 }
 }  // namespace
 
@@ -222,7 +227,6 @@ int __QAIC_HEADER(hexagon_rpc_open)(const char* uri, remote_handle64* handle) {
     return AEE_ENOMEMORY;
   }
   reset_device_api();
-  get_hexagon_rpc_server();
 
   return AEE_SUCCESS;
 }
@@ -235,6 +239,11 @@ int __QAIC_HEADER(hexagon_rpc_close)(remote_handle64 handle) {
   return AEE_SUCCESS;
 }
 
+int __QAIC_HEADER(hexagon_rpc_init)(remote_handle64 _h, uint32_t buff_size_bytes) {
+  get_hexagon_rpc_server(buff_size_bytes);
+  return AEE_SUCCESS;
+}
+
 /*!
  * \brief Send data from Host to Hexagon over RPCSession.
  * \param _handle The remote handle

src/runtime/hexagon/rpc/hexagon_rpc.idl

@@ -25,4 +25,5 @@ typedef sequence<octet> buffer;
 interface hexagon_rpc : remote_handle64 {
   AEEResult send(in buffer data);
   AEEResult receive(rout buffer buf, rout int64_t buf_written_size);
+  AEEResult init(in uint32_t buff_size_bytes);
 };

src/runtime/hexagon/rpc/simulator/session.cc

@@ -1311,6 +1311,8 @@ detail::Optional<HEXAPI_Nullptr> SimulatorRPCChannel::to_nullptr(const detail::M
 
 TVM_REGISTER_GLOBAL("tvm.contrib.hexagon.create_hexagon_session")
     .set_body([](TVMArgs args, TVMRetValue* rv) {
+      ICHECK(args.size() >= 4) << args.size() << " is less than 4";
+
       std::string session_name = args[0];
       // For target, the second parameter is remote_stack_size_bytes, ignore it.
       std::string sim_args = args[2];