[libc] Update RPC interface for system utilities on the GPU

[jhuber6]

Summary

This patch reworks the RPC interface to allow more generic memory
operations using the shared better. This patch decomposes the entire RPC
interface into opening a port and calling send or recv on it.

The send function sends a single packet of the length of the buffer.
The recv function is paired with the send call to then use the data.
So, any aribtrary combination of sending packets is possible. The only
restriction is that the client initiates the exchange with a send
while the server consumes it with a recv.

The operation of this is driven by two independent state machines that
tracks the buffer ownership during loads / stores. We keep track of two
so that we can transition between a send state and a recv state without
an extra wait. State transitions are observed via bit toggling, e.g.

This interface supports an efficient send -> ack -> send -> ack -> send
interface and allows for the last send to be ignored without checking
the ack.

A following patch will add some more comprehensive testing to this interface. I
I informally made an RPC call that simply incremented an integer and it took
roughly 10 microsends to complete an RPC call.

[goldstein.w.n]

libc/src/__support/RPC/rpc.h

57

IIUC this means it must be send/recv/senc/recv back-and-forth pattern? This seems contradictory to "mostly arbitrary combinations of 'send' and 'recv'" no?

[goldstein.w.n]

libc/src/__support/RPC/rpc.h

255

given that is only two states, exchange is enough, no need for CAS.

Likewise below.

[jhuber6]

Add sleeps and update recv_n.

[jhuber6]

libc/src/__support/RPC/rpc.h

57

Yes, poor choice of words. I meant arbitrary combinations on one side as long as it's mirrored on the other side. So, the following will work:

Client:  send -> send -> send -> recv -> send -> recv -> recv -> send -> recv
Server: recv -> recv -> recv -> send -> recv -> send -> send -> recv -> send

255

Good point, it's mostly a placeholder right now as we only support single threaded access. (At least this means if you call it with multiple threads it will wait until the other one is done).

[jhuber6]

Change atomic and adjust comments.

[jhuber6]

Remove unused variable.

[lntue]

libc/src/__support/OSUtil/gpu/io.cpp

19

Does send_n still expect msg.data() is null-terminated?

[jhuber6]

libc/src/__support/OSUtil/gpu/io.cpp

19

We copy the entire buffer now, including the terminating null byte. Previously we had to add a null byte because we printed it in chunks. it worked alright but would be sub-optimal when multiple threads are involved.

[lntue]

libc/src/__support/OSUtil/gpu/quick_exit.cpp

21

capture-by-value is better here.

libc/src/__support/RPC/rpc.h

33

Can this be enum class ?

74

Since this class is internally used, shall we make it enum class to be a bit safer? If you don't want to keep adding Flags::Data, you can add using Flags::Data and using Flags::Ack, or simply add using enum Flags with C++20 inside struct Process.

167

a != b is a safer option for boolean than bitwise-xor.

213

capture-by-value is better here.

230

capture-by-reference is not needed.

233

I think capture-by-value is good enough here.

250

a != b is a safer option for boolean than bitwise-xor.

262

a != b is a safer option for boolean than bitwise-xor.

265

nit: I actually prefer while (true) to for (;;)

287

a != b is actually a safer xor operator for boolean than bitwise-xor ^. So logically, this is equivalent to

bool(in & Data) == bool(out & Process::Ack)

Can you double check if it is correct?

299

nit: I actually prefer while (true) to for (;;)

libc/utils/gpu/loader/Server.h

41

nothing is captured with this lambda.

47

nothing is captured with this lambda.

[jhuber6]

Addressing comments

[jhuber6]

libc/src/__support/RPC/rpc.h

33

We could, but since this is passed over the wire I think it's okay to treat it as just an integer.

167

These are bits, we can't do things like & if it's an enum class without a ton of conversions. It should be uint32_t however.

230

This copies the size out.

265

I saw examples of both in the codebase, Maybe we should make a style rule for that.

287

The compiler will reduce it to an XOR AFAIK, but I explicitly use the language "bits are not equal" in the logic, so it's a lot clearer.

[jhuber6]

Make a few changes. Add a placeholder for an index for when multiple ports
will be supported. Also improve performance by removing the atomic load of the
write-only variable out.

[jhuber6]

Get rid of RAII interface for now. Was causing issues when using multiple threads.

[jdoerfert]

some comments, waiting on the next version now.

libc/src/__support/RPC/rpc.h

42

I would argue data should be int8_t[62] to get the size right.

74

Enum class or just constexpr constants.

164–166

Complex logic, please put it into a helper with comment. Probably the entire while loop.

173–174

reorder to simplify?

185–186

Same as above.

236

Here and above, move the cast out the loop, or at least stmt, to make it simpler (to read).

258–259

This should use the helper as well, and pass a retry flag as false, hence give up rather than loop. The scheme above otherwise exists 4 times.

[jhuber6]

Addressing comments.

[jhuber6]

Fixing.

[jhuber6]

libc/src/__support/RPC/rpc.h

164–166

Since the other use is an if I think it's just helpful to put the logic to a named function. I think showing an explicit busy-wait is good.

[jdoerfert]

Last real issue I have is the try open loop, see below.

libc/src/__support/RPC/rpc.h

43–44

static assert sizeof

303

I think {try,}open need a common helper impl. but should not use each other.
The problem right now is that we check the in/out multiple times for no good reason.
In fact, we might get stuck in a weird state which we could reasonably report as error.
Why not a single impl. with a flag, in case it is set you loop, otherwise you give up?

libc/utils/gpu/loader/Server.h

42

FWIW, later we want an easier way to get a single Ty value from the RPC. Seems like something common that should hide the casts and encoding.

[jhuber6]

libc/src/__support/RPC/rpc.h

43–44

My future plan is to collapse all this inbox / outbox stuff into a struct of arrays, but we should go for a static 64-byte buffer in any case.

303

The logic is that the client can only open a port capable of sending and the server can only open a port capable of receiving.

libc/utils/gpu/loader/Server.h

42

Definitely, I'm already tired of reinterpret cast.

[JonChesterfield]

This comment has been deleted.

libc/src/__support/RPC/rpc.h

253

This leaks ports if they're closed after a non-even number of send or recv calls

libc/utils/gpu/loader/Server.h

48

This leaks the port here^

[jhuber6]

libc/src/__support/RPC/rpc.h

253

The port can be opened in the Ack / Data state 0, 0 or 1, 1. it doesn't get stuck at 1, 1 because we treat 1, 1 the same as 0, 0 in the logic.

[JonChesterfield]

libc/src/__support/RPC/rpc.h

253

Uh, maybe. I'm still having trouble mapping between in/data in/ack out/data out/ack bits and client!=server and the two bit, client==server models.

[jdoerfert]

I'm fine with this. Any objections?

libc/utils/gpu/loader/Server.h

27

[JonChesterfield]

Main objection from me is it uses twice as much shared state as necessary. Using many more states than necessary makes the invariants difficult to follow.

Specifically, this implements the four states GPU owns, send to host, host owns, send to GPU, which are the right states for mutex here, but it does it using four bits of state instead of two. It misses the fundamental symmetry in the system - send and recv are the same thing.

Joseph and I are slowly working through that offline.

[jdoerfert]

We can reduce the bits from 4 to 2 later if that results in less complexity or better performance. I don't think that is in itself a blocker. When we see the patch to change it we can probably easily tell if it is simpler, and we hopefully have more stress testing then to trust whatever changes we make.

[JonChesterfield]

I think we should land D148191 and then adopt this send/recv scheme on top of it if preferred. That patch keeps the current semantics while extending to multiple warps, in a fashion which is known to support the interface proposed here.

What this patch does is simultaneously change the underlying state machine and the layers of code built on top of it, in a fashion which I think will need to be revised again when changing to support multiple warps.

[jhuber6]

Going to land this with the expectation that we can hack on it later and it will make subsequent patches easier to follow. Portions from D148191 will be added soon.

libc/src/__support/OSUtil/gpu/io.h

 
 #include "src/__support/CPP/string_view.h"
 #include "src/__support/macros/attributes.h" // LIBC_INLINE
 
 namespace __llvm_libc {
 
 void write_to_stderr(cpp::string_view msg);
 
-LIBC_INLINE void write_to_stderr(const char *msg) {
-  write_to_stderr(cpp::string_view(msg));
-}
-
 } // namespace __llvm_libc
 
 #endif // LLVM_LIBC_SRC_SUPPORT_OSUTIL_LINUX_IO_H

libc/src/__support/OSUtil/gpu/io.cpp

 //===-------------- GPU implementation of IO utils --------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "io.h"
 
 #include "src/__support/CPP/string_view.h"
 #include "src/__support/RPC/rpc_client.h"
 #include "src/string/memory_utils/memcpy_implementations.h"
 
 namespace __llvm_libc {
 
-namespace internal {
-
-static constexpr size_t BUFFER_SIZE = sizeof(rpc::Buffer) - sizeof(uint64_t);
-static constexpr size_t MAX_STRING_SIZE = BUFFER_SIZE;
-
-LIBC_INLINE void send_null_terminated(cpp::string_view src) {
-  rpc::client.run(
-      [&](rpc::Buffer *buffer) {
-        buffer->data[0] = rpc::Opcode::PRINT_TO_STDERR;
-        char *data = reinterpret_cast<char *>(&buffer->data[1]);
-        inline_memcpy(data, src.data(), src.size());
-        data[src.size()] = '\0';
-      },
-      [](rpc::Buffer *) { /* void */ });
-}
-
-} // namespace internal
-
 void write_to_stderr(cpp::string_view msg) {
-  bool send_empty_string = true;
-  for (; !msg.empty();) {
-    const auto chunk = msg.substr(0, internal::MAX_STRING_SIZE);
-    internal::send_null_terminated(chunk);
-    msg.remove_prefix(chunk.size());
-    send_empty_string = false;
-  }
-  if (send_empty_string)
-    internal::send_null_terminated("");
+  rpc::Port port = rpc::client.open(rpc::PRINT_TO_STDERR);
+  port.send_n(msg.data(), msg.size() + 1);
+  port.close();
 }
 
 } // namespace __llvm_libc

libc/src/__support/OSUtil/gpu/quick_exit.cpp

 #include "quick_exit.h"
 
 #include "src/__support/RPC/rpc_client.h"
 #include "src/__support/macros/properties/architectures.h"
 
 namespace __llvm_libc {
 
 void quick_exit(int status) {
-  // TODO: Support asynchronous calls so we don't wait and exit from the GPU
-  // immediately.
-  rpc::client.run(
-      [&](rpc::Buffer *buffer) {
-        buffer->data[0] = rpc::Opcode::EXIT;
-        buffer->data[1] = status;
-      },
-      [](rpc::Buffer *) { /* void */ });
+  rpc::Port port = rpc::client.open(rpc::EXIT);
+  port.send([&](rpc::Buffer *buffer) {
+    reinterpret_cast<uint32_t *>(buffer->data)[0] = status;
+  });
+  port.close();
 
 #if defined(LIBC_TARGET_ARCH_IS_NVPTX)
   asm("exit;" ::: "memory");
 #elif defined(LIBC_TARGET_ARCH_IS_AMDGPU)
   // This will terminate the entire wavefront, may not be valid with divergent
   // work items.
   asm("s_endpgm" ::: "memory");
 #endif
   __builtin_unreachable();
 }
 
 } // namespace __llvm_libc
 
 #endif // LLVM_LIBC_SRC_SUPPORT_OSUTIL_GPU_QUICK_EXIT_H

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/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. We
-/// reserve the first 255 values for internal libc usage.
-enum Opcode : uint64_t {
+/// 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,
-  LIBC_LAST = (1UL << 8) - 1,
 };
 
 /// A fixed size channel used to communicate between the RPC client and server.
-struct Buffer {
-  uint64_t data[8];
+struct alignas(64) Buffer {
+  uint8_t data[62];
+  uint16_t opcode;
 };
+static_assert(sizeof(Buffer) == 64, "Buffer size mismatch");
 
 /// 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.
+/// 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
+/// using this process are as follows:
+///   - The client will always start with a 'send' operation
+///   - The server will always start with a 'recv' operation
+///   - For every 'send' / 'recv' call on one side of the process there is a
+///     mirrored 'recv' / 'send' call.
+///
+/// The communication protocol is organized as a pair of two-state state
+/// machines. One state machine tracks outgoing sends and the other tracks
+/// incoming receives. For example, a 'send' operation uses its input 'Ack' bit
+/// and its output 'Data' bit. If these bits are equal the sender owns the
+/// buffer, otherwise the receiver owns the buffer and we wait. Similarly, a
+/// 'recv' operation uses its output 'Ack' bit and input 'Data' bit. If these
+/// 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;
+
+  cpp::Atomic<uint32_t> *lock;
   cpp::Atomic<uint32_t> *inbox;
   cpp::Atomic<uint32_t> *outbox;
   Buffer *buffer;
 
   /// Initialize the communication channels.
-  LIBC_INLINE void reset(void *inbox, void *outbox, void *buffer) {
+  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),
     };
   }
+
+  /// 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);
+  }
+};
+
+/// 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].opcode;
+  }
+
+  LIBC_INLINE void close() {
+    process.lock[index].store(0, cpp::MemoryOrder::RELAXED);
+  }
+
+private:
+  Process &process;
+  uint64_t index;
+  uint32_t out;
 };
 
 /// The RPC client used to make requests to the server.
 struct Client : public Process {
   LIBC_INLINE Client() = default;
   LIBC_INLINE Client(const Client &) = default;
   LIBC_INLINE Client &operator=(const Client &) = default;
   LIBC_INLINE ~Client() = default;
 
-  template <typename F, typename U> LIBC_INLINE void run(F fill, U use);
+  LIBC_INLINE cpp::optional<Port> try_open(uint16_t opcode);
+  LIBC_INLINE Port open(uint16_t opcode);
 };
 
 /// The RPC server used to respond to the client.
 struct Server : public Process {
   LIBC_INLINE Server() = default;
   LIBC_INLINE Server(const Server &) = default;
   LIBC_INLINE Server &operator=(const Server &) = default;
   LIBC_INLINE ~Server() = default;
 
-  template <typename W, typename C> LIBC_INLINE bool handle(W work, C clean);
+  LIBC_INLINE cpp::optional<Port> try_open();
+  LIBC_INLINE Port open();
 };
 
-/// Run the RPC client protocol to communicate with the server. We perform the
-/// following high level actions to complete a communication:
-///   - Apply \p fill to the shared buffer and write 1 to the outbox.
-///   - Wait until the inbox is 1.
-///   - Apply \p use to the shared buffer and write 0 to the outbox.
-///   - Wait until the inbox is 0.
-template <typename F, typename U> LIBC_INLINE void Client::run(F fill, U use) {
-  bool in = inbox->load(cpp::MemoryOrder::RELAXED);
-  bool out = outbox->load(cpp::MemoryOrder::RELAXED);
-  atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
-  // Apply the \p fill to the buffer and signal the server.
-  if (!in & !out) {
-    fill(buffer);
-    atomic_thread_fence(cpp::MemoryOrder::RELEASE);
-    outbox->store(1, cpp::MemoryOrder::RELAXED);
-    out = 1;
-  }
-  // Wait for the server to work on the buffer and respond.
-  if (!in & out) {
-    while (!in) {
-      sleep_briefly();
-      in = inbox->load(cpp::MemoryOrder::RELAXED);
-    }
-    atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
-  }
-  // Apply \p use to the buffer and signal the server.
-  if (in & out) {
-    use(buffer);
-    atomic_thread_fence(cpp::MemoryOrder::RELEASE);
-    outbox->store(0, cpp::MemoryOrder::RELAXED);
-    out = 0;
-  }
-  // Wait for the server to signal the end of the protocol.
-  if (in & !out) {
-    while (in) {
-      sleep_briefly();
-      in = inbox->load(cpp::MemoryOrder::RELAXED);
-    }
-    atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
+/// Applies \p fill to the shared buffer and initiates a send operation.
+template <typename F> LIBC_INLINE void Port::send(F fill) {
+  uint32_t in = process.inbox[index].load(cpp::MemoryOrder::RELAXED);
+
+  // We need to wait until we own the buffer before sending.
+  while (!Process::can_send_data(in, out)) {
+    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]);
+  out = out ^ Process::Data;
+  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)) {
+    in = process.inbox[index].load(cpp::MemoryOrder::RELAXED);
+
+  // Apply the \p use function to read the memory out of the buffer.
+  use(&process.buffer[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; });
+  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) {
+  uint64_t size = 0;
+  recv([&](Buffer *buffer) { size = buffer->data[0]; });
+  uint8_t *dst = reinterpret_cast<uint8_t *>(alloc(size));
+  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;
+      inline_memcpy(dst + idx, buffer->data, len);
+    });
+  }
+}
+
+/// 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) {
+  // Attempt to acquire the lock on this index.
+  if (lock->fetch_or(1, cpp::MemoryOrder::RELAXED))
+    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;
+  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();
   }
 }
 
-/// Run the RPC server protocol to communicate with the client. This is
-/// non-blocking and only checks the server a single time. We perform the
-/// following high level actions to complete a communication:
-///   - Query if the inbox is 1 and exit if there is no work to do.
-///   - Apply \p work to the shared buffer and write 1 to the outbox.
-///   - Wait until the inbox is 0.
-///   - Apply \p clean to the shared buffer and write 0 to the outbox.
-template <typename W, typename C>
-LIBC_INLINE bool Server::handle(W work, C clean) {
-  bool in = inbox->load(cpp::MemoryOrder::RELAXED);
-  bool out = outbox->load(cpp::MemoryOrder::RELAXED);
-  atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
-  // There is no work to do, exit early.
-  if (!in & !out)
-    return false;
-  // Apply \p work to the buffer and signal the client.
-  if (in & !out) {
-    work(buffer);
-    atomic_thread_fence(cpp::MemoryOrder::RELEASE);
-    outbox->store(1, cpp::MemoryOrder::RELAXED);
-    out = 1;
-  }
-  // Wait for the client to use the buffer and respond.
-  if (in & out) {
-    while (in)
-      in = inbox->load(cpp::MemoryOrder::RELAXED);
-    atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
+/// 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))
+    return cpp::nullopt;
+
+  out = outbox->load(cpp::MemoryOrder::RELAXED);
+
+  if (!can_recv_data(in, out)) {
+    lock->store(0, cpp::MemoryOrder::RELAXED);
+    return cpp::nullopt;
   }
-  // Clean up the buffer and signal the end of the protocol.
-  if (!in & out) {
-    clean(buffer);
-    atomic_thread_fence(cpp::MemoryOrder::RELEASE);
-    outbox->store(0, cpp::MemoryOrder::RELAXED);
-    out = 0;
-  }
+
+  return Port(*this, 0, out);
 
-  return true;
+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/startup/gpu/amdgpu/start.cpp

 //===-- Implementation of crt for amdgpu ----------------------------------===//
 //
 // 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/RPC/rpc_client.h"
 
+static __llvm_libc::cpp::Atomic<uint32_t> lock;
+
 extern "C" int main(int argc, char **argv, char **envp);
 
 extern "C" [[gnu::visibility("protected"), clang::amdgpu_kernel]] void
 _start(int argc, char **argv, char **envp, int *ret, void *in, void *out,
        void *buffer) {
-  __llvm_libc::rpc::client.reset(in, out, buffer);
+  __llvm_libc::rpc::client.reset(&lock, in, out, buffer);
 
   __atomic_fetch_or(ret, main(argc, argv, envp), __ATOMIC_RELAXED);
 }

libc/startup/gpu/nvptx/start.cpp

 //===-- Implementation of crt for nvptx -----------------------------------===//
 //
 // 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/RPC/rpc_client.h"
 
+static __llvm_libc::cpp::Atomic<uint32_t> lock;
+
 extern "C" int main(int argc, char **argv, char **envp);
 
 extern "C" [[gnu::visibility("protected")]] __attribute__((nvptx_kernel)) void
 _start(int argc, char **argv, char **envp, int *ret, void *in, void *out,
        void *buffer) {
-  __llvm_libc::rpc::client.reset(in, out, buffer);
+  __llvm_libc::rpc::client.reset(&lock, in, out, buffer);
 
   __atomic_fetch_or(ret, main(argc, argv, envp), __ATOMIC_RELAXED);
 }

libc/utils/gpu/loader/CMakeLists.txt

 add_library(gpu_loader OBJECT Main.cpp)
-target_include_directories(gpu_loader PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
+target_include_directories(gpu_loader PUBLIC
+  ${CMAKE_CURRENT_SOURCE_DIR}
+  ${LIBC_SOURCE_DIR}
+)
 
 find_package(hsa-runtime64 QUIET 1.2.0 HINTS ${CMAKE_INSTALL_PREFIX} PATHS /opt/rocm)
 if(hsa-runtime64_FOUND)
   add_subdirectory(amdgpu)
 else()
   message(STATUS "Skipping HSA loader for gpu target, no HSA was detected")
 endif()
 

libc/utils/gpu/loader/Server.h

+//===-- Generic RPC server interface --------------------------------------===//
+//
+// 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_UTILS_GPU_LOADER_RPC_H
+#define LLVM_LIBC_UTILS_GPU_LOADER_RPC_H
+
+#include <cstdint>
+#include <cstdio>
+#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;
+
+/// Queries the RPC client at least once and performs server-side work if there
+/// 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); });
+    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;
+  }
+  default:
+    port->recv([](__llvm_libc::rpc::Buffer *) { /* no-op */ });
+    return;
+  }
+  port->close();
+}
+#endif

libc/utils/gpu/loader/amdgpu/CMakeLists.txt

 add_executable(amdhsa_loader Loader.cpp)
 add_dependencies(amdhsa_loader libc.src.__support.RPC.rpc)
 
-target_include_directories(amdhsa_loader PRIVATE ${LIBC_SOURCE_DIR})
 target_link_libraries(amdhsa_loader
   PRIVATE
   hsa-runtime64::hsa-runtime64
   gpu_loader
 )

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

 //===-- Loader Implementation for AMDHSA devices --------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file impelements a simple loader to run images supporting the AMDHSA
 // architecture. The file launches the '_start' kernel which should be provided
 // by the device application start code and call ultimately call the 'main'
 // function.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Loader.h"
-
-#include "src/__support/RPC/rpc.h"
+#include "Server.h"
 
 #include <hsa/hsa.h>
 #include <hsa/hsa_ext_amd.h>
 
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
 #include <utility>
 
 /// The name of the kernel we will launch. All AMDHSA kernels end with '.kd'.
 constexpr const char *KERNEL_START = "_start.kd";
 
 /// The arguments to the '_start' kernel.
 struct kernel_args_t {
   int argc;
   void *argv;
   void *envp;
   void *ret;
   void *inbox;
   void *outbox;
   void *buffer;
 };
 
-static __llvm_libc::rpc::Server server;
-
-/// Queries the RPC client at least once and performs server-side work if there
-/// are any active requests.
-void handle_server() {
-  while (server.handle(
-      [&](__llvm_libc::rpc::Buffer *buffer) {
-        switch (static_cast<__llvm_libc::rpc::Opcode>(buffer->data[0])) {
-        case __llvm_libc::rpc::Opcode::PRINT_TO_STDERR: {
-          fputs(reinterpret_cast<const char *>(&buffer->data[1]), stderr);
-          break;
-        }
-        case __llvm_libc::rpc::Opcode::EXIT: {
-          exit(buffer->data[1]);
-          break;
-        }
-        default:
-          return;
-        };
-      },
-      [](__llvm_libc::rpc::Buffer *buffer) {}))
-    ;
-}
-
 /// Print the error code and exit if \p code indicates an error.
 static void handle_error(hsa_status_t code) {
   if (code == HSA_STATUS_SUCCESS || code == HSA_STATUS_INFO_BREAK)
     return;
 
   const char *desc;
   if (hsa_status_string(code, &desc) != HSA_STATUS_SUCCESS)
     desc = "Unknown error";
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(server_inbox, server_outbox, buffer);
+  server.reset(&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/CMakeLists.txt

 add_executable(nvptx_loader Loader.cpp)
 add_dependencies(nvptx_loader libc.src.__support.RPC.rpc)
 
-target_include_directories(nvptx_loader PRIVATE ${LIBC_SOURCE_DIR})
 target_link_libraries(nvptx_loader
   PRIVATE
   gpu_loader
   CUDA::cuda_driver
 )

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

 //===-- Loader Implementation for NVPTX devices --------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file impelements a simple loader to run images supporting the NVPTX
 // architecture. The file launches the '_start' kernel which should be provided
 // by the device application start code and call ultimately call the 'main'
 // function.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Loader.h"
-
-#include "src/__support/RPC/rpc.h"
+#include "Server.h"
 
 #include "cuda.h"
 #include <cstddef>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
 
 /// The arguments to the '_start' kernel.
 struct kernel_args_t {
   int argc;
   void *argv;
   void *envp;
   void *ret;
   void *inbox;
   void *outbox;
   void *buffer;
 };
 
-static __llvm_libc::rpc::Server server;
-
-/// Queries the RPC client at least once and performs server-side work if there
-/// are any active requests.
-void handle_server() {
-  while (server.handle(
-      [&](__llvm_libc::rpc::Buffer *buffer) {
-        switch (static_cast<__llvm_libc::rpc::Opcode>(buffer->data[0])) {
-        case __llvm_libc::rpc::Opcode::PRINT_TO_STDERR: {
-          fputs(reinterpret_cast<const char *>(&buffer->data[1]), stderr);
-          break;
-        }
-        case __llvm_libc::rpc::Opcode::EXIT: {
-          exit(buffer->data[1]);
-          break;
-        }
-        default:
-          return;
-        };
-      },
-      [](__llvm_libc::rpc::Buffer *buffer) {}))
-    ;
-}
-
 static void handle_error(CUresult err) {
   if (err == CUDA_SUCCESS)
     return;
 
   const char *err_str = nullptr;
   CUresult result = cuGetErrorString(err, &err_str);
   if (result != CUDA_SUCCESS)
     fprintf(stderr, "Unknown Error\n");
int load(int argc, char **argv, char **envp, void *image, size_t size,
   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(server_inbox, server_outbox, buffer);
+  server.reset(&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);