[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

Author: jhuber6

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);
 }

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);