[libc] Remove `MAX_LANE_SIZE` definition from the RPC server

[jhuber6]

Summary

This MAX_LANE_SIZE was a hack from the days when we used a single
instance of the server and had some GPU state handle it. Now that we
have everything templated this really shouldn't be used. This patch
removes its use and replaces it with template arguments.

[jhuber6]

Fix comment

[JonChesterfield]

I'm not sure about this being on the methods. Why not the class itself? We have a bad time if different places manage to disagree on the value

libc/utils/gpu/loader/Loader.h

114

Static assert that it's 32 or 64 here, or maybe go so far as a valid-for-target call?

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

425

handle_error if it isn't 32 or 64 please. Maybe a switch on size with handle_error in the default?

[jhuber6]

Address comments. The template is on the handle_server call because it's a variant class we use std::visit on.

[JonChesterfield]

libc/utils/gpu/loader/Loader.h

112

This will never pass

[jhuber6]

Oops, wrong conditional.

[JonChesterfield]

libc/utils/gpu/server/rpc_server.cpp

96

Was a functional change intended here?

[jhuber6]

libc/utils/gpu/server/rpc_server.cpp

96

This isn't technically a functional change, since this callback automatically uses the current thread mask I could put this part in here instead which serves the purpose of getting rid of the invocation of the lane size.

libc/src/__support/RPC/rpc_util.h

 #include "src/__support/CPP/type_traits.h"
 #include "src/__support/GPU/utils.h"
 #include "src/__support/macros/attributes.h" // LIBC_INLINE
 #include "src/__support/macros/properties/architectures.h"
 
 namespace __llvm_libc {
 namespace rpc {
 
-/// Maximum amount of data a single lane can use.
-constexpr uint64_t MAX_LANE_SIZE = 64;
-
 /// Suspend the thread briefly to assist the thread scheduler during busy loops.
 LIBC_INLINE void sleep_briefly() {
 #if defined(LIBC_TARGET_ARCH_IS_NVPTX) && __CUDA_ARCH__ >= 700
   LIBC_INLINE_ASM("nanosleep.u32 64;" ::: "memory");
 #elif defined(LIBC_TARGET_ARCH_IS_AMDGPU)
   __builtin_amdgcn_s_sleep(2);
 #elif defined(LIBC_TARGET_ARCH_IS_X86)
   __builtin_ia32_pause();

libc/utils/gpu/loader/Loader.h

inline void handle_error(const char *msg) {
   fprintf(stderr, "%s\n", msg);
   exit(EXIT_FAILURE);
 }
 
 inline void handle_error(rpc_status_t) {
   handle_error("Failure in the RPC server\n");
 }
 
+template <uint32_t lane_size>
 inline void register_rpc_callbacks(uint32_t device_id) {
+  static_assert(lane_size == 32 && lane_size == 64, "Invalid Lane size");
   // Register the ping test for the `libc` tests.
   rpc_register_callback(
       device_id, static_cast<rpc_opcode_t>(RPC_TEST_INCREMENT),
       [](rpc_port_t port, void *data) {
         rpc_recv_and_send(
             port,
             [](rpc_buffer_t *buffer, void *data) {
               reinterpret_cast<uint64_t *>(buffer->data)[0] += 1;
             },
             data);
rpc_register_callback(
               &cnt);
       },
       nullptr);
 
   // Register the stream test handler.
   rpc_register_callback(
       device_id, static_cast<rpc_opcode_t>(RPC_TEST_STREAM),
       [](rpc_port_t port, void *data) {
-        uint64_t sizes[RPC_MAXIMUM_LANE_SIZE] = {0};
-        void *dst[RPC_MAXIMUM_LANE_SIZE] = {nullptr};
+        uint64_t sizes[lane_size] = {0};
+        void *dst[lane_size] = {nullptr};
         rpc_recv_n(
             port, dst, sizes,
             [](uint64_t size, void *) -> void * { return new char[size]; },
             nullptr);
         rpc_send_n(port, dst, sizes);
-        for (uint64_t i = 0; i < RPC_MAXIMUM_LANE_SIZE; ++i) {
+        for (uint64_t i = 0; i < lane_size; ++i) {
           if (dst[i])
             delete[] reinterpret_cast<uint8_t *>(dst[i]);
         }
       },
       nullptr);
 }
 
 #endif

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

hsa_status_t launch_kernel(hsa_agent_t dev_agent, hsa_executable_t executable,
   // Look up the '_start' kernel in the loaded executable.
   hsa_executable_symbol_t symbol;
   if (hsa_status_t err = hsa_executable_get_symbol_by_name(
           executable, kernel_name, &dev_agent, &symbol))
     return err;
 
   // Register RPC callbacks for the malloc and free functions on HSA.
   uint32_t device_id = 0;
-  register_rpc_callbacks(device_id);
-
   auto tuple = std::make_tuple(dev_agent, coarsegrained_pool);
   rpc_register_callback(
       device_id, RPC_MALLOC,
       [](rpc_port_t port, void *data) {
         auto malloc_handler = [](rpc_buffer_t *buffer, void *data) -> void {
           auto &[dev_agent, pool] = *static_cast<decltype(tuple) *>(data);
           uint64_t size = buffer->data[0];
           void *dev_ptr = nullptr;
if (hsa_status_t err = hsa_amd_memory_pool_allocate(finegrained_pool, size,
       handle_error(err);
     hsa_amd_agents_allow_access(1, &dev_agent, nullptr, dev_ptr);
     return dev_ptr;
   };
   if (rpc_status_t err = rpc_server_init(device_id, RPC_MAXIMUM_PORT_COUNT,
                                          wavefront_size, rpc_alloc, &tuple))
     handle_error(err);
 
+  // Register callbacks for the RPC unit tests.
+  if (wavefront_size == 32)
+    register_rpc_callbacks<32>(device_id);
+  else if (wavefront_size == 64)
+    register_rpc_callbacks<64>(device_id);
+  else
+    handle_error("Invalid wavefront size");
+
   // Obtain the GPU's fixed-frequency clock rate and copy it to the GPU.
   // If the clock_freq symbol is missing, no work to do.
   hsa_executable_symbol_t freq_sym;
   if (HSA_STATUS_SUCCESS ==
       hsa_executable_get_symbol_by_name(executable, "__llvm_libc_clock_freq",
                                         &dev_agent, &freq_sym)) {
 
     void *host_clock_freq;

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

CUresult launch_kernel(CUmodule binary, CUstream stream,
   // needs to be done on a separate stream or else it will deadlock with the
   // executing kernel.
   CUstream memory_stream;
   if (CUresult err = cuStreamCreate(&memory_stream, CU_STREAM_NON_BLOCKING))
     handle_error(err);
 
   // Register RPC callbacks for the malloc and free functions on HSA.
   uint32_t device_id = 0;
-  register_rpc_callbacks(device_id);
+  register_rpc_callbacks<32>(device_id);
 
   rpc_register_callback(
       device_id, RPC_MALLOC,
       [](rpc_port_t port, void *data) {
         auto malloc_handler = [](rpc_buffer_t *buffer, void *data) -> void {
           CUstream memory_stream = *static_cast<CUstream *>(data);
           uint64_t size = buffer->data[0];
           CUdeviceptr dev_ptr;

libc/utils/gpu/server/rpc_server.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 = 512;
 
-/// The maximum number of parallel lanes that we can support.
-const uint64_t RPC_MAXIMUM_LANE_SIZE = 64;
-
 /// The symbol name associated with the client for use with the LLVM C library
 /// implementation.
 const char *const rpc_client_symbol_name = "__llvm_libc_rpc_client";
 
 /// status codes.
 typedef enum {
   RPC_STATUS_SUCCESS = 0x0,
   RPC_STATUS_CONTINUE = 0x1,

libc/utils/gpu/server/rpc_server.cpp

 using namespace __llvm_libc;
 
 static_assert(sizeof(rpc_buffer_t) == sizeof(rpc::Buffer),
               "Buffer size mismatch");
 
 static_assert(RPC_MAXIMUM_PORT_COUNT == rpc::MAX_PORT_COUNT,
               "Incorrect maximum port count");
 
-static_assert(RPC_MAXIMUM_LANE_SIZE == rpc::MAX_LANE_SIZE,
-              "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)) {}
 
   void reset(uint64_t port_count, void *buffer) {
rpc_status_t handle_server(
     auto port = server.try_open();
     if (!port)
       return RPC_STATUS_SUCCESS;
 
     switch (port->get_opcode()) {
     case RPC_WRITE_TO_STREAM:
     case RPC_WRITE_TO_STDERR:
     case RPC_WRITE_TO_STDOUT: {
-      uint64_t sizes[rpc::MAX_LANE_SIZE] = {0};
-      void *strs[rpc::MAX_LANE_SIZE] = {nullptr};
-      FILE *files[rpc::MAX_LANE_SIZE] = {nullptr};
+      uint64_t sizes[lane_size] = {0};
+      void *strs[lane_size] = {nullptr};
+      FILE *files[lane_size] = {nullptr};
       if (port->get_opcode() == RPC_WRITE_TO_STREAM)
         port->recv([&](rpc::Buffer *buffer, uint32_t id) {
           files[id] = reinterpret_cast<FILE *>(buffer->data[0]);
         });
       port->recv_n(strs, sizes, [&](uint64_t size) { return new char[size]; });
       port->send([&](rpc::Buffer *buffer, uint32_t id) {
         FILE *file =
             port->get_opcode() == RPC_WRITE_TO_STDOUT
                 ? stdout
                 : (port->get_opcode() == RPC_WRITE_TO_STDERR ? stderr
                                                              : files[id]);
         uint64_t ret = fwrite(strs[id], 1, sizes[id], file);
         std::memcpy(buffer->data, &ret, sizeof(uint64_t));
+        delete[] reinterpret_cast<uint8_t *>(strs[id]);
       });
-      for (uint64_t i = 0; i < rpc::MAX_LANE_SIZE; ++i) {
-        if (strs[i])
-          delete[] reinterpret_cast<uint8_t *>(strs[i]);
-      }
       break;
     }
     case RPC_READ_FROM_STREAM:
     case RPC_READ_FROM_STDIN: {
-      uint64_t sizes[rpc::MAX_LANE_SIZE] = {0};
-      void *data[rpc::MAX_LANE_SIZE] = {nullptr};
-      uint64_t rets[rpc::MAX_LANE_SIZE] = {0};
+      uint64_t sizes[lane_size] = {0};
+      void *data[lane_size] = {nullptr};
+      uint64_t rets[lane_size] = {0};
       port->recv([&](rpc::Buffer *buffer, uint32_t id) {
         sizes[id] = buffer->data[0];
         data[id] = new char[sizes[id]];
         FILE *file = port->get_opcode() == RPC_READ_FROM_STREAM
                          ? reinterpret_cast<FILE *>(buffer->data[1])
                          : stdin;
         rets[id] = fread(data[id], 1, sizes[id], file);
       });
       port->send_n(data, sizes);
       port->send([&](rpc::Buffer *buffer, uint32_t id) {
         delete[] reinterpret_cast<uint8_t *>(data[id]);
         std::memcpy(buffer->data, &rets[id], sizeof(uint64_t));
       });
       break;
     }
     case RPC_OPEN_FILE: {
-      uint64_t sizes[rpc::MAX_LANE_SIZE] = {0};
-      void *paths[rpc::MAX_LANE_SIZE] = {nullptr};
+      uint64_t sizes[lane_size] = {0};
+      void *paths[lane_size] = {nullptr};
       port->recv_n(paths, sizes, [&](uint64_t size) { return new char[size]; });
       port->recv_and_send([&](rpc::Buffer *buffer, uint32_t id) {
         FILE *file = fopen(reinterpret_cast<char *>(paths[id]),
                            reinterpret_cast<char *>(buffer->data));
         buffer->data[0] = reinterpret_cast<uintptr_t>(file);
       });
       break;
     }
case RPC_EXIT: {
       port->recv([](rpc::Buffer *buffer) {
         int status = 0;
         std::memcpy(&status, buffer->data, sizeof(int));
         exit(status);
       });
       break;
     }
     case RPC_HOST_CALL: {
-      uint64_t sizes[rpc::MAX_LANE_SIZE] = {0};
-      void *args[rpc::MAX_LANE_SIZE] = {nullptr};
+      uint64_t sizes[lane_size] = {0};
+      void *args[lane_size] = {nullptr};
       port->recv_n(args, sizes, [&](uint64_t size) { return new char[size]; });
       port->recv([&](rpc::Buffer *buffer, uint32_t id) {
         reinterpret_cast<void (*)(void *)>(buffer->data[0])(args[id]);
       });
       port->send([&](rpc::Buffer *, uint32_t id) {
         delete[] reinterpret_cast<uint8_t *>(args[id]);
       });
       break;