[libc] Treat the locks array as a bitfield

[jhuber6]

Summary

Currently we keep an internal buffer of device memory that is used to
indicate ownership of a port. Since we only use this as a single bit we
can simply turn this into a bitfield. I did this manually rather than
having a separate type as we need very special handling of the masks
used to interact with the locks.

[JonChesterfield]

Sad to see this isn't a type. Part of the point of encapsulating the indexing and mutation ops is the inbox outbox arrays could also be bitmaps.

Even without that, there's no reason to open code the division and modulo ops. Please put those behind named functions. Set nth bit or similar.

[JonChesterfield]

This comment has been deleted.

libc/src/__support/RPC/rpc.h

88

This isn't the size of the bitmap. I don't know what a bitfield is. Bits per word?

I'm not convinced we should pretend this stuff works for atomic types other than plain integers.

How about a static assert that atomic uint32 is the same size as a uint32 and write 32 here?

[jhuber6]

Addressing comments and moving logic to member functions.

[JonChesterfield]

Some nits. Still thinking about the tradeoffs between a type and functions but leaning in this direction

libc/src/__support/RPC/rpc.h

59–60

Maybe rename this while we're in the area, it's a maximum port count, not a default. Could static assert that it's a multiple of 32

161

available

248

the bitshifted boolean is probably ok - bool << unsigned probably promotes to an unsigned op - but could you either cast the cond or assign it to a uint32 before the expression to remove the ambiguity?

258

UINT32_MAX or ~0u please

[jhuber6]

Addressing comments

[JonChesterfield]

LG, thanks

[jplehr]

This may have broken the staging buildbot that tests libc on GPU https://lab.llvm.org/staging/#/builders/247/builds/3201

[jhuber6]

Hopefully fixing this for the buildbot

[JonChesterfield]

Working theory is the bot problem was miscompiling memcpy?

[jhuber6]

In D155511#4522907, @JonChesterfield wrote:

Working theory is the bot problem was miscompiling memcpy?

Unsure, did some basic checks and it seems to be working now on HPE. Going to do a git pull on both RONL and HPE systems to be doubly sure.

libc/src/__support/RPC/rpc.h

 
 /// A packet used to share data between the client and server across an entire
 /// lane. We use a lane as the minimum granularity for execution.
 template <uint32_t lane_size = gpu::LANE_SIZE> struct alignas(64) Packet {
   Header header;
   Payload<lane_size> payload;
 };
 
-// TODO: This should be configured by the server and passed in. The general rule
-//       of thumb is that you should have at least as many ports as possible
-//       concurrent work items on the GPU to mitigate the lack offorward
-//       progress guarantees on the GPU.
-constexpr uint64_t DEFAULT_PORT_COUNT = 64;
+/// The maximum number of parallel ports that the RPC interface can support.
+constexpr uint64_t MAX_PORT_COUNT = 512;
 
 /// A common process used to synchronize communication between a client and a
 /// server. The process contains a read-only inbox and a write-only outbox used
 /// for signaling ownership of the shared buffer between both sides. We assign
 /// ownership of the buffer to the client if the inbox and outbox bits match,
 /// otherwise it is owned by the server.
 ///
 /// This process is designed to allow the client and the server to exchange data
template <bool Invert, typename Packet> struct Process {
   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 *packet;
 
-  cpp::Atomic<uint32_t> lock[DEFAULT_PORT_COUNT] = {0};
+  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) {
     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)));
template <bool Invert, typename Packet> struct Process {
   /// 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 zeroth bit of the uint32_t at lock[index]
-  /// is set, and available when that bit is clear. Bits [1, 32) are zero.
-  /// Or with one is a no-op when the lock is already held.
+  /// single lock on success, e.g. the result of gpu::get_lane_mask().
+  /// The lock is held when the nth bit of the lock bitfield is set, otherwise
+  /// it is availible.
   [[clang::convergent]] LIBC_INLINE bool try_lock(uint64_t lane_mask,
                                                   uint64_t index) {
     // On amdgpu, test and set to lock[index] 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.
     uint32_t id = gpu::get_lane_id();
-    bool id_in_lane_mask = lane_mask & (1ul << id);
+    bool id_in_lane_mask = lane_mask & (1u << id);
 
     // All threads in the warp call fetch_or. Possibly at the same time.
-    bool before =
-        lock[index].fetch_or(id_in_lane_mask, cpp::MemoryOrder::RELAXED);
+    bool before = set_nth(lock, index, id_in_lane_mask);
     uint64_t packed = gpu::ballot(lane_mask, before);
 
     // If every bit set in lane_mask is also set in packed, every single thread
     // in the warp failed to get the lock. Ballot returns unset for threads not
     // in the lane mask.
     //
     // Cases, per thread:
     // mask==0 -> unspecified before, discarded by ballot -> 0
template <bool Invert, typename Packet> struct Process {
   [[clang::convergent]] LIBC_INLINE void unlock(uint64_t lane_mask,
                                                 uint64_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 bit at position 0 in lock[index]
-    // 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);
+    // Use exactly one thread to clear the associated lock bit. 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) {
     return port_count * sizeof(cpp::Atomic<uint32_t>);
   }
 
template <bool Invert, typename Packet> struct Process {
   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));
   }
+
+  /// 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) {
+    const uint32_t slot = index / NUM_BITS_IN_WORD;
+    const 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) {
+    const uint32_t slot = index / NUM_BITS_IN_WORD;
+    const uint32_t bit = index % NUM_BITS_IN_WORD;
+    return bits[slot].fetch_and(-1u ^ (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>
 static 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()]);

libc/startup/gpu/amdgpu/start.cpp

 }
 
 } // namespace __llvm_libc
 
 extern "C" [[gnu::visibility("protected"), clang::amdgpu_kernel]] void
 _begin(int argc, char **argv, char **env, void *rpc_shared_buffer) {
   // We need to set up the RPC client first in case any of the constructors
   // require it.
-  __llvm_libc::rpc::client.reset(__llvm_libc::rpc::DEFAULT_PORT_COUNT,
+  __llvm_libc::rpc::client.reset(__llvm_libc::rpc::MAX_PORT_COUNT,
                                  rpc_shared_buffer);
 
   // We want the fini array callbacks to be run after other atexit
   // callbacks are run. So, we register them before running the init
   // array callbacks as they can potentially register their own atexit
   // callbacks.
   __llvm_libc::atexit(&__llvm_libc::call_fini_array_callbacks);
   __llvm_libc::call_init_array_callbacks(argc, argv, env);

libc/startup/gpu/nvptx/start.cpp

 }
 
 } // namespace __llvm_libc
 
 extern "C" [[gnu::visibility("protected"), clang::nvptx_kernel]] void
 _begin(int argc, char **argv, char **env, void *rpc_shared_buffer) {
   // We need to set up the RPC client first in case any of the constructors
   // require it.
-  __llvm_libc::rpc::client.reset(__llvm_libc::rpc::DEFAULT_PORT_COUNT,
+  __llvm_libc::rpc::client.reset(__llvm_libc::rpc::MAX_PORT_COUNT,
                                  rpc_shared_buffer);
 
   // We want the fini array callbacks to be run after other atexit
   // callbacks are run. So, we register them before running the init
   // array callbacks as they can potentially register their own atexit
   // callbacks.
   __llvm_libc::atexit(&__llvm_libc::call_fini_array_callbacks);
   __llvm_libc::call_init_array_callbacks(argc, argv, env);

libc/utils/gpu/server/rpc_server.h

 
 #include "llvm-libc-types/rpc_opcodes_t.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 /// The maxium number of ports that can be opened for any server.
-const uint64_t RPC_MAXIMUM_PORT_COUNT = 64;
+const uint64_t RPC_MAXIMUM_PORT_COUNT = 512;
 
 /// The symbol name associated with the client for use with the LLVM C library
 /// implementation.
 inline const char *rpc_client_symbol_name = "__llvm_libc_rpc_client";
 
 /// status codes.
 typedef enum {
   RPC_STATUS_SUCCESS = 0x0,

libc/utils/gpu/server/rpc_server.cpp

 #include <variant>
 #include <vector>
 
 using namespace __llvm_libc;
 
 static_assert(sizeof(rpc_buffer_t) == sizeof(rpc::Buffer),
               "Buffer size mismatch");
 
-static_assert(RPC_MAXIMUM_PORT_COUNT == rpc::DEFAULT_PORT_COUNT,
+static_assert(RPC_MAXIMUM_PORT_COUNT == rpc::MAX_PORT_COUNT,
               "Incorrect maximum port count");
 
 // The client needs to support different lane sizes for the SIMT model. Because
 // of this we need to select between the possible sizes that the client can use.
 struct Server {
   template <uint32_t lane_size>
   Server(std::unique_ptr<rpc::Server<lane_size>> &&server)
       : server(std::move(server)) {}