[libc-commits] [libc] 62f57bc - [libc] Add other RPC callback methods to the RPC server

[Joseph Huber via libc-commits]

Author: Joseph Huber
Date: 2023-06-30T11:29:37-05:00
New Revision: 62f57bc9b0b03cdd8a2647e4960d7835b505429a

URL: https://github.com/llvm/llvm-project/commit/62f57bc9b0b03cdd8a2647e4960d7835b505429a
DIFF: https://github.com/llvm/llvm-project/commit/62f57bc9b0b03cdd8a2647e4960d7835b505429a.diff

LOG: [libc] Add other RPC callback methods to the RPC server

This patch adds the other two methods to the server so the external
users can use the interface through the obfuscated interface.

Reviewed By: JonChesterfield

Differential Revision: https://reviews.llvm.org/D154224

Added: 
    

Modified: 
    libc/utils/gpu/server/rpc_server.cpp
    libc/utils/gpu/server/rpc_server.h

Removed: 
    


################################################################################
diff  --git a/libc/utils/gpu/server/rpc_server.cpp b/libc/utils/gpu/server/rpc_server.cpp
index b5103f754392b4..ac7dca6232cebe 100644
--- a/libc/utils/gpu/server/rpc_server.cpp
+++ b/libc/utils/gpu/server/rpc_server.cpp
@@ -320,25 +320,50 @@ const void *rpc_get_client_buffer(uint32_t device_id) {
 
 uint64_t rpc_get_client_size() { return sizeof(rpc::Client); }
 
+using ServerPort = std::variant<rpc::Server<1>::Port *, rpc::Server<32>::Port *,
+                                rpc::Server<64>::Port *>;
+
+ServerPort getPort(rpc_port_t ref) {
+  if (ref.lane_size == 1)
+    return reinterpret_cast<rpc::Server<1>::Port *>(ref.handle);
+  else if (ref.lane_size == 32)
+    return reinterpret_cast<rpc::Server<32>::Port *>(ref.handle);
+  else if (ref.lane_size == 64)
+    return reinterpret_cast<rpc::Server<64>::Port *>(ref.handle);
+  else
+    __builtin_unreachable();
+}
+
+void rpc_send(rpc_port_t ref, rpc_port_callback_ty callback, void *data) {
+  auto port = getPort(ref);
+  std::visit(
+      [=](auto &port) {
+        port->send([=](rpc::Buffer *buffer) {
+          callback(reinterpret_cast<rpc_buffer_t *>(buffer), data);
+        });
+      },
+      port);
+}
+
+void rpc_recv(rpc_port_t ref, rpc_port_callback_ty callback, void *data) {
+  auto port = getPort(ref);
+  std::visit(
+      [=](auto &port) {
+        port->recv([=](rpc::Buffer *buffer) {
+          callback(reinterpret_cast<rpc_buffer_t *>(buffer), data);
+        });
+      },
+      port);
+}
+
 void rpc_recv_and_send(rpc_port_t ref, rpc_port_callback_ty callback,
                        void *data) {
-  if (ref.lane_size == 1) {
-    rpc::Server<1>::Port *port =
-        reinterpret_cast<rpc::Server<1>::Port *>(ref.handle);
-    port->recv_and_send([=](rpc::Buffer *buffer) {
-      callback(reinterpret_cast<rpc_buffer_t *>(buffer), data);
-    });
-  } else if (ref.lane_size == 32) {
-    rpc::Server<32>::Port *port =
-        reinterpret_cast<rpc::Server<32>::Port *>(ref.handle);
-    port->recv_and_send([=](rpc::Buffer *buffer) {
-      callback(reinterpret_cast<rpc_buffer_t *>(buffer), data);
-    });
-  } else if (ref.lane_size == 64) {
-    rpc::Server<64>::Port *port =
-        reinterpret_cast<rpc::Server<64>::Port *>(ref.handle);
-    port->recv_and_send([=](rpc::Buffer *buffer) {
-      callback(reinterpret_cast<rpc_buffer_t *>(buffer), data);
-    });
-  }
+  auto port = getPort(ref);
+  std::visit(
+      [=](auto &port) {
+        port->recv_and_send([=](rpc::Buffer *buffer) {
+          callback(reinterpret_cast<rpc_buffer_t *>(buffer), data);
+        });
+      },
+      port);
 }

diff  --git a/libc/utils/gpu/server/rpc_server.h b/libc/utils/gpu/server/rpc_server.h
index 1ac6a83f17660f..f4ba77fb906021 100644
--- a/libc/utils/gpu/server/rpc_server.h
+++ b/libc/utils/gpu/server/rpc_server.h
@@ -97,6 +97,12 @@ const void *rpc_get_client_buffer(uint32_t device_id);
 /// Returns the size of the client in bytes to be used for a memory copy.
 uint64_t rpc_get_client_size();
 
+/// Use the \p port to send a buffer using the \p callback.
+void rpc_send(rpc_port_t port, rpc_port_callback_ty callback, void *data);
+
+/// Use the \p port to recieve a buffer using the \p callback.
+void rpc_recv(rpc_port_t port, rpc_port_callback_ty callback, void *data);
+
 /// Use the \p port to receive and send a buffer using the \p callback.
 void rpc_recv_and_send(rpc_port_t port, rpc_port_callback_ty callback,
                        void *data);