[libc] Implement the RPC client / server for NVPTX

[jhuber6]

Summary

This patch adds the necessary code to impelement the existing RPC client
/ server interface when targeting NVPTX GPUs. This follows closely to
the implementation in the AMDGPU version. This does not yet enable unit
testing as the nvlink linker does not support static libraries. So
that will need to be worked around.

I am ignoring the RPC duplication between the AMDGPU and NVPTX loaders. This
will be changed completely later so there's no point unifying the code at this
stage. The implementation was tested manually with the following file and
compilation flags.

namespace __llvm_libc {
void write_to_stderr(const char *msg);
void quick_exit(int);
} // namespace __llvm_libc

using namespace __llvm_libc;

int main(int argc, char **argv, char **envp) {
  for (int i = 0; i < argc; ++i) {
    write_to_stderr(argv[i]);
    write_to_stderr("\n");
  }
  quick_exit(255);
}

$ clang++ crt1.o rpc_client.o quick_exit.o io.o main.cpp --target=nvptx64-nvidia-cuda -march=sm_70 -o image
$ ./nvptx_loader image 1 2 3
image
1
2
3
$ echo $?
255

Depends on D146681

[jdoerfert]

Matches the AMD side.

libc/startup/gpu/nvptx/CMakeLists.txt

 add_startup_object(
   crt1
   SRC
     start.cpp
+  DEPENDS
+    libc.src.__support.RPC.rpc_client
   COMPILE_OPTIONS
     -ffreestanding # To avoid compiler warnings about calling the main function.
     -fno-builtin
     -nogpulib # Do not include any GPU vendor libraries.
     -march=${LIBC_GPU_TARGET_ARCHITECTURE}
     --target=${LIBC_GPU_TARGET_TRIPLE}
   NO_GPU_BUNDLE # Compile this file directly without special GPU handling.
 )
 get_fq_target_name(crt1 fq_name)
 
 # Ensure that clang uses the correct linker for this object type.
 target_link_libraries(${fq_name}
   PUBLIC
   "-march=${LIBC_GPU_TARGET_ARCHITECTURE}"
   "--target=${LIBC_GPU_TARGET_TRIPLE}"
 )

libc/startup/gpu/nvptx/start.cpp

-//===-- Implementation of crt for amdgpu ----------------------------------===//
+//===-- Implementation of crt for nvptx -----------------------------------===//
 //
 // 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 "src/__support/RPC/rpc_client.h"
+
 extern "C" int main(int argc, char **argv, char **envp);
 
 extern "C" [[gnu::visibility("protected")]] __attribute__((nvptx_kernel)) void
 _start(int argc, char **argv, char **envp, int *ret, void *in, void *out,
        void *buffer) {
+  __llvm_libc::rpc::client.reset(in, out, buffer);
+
   __atomic_fetch_or(ret, main(argc, argv, envp), __ATOMIC_RELAXED);
 }

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 "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");
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.");
+
   // Set up the arguments to the '_start' kernel on the GPU.
-  // TODO: Setup RPC server implementation;
   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;
   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);
 
-  // TODO: Query the RPC server periodically while the kernel is running.
+  // 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();
 
   // 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;
 }