[libc-commits] [libc] 31c1548 - [libc] Allow the RPC client to be initialized via a H2D memcpy

[Joseph Huber via libc-commits]

Author: Joseph Huber
Date: 2023-06-26T10:41:32-05:00
New Revision: 31c154881c5f1e59f29443778cb9d583a10c9578

URL: https://github.com/llvm/llvm-project/commit/31c154881c5f1e59f29443778cb9d583a10c9578
DIFF: https://github.com/llvm/llvm-project/commit/31c154881c5f1e59f29443778cb9d583a10c9578.diff

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

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.

Reviewed By: JonChesterfield

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

Added: 
    

Modified: 
    libc/src/__support/RPC/rpc.h
    libc/src/__support/RPC/rpc_client.cpp
    libc/utils/gpu/server/rpc_server.cpp
    libc/utils/gpu/server/rpc_server.h

Removed: 
    


################################################################################
diff  --git a/libc/src/__support/RPC/rpc.h b/libc/src/__support/RPC/rpc.h
index fa727639d7b95..4e0b85b6da716 100644
--- a/libc/src/__support/RPC/rpc.h
+++ b/libc/src/__support/RPC/rpc.h
@@ -343,6 +343,10 @@ struct Client {
 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 {

diff  --git a/libc/src/__support/RPC/rpc_client.cpp b/libc/src/__support/RPC/rpc_client.cpp
index c03b75be1c01e..b25ff56410d7f 100644
--- a/libc/src/__support/RPC/rpc_client.cpp
+++ b/libc/src/__support/RPC/rpc_client.cpp
@@ -16,8 +16,10 @@ namespace rpc {
 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

diff  --git a/libc/utils/gpu/server/rpc_server.cpp b/libc/utils/gpu/server/rpc_server.cpp
index 644c19f99a052..a440efddeb8ab 100644
--- a/libc/utils/gpu/server/rpc_server.cpp
+++ b/libc/utils/gpu/server/rpc_server.cpp
@@ -179,6 +179,7 @@ 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;
 };
@@ -250,6 +251,7 @@ rpc_status_t rpc_server_init(uint32_t device_id, uint64_t num_ports,
     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;
 }
@@ -303,15 +305,19 @@ rpc_status_t rpc_register_callback(uint32_t device_id, rpc_opcode_t opcode,
 }
 
 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) {

diff  --git a/libc/utils/gpu/server/rpc_server.h b/libc/utils/gpu/server/rpc_server.h
index d27a0fe3bae41..1ac6a83f17660 100644
--- a/libc/utils/gpu/server/rpc_server.h
+++ b/libc/utils/gpu/server/rpc_server.h
@@ -20,6 +20,10 @@ extern "C" {
 /// 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,
@@ -86,6 +90,13 @@ rpc_status_t rpc_register_callback(uint32_t device_id, rpc_opcode_t opcode,
 /// 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);