[libc] Enable multiple threads to use RPC on the GPU

[jhuber6]

Summary

The execution model of the GPU expects that groups of threads will
execute in lock-step in SIMD fashion. It's both important for
performance and correctness that we treat this as the smallest possible
granularity for an RPC operation. Thus, we map multiple threads to a
single larger buffer and ship that across the wire.

This patch makes the necessary changes to support executing the RPC on
the GPU with multiple threads. This requires some workarounds to mimic
the model when handling the protocol from the CPU. I'm not completely
happy with some of the workarounds required, but I think it should work.

Uses some of the implementation details from D148191.

[jhuber6]

Cleanup leftover debugging code.

[jdoerfert]

Drive

libc/src/__support/RPC/rpc.h

48

Or

56

We should have a single generic warp size macro

115

why is data not typed here?

244

shouldn't 0 be something like gpu::get_id_in_lane()? Also below.

328

return;
}

[jhuber6]

libc/src/__support/RPC/rpc.h

56

True, we can put this in the GPU utils.

115

Just easier to manage since this needs to be copied from the host to the GPU and it's easier to model arguments to the kernel as void pointers.

244

This is basically just to let the CPU pretend like it has all the "threads" the GPU has. So given the GPU's lane size of 32. It'll execute this loop once. The CPU however will execute it 32 times because we set its lane size to 1. The idx is just an offset as you march in multiples of the lane size.

[jhuber6]

Address comments

[jhuber6]

Add MAX_LANE_SIZE for situations where we get the host to allocate enough to accomodate all sizes unconditionally.

[jhuber6]

This works as expected on my gfx1030, but seems to fail spuriously on sm_70. It's probably something to do with the Volta threading model, I'll need to investigate some more.

[jhuber6]

Updating to use an overloaded function type rather than make every function take uin32_t.

[jhuber6]

Add Volta warp syncs on divergent paths.

[jhuber6]

This should work as expected now on both AMDGPU and NVPTX.

[jhuber6]

Fix some iffy control flow that caused problems on AMDGPU sometimes. I've run
the tests 1000 times on both my gfx1030 and sm_70 and didn't observe any
failures, so I'm reasonably confident.

[jhuber6]

Update some names and allocate the proper amount of memory.

Also this adds a hack to work around a problem with the optimizer. The GPU
executes groups of threads in SIMD style. This means that basic blocks turn into
thread masks. THe problem was occuring when the optimizer would sink the common
'close()` call to release the thread lock until after another open call within
the same lane. This can be seen in this https://godbolt.org/z/cfEr4n5TE. There
is some ongoing work to make this actually work on AMD but for now it's a hack
to prevent these function calls from being merged. It has a performance penalty,
but it's minor.

[jhuber6]

Add a lane sync to address the convergence problem. Thanks to @nhaehnle for the suggestion.

[jhuber6]

Add acquire / release semantics on the device lock. The NVPTX backend doesn't
support release on stores so simply use an explicit fence and a relaxed
operation. The fence here will actually reduce to a full memory barrier, but
that should just be a slight performance hit because of Nvidia's lack of
support.

[JonChesterfield]

What's the reasoning behind the placement of sync_lanecalls? If it's meant to be one per basic block some are missing. The comment says it's added to close but the code has more calls than that.

If it's an amdgpu specific hack it should probably be named based on that and excluded on nvptx. Currently it reads like nvptx picks up a lot of synchronisation it doesn't need.

[jhuber6]

In D148943#4305191, @JonChesterfield wrote:

What's the reasoning behind the placement of sync_lanecalls? If it's meant to be one per basic block some are missing. The comment says it's added to close but the code has more calls than that.

If it's an amdgpu specific hack it should probably be named based on that and excluded on nvptx. Currently it reads like nvptx picks up a lot of synchronisation it doesn't need.

So, the convergence problem was present on Nvidia as well, it's just that its hardware model prevented it from deadlocking so it instead serialized the entire region. Right now we have a sync_lane in front of the open and the close mostly for safety to enforce the expected convergence of these operations. e.g. every thread that opens the port should close the port as well in the same context. The last sync is definitely necessary on the send_n call because the threads could copy variable amounts of data.

[jhuber6]

Add a check to detect divergence.

[jhuber6]

Updating to the new interface from @JonChesterfield

[jhuber6]

Rebase

[jhuber6]

Rebasing, the send and recv interface still needs to use the mask stored in the buffer because it's shared.

[jhuber6]

Remove file moved into another patch.

[jhuber6]

Rebasing, this currently deadlocks on the RPC test. Need to fix that.

[jhuber6]

Fixed deadlock, problem was the second lane sync wasn't rebased in.

[JonChesterfield]

I think this is OK. It seems likely that we'll be able to simplify parts of it later. I'll remove the outdated part of the commit message shortly, the state machine has changed since it was written

libc/src/__support/RPC/rpc.h

106–107

Not keen on this being a runtime value but it's not that clear how to avoid that on the host side so this seems OK

148

a bit frustrated that we need this sync for now

382–383

i think we're better off without the is_first_lane here, i.e. let all active lanes write the same value to buffer, but that should only be a codegen improvement

libc/src/__support/RPC/rpc_util.h

33

weird that these show up in the phab diff, they should already be in trunk

libc/src/__support/RPC/CMakeLists.txt

 if(NOT LIBC_TARGET_ARCHITECTURE_IS_GPU)
   return()
 endif()
 
 add_header_library(
   rpc
   HDRS
     rpc.h
     rpc_util.h
   DEPENDS
     libc.src.__support.common
     libc.src.__support.CPP.atomic
+    libc.src.__support.CPP.optional
+    libc.src.__support.CPP.functional
     libc.src.__support.GPU.utils
 )
 
 add_object_library(
   rpc_client
   SRCS
     rpc_client.cpp
   HDRS
     rpc_client.h
   DEPENDS
     libc.src.__support.GPU.utils
     .rpc
 )

libc/src/__support/RPC/rpc.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIBC_SRC_SUPPORT_RPC_RPC_H
 #define LLVM_LIBC_SRC_SUPPORT_RPC_RPC_H
 
 #include "rpc_util.h"
 #include "src/__support/CPP/atomic.h"
+#include "src/__support/CPP/functional.h"
 #include "src/__support/CPP/optional.h"
 #include "src/__support/GPU/utils.h"
 #include "src/string/memory_utils/memcpy_implementations.h"
 
 #include <stdint.h>
 
 namespace __llvm_libc {
 namespace rpc {
 
 /// A list of opcodes that we use to invoke certain actions on the server.
 enum Opcode : uint16_t {
   NOOP = 0,
   PRINT_TO_STDERR = 1,
   EXIT = 2,
   TEST_INCREMENT = 3,
 };
 
 /// A fixed size channel used to communicate between the RPC client and server.
-struct alignas(64) Buffer {
-  uint8_t data[62];
-  uint16_t opcode;
+struct Buffer {
+  uint64_t data[8];
 };
 static_assert(sizeof(Buffer) == 64, "Buffer size mismatch");
 
+/// The information associated with a packet. This indicates which operations to
+/// perform and which threads are active in the slots.
+struct Header {
+  uint64_t mask;
+  uint16_t opcode;
+};
+
+/// The data payload for the associated packet. We provide enough space for each
+/// thread in the cooperating lane to have a buffer.
+struct Payload {
+#if defined(LIBC_TARGET_ARCH_IS_GPU)
+  Buffer slot[gpu::LANE_SIZE];
+#else
+  // Flexible array size allocated at runtime to the appropriate size.
+  Buffer slot[];
+#endif
+};
+
+/// 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;
+};
+
 /// 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.
 ///
 /// No process writes to its inbox. Each toggles the bit in the outbox to pass
 /// ownership to the other process.
 /// When inbox == outbox, the current state machine owns the buffer.
 /// Initially the client is able to open any port as it will load 0 from both.

 ///     mirrored 'recv' / 'send' call.
 ///
 template <bool InvertInbox> struct Process {
   LIBC_INLINE Process() = default;
   LIBC_INLINE Process(const Process &) = default;
   LIBC_INLINE Process &operator=(const Process &) = default;
   LIBC_INLINE ~Process() = default;
 
+  uint32_t lane_size;
   cpp::Atomic<uint32_t> *lock;
   cpp::Atomic<uint32_t> *inbox;
   cpp::Atomic<uint32_t> *outbox;
-  Buffer *buffer;
+  Packet *buffer;
 
   /// Initialize the communication channels.
-  LIBC_INLINE void reset(void *lock, void *inbox, void *outbox, void *buffer) {
+  LIBC_INLINE void reset(uint32_t lane_size, void *lock, void *inbox,
+                         void *outbox, void *buffer) {
     *this = {
+        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<Buffer *>(buffer),
+        reinterpret_cast<Packet *>(buffer),
     };
   }
 
   /// Inverting the bits loaded from the inbox in exactly one of the pair of
   /// processes means that each can use the same state transitions.
   /// Whichever process has InvertInbox==false is the initial owner.
   /// Inbox equal Outbox => current process owns the buffer
   /// Inbox difer Outbox => current process does not own the buffer
   /// At startup, memory is zero initialised and raw loads of either mailbox
   /// would return zero. Thus both would succeed in opening a port and data
template <bool InvertInbox> struct Process {
   /// Attempt to claim the lock at index. Return true on lock taken.
   /// The lock is held when the zeroth bit of the uint32_t at lock[index]
   /// is set, and available when that bit is clear. Bits [1, 32) are zero.
   /// Or with one is a no-op when the lock is already held.
   LIBC_INLINE bool try_lock(uint64_t index) {
     return lock[index].fetch_or(1, cpp::MemoryOrder::RELAXED) == 0;
   }
 
-  // Unlock the lock at index.
+  /// Unlock at the specified index. Explicit lane syncs ensure convergent
+  /// behaviour, otherwise the compiler may sink the store and deadlock.
   LIBC_INLINE void unlock(uint64_t index) {
+    gpu::sync_lane(buffer[index].header.mask);
     lock[index].store(0, cpp::MemoryOrder::RELAXED);
+    gpu::sync_lane(buffer[index].header.mask);
+  }
+
+  /// 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) {
+    if constexpr (is_process_gpu()) {
+      fn(&buffer[index].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]);
+    }
+  }
+
+  /// 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) {
+    if constexpr (is_process_gpu()) {
+      fn(&buffer[index].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);
+    }
   }
 };
 
 /// 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 {
   // TODO: This should be move-only.
   LIBC_INLINE Port(Process<T> &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].opcode;
+    return process.buffer[index].header.opcode;
   }
 
   LIBC_INLINE void close() { process.unlock(index); }
 
 private:
   Process<T> &process;
   uint64_t index;
   uint32_t out;
template <bool T> template <typename F> LIBC_INLINE void Port<T>::send(F fill) {
 
   // We need to wait until we own the buffer before sending.
   while (Process<T>::buffer_unavailable(in, out)) {
     sleep_briefly();
     in = process.load_inbox(index);
   }
 
   // Apply the \p fill function to initialize the buffer and release the memory.
-  fill(&process.buffer[index]);
+  process.invoke_rpc(fill, index);
   out = !out;
   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 <bool T> template <typename U> LIBC_INLINE void Port<T>::recv(U use) {
   uint32_t in = process.load_inbox(index);
 
   // We need to wait until we own the buffer before receiving.
   while (Process<T>::buffer_unavailable(in, out)) {
     sleep_briefly();
     in = process.load_inbox(index);
   }
   atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
 
   // Apply the \p use function to read the memory out of the buffer.
-  use(&process.buffer[index]);
+  process.invoke_rpc(use, index);
   out = !out;
   process.outbox[index].store(out, cpp::MemoryOrder::RELAXED);
 }
 
 /// Combines a send and receive into a single function.
 template <bool T>
 template <typename F, typename U>
 LIBC_INLINE void Port<T>::send_and_recv(F fill, U use) {

 }
 
 /// Sends an arbitrarily sized data buffer \p src across the shared channel in
 /// multiples of the packet length.
 template <bool T>
 LIBC_INLINE void Port<T>::send_n(const void *src, uint64_t size) {
   // TODO: We could send the first bytes in this call and potentially save an
   // extra send operation.
-  send([=](Buffer *buffer) { buffer->data[0] = size; });
+  // TODO: We may need a way for the CPU to send different strings per thread.
+  send([=](Buffer *buffer) {
+    reinterpret_cast<uint64_t *>(buffer->data)[0] = 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);
 }
 
 /// 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 <bool T>
 template <typename A>
 LIBC_INLINE void Port<T>::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
+  // possibility that the sent data is of different length.
+  if constexpr (is_process_gpu()) {
     uint64_t size = 0;
-  recv([&](Buffer *buffer) { size = buffer->data[0]; });
-  uint8_t *dst = reinterpret_cast<uint8_t *>(alloc(size));
+    recv([&](Buffer *buffer) {
+      size = reinterpret_cast<uint64_t *>(buffer->data)[0];
+    });
+    uint8_t *dst = reinterpret_cast<uint8_t *>(alloc(size), gpu::get_lane_id());
     for (uint64_t idx = 0; idx < size; idx += sizeof(Buffer::data)) {
       recv([=](Buffer *buffer) {
-      uint64_t len =
-          size - idx > sizeof(Buffer::data) ? sizeof(Buffer::data) : size - idx;
+        uint64_t len = size - idx > sizeof(Buffer::data) ? sizeof(Buffer::data)
+                                                         : size - idx;
         inline_memcpy(dst + idx, buffer->data, len);
       });
     }
+    return;
+  } else {
+    uint64_t size[MAX_LANE_SIZE];
+    uint8_t *dst[MAX_LANE_SIZE];
+    uint64_t max = 0;
+    recv([&](Buffer *buffer, uint32_t id) {
+      size[id] = reinterpret_cast<uint64_t *>(buffer->data)[0];
+      dst[id] = reinterpret_cast<uint8_t *>(alloc(size[id], id));
+      max = size[id] > max ? size[id] : max;
+    });
+    for (uint64_t idx = 0; idx < max; idx += sizeof(Buffer::data)) {
+      recv([=](Buffer *buffer, uint32_t id) {
+        uint64_t len = size[id] - idx > sizeof(Buffer::data)
+                           ? sizeof(Buffer::data)
+                           : size[id] - idx;
+        if (idx < size[id])
+          inline_memcpy(dst[id] + idx, buffer->data, len);
+      });
+    }
+    return;
+  }
 }
 
 /// 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.
-LIBC_INLINE cpp::optional<Client::Port> Client::try_open(uint16_t opcode) {
+/// Opening a port is only valid if the `opcode` is the sam accross every
+/// participating thread.
+[[clang::convergent]] LIBC_INLINE cpp::optional<Client::Port>
+Client::try_open(uint16_t opcode) {
   constexpr uint64_t index = 0;
   // Attempt to acquire the lock on this index.
-  if (!try_lock(index))
+  uint64_t lane_mask = gpu::get_lane_mask();
+  bool claimed = false;
+  if (is_first_lane(lane_mask))
+    claimed = try_lock(index);
+  claimed = gpu::broadcast_value(claimed);
+
+  if (!claimed)
     return cpp::nullopt;
 
   // The mailbox state must be read with the lock held.
   atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
 
   uint32_t in = load_inbox(index);
   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 (buffer_unavailable(in, out)) {
     unlock(index);
     return cpp::nullopt;
   }
 
-  buffer->opcode = opcode;
-  return Port(*this, index, out);
+  // TODO: The index here will be non-zero when we support multiple slots.
+  if (is_first_lane(lane_mask)) {
+    buffer[index].header.opcode = opcode;
+    buffer[index].header.mask = lane_mask;
+  }
+  gpu::sync_lane(lane_mask);
+  return Port(*this, 0, out);
 }
 
 LIBC_INLINE Client::Port Client::open(uint16_t opcode) {
   for (;;) {
     if (cpp::optional<Client::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
-LIBC_INLINE cpp::optional<Server::Port> Server::try_open() {
+[[clang::convergent]] LIBC_INLINE cpp::optional<Server::Port>
+Server::try_open() {
   constexpr uint64_t index = 0;
   uint32_t in = load_inbox(index);
   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 (buffer_unavailable(in, out))
     return cpp::nullopt;
 
   // Attempt to acquire the lock on this index.
-  if (!try_lock(index))
+  uint64_t lane_mask = gpu::get_lane_mask();
+  bool claimed = false;
+  if (is_first_lane(lane_mask))
+    claimed = try_lock(index);
+  claimed = gpu::broadcast_value(claimed);
+
+  if (!claimed)
     return cpp::nullopt;
 
   // The mailbox state must be read with the lock held.
   atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
 
   in = load_inbox(index);
   out = outbox[index].load(cpp::MemoryOrder::RELAXED);
 

libc/src/__support/RPC/rpc_util.h

 //===-- Shared memory RPC client / server utilities -------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIBC_SRC_SUPPORT_RPC_RPC_UTILS_H
 #define LLVM_LIBC_SRC_SUPPORT_RPC_RPC_UTILS_H
 
+#include "src/__support/GPU/utils.h"
 #include "src/__support/macros/attributes.h"
 #include "src/__support/macros/properties/architectures.h"
 
 namespace __llvm_libc {
 namespace rpc {
 
+/// Maximum amount of data a single lane can use.
+constexpr uint64_t MAX_LANE_SIZE = 64;
+
 /// Suspend the thread briefly to assist the thread scheduler during busy loops.
 LIBC_INLINE void sleep_briefly() {
 #if defined(LIBC_TARGET_ARCH_IS_NVPTX) && __CUDA_ARCH__ >= 700
   asm("nanosleep.u32 64;" ::: "memory");
 #elif defined(LIBC_TARGET_ARCH_IS_AMDGPU)
   __builtin_amdgcn_s_sleep(2);
 #else
   // Simply do nothing if sleeping isn't supported on this platform.
 #endif
 }
 
+/// Get the first active thread inside the lane.
+LIBC_INLINE uint64_t get_first_lane_id(uint64_t lane_mask) {
+  return __builtin_ffsl(lane_mask) - 1;
+}
+
+/// Conditional that is only true for a single thread in a lane.
+LIBC_INLINE bool is_first_lane(uint64_t lane_mask) {
+  return gpu::get_lane_id() == get_first_lane_id(lane_mask);
+}
+
+/// Conditional to indicate if this process is running on the GPU.
+LIBC_INLINE constexpr bool is_process_gpu() {
+#if defined(LIBC_TARGET_ARCH_IS_GPU)
+  return true;
+#else
+  return false;
+#endif
+}
+
 } // namespace rpc
 } // namespace __llvm_libc
 
 #endif

libc/startup/gpu/amdgpu/start.cpp

void initialize(int argc, char **argv, char **env, void *in, void *out,
                 void *buffer) {
   // We need a single GPU thread to perform the initialization of the global
   // constructors and data. We simply mask off all but a single thread and
   // execute.
   count.fetch_add(1, cpp::MemoryOrder::RELAXED);
   if (gpu::get_thread_id() == 0 && gpu::get_block_id() == 0) {
     // We need to set up the RPC client first in case any of the constructors
     // require it.
-    rpc::client.reset(&lock, in, out, buffer);
+    rpc::client.reset(gpu::get_lane_size(), &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.
     atexit(&call_fini_array_callbacks);
     call_init_array_callbacks(argc, argv, env);
   }

libc/startup/gpu/nvptx/start.cpp

void initialize(int argc, char **argv, char **env, void *in, void *out,
                 void *buffer) {
   // We need a single GPU thread to perform the initialization of the global
   // constructors and data. We simply mask off all but a single thread and
   // execute.
   count.fetch_add(1, cpp::MemoryOrder::RELAXED);
   if (gpu::get_thread_id() == 0 && gpu::get_block_id() == 0) {
     // We need to set up the RPC client first in case any of the constructors
     // require it.
-    rpc::client.reset(&lock, in, out, buffer);
+    rpc::client.reset(gpu::get_lane_size(), &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.
     // FIXME: The function pointer escaping this TU causes warnings.
     __llvm_libc::atexit(&call_fini_array_callbacks);
     call_init_array_callbacks(argc, argv, env);

libc/test/integration/startup/gpu/CMakeLists.txt

add_integration_test(
   startup_rpc_test
   SUITE libc-startup-tests
   SRCS
     rpc_test.cpp
   DEPENDS
     libc.src.__support.RPC.rpc_client
     libc.src.__support.GPU.utils
   LOADER_ARGS
-    --blocks 16
-    --threads 1
+    --blocks-x 2
+    --blocks-y 2
+    --blocks-z 2
+    --threads-x 4
+    --threads-y 4
+    --threads-z 4
 )
 
 add_integration_test(
   init_fini_array_test
   SUITE libc-startup-tests
   SRCS
     init_fini_array_test.cpp
 )

libc/test/integration/startup/gpu/rpc_test.cpp

 //===-- Loader test to check the RPC interface with the loader ------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "src/__support/GPU/utils.h"
 #include "src/__support/RPC/rpc_client.h"
 #include "test/IntegrationTest/test.h"
 
 using namespace __llvm_libc;
 
 static void test_add_simple() {
-  uint32_t num_additions = 1000 + 10 * gpu::get_block_id_x();
+  uint32_t num_additions =
+      10 + 10 * gpu::get_thread_id() + 10 * gpu::get_block_id();
   uint64_t cnt = 0;
   for (uint32_t i = 0; i < num_additions; ++i) {
     rpc::Client::Port port = rpc::client.open(rpc::TEST_INCREMENT);
     port.send_and_recv(
         [=](rpc::Buffer *buffer) {
           reinterpret_cast<uint64_t *>(buffer->data)[0] = cnt;
         },
         [&](rpc::Buffer *buffer) {
           cnt = reinterpret_cast<uint64_t *>(buffer->data)[0];
         });
     port.close();
   }
   ASSERT_TRUE(cnt == num_additions && "Incorrect sum");
 }
 
+// Test to ensure that the RPC mechanism doesn't hang on divergence.
+static void test_noop(uint8_t data) {
+  rpc::Client::Port port = rpc::client.open(rpc::NOOP);
+  port.send([=](rpc::Buffer *buffer) { buffer->data[0] = data; });
+  port.close();
+}
+
 TEST_MAIN(int argc, char **argv, char **envp) {
   test_add_simple();
 
+  if (gpu::get_thread_id() % 2)
+    test_noop(1);
+  else
+    test_noop(2);
+
   return 0;
 }

libc/utils/gpu/loader/Loader.h

 };
 
 /// Generic interface to load the \p image and launch execution of the _start
 /// kernel on the target device. Copies \p argc and \p argv to the device.
 /// Returns the final value of the `main` function on the device.
 int load(int argc, char **argv, char **evnp, void *image, size_t size,
          const LaunchParameters &params);
 
+/// Return \p V aligned "upwards" according to \p Align.
+template <typename V, typename A> inline V align_up(V val, A align) {
+  return ((val + V(align) - 1) / V(align)) * V(align);
+}
+
 /// Copy the system's argument vector to GPU memory allocated using \p alloc.
 template <typename Allocator>
 void *copy_argument_vector(int argc, char **argv, Allocator alloc) {
   size_t argv_size = sizeof(char *) * (argc + 1);
   size_t str_size = 0;
   for (int i = 0; i < argc; ++i)
     str_size += strlen(argv[i]) + 1;
 

libc/utils/gpu/loader/Server.h

 /// are any active requests.
 void handle_server() {
   auto port = server.try_open();
   if (!port)
     return;
 
   switch (port->get_opcode()) {
   case __llvm_libc::rpc::Opcode::PRINT_TO_STDERR: {
-    uint64_t str_size;
-    char *str = nullptr;
-    port->recv_n([&](uint64_t size) {
-      str_size = size;
-      str = new char[size];
-      return str;
+    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];
     });
-    fwrite(str, str_size, 1, stderr);
-    delete[] str;
+    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 *) { /* no-op */ });
-    return;
+    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.
+  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>),
           /*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>),
           /*flags=*/0, &server_outbox))
     handle_error(err);
   if (hsa_status_t err = hsa_amd_memory_pool_allocate(
-          finegrained_pool, sizeof(__llvm_libc::rpc::Buffer),
+          finegrained_pool,
+          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);
 
   // Initialie all the arguments (explicit and implicit) to zero, then set the
   // explicit arguments to the values created above.
int load(int argc, char **argv, char **envp, void *image, size_t size,
   packet->kernarg_address = args;
 
   // Create a signal to indicate when this packet has been completed.
   if (hsa_status_t err =
           hsa_signal_create(1, 0, nullptr, &packet->completion_signal))
     handle_error(err);
 
   // Initialize the RPC server's buffer for host-device communication.
-  server.reset(&lock, server_inbox, server_outbox, buffer);
+  server.reset(wavefront_size, &lock, server_inbox, server_outbox, buffer);
 
   // Initialize the packet header and set the doorbell signal to begin execution
   // by the HSA runtime.
   uint16_t header =
       (HSA_PACKET_TYPE_KERNEL_DISPATCH << HSA_PACKET_HEADER_TYPE) |
       (HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_ACQUIRE_FENCE_SCOPE) |
       (HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_RELEASE_FENCE_SCOPE);
   __atomic_store_n(&packet->header, header | (packet->setup << 16),

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);
 
+  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(sizeof(__llvm_libc::rpc::Buffer));
+  void *buffer =
+      allocator(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.");
 
   // Set up the arguments to the '_start' kernel on the GPU.
   uint64_t args_size = sizeof(kernel_args_t);
   kernel_args_t args;
   std::memset(&args, 0, args_size);
   args.argc = argc;
   args.argv = dev_argv;
   args.envp = dev_envp;
   args.ret = reinterpret_cast<void *>(dev_ret);
   args.inbox = server_outbox;
   args.outbox = server_inbox;
   args.buffer = buffer;
   void *args_config[] = {CU_LAUNCH_PARAM_BUFFER_POINTER, &args,
                          CU_LAUNCH_PARAM_BUFFER_SIZE, &args_size,
                          CU_LAUNCH_PARAM_END};
 
   // Initialize the RPC server's buffer for host-device communication.
-  server.reset(&lock, server_inbox, server_outbox, buffer);
+  server.reset(warp_size, &lock, server_inbox, server_outbox, buffer);
 
   // Call the kernel with the given arguments.
   if (CUresult err = cuLaunchKernel(
           function, params.num_blocks_x, params.num_blocks_y,
           params.num_blocks_z, params.num_threads_x, params.num_threads_y,
           params.num_threads_z, 0, stream, nullptr, args_config))
     handle_error(err);