[libc][rpc] Simplify mailbox state tracking

[JonChesterfield]

Summary

Removes the redundant Ack/Data bit manipulation.

Represents the inbox/outbox state with one bit instead of two. This will
be useful if we change to a packed representation and otherwise cuts the
runtime state space from 16 to 4.

Further simplification is possible, this patch is intentionally minimal.

• can_{send,recv}_data are now in == out
• {client,server}::try_open can be factored into Process:try_open

This implements the state machine of D148191, modulo differences in atomic
ordering and fences.

[JonChesterfield]

• update header comment to remove references to ack/data bit

[jhuber6]

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

18

I guess since this is a templated type now we can't just have a generic port? Also what happened to the one in Server.h.

libc/src/__support/RPC/rpc.h

105

Can probably just call this can_use_buffer and get rid of the other one. I think keeping send and recv separate is more readable however.

[JonChesterfield]

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

18

Server.h used auto. I don't know what you have in mind for "generic port", but any user of this has a Client or Server instance available. Moving Port to a member of Process would remove some line noise, as would writing more of the methods on port in the struct,

libc/src/__support/RPC/rpc.h

105

I favour in != out at the call sites but didn't want to disturb them as part of this review. D149788, D149793 are separate to similarly avoid noise in the diff.

[jhuber6]

This worked on the unit tests in D149532 when I applied it so it looks good overall, thanks.

libc/src/__support/RPC/rpc.h

89–104

Header comment describing what this does and the inversion scheme please.

105

Yeah that's fair enough. I'll need to rebase those.

217–220

Put these and others in a single template.

[JonChesterfield]

• update header comment to remove references to ack/data bit
• header comment, fold template declarations

[JonChesterfield]

• words instead of symbols as requested

[JonChesterfield]

• back to working template declarations

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

 
 #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 {
 
 void write_to_stderr(cpp::string_view msg) {
-  rpc::Port port = rpc::client.open(rpc::PRINT_TO_STDERR);
+  rpc::Client::Port port = rpc::client.open(rpc::PRINT_TO_STDERR);
   port.send_n(msg.data(), msg.size());
   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) {
-  rpc::Port port = rpc::client.open(rpc::EXIT);
+  rpc::Client::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)

libc/src/__support/RPC/rpc.h

struct alignas(64) Buffer {
   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 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.
+/// The server inbox read is inverted, so it loads inbox==1, outbox==0 until
+/// the client has written to its outbox.
+///
 /// 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 {
+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;
 
-  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 *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.
+  /// 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
+  /// races result. If either inbox or outbox is inverted for one process, that
+  /// process interprets memory as Inbox!=Outbox and thus waits for the other.
+  /// It is simpler to invert reads from the inbox than writes to the outbox.
+  LIBC_INLINE uint32_t load_inbox(uint64_t index) {
+    uint32_t i = inbox[index].load(cpp::MemoryOrder::RELAXED);
+    return InvertInbox ? !i : i;
+  }
+
+  /// Determines if this process owns the buffer for a send.
   LIBC_INLINE static bool can_send_data(uint32_t in, uint32_t out) {
-    return bool(in & Process::Ack) == bool(out & Process::Data);
+    return in == out;
   }
 
-  /// 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.
+  /// Determines if this process owns the buffer for a receive.
   LIBC_INLINE static bool can_recv_data(uint32_t in, uint32_t out) {
-    return bool(in & Process::Data) != bool(out & Process::Ack);
+    return in == out;
   }
 };
 
 /// 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 {
+template <bool T> struct Port {
   // TODO: This should be move-only.
-  LIBC_INLINE Port(Process &process, uint64_t index, uint32_t out)
+  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;
   }
 
   LIBC_INLINE void close() {
     process.lock[index].store(0, cpp::MemoryOrder::RELAXED);
   }
 
 private:
-  Process &process;
+  Process<T> &process;
   uint64_t index;
   uint32_t out;
 };
 
 /// The RPC client used to make requests to the server.
-struct Client : public Process {
+struct Client : public Process<false> {
   LIBC_INLINE Client() = default;
   LIBC_INLINE Client(const Client &) = default;
   LIBC_INLINE Client &operator=(const Client &) = default;
   LIBC_INLINE ~Client() = default;
 
+  using Port = rpc::Port<false>;
   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 {
+struct Server : public Process<true> {
   LIBC_INLINE Server() = default;
   LIBC_INLINE Server(const Server &) = default;
   LIBC_INLINE Server &operator=(const Server &) = default;
   LIBC_INLINE ~Server() = default;
 
+  using Port = rpc::Port<true>;
   LIBC_INLINE cpp::optional<Port> try_open();
   LIBC_INLINE Port open();
 };
 
 /// 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);
+template <bool T> template <typename F> LIBC_INLINE void Port<T>::send(F fill) {
+  uint32_t in = process.load_inbox(index);
 
   // We need to wait until we own the buffer before sending.
-  while (!Process::can_send_data(in, out)) {
+  while (!Process<T>::can_send_data(in, out)) {
     sleep_briefly();
-    in = process.inbox[index].load(cpp::MemoryOrder::RELAXED);
+    in = process.load_inbox(index);
   }
 
   // Apply the \p fill function to initialize the buffer and release the memory.
   fill(&process.buffer[index]);
-  out = out ^ Process::Data;
+  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 <typename U> LIBC_INLINE void Port::recv(U use) {
-  uint32_t in = process.inbox[index].load(cpp::MemoryOrder::RELAXED);
+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::can_recv_data(in, out)) {
+  while (!Process<T>::can_recv_data(in, out)) {
     sleep_briefly();
-    in = process.inbox[index].load(cpp::MemoryOrder::RELAXED);
+    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]);
-  out = out ^ Process::Ack;
+  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::send_and_recv(F fill, U use) {
+LIBC_INLINE void Port<T>::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) {
+template <bool T>
+template <typename W>
+LIBC_INLINE void Port<T>::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) {
+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; });
   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) {
+template <bool T>
+template <typename A>
+LIBC_INLINE void Port<T>::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) {
+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 (lock->fetch_or(1, cpp::MemoryOrder::RELAXED))
     return cpp::nullopt;
 
   // The mailbox state must be read with the lock held.
   atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
 
-  uint32_t in = inbox->load(cpp::MemoryOrder::RELAXED);
+  uint32_t in = load_inbox(index);
   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);
+  return Port(*this, index, out);
 }
 
-LIBC_INLINE Port Client::open(uint16_t opcode) {
+LIBC_INLINE Client::Port Client::open(uint16_t opcode) {
   for (;;) {
-    if (cpp::optional<Port> p = try_open(opcode))
+    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<Port> Server::try_open() {
-  uint32_t in = inbox->load(cpp::MemoryOrder::RELAXED);
+LIBC_INLINE cpp::optional<Server::Port> Server::try_open() {
+  constexpr uint64_t index = 0;
+  uint32_t in = load_inbox(index);
   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;
 
   // The mailbox state must be read with the lock held.
   atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
 
-  in = inbox->load(cpp::MemoryOrder::RELAXED);
+  in = load_inbox(index);
   out = outbox->load(cpp::MemoryOrder::RELAXED);
 
   if (!can_recv_data(in, out)) {
     lock->store(0, cpp::MemoryOrder::RELAXED);
     return cpp::nullopt;
   }
 
-  return Port(*this, 0, out);
+  return Port(*this, index, out);
 }
 
-LIBC_INLINE Port Server::open() {
+LIBC_INLINE Server::Port Server::open() {
   for (;;) {
-    if (cpp::optional<Port> p = try_open())
+    if (cpp::optional<Server::Port> p = try_open())
       return p.value();
     sleep_briefly();
   }
 }
 
 } // namespace rpc
 } // namespace __llvm_libc
 
 #endif

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

 #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();
   uint64_t cnt = 0;
   for (uint32_t i = 0; i < num_additions; ++i) {
-    rpc::Port port = rpc::client.open(rpc::TEST_INCREMENT);
+    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();