[libc][rpc] Allocate a single block of shared memory instead of three

[JonChesterfield]

Summary

Allows moving the pointer swap between server and client into reset.
Single allocation simplifies whatever allocates the client/server, currently
the libc loaders.

[JonChesterfield]

constexpr on the functions allowed ad hoc testing along the lines of:

enum {N = 128};
struct equiv {
  atomic<uint32_t> primary[N];
  atomic<uint32_t> secondary[N];
  Packet buffer[N];
};
static_assert(Process<false>::memory_offset_primary_mailbox(N) == __builtin_offset(equiv, primary));

but I can't see an obvious place to check in tests like that

[jhuber6]

libc/src/__support/RPC/rpc.h

123

Static cast

123

This is unrelated, but I think we might want to call InvertInbox something more descriptive in the fact that it's the server / client.

131

274

Not sure we need to bother with the macro, I'm sure the optimizer can handle forwarding 32 or __AMDGCN_WAVEFRONTSIZE.

libc/src/__support/RPC/rpc_util.h

54 ↗

(On Diff #521189)

Did you format this with a different style?

libc/startup/gpu/amdgpu/start.cpp

41

I'd like something more like rpc_shared_buffer or something so we know what it's for.

[JonChesterfield]

libc/src/__support/RPC/rpc.h

274

gone with this so it matches the definition of the payload type. alternative is something like sizeof(header) + lane_size*N which seems more obscure

libc/src/__support/RPC/rpc_util.h

54 ↗

(On Diff #521189)

no, adding constexpr hit a line limit. can drop constexpr across the board from this patch if necessary

[JonChesterfield]

• rename variable, drop constexpr in various places

[JonChesterfield]

• comma requested

[jhuber6]

We can clean some of this stuff up later. Thinking about doing a big NFC for comments and other things.

libc/src/__support/RPC/rpc.h

template <bool InvertInbox> struct Process {
   uint32_t lane_size;
   cpp::Atomic<uint32_t> *inbox;
   cpp::Atomic<uint32_t> *outbox;
   Packet *packet;
 
   cpp::Atomic<uint32_t> lock[default_port_count] = {0};
 
   /// Initialize the communication channels.
-  LIBC_INLINE void reset(uint64_t port_count, uint32_t lane_size, void *inbox,
-                         void *outbox, void *packet) {
+  LIBC_INLINE void reset(uint64_t port_count, uint32_t lane_size, void *state) {
+    uint64_t p = memory_offset_primary_mailbox(port_count);
+    uint64_t s = memory_offset_secondary_mailbox(port_count);
     this->port_count = port_count;
     this->lane_size = lane_size;
-    this->inbox = reinterpret_cast<cpp::Atomic<uint32_t> *>(inbox);
-    this->outbox = reinterpret_cast<cpp::Atomic<uint32_t> *>(outbox);
-    this->packet = reinterpret_cast<Packet *>(packet);
+    this->inbox = reinterpret_cast<cpp::Atomic<uint32_t> *>(
+        static_cast<char *>(state) + (InvertInbox ? s : p));
+    this->outbox = reinterpret_cast<cpp::Atomic<uint32_t> *>(
+        static_cast<char *>(state) + (InvertInbox ? p : s));
+    this->packet = reinterpret_cast<Packet *>(static_cast<char *>(state) +
+                                              memory_offset_buffer(port_count));
+  }
+
+  /// Allocate a single block of memory for use by client and server
+  /// template<size_t N>, N is generally a runtime value
+  /// struct equivalent {
+  ///   atomic<uint32_t> primary[N];
+  ///   atomic<uint32_t> secondary[N];
+  ///   Packet buffer[N];
+  /// };
+  LIBC_INLINE static uint64_t allocation_size(uint64_t port_count,
+                                              uint32_t lane_size) {
+    return memory_offset_buffer(port_count) +
+           memory_allocated_buffer(port_count, lane_size);
   }
 
   /// The length of the packet is flexible because the server needs to look up
   /// the lane size at runtime. This helper indexes at the proper offset.
   LIBC_INLINE Packet &get_packet(uint64_t index) {
     return *reinterpret_cast<Packet *>(
         reinterpret_cast<uint8_t *>(packet) +
         index * align_up(sizeof(Header) + lane_size * sizeof(Buffer),
LIBC_INLINE void invoke_rpc(cpp::function<void(Buffer *, uint32_t)> fn,
     if constexpr (is_process_gpu()) {
       fn(&packet.payload.slot[gpu::get_lane_id()], gpu::get_lane_id());
     } else {
       for (uint32_t i = 0; i < lane_size; i += gpu::get_lane_size())
         if (packet.header.mask & 1ul << i)
           fn(&packet.payload.slot[i], i);
     }
   }
+
+  /// Number of bytes allocated for mailbox or buffer
+  LIBC_INLINE static uint64_t memory_allocated_mailbox(uint64_t port_count) {
+    return port_count * sizeof(cpp::Atomic<uint32_t>);
+  }
+
+  LIBC_INLINE static uint64_t memory_allocated_buffer(uint64_t port_count,
+                                                      uint32_t lane_size) {
+#if defined(LIBC_TARGET_ARCH_IS_GPU)
+    (void)lane_size;
+    return port_count * sizeof(Packet);
+#else
+    return port_count * (sizeof(Packet) + sizeof(Buffer) * lane_size);
+#endif
+  }
+
+  /// Offset of mailbox/buffer in single allocation
+  LIBC_INLINE static uint64_t
+  memory_offset_primary_mailbox(uint64_t /*port_count*/) {
+    return 0;
+  }
+  LIBC_INLINE static uint64_t
+  memory_offset_secondary_mailbox(uint64_t port_count) {
+    return memory_allocated_mailbox(port_count);
+  }
+  LIBC_INLINE static uint64_t memory_offset_buffer(uint64_t port_count) {
+    return align_up(2 * memory_allocated_mailbox(port_count), alignof(Packet));
+  }
 };
 
 /// The port provides the interface to communicate between the multiple
 /// processes. A port is conceptually an index into the memory provided by the
 /// underlying process that is guarded by a lock bit.
 template <bool T> struct Port {
   LIBC_INLINE Port(Process<T> &process, uint64_t lane_mask, uint64_t index,
                    uint32_t out)

libc/startup/gpu/amdgpu/start.cpp

static void call_fini_array_callbacks() {
   size_t fini_array_size = __fini_array_end - __fini_array_start;
   for (size_t i = 0; i < fini_array_size; ++i)
     reinterpret_cast<FiniCallback *>(__fini_array_start[i])();
 }
 
 } // namespace __llvm_libc
 
 extern "C" [[gnu::visibility("protected"), clang::amdgpu_kernel]] void
-_begin(int argc, char **argv, char **env, void *in, void *out, void *buffer) {
+_begin(int argc, char **argv, char **env, void *rpc_shared_buffer) {
   // We need to set up the RPC client first in case any of the constructors
   // require it.
   __llvm_libc::rpc::client.reset(__llvm_libc::rpc::default_port_count,
-                                 __llvm_libc::gpu::get_lane_size(), in, out,
-                                 buffer);
+                                 __llvm_libc::gpu::get_lane_size(),
+                                 rpc_shared_buffer);
 
   // We want the fini array callbacks to be run after other atexit
   // callbacks are run. So, we register them before running the init
   // array callbacks as they can potentially register their own atexit
   // callbacks.
   __llvm_libc::atexit(&__llvm_libc::call_fini_array_callbacks);
   __llvm_libc::call_init_array_callbacks(argc, argv, env);
 }

libc/startup/gpu/nvptx/start.cpp

static void call_fini_array_callbacks() {
   size_t fini_array_size = __fini_array_end - __fini_array_start;
   for (size_t i = 0; i < fini_array_size; ++i)
     reinterpret_cast<FiniCallback *>(__fini_array_start[i])();
 }
 
 } // namespace __llvm_libc
 
 extern "C" [[gnu::visibility("protected"), clang::nvptx_kernel]] void
-_begin(int argc, char **argv, char **env, void *in, void *out, void *buffer) {
+_begin(int argc, char **argv, char **env, void *rpc_shared_buffer) {
   // We need to set up the RPC client first in case any of the constructors
   // require it.
   __llvm_libc::rpc::client.reset(__llvm_libc::rpc::default_port_count,
-                                 __llvm_libc::gpu::get_lane_size(), in, out,
-                                 buffer);
+                                 __llvm_libc::gpu::get_lane_size(),
+                                 rpc_shared_buffer);
 
   // We want the fini array callbacks to be run after other atexit
   // callbacks are run. So, we register them before running the init
   // array callbacks as they can potentially register their own atexit
   // callbacks.
   __llvm_libc::atexit(&__llvm_libc::call_fini_array_callbacks);
   __llvm_libc::call_init_array_callbacks(argc, argv, env);
 }

libc/utils/gpu/loader/Loader.h

struct LaunchParameters {
   uint32_t num_blocks_z;
 };
 
 /// The arguments to the '_begin' kernel.
 struct begin_args_t {
   int argc;
   void *argv;
   void *envp;
-  void *inbox;
-  void *outbox;
-  void *buffer;
+  void *rpc_shared_buffer;
 };
 
 /// The arguments to the '_start' kernel.
 struct start_args_t {
   int argc;
   void *argv;
   void *envp;
   void *ret;

libc/utils/gpu/loader/amdgpu/Loader.cpp

int load(int argc, char **argv, char **envp, void *image, size_t size,
   hsa_amd_memory_fill(dev_ret, 0, sizeof(int));
 
   // Allocate finegrained memory for the RPC server and client to share.
   uint64_t port_size = __llvm_libc::rpc::default_port_count;
   uint32_t wavefront_size = 0;
   if (hsa_status_t err = hsa_agent_get_info(
           dev_agent, HSA_AGENT_INFO_WAVEFRONT_SIZE, &wavefront_size))
     handle_error(err);
-  void *server_inbox;
-  void *server_outbox;
-  void *buffer;
-  if (hsa_status_t err = hsa_amd_memory_pool_allocate(
-          finegrained_pool, port_size * sizeof(__llvm_libc::cpp::Atomic<int>),
-          /*flags=*/0, &server_inbox))
-    handle_error(err);
-  if (hsa_status_t err = hsa_amd_memory_pool_allocate(
-          finegrained_pool, port_size * sizeof(__llvm_libc::cpp::Atomic<int>),
-          /*flags=*/0, &server_outbox))
-    handle_error(err);
-  if (hsa_status_t err = hsa_amd_memory_pool_allocate(
-          finegrained_pool,
-          port_size *
-              align_up(sizeof(__llvm_libc::rpc::Header) +
-                           (wavefront_size * sizeof(__llvm_libc::rpc::Buffer)),
-                       alignof(__llvm_libc::rpc::Packet)),
-          /*flags=*/0, &buffer))
-    handle_error(err);
-  hsa_amd_agents_allow_access(1, &dev_agent, nullptr, server_inbox);
-  hsa_amd_agents_allow_access(1, &dev_agent, nullptr, server_outbox);
-  hsa_amd_agents_allow_access(1, &dev_agent, nullptr, buffer);
+
+  uint64_t rpc_shared_buffer_size =
+      __llvm_libc::rpc::Server::allocation_size(port_size, wavefront_size);
+  void *rpc_shared_buffer;
+  if (hsa_status_t err =
+          hsa_amd_memory_pool_allocate(finegrained_pool, rpc_shared_buffer_size,
+                                       /*flags=*/0, &rpc_shared_buffer))
+    handle_error(err);
+  hsa_amd_agents_allow_access(1, &dev_agent, nullptr, rpc_shared_buffer);
 
   // Initialize the RPC server's buffer for host-device communication.
-  server.reset(port_size, wavefront_size, server_inbox, server_outbox, buffer);
+  server.reset(port_size, wavefront_size, rpc_shared_buffer);
 
   // Obtain a queue with the minimum (power of two) size, used to send commands
   // to the HSA runtime and launch execution on the device.
   uint64_t queue_size;
   if (hsa_status_t err = hsa_agent_get_info(
           dev_agent, HSA_AGENT_INFO_QUEUE_MIN_SIZE, &queue_size))
     handle_error(err);
   hsa_queue_t *queue = nullptr;
   if (hsa_status_t err =
           hsa_queue_create(dev_agent, queue_size, HSA_QUEUE_TYPE_MULTI, nullptr,
                            nullptr, UINT32_MAX, UINT32_MAX, &queue))
     handle_error(err);
 
   LaunchParameters single_threaded_params = {1, 1, 1, 1, 1, 1};
-  begin_args_t init_args = {argc,          dev_argv,     dev_envp,
-                            server_outbox, server_inbox, buffer};
+  begin_args_t init_args = {argc, dev_argv, dev_envp, rpc_shared_buffer};
   if (hsa_status_t err =
           launch_kernel(dev_agent, executable, kernargs_pool, queue,
                         single_threaded_params, "_begin.kd", init_args))
     handle_error(err);
 
   start_args_t args = {argc, dev_argv, dev_envp, dev_ret};
   if (hsa_status_t err = launch_kernel(dev_agent, executable, kernargs_pool,
                                        queue, params, "_start.kd", args))
if (hsa_status_t err =
                         single_threaded_params, "_end.kd", fini_args))
     handle_error(err);
 
   // Free the memory allocated for the device.
   if (hsa_status_t err = hsa_amd_memory_pool_free(dev_argv))
     handle_error(err);
   if (hsa_status_t err = hsa_amd_memory_pool_free(dev_ret))
     handle_error(err);
-  if (hsa_status_t err = hsa_amd_memory_pool_free(server_inbox))
-    handle_error(err);
-  if (hsa_status_t err = hsa_amd_memory_pool_free(server_outbox))
-    handle_error(err);
-  if (hsa_status_t err = hsa_amd_memory_pool_free(buffer))
+  if (hsa_status_t err = hsa_amd_memory_pool_free(rpc_shared_buffer))
     handle_error(err);
   if (hsa_status_t err = hsa_amd_memory_pool_free(host_ret))
     handle_error(err);
 
   if (hsa_status_t err = hsa_signal_destroy(memory_signal))
     handle_error(err);
   if (hsa_status_t err = hsa_queue_destroy(queue))
     handle_error(err);

libc/utils/gpu/loader/nvptx/Loader.cpp

int load(int argc, char **argv, char **envp, void *image, size_t size,
   CUdeviceptr dev_ret;
   if (CUresult err = cuMemAlloc(&dev_ret, sizeof(int)))
     handle_error(err);
   if (CUresult err = cuMemsetD32(dev_ret, 0, 1))
     handle_error(err);
 
   uint64_t port_size = __llvm_libc::rpc::default_port_count;
   uint32_t warp_size = 32;
-  void *server_inbox =
-      allocator(port_size * sizeof(__llvm_libc::cpp::Atomic<int>));
-  void *server_outbox =
-      allocator(port_size * sizeof(__llvm_libc::cpp::Atomic<int>));
-  void *buffer = allocator(
-      port_size * align_up(sizeof(__llvm_libc::rpc::Header) +
-                               (warp_size * sizeof(__llvm_libc::rpc::Buffer)),
-                           alignof(__llvm_libc::rpc::Packet)));
-  if (!server_inbox || !server_outbox || !buffer)
+
+  uint64_t rpc_shared_buffer_size =
+      __llvm_libc::rpc::Server::allocation_size(port_size, warp_size);
+  void *rpc_shared_buffer = allocator(rpc_shared_buffer_size);
+
+  if (!rpc_shared_buffer)
     handle_error("Failed to allocate memory the RPC client / server.");
 
   // Initialize the RPC server's buffer for host-device communication.
-  server.reset(port_size, warp_size, server_inbox, server_outbox, buffer);
+  server.reset(port_size, warp_size, rpc_shared_buffer);
 
   LaunchParameters single_threaded_params = {1, 1, 1, 1, 1, 1};
   // Call the kernel to
-  begin_args_t init_args = {argc,          dev_argv,     dev_envp,
-                            server_outbox, server_inbox, buffer};
+  begin_args_t init_args = {argc, dev_argv, dev_envp, rpc_shared_buffer};
   if (CUresult err = launch_kernel(binary, stream, single_threaded_params,
                                    "_begin", init_args))
     handle_error(err);
 
   start_args_t args = {argc, dev_argv, dev_envp,
                        reinterpret_cast<void *>(dev_ret)};
   if (CUresult err = launch_kernel(binary, stream, params, "_start", args))
     handle_error(err);
int load(int argc, char **argv, char **envp, void *image, size_t size,
 
   // Free the memory allocated for the device.
   if (CUresult err = cuMemFreeHost(*memory_or_err))
     handle_error(err);
   if (CUresult err = cuMemFree(dev_ret))
     handle_error(err);
   if (CUresult err = cuMemFreeHost(dev_argv))
     handle_error(err);
-  if (CUresult err = cuMemFreeHost(server_inbox))
-    handle_error(err);
-  if (CUresult err = cuMemFreeHost(server_outbox))
-    handle_error(err);
-  if (CUresult err = cuMemFreeHost(buffer))
+  if (CUresult err = cuMemFreeHost(rpc_shared_buffer))
     handle_error(err);
 
   // Destroy the context and the loaded binary.
   if (CUresult err = cuModuleUnload(binary))
     handle_error(err);
   if (CUresult err = cuDevicePrimaryCtxRelease(device))
     handle_error(err);
   return host_ret;
 }