[libc][NFC] Clean up the memory buffer handling for RPC

[jhuber6]

Summary

We do a lot of arithmetic on void pointers here, so include a helper and
make some more consistent names. Changes no functionality.

[JonChesterfield]

Size is shorter than the existing names but more likely to be conflated with the array length or number of ports. How about inbox_bytes or similar?

libc/src/__support/RPC/rpc.h

287

Use the template parameter directly please. Memory layout is a compile time property, routing it through a function parameter obfuscates that.

[jhuber6]

Suggested changes.

libc/src/__support/RPC/rpc.h

template <bool InvertInbox> struct Process {
   uint32_t lane_size;
   cpp::Atomic<uint32_t> *inbox;
   cpp::Atomic<uint32_t> *outbox;
   Packet *packet;
 
   cpp::Atomic<uint32_t> lock[DEFAULT_PORT_COUNT] = {0};
 
   /// Initialize the communication channels.
-  LIBC_INLINE void reset(uint64_t port_count, uint32_t lane_size, void *state) {
-    uint64_t p = memory_offset_primary_mailbox(port_count);
-    uint64_t s = memory_offset_secondary_mailbox(port_count);
+  LIBC_INLINE void reset(uint64_t port_count, uint32_t lane_size,
+                         void *buffer) {
     this->port_count = port_count;
     this->lane_size = lane_size;
     this->inbox = reinterpret_cast<cpp::Atomic<uint32_t> *>(
-        static_cast<char *>(state) + (InvertInbox ? s : p));
+        advance(buffer, inbox_offset(port_count)));
     this->outbox = reinterpret_cast<cpp::Atomic<uint32_t> *>(
-        static_cast<char *>(state) + (InvertInbox ? p : s));
-    this->packet = reinterpret_cast<Packet *>(static_cast<char *>(state) +
-                                              memory_offset_buffer(port_count));
+        advance(buffer, outbox_offset(port_count)));
+    this->packet =
+        reinterpret_cast<Packet *>(advance(buffer, buffer_offset(port_count)));
   }
 
-  /// Allocate a single block of memory for use by client and server
-  /// template<size_t N>, N is generally a runtime value
-  /// struct equivalent {
-  ///   atomic<uint32_t> primary[N];
-  ///   atomic<uint32_t> secondary[N];
-  ///   Packet buffer[N];
+  /// Allocate a memory buffer sufficient to store the following equivalent
+  /// representation in memory.
+  ///
+  /// struct Equivalent {
+  ///   Atomic<uint32_t> primary[port_count];
+  ///   Atomic<uint32_t> secondary[port_count];
+  ///   Packet buffer[port_count];
   /// };
   LIBC_INLINE static uint64_t allocation_size(uint64_t port_count,
                                               uint32_t lane_size) {
-    return memory_offset_buffer(port_count) +
-           memory_allocated_buffer(port_count, lane_size);
+    return buffer_offset(port_count) + buffer_bytes(port_count, lane_size);
   }
 
   /// The length of the packet is flexible because the server needs to look up
   /// the lane size at runtime. This helper indexes at the proper offset.
   LIBC_INLINE Packet &get_packet(uint64_t index) {
-    return *reinterpret_cast<Packet *>(
-        reinterpret_cast<uint8_t *>(packet) +
-        index * align_up(sizeof(Header) + lane_size * sizeof(Buffer),
-                         alignof(Packet)));
+    return *reinterpret_cast<Packet *>(advance(
+        packet, index * align_up(sizeof(Header) + lane_size * sizeof(Buffer),
+                                 alignof(Packet))));
   }
 
   /// 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
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 allocated for mailbox or buffer
-  LIBC_INLINE static uint64_t memory_allocated_mailbox(uint64_t port_count) {
+  /// Number of bytes to allocate for an inbox or outbox.
+  LIBC_INLINE static uint64_t mailbox_bytes(uint64_t port_count) {
     return port_count * sizeof(cpp::Atomic<uint32_t>);
   }
 
-  LIBC_INLINE static uint64_t memory_allocated_buffer(uint64_t port_count,
-                                                      uint32_t lane_size) {
-#if defined(LIBC_TARGET_ARCH_IS_GPU)
-    (void)lane_size;
-    return port_count * sizeof(Packet);
-#else
-    return port_count * (sizeof(Packet) + sizeof(Buffer) * lane_size);
-#endif
+  /// Number of bytes to allocate for the buffer containing the packets.
+  LIBC_INLINE static uint64_t buffer_bytes(uint64_t port_count,
+    return is_process_gpu()
+               ? port_count * sizeof(Packet)
+               : port_count *
+                     align_up(sizeof(Header) + (lane_size * sizeof(Buffer)),
+                              alignof(Packet));
   }
 
-  /// Offset of mailbox/buffer in single allocation
-  LIBC_INLINE static uint64_t
-  memory_offset_primary_mailbox(uint64_t /*port_count*/) {
-    return 0;
-  }
-  LIBC_INLINE static uint64_t
-  memory_offset_secondary_mailbox(uint64_t port_count) {
-    return memory_allocated_mailbox(port_count);
+  /// Offset of the inbox in memory. This is the same as the outbox if inverted.
+  LIBC_INLINE static uint64_t inbox_offset(uint64_t port_count) {
+    return InvertInbox ? mailbox_bytes(port_count) : 0;
   }
-  LIBC_INLINE static uint64_t memory_offset_buffer(uint64_t port_count) {
-    return align_up(2 * memory_allocated_mailbox(port_count), alignof(Packet));
+
+  /// Offset of the outbox in memory. This is the same as the inbox if inverted.
+  LIBC_INLINE static uint64_t outbox_offset(uint64_t port_count) {
+    return InvertInbox ? 0 : mailbox_bytes(port_count);
+  }
+
+  /// Offset of the buffer containing the packets after the inbox and outbox.
+  LIBC_INLINE static uint64_t buffer_offset(uint64_t port_count) {
+    return align_up(2 * mailbox_bytes(port_count), alignof(Packet));
   }
 };
 
 /// The port provides the interface to communicate between the multiple
 /// processes. A port is conceptually an index into the memory provided by the
 /// underlying process that is guarded by a lock bit.
 template <bool T> struct Port {
   LIBC_INLINE Port(Process<T> &process, uint64_t lane_mask, uint64_t index,

libc/src/__support/RPC/rpc_util.h

template <typename V> LIBC_INLINE V &lane_value(V *val, uint32_t id) {
   return val[id];
 }
 
 /// Helper to get the maximum value.
 template <typename T> LIBC_INLINE const T &max(const T &x, const T &y) {
   return x < y ? y : x;
 }
 
+/// Advance the \p ptr by \p bytes.
+template <typename T, typename U> LIBC_INLINE T *advance(T ptr, U bytes) {
+  return reinterpret_cast<T *>(reinterpret_cast<uint8_t *>(ptr) + bytes);
+}
+
 } // namespace rpc
 } // namespace __llvm_libc
 
 #endif