[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

100

why is data not typed here?

203

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

280

return;
}

[jhuber6]

libc/src/__support/RPC/rpc.h

56

True, we can put this in the GPU utils.

100

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.

203

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

98–99

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

109

a bit frustrated that we need this sync for now

330–333

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

 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;
+  uint8_t data[64];
 };
 static_assert(sizeof(Buffer) == 64, "Buffer size mismatch");
 
+/// A buffer used to swap data between the client and server across an entire
+/// lane. The GPU executes threads in lock-step so we use the size of a 'warp'
+/// or 'wavefront' as our minimum granularity.
+struct LaneBuffer {
+  uint16_t opcode;
+  uint64_t mask;
+#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 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
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;
 
+  uint32_t lane_size;
   cpp::Atomic<uint32_t> *lock;
   cpp::Atomic<uint32_t> *inbox;
   cpp::Atomic<uint32_t> *outbox;
-  Buffer *buffer;
+  LaneBuffer *buffer;
 
   /// Initialize the communication channels.
-  LIBC_INLINE void reset(void *lock, void *inbox, void *outbox, void *buffer) {
-    *this = {
-        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),
-    };
+  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<LaneBuffer *>(data);
   }
 
   /// 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) {
+    for (uint32_t i = 0; i < lane_size; i += gpu::get_lane_size())
+      if (is_process_gpu() || buffer[index].mask & 1ul << i)
+        fn(&buffer[index].slot[gpu::get_lane_id() + 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) {
+    for (uint32_t i = 0; i < lane_size; i += gpu::get_lane_size())
+      if (is_process_gpu() || buffer[index].mask & 1ul << i)
+        fn(&buffer[index].slot[gpu::get_lane_id() + i], gpu::get_lane_id() + 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)
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.
-  fill(&process.buffer[index]);
+  process.invoke_rpc(fill, 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.
-  use(&process.buffer[index]);
+  process.invoke_rpc(use, 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);
   recv(use);
 }
 
 /// Combines a receive and send operation into a single function. The \p work
 /// function modifies the buffer in-place and the send is only used to initiate
 /// the copy back.
 template <typename W> LIBC_INLINE void Port::recv_and_send(W work) {
   recv(work);
   send([](Buffer *) { /* no-op */ });
 }
 
 /// Sends an arbitrarily sized data buffer \p src across the shared channel in
 /// multiples of the packet length.
 LIBC_INLINE void Port::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);
     });
   }
 }
 
 /// 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
+  // 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<Port> Client::try_open(uint16_t opcode) {
+/// Opening a port is only valid if the `opcode` is the sam accross every
+/// participating thread.
+LIBC_INLINE cpp::optional<Port> Client::try_open(const uint16_t opcode) {
   // Attempt to acquire the lock on this index.
-  if (lock->fetch_or(1, cpp::MemoryOrder::RELAXED))
+  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::RELAXED);
+  previous = gpu::broadcast_value(previous);
+
+  if (previous != 0)
     return cpp::nullopt;
 
   uint32_t in = inbox->load(cpp::MemoryOrder::RELAXED);
   uint32_t out = outbox->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);
     return cpp::nullopt;
   }
 
-  buffer->opcode = opcode;
+  if (is_first_lane(lane_mask)) {
+    buffer[0].opcode = opcode;
+    buffer[0].mask = lane_mask;
+  }
   return Port(*this, 0, out);
 }
 
 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
 LIBC_INLINE cpp::optional<Port> Server::try_open() {
   uint32_t in = inbox->load(cpp::MemoryOrder::RELAXED);
   uint32_t out = outbox->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;
 
   // Attempt to acquire the lock on this index.
-  if (lock->fetch_or(1, cpp::MemoryOrder::RELAXED))
+  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::RELAXED);
+  previous = gpu::broadcast_value(previous);
+
+  if (previous != 0)
     return cpp::nullopt;
 
   in = inbox->load(cpp::MemoryOrder::RELAXED);
   out = outbox->load(cpp::MemoryOrder::RELAXED);
 
   if (!can_recv_data(in, out)) {
     lock->store(0, cpp::MemoryOrder::RELAXED);
     return cpp::nullopt;

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

 
 static cpp::Atomic<uint32_t> lock = 0;
 
 static cpp::Atomic<uint32_t> init = 0;
 
 void init_rpc(void *in, void *out, void *buffer) {
   // Only a single thread should update the RPC data.
   if (gpu::get_thread_id() == 0 && gpu::get_block_id() == 0) {
-    rpc::client.reset(&lock, in, out, buffer);
+    rpc::client.reset(gpu::get_lane_size(), &lock, in, out, buffer);
     init.store(1, cpp::MemoryOrder::RELAXED);
   }
 
   // Wait until the previous thread signals that the data has been written.
   while (!init.load(cpp::MemoryOrder::RELAXED))
     rpc::sleep_briefly();
 
   // Wait for the threads in the block to converge and fence the write.

libc/startup/gpu/nvptx/start.cpp

 
 static cpp::Atomic<uint32_t> lock = 0;
 
 static cpp::Atomic<uint32_t> init = 0;
 
 void init_rpc(void *in, void *out, void *buffer) {
   // Only a single thread should update the RPC data.
   if (gpu::get_thread_id() == 0 && gpu::get_block_id() == 0) {
-    rpc::client.reset(&lock, in, out, buffer);
+    rpc::client.reset(gpu::get_lane_size(), &lock, in, out, buffer);
     init.store(1, cpp::MemoryOrder::RELAXED);
   }
 
   // Wait until the previous thread signals that the data has been written.
   while (!init.load(cpp::MemoryOrder::RELAXED))
     rpc::sleep_briefly();
 
   // Wait for the threads in the block to converge and fence the write.

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
 )

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::Port port = rpc::client.open(rpc::TEST_INCREMENT);
     port.send_and_recv(
-        [=](rpc::Buffer *buffer) {
+        [=](rpc::Buffer *buffer, uint32_t) {
           reinterpret_cast<uint64_t *>(buffer->data)[0] = cnt;
         },
-        [&](rpc::Buffer *buffer) {
+        [&](rpc::Buffer *buffer, uint32_t) {
           cnt = reinterpret_cast<uint64_t *>(buffer->data)[0];
         });
     port.close();
   }
   ASSERT_TRUE(cnt == num_additions && "Incorrect sum");
 }
 
 TEST_MAIN(int argc, char **argv, char **envp) {
   test_add_simple();
 
   return 0;
 }

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: {
-    void *str = nullptr;
-    port->recv_n([&](uint64_t size) { return str = malloc(size); });
+    void *strs[__llvm_libc::rpc::MAX_LANE_SIZE] = {nullptr};
+    port->recv_n(
+        [&](uint64_t size, uint32_t id) { return strs[id] = malloc(size); });
+    for (void *str : strs) {
+      if (str) {
         fputs(reinterpret_cast<char *>(str), stderr);
         free(str);
+      }
+    }
     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,
+          sizeof(__llvm_libc::rpc::LaneBuffer) +
+              (wavefront_size * sizeof(__llvm_libc::rpc::Buffer)),
           /*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(sizeof(__llvm_libc::rpc::LaneBuffer) +
+                           (warp_size * sizeof(__llvm_libc::rpc::Buffer)));
   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);