[libc] Support concurrent RPC port access on the GPU

[jhuber6]

Summary

Previously we used a single port to implement the RPC. This was
sufficient for single threaded tests but can potentially cause deadlocks
when using multiple threads. The reason for this is that GPUs make no
forward progress guarantees. Therefore one group of threads waiting on
another group of threads can spin forever because there is no guarantee
that the other threads will continue executing. The typical workaround
for this is to allocate enough memory that a sufficiently large number
of work groups can make progress. As long as this number is somewhat
close to the amount of total concurrency we can obtain reliable
execution around a shared resource.

This patch enables using multiple ports by widening the arrays to a
predetermined size and indexes into them. Empty ports are currently
obtained via a trivial linker scan. This should be imporoved in the
future for performance reasons. Portions of D148191 were applied to
achieve parallel support.

Depends on D149581

[JonChesterfield]

libc/src/__support/RPC/rpc.h

353

Your data representation induces overhead here.

inbox.load hits the bus to retrieve one bit of information. If it fails, next step through the loop starts over from scratch.

If you change to packed bits, the first atomic load gets you the state of ~32 mailboxes in one tick. If your current index doesn't work out, you now have enough information to pick the next credible candidate without hitting the bus at all.

Also you're loading from the outbox despite the value of that being known already. Maybe the latency on that redundant load is hidden by the latency on loading from the inbox.

[JonChesterfield]

libc/src/__support/RPC/rpc.h

355

You could make this less vulnerable to data races on volta by getting the appropriate thread mask outside of the loop. That's also likely to be faster on amdgpu.

[jhuber6]

libc/src/__support/RPC/rpc.h

353

I can see the performance benefits of bit-packing here I'd assume unless the RPC client is mostly unused the performance penalty of false sharing more atomic indices will be small compared to saving several atomic checks in the client / server. We can also definitely hoist the outbox load.

355

Good point, since this is marked convergent it's unlikely that the compiler will remove this from the loop itself.

[JonChesterfield]

libc/src/__support/RPC/rpc.h

398

Is all this a straight copy/paste from the above? Why are client/server open even different functions, they could be identical

[jhuber6]

libc/src/__support/RPC/rpc.h

398

They're slightly different. The client open function only works if the port is capable of sending data while the server only works if it's capable of receiving it. I could probably template that somehow, but I didn't see the need.

[JonChesterfield]

libc/src/__support/RPC/rpc.h

398

Refer to D148191 for an example of using a single function for both. Traditional arguments for not copy/paste/slightly-change functions which pick up performance related changes over time apply really.

[jhuber6]

Update to new interface after @JonChesterfield's merge

[jhuber6]

Fix again.

[jhuber6]

Rebase

[JonChesterfield]

OK, much easier to parse than the earlier versions. Thanks! There's some cleanup to do post commit I think - in particular the try_open methods are messy and really similar, and it's probably worth having the inbox/outbox/locks pointers private to process.

libc/src/__support/RPC/rpc.h

106

port_count? - its how many ports are allocated, not how big a port is

149

I don't see a remaining use of *this, could be a free function (or at least static)

libc/src/__support/RPC/rpc.h

 
 /// A packet used to share data between the client and server across an entire
 /// lane. We use a lane as the minimum granularity for execution.
 struct alignas(64) Packet {
   Header header;
   Payload payload;
 };
 
+// TODO: This should be configured by the server and passed in. The general rule
+//       of thumb is that you should have at least as many ports as possible
+//       concurrent work items on the GPU to mitigate the lack offorward
+//       progress guarantees on the GPU.
+constexpr uint64_t default_port_size = 64;
+
 /// A common process used to synchronize communication between a client and a
 /// server. The process contains an inbox and an outbox used for signaling
 /// ownership of the shared buffer between both sides.
 ///
 /// This process is designed to support mostly arbitrary combinations of 'send'
 /// and 'recv' operations on the shared buffer as long as these operations are
 /// mirrored by the other process. These operations exchange ownership of the
 /// fixed-size buffer between the users of the protocol. The assumptions when

 /// bits are not equal the receiver owns the buffer, otherwise the sender owns
 /// the buffer.
 struct Process {
   LIBC_INLINE Process() = default;
   LIBC_INLINE Process(const Process &) = default;
   LIBC_INLINE Process &operator=(const Process &) = default;
   LIBC_INLINE ~Process() = default;
 
   static constexpr uint32_t Data = 0b01;
   static constexpr uint32_t Ack = 0b10;
 
+  uint64_t port_size;
   uint32_t lane_size;
   cpp::Atomic<uint32_t> *lock;
   cpp::Atomic<uint32_t> *inbox;
   cpp::Atomic<uint32_t> *outbox;
-  Packet *buffer;
+  Packet *packet;
 
   /// Initialize the communication channels.
-  LIBC_INLINE void reset(uint32_t size, void *mtx, void *in, void *out,
-                         void *data) {
-    lane_size = size;
-    lock = reinterpret_cast<cpp::Atomic<uint32_t> *>(mtx);
-    inbox = reinterpret_cast<cpp::Atomic<uint32_t> *>(in);
-    outbox = reinterpret_cast<cpp::Atomic<uint32_t> *>(out);
-    buffer = reinterpret_cast<Packet *>(data);
+  LIBC_INLINE void reset(uint64_t port_size, uint32_t lane_size, void *lock,
+                         void *inbox, void *outbox, void *packet) {
+    *this = {port_size,
+             lane_size,
+             reinterpret_cast<cpp::Atomic<uint32_t> *>(lock),
+             reinterpret_cast<cpp::Atomic<uint32_t> *>(inbox),
+             reinterpret_cast<cpp::Atomic<uint32_t> *>(outbox),
+             reinterpret_cast<Packet *>(packet)};
+  }
+
+  /// 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),
+                         alignof(Packet)));
   }
 
   /// Determines if this process owns the buffer for a send. We can send data if
   /// the output data bit matches the input acknowledge bit.
   LIBC_INLINE static bool can_send_data(uint32_t in, uint32_t out) {
     return bool(in & Process::Ack) == bool(out & Process::Data);
   }
 
   /// Determines if this process owns the buffer for a receive. We can send data
   /// if the output acknowledge bit does not match the input data bit.
   LIBC_INLINE static bool can_recv_data(uint32_t in, uint32_t out) {
     return bool(in & Process::Data) != bool(out & Process::Ack);
   }
 
   /// Invokes a function accross every active buffer across the total lane size.
   LIBC_INLINE void invoke_rpc(cpp::function<void(Buffer *)> fn,
-                              uint32_t index) {
+                              Packet &packet) {
     if constexpr (is_process_gpu()) {
-      fn(&buffer[index].payload.slot[gpu::get_lane_id()]);
+      fn(&packet.payload.slot[gpu::get_lane_id()]);
     } else {
       for (uint32_t i = 0; i < lane_size; i += gpu::get_lane_size())
-        if (buffer[index].header.mask & 1ul << i)
-          fn(&buffer[index].payload.slot[i]);
+        if (packet.header.mask & 1ul << i)
+          fn(&packet.payload.slot[i]);
     }
   }
 
   /// Alternate version that also provides the index of the current lane.
   LIBC_INLINE void invoke_rpc(cpp::function<void(Buffer *, uint32_t)> fn,
-                              uint32_t index) {
+                              Packet &packet) {
     if constexpr (is_process_gpu()) {
-      fn(&buffer[index].payload.slot[gpu::get_lane_id()], gpu::get_lane_id());
+      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 (buffer[index].header.mask & 1ul << i)
-          fn(&buffer[index].payload.slot[i], i);
+        if (packet.header.mask & 1ul << i)
+          fn(&packet.payload.slot[i], i);
     }
   }
 };
 
 /// 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.
 struct Port {
   // TODO: This should be move-only.
   LIBC_INLINE Port(Process &process, uint64_t index, uint32_t out)
       : process(process), index(index), out(out) {}
   LIBC_INLINE Port(const Port &) = default;
   LIBC_INLINE Port &operator=(const Port &) = delete;
   LIBC_INLINE ~Port() = default;
 
   template <typename U> LIBC_INLINE void recv(U use);
   template <typename F> LIBC_INLINE void send(F fill);
   template <typename F, typename U>
   LIBC_INLINE void send_and_recv(F fill, U use);
   template <typename W> LIBC_INLINE void recv_and_send(W work);
   LIBC_INLINE void send_n(const void *src, uint64_t size);
   template <typename A> LIBC_INLINE void recv_n(A alloc);
 
   LIBC_INLINE uint16_t get_opcode() const {
-    return process.buffer[index].header.opcode;
+    return process.get_packet(index).header.opcode;
   }
 
   [[clang::convergent]] LIBC_INLINE void close() {
-    uint64_t mask = process.buffer[index].header.mask;
+    uint64_t mask = process.get_packet(index).header.mask;
     process.lock[index].store(0, cpp::MemoryOrder::RELAXED);
     cpp::atomic_thread_fence(cpp::MemoryOrder::RELEASE);
     // This lane sync forces explicit ordering of the open and close calls.
     gpu::sync_lane(mask);
   }
 
 private:
   Process &process;
template <typename F> LIBC_INLINE void Port::send(F fill) {
 
   // We need to wait until we own the buffer before sending.
   while (!Process::can_send_data(in, out)) {
     sleep_briefly();
     in = process.inbox[index].load(cpp::MemoryOrder::RELAXED);
   }
 
   // Apply the \p fill function to initialize the buffer and release the memory.
-  process.invoke_rpc(fill, index);
+  process.invoke_rpc(fill, process.get_packet(index));
   out = out ^ Process::Data;
   atomic_thread_fence(cpp::MemoryOrder::RELEASE);
   process.outbox[index].store(out, cpp::MemoryOrder::RELAXED);
 }
 
 /// Applies \p use to the shared buffer and acknowledges the send.
 template <typename U> LIBC_INLINE void Port::recv(U use) {
   uint32_t in = process.inbox[index].load(cpp::MemoryOrder::RELAXED);
 
   // We need to wait until we own the buffer before receiving.
   while (!Process::can_recv_data(in, out)) {
     sleep_briefly();
     in = process.inbox[index].load(cpp::MemoryOrder::RELAXED);
   }
   atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
 
   // Apply the \p use function to read the memory out of the buffer.
-  process.invoke_rpc(use, index);
+  process.invoke_rpc(use, process.get_packet(index));
   out = out ^ Process::Ack;
   process.outbox[index].store(out, cpp::MemoryOrder::RELAXED);
 }
 
 /// Combines a send and receive into a single function.
 template <typename F, typename U>
 LIBC_INLINE void Port::send_and_recv(F fill, U use) {
   send(fill);
LIBC_INLINE void Port::send_n(const void *src, uint64_t size) {
   const uint8_t *ptr = reinterpret_cast<const uint8_t *>(src);
   for (uint64_t idx = 0; idx < size; idx += sizeof(Buffer::data)) {
     send([=](Buffer *buffer) {
       const uint64_t len =
           size - idx > sizeof(Buffer::data) ? sizeof(Buffer::data) : size - idx;
       inline_memcpy(buffer->data, ptr + idx, len);
     });
   }
-  gpu::sync_lane(process.buffer[index].header.mask);
+  gpu::sync_lane(process.get_packet(index).header.mask);
 }
 
 /// Receives an arbitrarily sized data buffer across the shared channel in
 /// multiples of the packet length. The \p alloc function is called with the
 /// size of the data so that we can initialize the size of the \p dst buffer.
 template <typename A> LIBC_INLINE void Port::recv_n(A alloc) {
   // The GPU handles thread private variables and masking implicitly through its
   // execution model. If this is the CPU we need to manually handle the

 /// Attempts to open a port to use as the client. The client can only open a
 /// port if we find an index that is in a valid sending state. That is, there
 /// are send operations pending that haven't been serviced on this port. Each
 /// port instance uses an associated \p opcode to tell the server what to do.
 /// Opening a port is only valid if the `opcode` is the sam accross every
 /// participating thread.
 [[clang::convergent]] LIBC_INLINE cpp::optional<Port>
 Client::try_open(uint16_t opcode) {
+  // Perform a naive linear scan for a port that can be opened to send data.
+  for (uint64_t index = 0; index < port_size; ++index) {
     // Attempt to acquire the lock on this index.
     uint64_t lane_mask = gpu::get_lane_mask();
     uint32_t previous = 0;
     if (is_first_lane(lane_mask))
-    previous = lock->fetch_or(1, cpp::MemoryOrder::ACQ_REL);
+      previous = lock[index].fetch_or(1, cpp::MemoryOrder::ACQ_REL);
     previous = gpu::broadcast_value(previous);
 
     if (previous != 0)
-    return cpp::nullopt;
+      continue;
 
-  uint32_t in = inbox->load(cpp::MemoryOrder::RELAXED);
-  uint32_t out = outbox->load(cpp::MemoryOrder::RELAXED);
+    uint32_t in = inbox[index].load(cpp::MemoryOrder::RELAXED);
+    uint32_t out = outbox[index].load(cpp::MemoryOrder::RELAXED);
 
     // Once we acquire the index we need to check if we are in a valid sending
     // state.
     if (!can_send_data(in, out)) {
-    lock->store(0, cpp::MemoryOrder::RELAXED);
+      lock[index].store(0, cpp::MemoryOrder::RELAXED);
       cpp::atomic_thread_fence(cpp::MemoryOrder::RELEASE);
-    return cpp::nullopt;
+      continue;
     }
 
-  // TODO: The index here will be non-zero when we support multiple slots.
     if (is_first_lane(lane_mask)) {
-    buffer[0].header.opcode = opcode;
-    buffer[0].header.mask = lane_mask;
+      get_packet(index).header.opcode = opcode;
+      get_packet(index).header.mask = lane_mask;
     }
     gpu::sync_lane(lane_mask);
-  return Port(*this, 0, out);
+    return Port(*this, index, out);
+  }
+  return cpp::nullopt;
 }
 
 LIBC_INLINE Port Client::open(uint16_t opcode) {
   for (;;) {
     if (cpp::optional<Port> p = try_open(opcode))
       return p.value();
     sleep_briefly();
   }
 }
 
 /// Attempts to open a port to use as the server. The server can only open a
 /// port if it has a pending receive operation
 [[clang::convergent]] LIBC_INLINE cpp::optional<Port> Server::try_open() {
-  uint32_t in = inbox->load(cpp::MemoryOrder::RELAXED);
-  uint32_t out = outbox->load(cpp::MemoryOrder::RELAXED);
+  // Perform a naive linear scan for a port that has a pending request.
+  for (uint64_t index = 0; index < port_size; ++index) {
+    uint32_t in = inbox[index].load(cpp::MemoryOrder::RELAXED);
+    uint32_t out = outbox[index].load(cpp::MemoryOrder::RELAXED);
 
     // The server is passive, if there is no work pending don't bother
     // opening a port.
     if (!can_recv_data(in, out))
-    return cpp::nullopt;
+      continue;
 
     // Attempt to acquire the lock on this index.
     uint64_t lane_mask = gpu::get_lane_mask();
     uint32_t previous = 0;
     if (is_first_lane(lane_mask))
-    previous = lock->fetch_or(1, cpp::MemoryOrder::ACQ_REL);
+      previous = lock[index].fetch_or(1, cpp::MemoryOrder::ACQ_REL);
     previous = gpu::broadcast_value(previous);
 
     if (previous != 0)
-    return cpp::nullopt;
+      continue;
 
-  in = inbox->load(cpp::MemoryOrder::RELAXED);
-  out = outbox->load(cpp::MemoryOrder::RELAXED);
+    in = inbox[index].load(cpp::MemoryOrder::RELAXED);
+    out = outbox[index].load(cpp::MemoryOrder::RELAXED);
 
     if (!can_recv_data(in, out)) {
       lock->store(0, cpp::MemoryOrder::RELAXED);
       cpp::atomic_thread_fence(cpp::MemoryOrder::RELEASE);
-    return cpp::nullopt;
+      continue;
     }
 
-  return Port(*this, 0, out);
+    return Port(*this, index, out);
+  }
+  return cpp::nullopt;
 }
 
 LIBC_INLINE Port Server::open() {
   for (;;) {
     if (cpp::optional<Port> p = try_open())
       return p.value();
     sleep_briefly();
   }
 }
 
 } // namespace rpc
 } // namespace __llvm_libc
 
 #endif

libc/src/__support/RPC/rpc_util.h

 LIBC_INLINE constexpr bool is_process_gpu() {
 #if defined(LIBC_TARGET_ARCH_IS_GPU)
   return true;
 #else
   return false;
 #endif
 }
 
+/// Return \p val aligned "upwards" according to \p align.
+template <typename V, typename A> LIBC_INLINE V align_up(V val, A align) {
+  return ((val + V(align) - 1) / V(align)) * V(align);
+}
+
 } // namespace rpc
 } // namespace __llvm_libc
 
 #endif

libc/startup/gpu/amdgpu/start.cpp

 #include "src/__support/RPC/rpc_client.h"
 #include "src/stdlib/atexit.h"
 #include "src/stdlib/exit.h"
 
 extern "C" int main(int argc, char **argv, char **envp);
 
 namespace __llvm_libc {
 
-static cpp::Atomic<uint32_t> lock = 0;
+static cpp::Atomic<uint32_t> lock[rpc::default_port_size] = {0};
 
 extern "C" uintptr_t __init_array_start[];
 extern "C" uintptr_t __init_array_end[];
 extern "C" uintptr_t __fini_array_start[];
 extern "C" uintptr_t __fini_array_end[];
 
 using InitCallback = void(int, char **, char **);
 using FiniCallback = void(void);

 }
 
 } // 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) {
   // We need to set up the RPC client first in case any of the constructors
   // require it.
-  __llvm_libc::rpc::client.reset(__llvm_libc::gpu::get_lane_size(),
+  __llvm_libc::rpc::client.reset(__llvm_libc::rpc::default_port_size,
+                                 __llvm_libc::gpu::get_lane_size(),
                                  &__llvm_libc::lock, in, out, 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

 #include "src/__support/RPC/rpc_client.h"
 #include "src/stdlib/atexit.h"
 #include "src/stdlib/exit.h"
 
 extern "C" int main(int argc, char **argv, char **envp);
 
 namespace __llvm_libc {
 
-static cpp::Atomic<uint32_t> lock = 0;
+static cpp::Atomic<uint32_t> lock[rpc::default_port_size] = {0};
 
 extern "C" {
 // Nvidia's 'nvlink' linker does not provide these sections. We instead need
 // to manually create them and update the globals in the loader implememtation.
 uintptr_t *__init_array_start;
 uintptr_t *__init_array_end;
 uintptr_t *__fini_array_start;
 uintptr_t *__fini_array_end;

 }
 
 } // 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) {
   // We need to set up the RPC client first in case any of the constructors
   // require it.
-  __llvm_libc::rpc::client.reset(__llvm_libc::gpu::get_lane_size(),
+  __llvm_libc::rpc::client.reset(__llvm_libc::rpc::default_port_size,
+                                 __llvm_libc::gpu::get_lane_size(),
                                  &__llvm_libc::lock, in, out, 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/Server.h

 #include <cstdlib>
 #include <cstring>
 #include <stddef.h>
 
 #include "src/__support/RPC/rpc.h"
 
 static __llvm_libc::rpc::Server server;
 
-static __llvm_libc::cpp::Atomic<uint32_t> lock;
+static __llvm_libc::cpp::Atomic<uint32_t>
+    lock[__llvm_libc::rpc::default_port_size] = {0};
 
 /// Queries the RPC client at least once and performs server-side work if there
 /// are any active requests.
 void handle_server() {
+  // Continue servicing the client until there is no work left and we return.
+  for (;;) {
     auto port = server.try_open();
     if (!port)
       return;
 
     switch (port->get_opcode()) {
     case __llvm_libc::rpc::Opcode::PRINT_TO_STDERR: {
       uint64_t str_size[__llvm_libc::rpc::MAX_LANE_SIZE] = {0};
       char *strs[__llvm_libc::rpc::MAX_LANE_SIZE] = {nullptr};
       port->recv_n([&](uint64_t size, uint32_t id) {
         str_size[id] = size;
         strs[id] = new char[size];
         return strs[id];
       });
       for (uint64_t i = 0; i < __llvm_libc::rpc::MAX_LANE_SIZE; ++i) {
         if (strs[i]) {
           fwrite(strs[i], str_size[i], 1, stderr);
           delete[] strs[i];
         }
       }
       break;
     }
     case __llvm_libc::rpc::Opcode::EXIT: {
       port->recv([](__llvm_libc::rpc::Buffer *buffer) {
         exit(reinterpret_cast<uint32_t *>(buffer->data)[0]);
       });
       break;
     }
     case __llvm_libc::rpc::Opcode::TEST_INCREMENT: {
       port->recv_and_send([](__llvm_libc::rpc::Buffer *buffer) {
         reinterpret_cast<uint64_t *>(buffer->data)[0] += 1;
       });
       break;
     }
     default:
       port->recv([](__llvm_libc::rpc::Buffer *buffer) {});
     }
     port->close();
   }
+}
 #endif

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

int load(int argc, char **argv, char **envp, void *image, size_t size,
   void *dev_ret;
   if (hsa_status_t err =
           hsa_amd_memory_pool_allocate(coarsegrained_pool, sizeof(int),
                                        /*flags=*/0, &dev_ret))
     handle_error(err);
   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_size;
   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, sizeof(__llvm_libc::cpp::Atomic<int>),
+          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, sizeof(__llvm_libc::cpp::Atomic<int>),
+          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);
 
   // Initialize the RPC server's buffer for host-device communication.
-  server.reset(wavefront_size, &lock, server_inbox, server_outbox, buffer);
+  server.reset(port_size, wavefront_size, &lock, server_inbox, server_outbox,
+               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;

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

int load(int argc, char **argv, char **envp, void *image, size_t size,
 
   // Allocate space for the return pointer and initialize it to zero.
   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_size;
   uint32_t warp_size = 32;
-  void *server_inbox = allocator(sizeof(__llvm_libc::cpp::Atomic<int>));
-  void *server_outbox = allocator(sizeof(__llvm_libc::cpp::Atomic<int>));
-  void *buffer =
-      allocator(align_up(sizeof(__llvm_libc::rpc::Header) +
+  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)
     handle_error("Failed to allocate memory the RPC client / server.");
 
   // Initialize the RPC server's buffer for host-device communication.
-  server.reset(warp_size, &lock, server_inbox, server_outbox, buffer);
+  server.reset(port_size, warp_size, &lock, server_inbox, server_outbox,
+               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};
   if (CUresult err = launch_kernel(binary, stream, single_threaded_params,
                                    "_begin", init_args))
     handle_error(err);