[libc] Begin implementing a library for the RPC server

[jhuber6]

Summary

This patch begins providing a generic static library that wraps around
the raw rpc.h interface. As discussed in the corresponding RFC,
https://discourse.llvm.org/t/rfc-libc-exporting-the-rpc-interface-for-the-gpu-libc/71030,
we want to begin exporting RPC services to external users. In order to
do this we decided to not expose the rpc.h header by wrapping around
its functionality. This is done with a C-interface as we make heavy use
of callbacks and allows us to provide a predictable interface.

[jhuber6]

Make include directory public to simplify usage.

[jhuber6]

Changing implementation to use a C-based interface. This is primarily because
it's required to actually hide the implementation details of the RPC server. The
rpc.h header should not be visible outside of internal libc projects. This
also makes it easier to provide as a simple library with an expected symbol
name.

[jhuber6]

Forgot to reset memory, and ping.

[jhuber6]

Updating and rebasing.

I'm not completely happy with the interface but we can always modify it later.
This should provide what the libc developers wanted and separate the
implementation of rpc.h and the provided server. This has some downsides
compared to just exporting the full header somehow, but also has some
advantages, considering that we may not need to really provde that much custom
utility outside of libc for users, so it's easier to consoliate the
functionality in a defined interface.

Let me know what should be changed, there's a lot of cruft associated with
registering the custom handlers for malloc / free but I don't think there's
another way to get around it.

[sivachandra]

Answering the question from discord: Normally, a single header file named <some>_service.h is to be used by both the client and the server. Client's will use the client API from that header file, where as servers will use the server API. That way, the message types (or opcodes as you are calling) will be in a single shared header file.

libc/utils/gpu/server/Server.h

46 ↗

(On Diff #528564)

Why are there separate functions allocation and deallocation?

51 ↗

(On Diff #528564)

Update comment.

54 ↗

(On Diff #528564)

Ditto.

61 ↗

(On Diff #528564)

Why not just rpc_shutdown?

[jhuber6]

libc/utils/gpu/server/Server.h

46 ↗

(On Diff #528564)

I'll change the name to free but we need both to adequately de allocate the shared memory.

[jhuber6]

Making suggested changes

[jhuber6]

Ping

[sivachandra]

OK from my side but GPU review should be done by a GPU expert.

[JonChesterfield]

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

55

Call into the other one,

static void handle_error(rpc_status_t) {
  handle_error("Failure in the RPC server");
}

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

50

i still really dislike the copy/paste going on here

72

Does this work? It looks like the same stream running the kernel is being used to provide malloc/free, and I'd expect that to deadlock

libc/utils/gpu/server/Server.cpp

35 ↗

(On Diff #528799)

Could we go with a vector of Device instead of the new+array construct?

44 ↗

(On Diff #528799)

Why is this a heap allocated thing, as opposed to static State state; ?

51 ↗

(On Diff #528799)

Could make the counter 64 bit and delete the test against max as a counter >= address space size can't overflow

In general the DIY reference counting is a bit odd - is there a reason this isn't a shared_ptr?

107 ↗

(On Diff #528799)

I'm still hopeful that we'll come up with a better idea than rpc::MAX_LANE_SIZE

libc/utils/gpu/server/Server.h

57 ↗

(On Diff #528799)

This looks like the type of a function, not the type of a function pointer. It's used as an argument to functions where it'll decay to the pointer type. More conventionally written with an extra *

typedef void(*rpc_free_ty)(void *ptr, void *data);

Is there a benefit to declaring this as the function type as opposed to the function pointer type?

70 ↗

(On Diff #528799)

Want (void) if this is meant to be usable from C (the guards about suggest it is)

C++ thinks foo() is a function of no arguments. C thinks it's some aberration from the past (though that might have been dropped in the last standard).

[jhuber6]

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

72

There's a test for this that's been running on https://lab.llvm.org/buildbot/#/builders/46 for a few weeks now and it hasn't deadlocked as far as I can tell. It's a completely separate stream called memory_stream that's just created here. The one running the kernel is just called stream. This requires CUDA 11.2 IIRC.

libc/utils/gpu/server/Server.cpp

35 ↗

(On Diff #528799)

It's a static size so a constant sized array should be more correct.

44 ↗

(On Diff #528799)

It's just easier to check if it's been initialized because the pointer is nullable. We coiuld probably make it a static thing and have a flag instead if you'd like.

51 ↗

(On Diff #528799)

Would a shared pointer give us the same semantics? We would be allocating it multiple times and not copying it.

107 ↗

(On Diff #528799)

We could use a vector and push back into it instead, or preallocate according to the size above, but this was the easiest solution.

[jhuber6]

libc/utils/gpu/server/Server.h

57 ↗

(On Diff #528799)

I just forgot to add it and my IDE took care of the conversions when I got an error, will change.

70 ↗

(On Diff #528799)

Forgot about that quirk of C, thanks.

[jhuber6]

Addressing comments

[JonChesterfield]

Interesting that the cuda path doesn't deadlock, I wonder if that's a fix from previous cuda revisions. Thanks

[jhuber6]

In D147054#4422839, @JonChesterfield wrote:

Interesting that the cuda path doesn't deadlock, I wonder if that's a fix from previous cuda revisions. Thanks

Appreciate the thorough review.

Yeah I heard about that problem as well which made me nervous about how well we could support a real malloc in the future. But pleasingly it's been running on the sm_70 and sm_60 testers for awhile and it seems good as long as your CUDA is new enough. As it stands we'll probably need to turn this function into a hard error if we integrate it into OpenMP and CUDA is too old.

libc/utils/gpu/CMakeLists.txt

+add_subdirectory(server)
 add_subdirectory(loader)

libc/utils/gpu/loader/amdgpu/CMakeLists.txt

 add_executable(amdhsa_loader Loader.cpp)
-add_dependencies(amdhsa_loader libc.src.__support.RPC.rpc)
 
-target_include_directories(amdhsa_loader PRIVATE ${LIBC_SOURCE_DIR})
 target_link_libraries(amdhsa_loader
   PRIVATE
   hsa-runtime64::hsa-runtime64
+  rpc_server
   gpu_loader
 )

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

 // architecture. The file launches the '_start' kernel which should be provided
 // by the device application start code and call ultimately call the 'main'
 // function.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Loader.h"
 
-#include "src/__support/RPC/rpc.h"
+#include "rpc_server.h"
 
 #include <hsa/hsa.h>
 #include <hsa/hsa_ext_amd.h>
 
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
+#include <tuple>
 #include <utility>
 
 /// The name of the kernel we will launch. All AMDHSA kernels end with '.kd'.
 constexpr const char *KERNEL_START = "_start.kd";
 
 /// The arguments to the '_start' kernel.
 struct kernel_args_t {
   int argc;
   void *argv;
   void *envp;
   void *ret;
   void *inbox;
   void *outbox;
   void *buffer;
 };
 
-static __llvm_libc::rpc::Server server;
-
-/// Queries the RPC client at least once and performs server-side work if there
-/// are any active requests.
-void handle_server() {
-  while (server.handle(
-      [&](__llvm_libc::rpc::Buffer *buffer) {
-        switch (static_cast<__llvm_libc::rpc::Opcode>(buffer->data[0])) {
-        case __llvm_libc::rpc::Opcode::PRINT_TO_STDERR: {
-          fputs(reinterpret_cast<const char *>(&buffer->data[1]), stderr);
-          break;
-        }
-        case __llvm_libc::rpc::Opcode::EXIT: {
-          exit(buffer->data[1]);
-          break;
-        }
-        default:
-          return;
-        };
-      },
-      [](__llvm_libc::rpc::Buffer *buffer) {}))
-    ;
-}
-
 /// Print the error code and exit if \p code indicates an error.
 static void handle_error(hsa_status_t code) {
   if (code == HSA_STATUS_SUCCESS || code == HSA_STATUS_INFO_BREAK)
     return;
 
   const char *desc;
   if (hsa_status_string(code, &desc) != HSA_STATUS_SUCCESS)
     desc = "Unknown error";
   fprintf(stderr, "%s\n", desc);
   exit(EXIT_FAILURE);
 }
 
 static void handle_error(const char *msg) {
   fprintf(stderr, "%s\n", msg);
   exit(EXIT_FAILURE);
 }
 
 /// Generic interface for iterating using the HSA callbacks.
 template <typename elem_ty, typename func_ty, typename callback_ty>
 hsa_status_t iterate(func_ty func, callback_ty cb) {
   auto l = [](elem_ty elem, void *data) -> hsa_status_t {
int load(int argc, char **argv, char **envp, void *image, size_t size) {
   void *dev_ret;
   if (hsa_status_t err =
           hsa_amd_memory_pool_allocate(coarsegrained_pool, sizeof(int),
                                        /*flags=*/0, &dev_ret))
     handle_error(err);
   hsa_amd_memory_fill(dev_ret, 0, sizeof(int));
 
   // Allocate finegrained memory for the RPC server and client to share.
-  void *server_inbox;
-  void *server_outbox;
-  void *buffer;
-  if (hsa_status_t err = hsa_amd_memory_pool_allocate(
-          finegrained_pool, sizeof(__llvm_libc::cpp::Atomic<int>),
-          /*flags=*/0, &server_inbox))
-    handle_error(err);
-  if (hsa_status_t err = hsa_amd_memory_pool_allocate(
-          finegrained_pool, sizeof(__llvm_libc::cpp::Atomic<int>),
-          /*flags=*/0, &server_outbox))
-    handle_error(err);
-  if (hsa_status_t err = hsa_amd_memory_pool_allocate(
-          finegrained_pool, sizeof(__llvm_libc::rpc::Buffer),
-          /*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);
+  auto rpc_data = std::make_tuple(finegrained_pool, dev_agent);
+  auto rpc_allocator = [](uint64_t size, void *data) -> void * {
+    auto &[finegrained_pool, dev_agent] =
+        *reinterpret_cast<decltype(rpc_data) *>(data);
+    void *dev_ptr = nullptr;
+    if (hsa_status_t err = hsa_amd_memory_pool_allocate(finegrained_pool, size,
+                                                        /*flags=*/0, &dev_ptr))
+      handle_error(err);
+    hsa_amd_agents_allow_access(1, &dev_agent, nullptr, dev_ptr);
+    return dev_ptr;
+  };
+  rpc_init(rpc_allocator, &rpc_data);
 
   // Initialie all the arguments (explicit and implicit) to zero, then set the
   // explicit arguments to the values created above.
   std::memset(args, 0, args_size);
   kernel_args_t *kernel_args = reinterpret_cast<kernel_args_t *>(args);
   kernel_args->argc = argc;
   kernel_args->argv = dev_argv;
   kernel_args->envp = dev_envp;
   kernel_args->ret = dev_ret;
-  kernel_args->inbox = server_outbox;
-  kernel_args->outbox = server_inbox;
-  kernel_args->buffer = buffer;
+  kernel_args->inbox = rpc_get_outbox();
+  kernel_args->outbox = rpc_get_inbox();
+  kernel_args->buffer = rpc_get_buffer();
 
   // Obtain a packet from the queue.
   uint64_t packet_id = hsa_queue_add_write_index_relaxed(queue, 1);
   while (packet_id - hsa_queue_load_read_index_scacquire(queue) >= queue_size)
     ;
 
   const uint32_t mask = queue_size - 1;
   hsa_kernel_dispatch_packet_t *packet =
int load(int argc, char **argv, char **envp, void *image, size_t size) {
   packet->kernel_object = kernel;
   packet->kernarg_address = args;
 
   // Create a signal to indicate when this packet has been completed.
   if (hsa_status_t err =
           hsa_signal_create(1, 0, nullptr, &packet->completion_signal))
     handle_error(err);
 
-  // Initialize the RPC server's buffer for host-device communication.
-  server.reset(server_inbox, server_outbox, buffer);
-
   // Initialize the packet header and set the doorbell signal to begin execution
   // by the HSA runtime.
   uint16_t header =
       (HSA_PACKET_TYPE_KERNEL_DISPATCH << HSA_PACKET_HEADER_TYPE) |
       (HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_ACQUIRE_FENCE_SCOPE) |
       (HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_RELEASE_FENCE_SCOPE);
   __atomic_store_n(&packet->header, header | (packet->setup << 16),
                    __ATOMIC_RELEASE);
   hsa_signal_store_relaxed(queue->doorbell_signal, packet_id);
 
   // Wait until the kernel has completed execution on the device. Periodically
   // check the RPC client for work to be performed on the server.
   while (hsa_signal_wait_scacquire(
              packet->completion_signal, HSA_SIGNAL_CONDITION_EQ, 0,
              /*timeout_hint=*/1024, HSA_WAIT_STATE_ACTIVE) != 0)
-    handle_server();
+    rpc_handle();
 
   // Create a memory signal and copy the return value back from the device into
   // a new buffer.
   hsa_signal_t memory_signal;
   if (hsa_status_t err = hsa_signal_create(1, 0, nullptr, &memory_signal))
     handle_error(err);
 
   void *host_ret;

libc/utils/gpu/loader/nvptx/CMakeLists.txt

 add_executable(nvptx_loader Loader.cpp)
-add_dependencies(nvptx_loader libc.src.__support.RPC.rpc)
 
-target_include_directories(nvptx_loader PRIVATE ${LIBC_SOURCE_DIR})
 target_link_libraries(nvptx_loader
   PRIVATE
   gpu_loader
+  rpc_server
   CUDA::cuda_driver
 )

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

 // architecture. The file launches the '_start' kernel which should be provided
 // by the device application start code and call ultimately call the 'main'
 // function.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Loader.h"
 
-#include "src/__support/RPC/rpc.h"
+#include "rpc_server.h"
 
 #include "cuda.h"
 #include <cstddef>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
 
 /// The arguments to the '_start' kernel.
 struct kernel_args_t {
   int argc;
   void *argv;
   void *envp;
   void *ret;
   void *inbox;
   void *outbox;
   void *buffer;
 };
 
-static __llvm_libc::rpc::Server server;
-
-/// Queries the RPC client at least once and performs server-side work if there
-/// are any active requests.
-void handle_server() {
-  while (server.handle(
-      [&](__llvm_libc::rpc::Buffer *buffer) {
-        switch (static_cast<__llvm_libc::rpc::Opcode>(buffer->data[0])) {
-        case __llvm_libc::rpc::Opcode::PRINT_TO_STDERR: {
-          fputs(reinterpret_cast<const char *>(&buffer->data[1]), stderr);
-          break;
-        }
-        case __llvm_libc::rpc::Opcode::EXIT: {
-          exit(buffer->data[1]);
-          break;
-        }
-        default:
-          return;
-        };
-      },
-      [](__llvm_libc::rpc::Buffer *buffer) {}))
-    ;
-}
-
 static void handle_error(CUresult err) {
   if (err == CUDA_SUCCESS)
     return;
 
   const char *err_str = nullptr;
   CUresult result = cuGetErrorString(err, &err_str);
   if (result != CUDA_SUCCESS)
     fprintf(stderr, "Unknown Error\n");
   else
     fprintf(stderr, "%s\n", err_str);
   exit(1);
 }
 
 static void handle_error(const char *msg) {
   fprintf(stderr, "%s\n", msg);
   exit(EXIT_FAILURE);
 }
 
 int load(int argc, char **argv, char **envp, void *image, size_t size) {
   if (CUresult err = cuInit(0))
     handle_error(err);
 
   // Obtain the first device found on the system.
   CUdevice device;
   if (CUresult err = cuDeviceGet(&device, 0))
     handle_error(err);
 
   // Initialize the CUDA context and claim it for this execution.
   CUcontext context;
   if (CUresult err = cuDevicePrimaryCtxRetain(&context, device))
     handle_error(err);
   if (CUresult err = cuCtxSetCurrent(context))
     handle_error(err);
 
   // Initialize a non-blocking CUDA stream to execute the kernel.
   CUstream stream;
   if (CUresult err = cuStreamCreate(&stream, CU_STREAM_NON_BLOCKING))
     handle_error(err);
 
   // Load the image into a CUDA module.
   CUmodule binary;
   if (CUresult err = cuModuleLoadDataEx(&binary, image, 0, nullptr, nullptr))
     handle_error(err);
 
int load(int argc, char **argv, char **envp, void *image, size_t size) {
 
   // Allocate space for the return pointer and initialize it to zero.
   CUdeviceptr dev_ret;
   if (CUresult err = cuMemAlloc(&dev_ret, sizeof(int)))
     handle_error(err);
   if (CUresult err = cuMemsetD32(dev_ret, 0, 1))
     handle_error(err);
 
-  void *server_inbox = allocator(sizeof(__llvm_libc::cpp::Atomic<int>));
-  void *server_outbox = allocator(sizeof(__llvm_libc::cpp::Atomic<int>));
-  void *buffer = allocator(sizeof(__llvm_libc::rpc::Buffer));
-  if (!server_inbox || !server_outbox || !buffer)
-    handle_error("Failed to allocate memory the RPC client / server.");
+  // Allocate finegrained memory for the RPC server and client to share.
+  auto rpc_allocator = [](uint64_t size, void *) -> void * {
+    void *dev_ptr;
+    if (CUresult err = cuMemAllocHost(&dev_ptr, size))
+      handle_error(err);
+    return dev_ptr;
+  };
+  rpc_init(rpc_allocator, nullptr);
 
   // Set up the arguments to the '_start' kernel on the GPU.
   uint64_t args_size = sizeof(kernel_args_t);
   kernel_args_t args;
   std::memset(&args, 0, args_size);
   args.argc = argc;
   args.argv = dev_argv;
   args.envp = dev_envp;
   args.ret = reinterpret_cast<void *>(dev_ret);
-  args.inbox = server_outbox;
-  args.outbox = server_inbox;
-  args.buffer = buffer;
+  args.inbox = rpc_get_outbox();
+  args.outbox = rpc_get_inbox();
+  args.buffer = rpc_get_buffer();
   void *args_config[] = {CU_LAUNCH_PARAM_BUFFER_POINTER, &args,
                          CU_LAUNCH_PARAM_BUFFER_SIZE, &args_size,
                          CU_LAUNCH_PARAM_END};
 
-  // Initialize the RPC server's buffer for host-device communication.
-  server.reset(server_inbox, server_outbox, buffer);
-
   // Call the kernel with the given arguments.
   if (CUresult err =
           cuLaunchKernel(function, /*gridDimX=*/1, /*gridDimY=*/1,
                          /*gridDimZ=*/1, /*blockDimX=*/1, /*blockDimY=*/1,
                          /*bloackDimZ=*/1, 0, stream, nullptr, args_config))
     handle_error(err);
 
   // Wait until the kernel has completed execution on the device. Periodically
   // check the RPC client for work to be performed on the server.
   while (cuStreamQuery(stream) == CUDA_ERROR_NOT_READY)
-    handle_server();
+    rpc_handle();
 
   // Copy the return value back from the kernel and wait.
   int host_ret = 0;
   if (CUresult err = cuMemcpyDtoH(&host_ret, dev_ret, sizeof(int)))
     handle_error(err);
 
   if (CUresult err = cuStreamSynchronize(stream))
     handle_error(err);
 
   // Destroy the context and the loaded binary.
   if (CUresult err = cuModuleUnload(binary))
     handle_error(err);
   if (CUresult err = cuDevicePrimaryCtxRelease(device))
     handle_error(err);
   return host_ret;
 }

libc/utils/gpu/server/CMakeLists.txt

+add_library(rpc_server STATIC rpc_server.h rpc_server.cpp)
+
+# Include the RPC implemenation from libc.
+add_dependencies(rpc_server libc.src.__support.RPC.rpc)
+target_include_directories(rpc_server PRIVATE ${LIBC_SOURCE_DIR})
+target_include_directories(rpc_server PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})

libc/utils/gpu/server/rpc_server.h

+//===-- Shared memory RPC server instantiation ------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIBC_UTILS_GPU_SERVER_RPC_SERVER_H
+#define LLVM_LIBC_UTILS_GPU_SERVER_RPC_SERVER_H
+
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+typedef void *(rpc_alloc_ty)(uint64_t size, void *data);
+
+typedef void(rpc_dealloc_ty)(void *ptr, void *data);
+
+/// Initialize the server with unified memory to communicate with the client.
+void rpc_init(rpc_alloc_ty alloc, void *data);
+
+/// Deallocate the memory associated with the server.
+void rpc_deinit(rpc_dealloc_ty, void *data);
+
+/// Queries the RPC client at least once and performs server-side work if there
+/// are any active requests.
+void rpc_handle();
+
+/// Get the pointer to the data inbox.
+/// TODO: We should try to compress this into a single buffer.
+void *rpc_get_inbox();
+
+/// Get the pointer to the data outbox.
+void *rpc_get_outbox();
+
+/// Get the pointer to the data buffer.
+void *rpc_get_buffer();
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif

libc/utils/gpu/server/rpc_server.cpp

+//===-- Shared memory RPC server instantiation ------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "rpc_server.h"
+
+#include "src/__support/RPC/rpc.h"
+
+#include <cstdio>
+#include <cstdlib>
+
+/// The server instance used to communicate with the libc client.
+__llvm_libc::rpc::Server server;
+
+void rpc_init(rpc_alloc_ty alloc, void *data) {
+  void *inbox = alloc(sizeof(__llvm_libc::cpp::Atomic<int>), data);
+  void *outbox = alloc(sizeof(__llvm_libc::cpp::Atomic<int>), data);
+  void *buffer = alloc(sizeof(__llvm_libc::rpc::Buffer), data);
+  server.reset(inbox, outbox, buffer);
+}
+
+void rpc_deinit(rpc_dealloc_ty dealloc, void *data) {
+  dealloc(server.inbox, data);
+  dealloc(server.outbox, data);
+  dealloc(server.buffer, data);
+}
+
+void rpc_handle() {
+  while (server.handle(
+      [&](__llvm_libc::rpc::Buffer *buffer) {
+        switch (static_cast<__llvm_libc::rpc::Opcode>(buffer->data[0])) {
+        case __llvm_libc::rpc::Opcode::PRINT_TO_STDERR: {
+          fputs(reinterpret_cast<const char *>(&buffer->data[1]), stderr);
+          break;
+        }
+        case __llvm_libc::rpc::Opcode::EXIT: {
+          exit(buffer->data[1]);
+          break;
+        }
+        default:
+          return;
+        };
+      },
+      [](__llvm_libc::rpc::Buffer *buffer) {}))
+    ;
+}
+
+void *rpc_get_inbox() { return server.inbox; }
+
+void *rpc_get_outbox() { return server.outbox; }
+
+void *rpc_get_buffer() { return server.buffer; }