[libc] Drop RPC port count and index to 32-bit numbers

[jhuber6]

Summary

The port count and index into the ports was originally written as a
64-bit number. This was with an abundance of caution, however it's
highly unlikely that any configuration will excede a 32-bit number as
most machines will require something in the low-thousands. Because GPUs
are functionally 32-bit in many of their operations this costs us some
extra time and registers to do the 64-bit operations. Doing this saves
us about four registers in most tests.

[JonChesterfield]

It would probably be sensible to type some of these integers to detect passing them in the wrong order. Enum class.

libc/src/__support/RPC/rpc.h

 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 = 0;
+  uint32_t port_count = 0;
   cpp::Atomic<uint32_t> *inbox = nullptr;
   cpp::Atomic<uint32_t> *outbox = nullptr;
   Packet *packet = nullptr;
 
   static constexpr uint64_t NUM_BITS_IN_WORD = sizeof(uint32_t) * 8;
   cpp::Atomic<uint32_t> lock[MAX_PORT_COUNT / NUM_BITS_IN_WORD] = {0};
 
   /// Initialize the communication channels.
-  LIBC_INLINE void reset(uint64_t port_count, void *buffer) {
+  LIBC_INLINE void reset(uint32_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 *>(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.
   ///
   /// struct Equivalent {
   ///   Atomic<uint32_t> primary[port_count];
   ///   Atomic<uint32_t> secondary[port_count];
   ///   Packet buffer[port_count];
   /// };
-  LIBC_INLINE static constexpr uint64_t allocation_size(uint64_t port_count) {
+  LIBC_INLINE static constexpr uint64_t allocation_size(uint32_t port_count) {
     return buffer_offset(port_count) + buffer_bytes(port_count);
   }
 
   /// Retrieve the inbox state from memory shared between processes.
-  LIBC_INLINE uint32_t load_inbox(uint64_t index) {
+  LIBC_INLINE uint32_t load_inbox(uint32_t index) {
     return inbox[index].load(cpp::MemoryOrder::RELAXED);
   }
 
   /// Retrieve the outbox state from memory shared between processes.
-  LIBC_INLINE uint32_t load_outbox(uint64_t index) {
+  LIBC_INLINE uint32_t load_outbox(uint32_t index) {
     return outbox[index].load(cpp::MemoryOrder::RELAXED);
   }
 
   /// Signal to the other process that this one is finished with the buffer.
   /// Equivalent to loading outbox followed by store of the inverted value
   /// The outbox is write only by this warp and tracking the value locally is
   /// cheaper than calling load_outbox to get the value to store.
-  LIBC_INLINE uint32_t invert_outbox(uint64_t index, uint32_t current_outbox) {
+  LIBC_INLINE uint32_t invert_outbox(uint32_t index, uint32_t current_outbox) {
     uint32_t inverted_outbox = !current_outbox;
     atomic_thread_fence(cpp::MemoryOrder::RELEASE);
     outbox[index].store(inverted_outbox, cpp::MemoryOrder::RELAXED);
     return inverted_outbox;
   }
 
   // Given the current outbox and inbox values, wait until the inbox changes
   // to indicate that this thread owns the buffer element.
-  LIBC_INLINE void wait_for_ownership(uint64_t index, uint32_t outbox,
+  LIBC_INLINE void wait_for_ownership(uint32_t index, uint32_t outbox,
                                       uint32_t in) {
     while (buffer_unavailable(in, outbox)) {
       sleep_briefly();
       in = load_inbox(index);
     }
     atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
   }
 
   /// Determines if this process needs to wait for ownership of the buffer. We
   /// invert the condition on one of the processes to indicate that if one
   /// process owns the buffer then the other does not.
   LIBC_INLINE static bool buffer_unavailable(uint32_t in, uint32_t out) {
     bool cond = in != out;
     return Invert ? !cond : cond;
   }
 
   /// Attempt to claim the lock at index. Return true on lock taken.
   /// lane_mask is a bitmap of the threads in the warp that would hold the
   /// single lock on success, e.g. the result of gpu::get_lane_mask()
   /// The lock is held when the n-th bit of the lock bitfield is set.
   [[clang::convergent]] LIBC_INLINE bool try_lock(uint64_t lane_mask,
-                                                  uint64_t index) {
+                                                  uint32_t index) {
     // On amdgpu, test and set to the nth lock bit and a sync_lane would suffice
     // On volta, need to handle differences between the threads running and
     // the threads that were detected in the previous call to get_lane_mask()
     //
     // All threads in lane_mask try to claim the lock. At most one can succeed.
     // There may be threads active which are not in lane mask which must not
     // succeed in taking the lock, as otherwise it will leak. This is handled
     // by making threads which are not in lane_mask or with 0, a no-op.
[[clang::convergent]] LIBC_INLINE bool try_lock(uint64_t lane_mask,
     if (holding_lock)
       atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
     return holding_lock;
   }
 
   /// Unlock the lock at index. We need a lane sync to keep this function
   /// convergent, otherwise the compiler will sink the store and deadlock.
   [[clang::convergent]] LIBC_INLINE void unlock(uint64_t lane_mask,
-                                                uint64_t index) {
+                                                uint32_t index) {
     // Do not move any writes past the unlock
     atomic_thread_fence(cpp::MemoryOrder::RELEASE);
 
     // Wait for other threads in the warp to finish using the lock
     gpu::sync_lane(lane_mask);
 
     // Use exactly one thread to clear the nth bit in the lock array 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.
     clear_nth(lock, index, rpc::is_first_lane(lane_mask));
     gpu::sync_lane(lane_mask);
   }
 
   /// Number of bytes to allocate for an inbox or outbox.
-  LIBC_INLINE static constexpr uint64_t mailbox_bytes(uint64_t port_count) {
+  LIBC_INLINE static constexpr uint64_t mailbox_bytes(uint32_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) {
+  LIBC_INLINE static constexpr uint64_t buffer_bytes(uint32_t port_count) {
     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) {
+  LIBC_INLINE static constexpr uint64_t inbox_offset(uint32_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) {
+  LIBC_INLINE static constexpr uint64_t outbox_offset(uint32_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) {
+  LIBC_INLINE static constexpr uint64_t buffer_offset(uint32_t port_count) {
     return align_up(2 * mailbox_bytes(port_count), alignof(Packet));
   }
 
   /// Conditionally set the n-th bit in the atomic bitfield.
   LIBC_INLINE static constexpr uint32_t set_nth(cpp::Atomic<uint32_t> *bits,
-                                                uint64_t index, bool cond) {
-    uint64_t slot = index / NUM_BITS_IN_WORD;
-    uint64_t bit = index % NUM_BITS_IN_WORD;
+                                                uint32_t index, bool cond) {
+    uint32_t slot = index / NUM_BITS_IN_WORD;
+    uint32_t bit = index % NUM_BITS_IN_WORD;
     return bits[slot].fetch_or(static_cast<uint32_t>(cond) << bit,
                                cpp::MemoryOrder::RELAXED) &
            (1u << bit);
   }
 
   /// Conditionally clear the n-th bit in the atomic bitfield.
   LIBC_INLINE static constexpr uint32_t clear_nth(cpp::Atomic<uint32_t> *bits,
-                                                  uint64_t index, bool cond) {
-    uint64_t slot = index / NUM_BITS_IN_WORD;
-    uint64_t bit = index % NUM_BITS_IN_WORD;
+                                                  uint32_t index, bool cond) {
+    uint32_t slot = index / NUM_BITS_IN_WORD;
+    uint32_t bit = index % NUM_BITS_IN_WORD;
     return bits[slot].fetch_and(~0u ^ (static_cast<uint32_t>(cond) << bit),
                                 cpp::MemoryOrder::RELAXED) &
            (1u << bit);
   }
 };
 
 /// Invokes a function accross every active buffer across the total lane size.
 template <uint32_t lane_size>
for (uint32_t i = 0; i < lane_size; i += gpu::get_lane_size())
         fn(&packet.payload.slot[i], 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.
 template <bool T, typename S> struct Port {
-  LIBC_INLINE Port(Process<T, S> &process, uint64_t lane_mask, uint64_t index,
+  LIBC_INLINE Port(Process<T, S> &process, uint64_t lane_mask, uint32_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;
   LIBC_INLINE Port &operator=(const Port &) = delete;
LIBC_INLINE void close() {
     if (owns_buffer && T)
       out = process.invert_outbox(index, out);
     process.unlock(lane_mask, index);
   }
 
 private:
   Process<T, S> &process;
   uint64_t lane_mask;
-  uint64_t index;
+  uint32_t index;
   uint32_t out;
   bool receive;
   bool owns_buffer;
 };
 
 /// 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, 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) {
+  LIBC_INLINE void reset(uint32_t port_count, void *buffer) {
     process.reset(port_count, buffer);
   }
 
 private:
   Process<false, Packet<gpu::LANE_SIZE>> process;
 };
 static_assert(cpp::is_trivially_copyable<Client>::value &&
                   sizeof(Process<false, Packet<1>>) ==
                       sizeof(Process<false, Packet<32>>),
               "The client is not trivially copyable from the server");
 
 /// 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, Packet<lane_size>>;
   LIBC_INLINE cpp::optional<Port> try_open();
   LIBC_INLINE Port open();
 
-  LIBC_INLINE void reset(uint64_t port_count, void *buffer) {
+  LIBC_INLINE void reset(uint32_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) {
+  LIBC_INLINE static uint64_t allocation_size(uint32_t port_count) {
     return Process<true, Packet<lane_size>>::allocation_size(port_count);
   }
 
 private:
   Process<true, Packet<lane_size>> process;
 };
 
 /// Applies \p fill to the shared buffer and initiates a send operation.

 /// 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.
 template <uint16_t opcode>
 [[clang::convergent]] LIBC_INLINE cpp::optional<Client::Port>
 Client::try_open() {
   // Perform a naive linear scan for a port that can be opened to send data.
-  for (uint64_t index = 0; index < process.port_count; ++index) {
+  for (uint32_t index = 0; index < process.port_count; ++index) {
     // Attempt to acquire the lock on this index.
     uint64_t lane_mask = gpu::get_lane_mask();
     if (!process.try_lock(lane_mask, index))
       continue;
 
     uint32_t in = process.load_inbox(index);
     uint32_t out = process.load_outbox(index);
 

 
 /// Attempts to open a port to use as the server. The server can only open a
 /// port if it has a pending receive operation
 template <uint32_t lane_size>
 [[clang::convergent]] LIBC_INLINE
     cpp::optional<typename Server<lane_size>::Port>
     Server<lane_size>::try_open() {
   // Perform a naive linear scan for a port that has a pending request.
-  for (uint64_t index = 0; index < process.port_count; ++index) {
+  for (uint32_t index = 0; index < process.port_count; ++index) {
     uint32_t in = process.load_inbox(index);
     uint32_t out = process.load_outbox(index);
 
     // The server is passive, if there is no work pending don't bother
     // opening a port.
     if (process.buffer_unavailable(in, out))
       continue;