[libc] Factor specifics of packet type out of process

[JonChesterfield]

Summary

NFC. Simplifies process slightly, gives more options for testing it.

[JonChesterfield]

libc/src/__support/RPC/rpc.h

255

diff hasn't gone well here - moving the methods into free functions changes where they are in the file. The indexing functions changed insofar as Packet<lane_size> is now Packet

[JonChesterfield]

• avoid moving functions in file to decrease diff size

[JonChesterfield]

• now compiles, moved invoke back to free functions

[jhuber6]

LG, thanks

libc/src/__support/RPC/rpc.h

253

LLVM style, static over anonymous namespace. ditto below.

[JonChesterfield]

• namespace => static

libc/src/__support/RPC/rpc.h

 /// otherwise it is owned by the server.
 ///
 /// This process is designed to allow the client and the server to exchange data
 /// using a fixed size packet in a mostly arbitrary order using the 'send' and
 /// 'recv' operations. The following restrictions to this scheme apply:
 ///   - The client will always start with a 'send' operation.
 ///   - The server will always start with a 'recv' operation.
 ///   - Every 'send' or 'recv' call is mirrored by the other process.
-template <bool Invert, uint32_t lane_size> struct Process {
+template <bool Invert, typename Packet> struct Process {
   LIBC_INLINE Process() = default;
   LIBC_INLINE Process(const Process &) = delete;
   LIBC_INLINE Process &operator=(const Process &) = delete;
   LIBC_INLINE Process(Process &&) = default;
   LIBC_INLINE Process &operator=(Process &&) = default;
   LIBC_INLINE ~Process() = default;
 
   uint64_t port_count;
   cpp::Atomic<uint32_t> *inbox;
   cpp::Atomic<uint32_t> *outbox;
-  Packet<lane_size> *packet;
+  Packet *packet;
 
   cpp::Atomic<uint32_t> lock[DEFAULT_PORT_COUNT] = {0};
 
   /// Initialize the communication channels.
   LIBC_INLINE void reset(uint64_t port_count, void *buffer) {
     this->port_count = port_count;
     this->inbox = reinterpret_cast<cpp::Atomic<uint32_t> *>(
         advance(buffer, inbox_offset(port_count)));
     this->outbox = reinterpret_cast<cpp::Atomic<uint32_t> *>(
         advance(buffer, outbox_offset(port_count)));
-    this->packet = reinterpret_cast<Packet<lane_size> *>(
-        advance(buffer, buffer_offset(port_count)));
+    this->packet =
+        reinterpret_cast<Packet *>(advance(buffer, buffer_offset(port_count)));
   }
 
   /// Returns the beginning of the unified buffer. Intended for initializing the
   /// client after the server has been started.
   LIBC_INLINE void *get_buffer_start() const { return Invert ? outbox : inbox; }
 
   /// Allocate a memory buffer sufficient to store the following equivalent
   /// representation in memory.
[[clang::convergent]] LIBC_INLINE void unlock(uint64_t lane_mask,
     // Must restrict to a single thread to avoid one thread dropping the lock,
     // then an unrelated warp claiming the lock, then a second thread in this
     // warp dropping the lock again.
     uint32_t and_mask = ~(rpc::is_first_lane(lane_mask) ? 1 : 0);
     lock[index].fetch_and(and_mask, cpp::MemoryOrder::RELAXED);
     gpu::sync_lane(lane_mask);
   }
 
-  /// Invokes a function accross every active buffer across the total lane size.
-  LIBC_INLINE void invoke_rpc(cpp::function<void(Buffer *)> fn,
-                              Packet<lane_size> &packet) {
-    if constexpr (is_process_gpu()) {
-      fn(&packet.payload.slot[gpu::get_lane_id()]);
-    } else {
-      for (uint32_t i = 0; i < lane_size; i += gpu::get_lane_size())
-        if (packet.header.mask & 1ul << i)
-          fn(&packet.payload.slot[i]);
-    }
-  }
-
-  /// Alternate version that also provides the index of the current lane.
-  LIBC_INLINE void invoke_rpc(cpp::function<void(Buffer *, uint32_t)> fn,
-                              Packet<lane_size> &packet) {
-    if constexpr (is_process_gpu()) {
-      fn(&packet.payload.slot[gpu::get_lane_id()], gpu::get_lane_id());
-    } else {
-      for (uint32_t i = 0; i < lane_size; i += gpu::get_lane_size())
-        if (packet.header.mask & 1ul << i)
-          fn(&packet.payload.slot[i], i);
-    }
-  }
-
   /// Number of bytes to allocate for an inbox or outbox.
   LIBC_INLINE static constexpr uint64_t mailbox_bytes(uint64_t port_count) {
     return port_count * sizeof(cpp::Atomic<uint32_t>);
   }
 
   /// Number of bytes to allocate for the buffer containing the packets.
   LIBC_INLINE static constexpr uint64_t buffer_bytes(uint64_t port_count) {
-    return port_count * sizeof(Packet<lane_size>);
+    return port_count * sizeof(Packet);
   }
 
   /// Offset of the inbox in memory. This is the same as the outbox if inverted.
   LIBC_INLINE static constexpr uint64_t inbox_offset(uint64_t port_count) {
     return Invert ? mailbox_bytes(port_count) : 0;
   }
 
   /// Offset of the outbox in memory. This is the same as the inbox if inverted.
   LIBC_INLINE static constexpr uint64_t outbox_offset(uint64_t port_count) {
     return Invert ? 0 : mailbox_bytes(port_count);
   }
 
   /// Offset of the buffer containing the packets after the inbox and outbox.
   LIBC_INLINE static constexpr uint64_t buffer_offset(uint64_t port_count) {
-    return align_up(2 * mailbox_bytes(port_count), alignof(Packet<lane_size>));
+    return align_up(2 * mailbox_bytes(port_count), alignof(Packet));
   }
 };
 
+namespace {
+/// Invokes a function accross every active buffer across the total lane size.
+template <uint32_t lane_size>
+LIBC_INLINE void invoke_rpc(cpp::function<void(Buffer *)> fn,
+                            Packet<lane_size> &packet) {
+  if constexpr (is_process_gpu()) {
+    fn(&packet.payload.slot[gpu::get_lane_id()]);
+  } else {
+    for (uint32_t i = 0; i < lane_size; i += gpu::get_lane_size())
+      if (packet.header.mask & 1ul << i)
+        fn(&packet.payload.slot[i]);
+  }
+}
+
+/// Alternate version that also provides the index of the current lane.
+template <uint32_t lane_size>
+LIBC_INLINE void invoke_rpc(cpp::function<void(Buffer *, uint32_t)> fn,
+                            Packet<lane_size> &packet) {
+  if constexpr (is_process_gpu()) {
+    fn(&packet.payload.slot[gpu::get_lane_id()], gpu::get_lane_id());
+  } else {
+    for (uint32_t i = 0; i < lane_size; i += gpu::get_lane_size())
+      if (packet.header.mask & 1ul << i)
+        fn(&packet.payload.slot[i], i);
+  }
+}
+} // namespace
+
 /// 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, uint32_t S> struct Port {
+template <bool T, typename S> struct Port {
   LIBC_INLINE Port(Process<T, S> &process, uint64_t lane_mask, uint64_t index,
                    uint32_t out)
       : process(process), lane_mask(lane_mask), index(index), out(out),
         receive(false), owns_buffer(true) {}
   LIBC_INLINE ~Port() = default;
 
 private:
   LIBC_INLINE Port(const Port &) = delete;

 
 /// The RPC client used to make requests to the server.
 struct Client {
   LIBC_INLINE Client() = default;
   LIBC_INLINE Client(const Client &) = delete;
   LIBC_INLINE Client &operator=(const Client &) = delete;
   LIBC_INLINE ~Client() = default;
 
-  using Port = rpc::Port<false, gpu::LANE_SIZE>;
+  using Port = rpc::Port<false, Packet<gpu::LANE_SIZE>>;
   template <uint16_t opcode> LIBC_INLINE cpp::optional<Port> try_open();
   template <uint16_t opcode> LIBC_INLINE Port open();
 
   LIBC_INLINE void reset(uint64_t port_count, void *buffer) {
     process.reset(port_count, buffer);
   }
 
 private:
-  Process<false, gpu::LANE_SIZE> process;
+  Process<false, Packet<gpu::LANE_SIZE>> process;
 };
 
 /// The RPC server used to respond to the client.
 template <uint32_t lane_size> struct Server {
   LIBC_INLINE Server() = default;
   LIBC_INLINE Server(const Server &) = delete;
   LIBC_INLINE Server &operator=(const Server &) = delete;
   LIBC_INLINE ~Server() = default;
 
-  using Port = rpc::Port<true, lane_size>;
+  using Port = rpc::Port<true, Packet<lane_size>>;
   LIBC_INLINE cpp::optional<Port> try_open();
   LIBC_INLINE Port open();
 
   LIBC_INLINE void reset(uint64_t port_count, void *buffer) {
     process.reset(port_count, buffer);
   }
 
   LIBC_INLINE void *get_buffer_start() const {
     return process.get_buffer_start();
   }
 
   LIBC_INLINE static uint64_t allocation_size(uint64_t port_count) {
-    return Process<true, lane_size>::allocation_size(port_count);
+    return Process<true, Packet<lane_size>>::allocation_size(port_count);
   }
 
 private:
-  Process<true, lane_size> process;
+  Process<true, Packet<lane_size>> process;
 };
 
 /// Applies \p fill to the shared buffer and initiates a send operation.
-template <bool T, uint32_t S>
+template <bool T, typename S>
 template <typename F>
 LIBC_INLINE void Port<T, S>::send(F fill) {
   uint32_t in = owns_buffer ? out ^ T : process.load_inbox(index);
 
   // We need to wait until we own the buffer before sending.
   process.wait_for_ownership(index, out, in);
 
   // Apply the \p fill function to initialize the buffer and release the memory.
-  process.invoke_rpc(fill, process.packet[index]);
+  invoke_rpc(fill, process.packet[index]);
   out = process.invert_outbox(index, out);
   owns_buffer = false;
   receive = false;
 }
 
 /// Applies \p use to the shared buffer and acknowledges the send.
-template <bool T, uint32_t S>
+template <bool T, typename S>
 template <typename U>
 LIBC_INLINE void Port<T, S>::recv(U use) {
   // We only exchange ownership of the buffer during a receive if we are waiting
   // for a previous receive to finish.
   if (receive) {
     out = process.invert_outbox(index, out);
     owns_buffer = false;
   }
 
   uint32_t in = owns_buffer ? out ^ T : process.load_inbox(index);
 
   // We need to wait until we own the buffer before receiving.
   process.wait_for_ownership(index, out, in);
 
   // Apply the \p use function to read the memory out of the buffer.
-  process.invoke_rpc(use, process.packet[index]);
+  invoke_rpc(use, process.packet[index]);
   receive = true;
   owns_buffer = true;
 }
 
 /// Combines a send and receive into a single function.
-template <bool T, uint32_t S>
+template <bool T, typename S>
 template <typename F, typename U>
 LIBC_INLINE void Port<T, S>::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 <bool T, uint32_t S>
+template <bool T, typename S>
 template <typename W>
 LIBC_INLINE void Port<T, S>::recv_and_send(W work) {
   recv(work);
   send([](Buffer *) { /* no-op */ });
 }
 
 /// Helper routine to simplify the interface when sending from the GPU using
 /// thread private pointers to the underlying value.
-template <bool T, uint32_t S>
+template <bool T, typename S>
 LIBC_INLINE void Port<T, S>::send_n(const void *src, uint64_t size) {
   static_assert(is_process_gpu(), "Only valid when running on the GPU");
   const void **src_ptr = &src;
   uint64_t *size_ptr = &size;
   send_n(src_ptr, size_ptr);
 }
 
 /// Sends an arbitrarily sized data buffer \p src across the shared channel in
 /// multiples of the packet length.
-template <bool T, uint32_t S>
+template <bool T, typename S>
 LIBC_INLINE void Port<T, S>::send_n(const void *const *src, uint64_t *size) {
   uint64_t num_sends = 0;
   send([&](Buffer *buffer, uint32_t id) {
     reinterpret_cast<uint64_t *>(buffer->data)[0] = lane_value(size, id);
     num_sends = is_process_gpu() ? lane_value(size, id)
                                  : max(lane_value(size, id), num_sends);
     uint64_t len =
         lane_value(size, id) > sizeof(Buffer::data) - sizeof(uint64_t)
while (gpu::ballot(mask, idx < num_sends)) {
     });
     idx += sizeof(Buffer::data);
   }
 }
 
 /// 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, uint32_t S>
+template <bool T, typename S>
 template <typename A>
 LIBC_INLINE void Port<T, S>::recv_n(void **dst, uint64_t *size, A &&alloc) {
   uint64_t num_recvs = 0;
   recv([&](Buffer *buffer, uint32_t id) {
     lane_value(size, id) = reinterpret_cast<uint64_t *>(buffer->data)[0];
     lane_value(dst, id) =
         reinterpret_cast<uint8_t *>(alloc(lane_value(size, id)));
     num_recvs = is_process_gpu() ? lane_value(size, id)

libc/test/src/__support/RPC/rpc_smoke_test.cpp

 //===-- smoke tests for RPC -----------------------------------------------===//
 //
 // 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.h"
 
 #include "test/UnitTest/Test.h"
 
 namespace {
 enum { lane_size = 8, port_count = 4 };
 
-using ProcAType = __llvm_libc::rpc::Process<false, lane_size>;
-using ProcBType = __llvm_libc::rpc::Process<true, lane_size>;
+struct Packet {
+  uint64_t unused;
+};
+
+using ProcAType = __llvm_libc::rpc::Process<false, Packet>;
+using ProcBType = __llvm_libc::rpc::Process<true, Packet>;
 
 static_assert(ProcAType::inbox_offset(port_count) ==
               ProcBType::outbox_offset(port_count));
 
 static_assert(ProcAType::outbox_offset(port_count) ==
               ProcBType::inbox_offset(port_count));
 
 enum { alloc_size = ProcAType::allocation_size(port_count) };