[libc] Allow the RPC client to be initialized via a H2D memcpy

[jhuber6]

Summary

The RPC client must be initialized to set a pointer to the underlying
buffer. This is currently done with the reset method which may not be
ideal for the use-case. We want runtimes to be able to initialize this
without needing to call a kernel. Recent changes allowed the Client
type to be trivially copyable. That means we can create a client on the
server side and then copy it over. To that end we take the existing
externally visible symbol and initialize it to the client's pointer.
Therefore we can look up the symbol and copy it over once loaded.

No test currently, I tested with a demo OpenMP application but couldn't think of
how to put that in-tree.

[jhuber6]

Add the symbol's name to the RPC header so users can look it up in the runtime.

[jhuber6]

Move string constant to the server and make it static, since only the server should care about it.

[jhuber6]

Simplify

[JonChesterfield]

Thanks for the update, LG

libc/src/__support/RPC/rpc.h

struct Client {
 
   LIBC_INLINE void reset(uint64_t port_count, void *buffer) {
     process.reset(port_count, buffer);
   }
 
 private:
   Process<false, Packet<gpu::LANE_SIZE>> process;
 };
+static_assert(cpp::is_trivially_copyable<Client>::value &&
+                  sizeof(Process<false, Packet<1>>) ==
+                      sizeof(Process<false, Packet<32>>),
+              "The client is not trivially copyable from the server");
 
 /// The RPC server used to respond to the client.
 template <uint32_t lane_size> struct Server {
   LIBC_INLINE Server() = default;
   LIBC_INLINE Server(const Server &) = delete;
   LIBC_INLINE Server &operator=(const Server &) = delete;
   LIBC_INLINE ~Server() = default;
 

libc/src/__support/RPC/rpc_client.cpp

 
 namespace __llvm_libc {
 namespace rpc {
 
 /// The libc client instance used to communicate with the server.
 Client client;
 
 /// Externally visible symbol to signify the usage of an RPC client to
-/// whomever needs to run the server.
-extern "C" [[gnu::visibility("protected")]] const bool __llvm_libc_rpc = false;
+/// whomever needs to run the server as well as provide a way to initialize
+/// the client with a copy..
+extern "C" [[gnu::visibility("protected")]] void *__llvm_libc_rpc_client =
+    &client;
 
 } // namespace rpc
 } // namespace __llvm_libc

libc/utils/gpu/server/rpc_server.h

 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 /// The maxium number of ports that can be opened for any server.
 const uint64_t RPC_MAXIMUM_PORT_COUNT = 64;
 
+/// The symbol name associated with the client for use with the LLVM C library
+/// implementation.
+static const char *rpc_client_symbol_name = "__llvm_libc_rpc_client";
+
 /// status codes.
 typedef enum {
   RPC_STATUS_SUCCESS = 0x0,
   RPC_STATUS_CONTINUE = 0x1,
   RPC_STATUS_ERROR = 0x1000,
   RPC_STATUS_OUT_OF_RANGE = 0x1001,
   RPC_STATUS_UNHANDLED_OPCODE = 0x1002,
   RPC_STATUS_INVALID_LANE_SIZE = 0x1003,

 /// data must remain accessible as long as the user intends to handle the server
 /// with this callback.
 rpc_status_t rpc_register_callback(uint32_t device_id, rpc_opcode_t opcode,
                                    rpc_opcode_callback_ty callback, void *data);
 
 /// Obtain a pointer to the memory buffer used to run the RPC client and server.
 void *rpc_get_buffer(uint32_t device_id);
 
+/// 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 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

std::variant<std::unique_ptr<rpc::Server<1>>,
                std::unique_ptr<rpc::Server<64>>>
       server;
 };
 
 struct Device {
   template <typename T>
   Device(std::unique_ptr<T> &&server) : server(std::move(server)) {}
   Server server;
+  rpc::Client client;
   std::unordered_map<rpc_opcode_t, rpc_opcode_callback_ty> callbacks;
   std::unordered_map<rpc_opcode_t, void *> callback_data;
 };
 
 // A struct containing all the runtime state required to run the RPC server.
 struct State {
   State(uint32_t num_devices)
       : num_devices(num_devices), devices(num_devices), reference_count(0u) {}
rpc_status_t rpc_server_init(uint32_t device_id, uint64_t num_ports,
 
   uint64_t size = state->devices[device_id]->server.allocation_size(num_ports);
   void *buffer = alloc(size, data);
 
   if (!buffer)
     return RPC_STATUS_ERROR;
 
   state->devices[device_id]->server.reset(num_ports, buffer);
+  state->devices[device_id]->client.reset(num_ports, buffer);
 
   return RPC_STATUS_SUCCESS;
 }
 
 rpc_status_t rpc_server_shutdown(uint32_t device_id, rpc_free_ty dealloc,
                                  void *data) {
   if (!state)
     return RPC_STATUS_NOT_INITIALIZED;
if (!state->devices[device_id])
     return RPC_STATUS_ERROR;
 
   state->devices[device_id]->callbacks[opcode] = callback;
   state->devices[device_id]->callback_data[opcode] = data;
   return RPC_STATUS_SUCCESS;
 }
 
 void *rpc_get_buffer(uint32_t device_id) {
-  if (!state)
-    return nullptr;
-  if (device_id >= state->num_devices)
-    return nullptr;
-  if (!state->devices[device_id])
+  if (!state || device_id >= state->num_devices || !state->devices[device_id])
     return nullptr;
   return state->devices[device_id]->server.get_buffer_start();
 }
 
+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) {
       callback(reinterpret_cast<rpc_buffer_t *>(buffer), data);
     });