[libc][rpc] Allocate locks array within process

[JonChesterfield]

Summary

Replaces the globals currently used. Worth changing to a bitmap
before allowing runtime number of ports >> 64. One bit per port is likely
to be cheap enough that sizing for the worst case is always fine, otherwise
in the future we can change to dynamically allocating it.

[JonChesterfield]

Main motivation is to make the client/server allocation simpler before extending the test coverage for these classes

libc/src/__support/RPC/rpc.h

107

Number doesn't matter hugely at this point, have gone with deliberately != default_port_count.

[jhuber6]

libc/src/__support/RPC/rpc.h

107

What's the point of keeping both around?

libc/utils/gpu/loader/amdgpu/Loader.cpp

362

Needs the same in the nvptx loader.

[JonChesterfield]

libc/src/__support/RPC/rpc.h

107

I'm expecting the number of ports to be chosen at runtime based on the GPU and the number of locks to be a fixed worst case thing using a couple of kb of device memory

libc/utils/gpu/loader/amdgpu/Loader.cpp

362

Ah yes, will fix. Thanks

[JonChesterfield]

• also nvptx loader

[JonChesterfield]

• fold counts together for now

libc/src/__support/RPC/rpc.h

 template <bool InvertInbox> 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;
   uint32_t lane_size;
-  cpp::Atomic<uint32_t> *lock;
   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 *lock,
-                         void *inbox, void *outbox, void *packet) {
-    *this = {port_count,
-             lane_size,
-             reinterpret_cast<cpp::Atomic<uint32_t> *>(lock),
-             reinterpret_cast<cpp::Atomic<uint32_t> *>(inbox),
-             reinterpret_cast<cpp::Atomic<uint32_t> *>(outbox),
-             reinterpret_cast<Packet *>(packet)};
+  LIBC_INLINE void reset(uint64_t port_count, uint32_t lane_size, void *inbox,
+                         void *outbox, void *packet) {
+    this->port_count = port_count;
+    this->lane_size = lane_size;
+    this->inbox = reinterpret_cast<cpp::Atomic<uint32_t> *>(inbox);
+    this->outbox = reinterpret_cast<cpp::Atomic<uint32_t> *>(outbox);
+    this->packet = reinterpret_cast<Packet *>(packet);
   }
 
   /// 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),

libc/startup/gpu/amdgpu/start.cpp

 #include "src/__support/RPC/rpc_client.h"
 #include "src/stdlib/atexit.h"
 #include "src/stdlib/exit.h"
 
 extern "C" int main(int argc, char **argv, char **envp);
 
 namespace __llvm_libc {
 
-static cpp::Atomic<uint32_t> lock[rpc::default_port_count] = {0};
-
 extern "C" uintptr_t __init_array_start[];
 extern "C" uintptr_t __init_array_end[];
 extern "C" uintptr_t __fini_array_start[];
 extern "C" uintptr_t __fini_array_end[];
 
 using InitCallback = void(int, char **, char **);
 using FiniCallback = void(void);
 

 
 } // namespace __llvm_libc
 
 extern "C" [[gnu::visibility("protected"), clang::amdgpu_kernel]] void
 _begin(int argc, char **argv, char **env, void *in, void *out, void *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::gpu::get_lane_size(),
-                                 &__llvm_libc::lock, in, out, buffer);
+                                 __llvm_libc::gpu::get_lane_size(), in, out,
+                                 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

 #include "src/__support/RPC/rpc_client.h"
 #include "src/stdlib/atexit.h"
 #include "src/stdlib/exit.h"
 
 extern "C" int main(int argc, char **argv, char **envp);
 
 namespace __llvm_libc {
 
-static cpp::Atomic<uint32_t> lock[rpc::default_port_count] = {0};
-
 extern "C" {
 // Nvidia's 'nvlink' linker does not provide these symbols. We instead need
 // to manually create them and update the globals in the loader implememtation.
 uintptr_t *__init_array_start [[gnu::visibility("protected")]];
 uintptr_t *__init_array_end [[gnu::visibility("protected")]];
 uintptr_t *__fini_array_start [[gnu::visibility("protected")]];
 uintptr_t *__fini_array_end [[gnu::visibility("protected")]];
 }

 
 } // namespace __llvm_libc
 
 extern "C" [[gnu::visibility("protected"), clang::nvptx_kernel]] void
 _begin(int argc, char **argv, char **env, void *in, void *out, void *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::gpu::get_lane_size(),
-                                 &__llvm_libc::lock, in, out, buffer);
+                                 __llvm_libc::gpu::get_lane_size(), in, out,
+                                 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/loader/Server.h

 #include <cstdlib>
 #include <cstring>
 #include <stddef.h>
 
 #include "src/__support/RPC/rpc.h"
 
 static __llvm_libc::rpc::Server server;
 
-static __llvm_libc::cpp::Atomic<uint32_t>
-    lock[__llvm_libc::rpc::default_port_count] = {0};
-
 /// Queries the RPC client at least once and performs server-side work if there
 /// are any active requests.
 void handle_server() {
   // Continue servicing the client until there is no work left and we return.
   for (;;) {
     auto port = server.try_open();
     if (!port)
       return;

libc/utils/gpu/loader/amdgpu/Loader.cpp

if (hsa_status_t err = hsa_amd_memory_pool_allocate(
                        alignof(__llvm_libc::rpc::Packet)),
           /*flags=*/0, &buffer))
     handle_error(err);
   hsa_amd_agents_allow_access(1, &dev_agent, nullptr, server_inbox);
   hsa_amd_agents_allow_access(1, &dev_agent, nullptr, server_outbox);
   hsa_amd_agents_allow_access(1, &dev_agent, nullptr, buffer);
 
   // Initialize the RPC server's buffer for host-device communication.
-  server.reset(port_size, wavefront_size, &lock, server_inbox, server_outbox,
-               buffer);
+  server.reset(port_size, wavefront_size, server_inbox, server_outbox, buffer);
 
   // Obtain a queue with the minimum (power of two) size, used to send commands
   // to the HSA runtime and launch execution on the device.
   uint64_t queue_size;
   if (hsa_status_t err = hsa_agent_get_info(
           dev_agent, HSA_AGENT_INFO_QUEUE_MIN_SIZE, &queue_size))
     handle_error(err);
   hsa_queue_t *queue = nullptr;

libc/utils/gpu/loader/nvptx/Loader.cpp

int load(int argc, char **argv, char **envp, void *image, size_t size,
   void *buffer = allocator(
       port_size * align_up(sizeof(__llvm_libc::rpc::Header) +
                                (warp_size * sizeof(__llvm_libc::rpc::Buffer)),
                            alignof(__llvm_libc::rpc::Packet)));
   if (!server_inbox || !server_outbox || !buffer)
     handle_error("Failed to allocate memory the RPC client / server.");
 
   // Initialize the RPC server's buffer for host-device communication.
-  server.reset(port_size, warp_size, &lock, server_inbox, server_outbox,
-               buffer);
+  server.reset(port_size, warp_size, server_inbox, server_outbox, buffer);
 
   LaunchParameters single_threaded_params = {1, 1, 1, 1, 1, 1};
   // Call the kernel to
   begin_args_t init_args = {argc,          dev_argv,     dev_envp,
                             server_outbox, server_inbox, buffer};
   if (CUresult err = launch_kernel(binary, stream, single_threaded_params,
                                    "_begin", init_args))
     handle_error(err);