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

[jhuber6]

Summary

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

libc/utils/gpu/server/rpc_server.h

 
 /// Obtain a pointer to a local client buffer that can be copied directly to the
 /// other process.
 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);
 
 #ifdef __cplusplus
 }
 #endif
 
 #endif

libc/utils/gpu/server/rpc_server.cpp

 const void *rpc_get_client_buffer(uint32_t device_id) {
   if (!state || device_id >= state->num_devices || !state->devices[device_id])
     return nullptr;
   return &state->devices[device_id]->client;
 }
 
 uint64_t rpc_get_client_size() { return sizeof(rpc::Client); }
 
-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) {
+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);
         });
-  } 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) {
+      },
+      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);
         });
-  } else if (ref.lane_size == 64) {
-    rpc::Server<64>::Port *port =
-        reinterpret_cast<rpc::Server<64>::Port *>(ref.handle);
+      },
+      port);
+}
+
+void rpc_recv_and_send(rpc_port_t ref, rpc_port_callback_ty callback,
+                       void *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);
 }