[libc-commits] [PATCH] D153633: [libc] Allow the RPC client to be initialized via a H2D memcpy

[Joseph Huber via Phabricator via libc-commits]

jhuber6 created this revision.
jhuber6 added reviewers: jdoerfert, tianshilei1992, JonChesterfield, sivachandra.
Herald added projects: libc-project, All.
Herald added a subscriber: libc-commits.
jhuber6 requested review of this revision.
Herald added subscribers: jplehr, sstefan1.

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.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D153633

Files:
  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


Index: libc/utils/gpu/server/rpc_server.h
===================================================================
--- libc/utils/gpu/server/rpc_server.h
+++ libc/utils/gpu/server/rpc_server.h
@@ -86,6 +86,13 @@
 /// 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);
Index: libc/utils/gpu/server/rpc_server.cpp
===================================================================
--- libc/utils/gpu/server/rpc_server.cpp
+++ libc/utils/gpu/server/rpc_server.cpp
@@ -36,6 +36,7 @@
   void reset(uint64_t port_count, void *buffer) {
     std::visit([&](auto &server) { server->reset(port_count, buffer); },
                server);
+    client.reset(port_count, buffer);
   }
 
   uint64_t allocation_size(uint64_t port_count) {
@@ -51,6 +52,8 @@
     return ret;
   }
 
+  const void *get_client_data() const { return &client; }
+
   rpc_status_t handle_server(
       std::unordered_map<rpc_opcode_t, rpc_opcode_callback_ty> &callbacks,
       std::unordered_map<rpc_opcode_t, void *> &callback_data) {
@@ -173,6 +176,7 @@
                std::unique_ptr<rpc::Server<32>>,
                std::unique_ptr<rpc::Server<64>>>
       server;
+  rpc::Client client;
 };
 
 struct Device {
@@ -312,6 +316,18 @@
   return state->devices[device_id]->server.get_buffer_start();
 }
 
+const void *rpc_get_client_buffer(uint32_t device_id) {
+  if (!state)
+    return nullptr;
+  if (device_id >= state->num_devices)
+    return nullptr;
+  if (!state->devices[device_id])
+    return nullptr;
+  return state->devices[device_id]->server.get_client_data();
+}
+
+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) {
Index: libc/src/__support/RPC/rpc_client.cpp
===================================================================
--- libc/src/__support/RPC/rpc_client.cpp
+++ libc/src/__support/RPC/rpc_client.cpp
@@ -16,8 +16,10 @@
 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
Index: libc/src/__support/RPC/rpc.h
===================================================================
--- libc/src/__support/RPC/rpc.h
+++ libc/src/__support/RPC/rpc.h
@@ -343,6 +343,10 @@
 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 {


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D153633.533955.patch
Type: text/x-patch
Size: 3550 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/libc-commits/attachments/20230623/da78230e/attachment-0001.bin>