[libc][rpc] Pass lane_mask into lock functions

[JonChesterfield]

Summary

Will be necessary for correct locking on volta.

API-only change to help with rebasing the rest of the stack on this.

libc/src/__support/RPC/rpc.h

template <bool InvertInbox> struct Process {
   LIBC_INLINE static bool buffer_unavailable(uint32_t in, uint32_t out) {
     return in != out;
   }
 
   /// Attempt to claim the lock at index. Return true on lock taken.
   /// 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.
-  LIBC_INLINE bool try_lock(uint64_t index) {
+  LIBC_INLINE bool try_lock(uint64_t, uint64_t index) {
     return lock[index].fetch_or(1, cpp::MemoryOrder::RELAXED) == 0;
   }
 
   // Unlock the lock at index.
-  LIBC_INLINE void unlock(uint64_t index) {
+  LIBC_INLINE void unlock(uint64_t, uint64_t index) {
     lock[index].store(0, cpp::MemoryOrder::RELAXED);
   }
 };
 
 /// 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 {
   // TODO: This should be move-only.
-  LIBC_INLINE Port(Process<T> &process, uint64_t index, uint32_t out)
-      : process(process), index(index), out(out) {}
+  LIBC_INLINE Port(Process<T> &process, uint64_t lane_mask, uint64_t index,
+                   uint32_t out)
+      : process(process), lane_mask(lane_mask), 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.unlock(index); }
+  LIBC_INLINE void close() { process.unlock(lane_mask, index); }
 
 private:
   Process<T> &process;
+  uint64_t lane_mask;
   uint64_t index;
   uint32_t out;
 };
 
 /// The RPC client used to make requests to the server.
 struct Client : public Process<false> {
   LIBC_INLINE Client() = default;
   LIBC_INLINE Client(const Client &) = default;

 }
 
 /// 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<Client::Port> Client::try_open(uint16_t opcode) {
   constexpr uint64_t index = 0;
+  const uint64_t lane_mask = gpu::get_lane_mask();
+
   // Attempt to acquire the lock on this index.
-  if (!try_lock(index))
+  if (!try_lock(lane_mask, index))
     return cpp::nullopt;
 
   // The mailbox state must be read with the lock held.
   atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
 
   uint32_t in = load_inbox(index);
   uint32_t out = outbox[index].load(cpp::MemoryOrder::RELAXED);
 
   // Once we acquire the index we need to check if we are in a valid sending
   // state.
 
   if (buffer_unavailable(in, out)) {
-    unlock(index);
+    unlock(lane_mask, index);
     return cpp::nullopt;
   }
 
   buffer->opcode = opcode;
-  return Port(*this, index, out);
+  return Port(*this, lane_mask, index, out);
 }
 
 LIBC_INLINE Client::Port Client::open(uint16_t opcode) {
   for (;;) {
     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<Server::Port> Server::try_open() {
   constexpr uint64_t index = 0;
+  const uint64_t lane_mask = gpu::get_lane_mask();
+
   uint32_t in = load_inbox(index);
   uint32_t out = outbox[index].load(cpp::MemoryOrder::RELAXED);
 
   // The server is passive, if there is no work pending don't bother
   // opening a port.
   if (buffer_unavailable(in, out))
     return cpp::nullopt;
 
   // Attempt to acquire the lock on this index.
-  if (!try_lock(index))
+  if (!try_lock(lane_mask, index))
     return cpp::nullopt;
 
   // The mailbox state must be read with the lock held.
   atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
 
   in = load_inbox(index);
   out = outbox[index].load(cpp::MemoryOrder::RELAXED);
 
   if (buffer_unavailable(in, out)) {
-    unlock(index);
+    unlock(lane_mask, index);
     return cpp::nullopt;
   }
 
-  return Port(*this, index, out);
+  return Port(*this, lane_mask, index, out);
 }
 
 LIBC_INLINE Server::Port Server::open() {
   for (;;) {
     if (cpp::optional<Server::Port> p = try_open())
       return p.value();
     sleep_briefly();
   }
 }
 
 } // namespace rpc
 } // namespace __llvm_libc
 
 #endif