[libc] Add initial support for an RPC mechanism for the GPU

[jhuber6]

Summary

This patch adds initial support for an RPC client / server architecture.
The GPU is unable to perform several system utilities on its own, so in
order to implement features like printing or memory allocation we need
to be able to communicate with the executing process. This is done via a
buffer of "sharable" memory. That is, a buffer with a unified pointer
that both the client and server can use to communicate.

The implementation here is based off of Jon Chesterfields minimal RPC
example in his work. We use an inbox and outbox to communicate
between if there is an RPC request and to signify when work is done.
We use a fixed-size buffer for the communication channel. This is fixed
size so that we can ensure that there is enough space for all
compute-units on the GPU to issue work to any of the ports. Right now
the implementation is single threaded so there is only a single buffer
that is not shared.

This implementation still has several features missing to be complete.
Such as multi-threaded support and asynchrnonous calls.

Depends on D145912

[jhuber6]

Few fixes and more comments.

[jdoerfert]

libc/src/__support/OSUtil/gpu/io.cpp

26

Nit: braces

libc/src/__support/RPC/rpc.h

7

file comment how this is going to work.

Also add a web documentation document

20

reserve it, don't call it noon now.

50

Comment for all classes and enums and functions.

75

Commented with full sentences please.

83

maybe handle, or query?

[jhuber6]

libc/src/__support/RPC/rpc.h

7

File comment is good. I'll documentation online in a later patch.

20

How about null? I think it's important that 0 doesn't do anything as we expect that to be the default clean state of this buffer.

83

Handle would probably be good, since I decided to make this non-blocking to make it more versatile on the server-side.

[jdoerfert]

libc/src/__support/OSUtil/gpu/io.cpp

28

Why is this function unused? (also braces for the loop).
Let's avoid "print" for now and go with fputs and exit first.
We certainly want fputs to be the fast path for fprintf, and the non-f versions can be resolved on the device.

libc/src/__support/RPC/rpc.h

20

OK, null.

31

Since data[0] is supposed to be the Opcode, should we not make it a opcode member and then 7 * 8 bytes payload?

[jhuber6]

libc/src/__support/OSUtil/gpu/io.cpp

28

It's the function to use the return value from the server. Since this is void there's nothing to use. In the future this should be "asynchonrous" where we don't wait for the return value.

libc/src/__support/RPC/rpc.h

31

I'm wondering if we should provide a bunch of structs that outline the structure of the arguments and just reinterpret case the underlying buffer to that struct type.

[jhuber6]

Addressing more comments.

[JonChesterfield]

The GPU plumbing looks ok to me. Implementing the minimal can-write-stderr version first is good too, details can be revised in tree as needed.

[michaelrj]

libc/src/__support/RPC/rpc.h

59

nit: communication

[jhuber6]

There is no usage of these functions currently. Adding test support is WIP. Right now the following works if I use the tools manually

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) {
  for (int i = 0; i < argc; ++i) {
    write_to_stderr(argv[i]);
    write_to_stderr("\n");
  }
  quick_exit(127);
}

$ clang++ crt1.o rpc_client.o io.o quick_exit.o main.cpp -flto --target=amdgcn-amd-amdhsa -mcpu=gfx1030 -o image
$ ./amdhsa_loader image args to the main function
image
args
to
the
main
function
$ echo $?
127

[sivachandra]

I have no major comments so stepping aside.

libc/src/__support/RPC/rpc.h

2

Should this RPC directory be nested under OSUtil/gpu as it is a GPU specific? AFAICT there is nothing GPU specific though.

40

Constants are in UPPER_CASE style.

libc/src/__support/RPC/rpc_client.cpp

23

false ? Also, is this being used anywhere?

[jhuber6]

libc/src/__support/RPC/rpc.h

2

Right now it's "gpu specific" but theoretically it could be used by any heterogeneous system that can share a memory region atomically. Maybe we could port it to run on FPGAs some day. But no other targets will use this right now, so it would be fine to put it under a GPU only directory if you want.

libc/src/__support/RPC/rpc_client.cpp

23

This isn't used right now. The idea is that when we provide this as a static library, any entrypoint that requires the RPC will pull in this file to resolve the symbol. That will in turn pull in this externally visible symbol which we can then read from the image directly. Therefore, if the GPU image does not contain __llvm_libc_rpc we don't need to bother with the runtime cost of spinning up a server on the host CPU.

It's set to false because constants don't get emitted as symbols if they are undefined, and we want to use the default initializer. It might be fine to make it true, but the presence of the symbol is the boolean here.

[sivachandra]

libc/src/__support/RPC/rpc_client.cpp

23

It's set to false because constants don't get emitted as symbols if they are undefined, and we want to use the default initializer. It might be fine to make it true, but the presence of the symbol is the boolean here.

My comment was about changing it to false instead of 0.

[jhuber6]

Addressing comments

[gchatelet]

I stumbled upon this code while looking at write_to_stderr and I think that the PRINT_TO_STDERR opcode is buggy (see comments). Also since we're exchanging messages of 64B, why waste 8B on the Opcode? Do we need such a huge opcode space ?
I'm working on a patch to try to fix these issues, I'll ping back when it's ready for discussion.

libc/src/__support/OSUtil/gpu/io.cpp

25

The number of bytes copied at each round are either buffer_len (i.e. 56B) or the total length of the string, it should be the remainder.

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

50

If the message's length is greater than 56B, &buffer->data[1] contains only a fragment of the message and is not null-terminated, leading fputs to read past it's allocated buffer.

314

Are there any guarantees that this piece of memory has the same alignment that __llvm_libc::cpp::Atomic<int>.
Same for the Cuda loader.

[jhuber6]

In D145913#4236641, @gchatelet wrote:

I stumbled upon this code while looking at write_to_stderr and I think that the PRINT_TO_STDERR opcode is buggy (see comments). Also since we're exchanging messages of 64B, why waste 8B on the Opcode? Do we need such a huge opcode space ?
I'm working on a patch to try to fix these issues, I'll ping back when it's ready for discussion.

This code is heavily WIP. Almost all of this is going to change at some point in the future. This exists mainly to let us stand up unit tests. The reason for the opcode's size right now was mainly just alignment. We want the actual arguments to at least have 8B alignment. I'm going to change this in the future. Something like this, but I'm still thinking about it.

struct Port {
  cpp::Atomic<uint32_t> flags;
  uint32_t opcode[WARP_SIZE];
  uint64_t activemask;
  Buffer data[WARP_SIZE];
}

libc/src/__support/OSUtil/gpu/io.cpp

25

Yeah, don't know why this one works. Probably more accidental behavior given that the memory allocation is a lot larger than the size of the buffer I'm using.

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

50

You're right. This worked in practice so I never noticed. Most likely because allocating fine-grained memory most likely gives you a full page at a minimum. I never write to the rest of that data and it's implicitly initialized to zero.

314

Fine-grained memory in this context is always going to be aligned to a page as far as I know. So that'll usually be aligned on a 4096 byte boundary.

[gchatelet]

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

314

Does that mean that you're reserving one page for just a few bytes? Maybe it would make more sense to reserve a larger chunk of memory and place the objects ourselves (not sure what this implies for the GPU).

Regarding the rpc mechanism, should the two atomics share the same cache line or would it make sense to have them in separate cache lines to prevent false sharing? It's probably OK to have them on the same cache line if there is only one client and one server.

Performance wise, it seems important that the atomic don't cross cache line boundaries though. And since placement is important (alignof(cpp::atomic<T>) should be honored for proper codegen) maybe the server should just allocate a sufficiently large chunk of memory and let a common function do the actual placement. This would prevent duplicate logic in each server (loader).

[JonChesterfield]

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

314

The two atomics are single-writer. Not only should they be on different cache lines, they should probably be on different pages (so that a given page is only every written by one of the agents involved).

Atomic variables definitely shouldn't cross cache lines. I don't know whether they'd work on GPUs if they did, that seems likely to be a correctness issue. But as they're naturally aligned and smaller than cache lines we're fine.

[jhuber6]

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

314

Does that mean that you're reserving one page for just a few bytes? Maybe it would make more sense to reserve a larger chunk of memory and place the objects ourselves (not sure what this implies for the GPU).

Yes, this is really wasteful but I'm ignoring it for now since the resource usage of these tests is quite low. This will change when we need to support multiple client calls from GPU threads. To meet hardware requirements the size of the buffer will probably be about 2 MB.

Regarding the rpc mechanism, should the two atomics share the same cache line or would it make sense to have them in separate cache lines to prevent false sharing? It's probably OK to have them on the same cache line if there is only one client and one server.

I'm probably going to merge these into a single atomic that we use as a bit-field. This is mostly because the underlying RPC state machine will need to support more than four states in the future (client send, server reply, client use, server clean). I'm actually not entirely sure what false sharing would mean in this context. This memory is basically a whole page that's shared with the GPU via the PCI(e) bus. I'm not privy enough to the internals to know if this can cause issues.

Performance wise, it seems important that the atomic don't cross cache line boundaries though. And since placement is important (alignof(cpp::atomic<T>) should be honored for proper codegen) maybe the server should just allocate a sufficiently large chunk of memory and let a common function do the actual placement. This would prevent duplicate logic in each server (loader).

We should definitely put alignas on the struct. I was thinking about doing the above. Having a single buffer makes it much easier to write to the GPU as well. I'm planning on just making this an array of some struct type.

libc/src/__support/CMakeLists.txt

add_header_library(
   DEPENDS
     .integer_to_string
     libc.src.__support.OSUtil.osutil
 )
 
 add_subdirectory(FPUtil)
 add_subdirectory(OSUtil)
 add_subdirectory(StringUtil)
+add_subdirectory(RPC)
 
 # Thread support is used by other "File". So, we add the "threads"
 # before "File".
 add_subdirectory(threads)
 
 add_subdirectory(File)

libc/src/__support/OSUtil/CMakeLists.txt

 if(NOT EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${LIBC_TARGET_OS})
   return()
 endif()
 
 add_subdirectory(${LIBC_TARGET_OS})
 set(target_os_util libc.src.__support.OSUtil.${LIBC_TARGET_OS}.${LIBC_TARGET_OS}_util)
 if(NOT TARGET ${target_os_util})
   return()
 endif()
 
+# The OSUtil is an object library in GPU mode.
+if(NOT LIBC_TARGET_ARCHITECTURE_IS_GPU)
   add_header_library(
     osutil
     HDRS
       io.h
       quick_exit.h
       syscall.h
     DEPENDS
       ${target_os_util}
   )
+else()
+  add_object_library(
+    osutil
+    ALIAS
+      ${target_os_util}
+    DEPENDS
+      ${target_os_util}
+  )
+endif()

libc/src/__support/OSUtil/gpu/CMakeLists.txt

 add_object_library(
   gpu_util
   SRCS
     quick_exit.cpp
+    io.cpp
   HDRS
     quick_exit.h
+    io.h
   DEPENDS
     libc.src.__support.common
+    libc.src.__support.RPC.rpc_client
 )

libc/src/__support/OSUtil/gpu/io.h

-//===---------------- Implementation of IO utils ----------------*- C++ -*-===//
+//===-------------- GPU implementation of IO utils --------------*- 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_SRC_SUPPORT_OSUTIL_IO_H
-#define LLVM_LIBC_SRC_SUPPORT_OSUTIL_IO_H
+#ifndef LLVM_LIBC_SRC_SUPPORT_OSUTIL_GPU_IO_H
+#define LLVM_LIBC_SRC_SUPPORT_OSUTIL_GPU_IO_H
 
-#ifdef __unix__
-#include "linux/io.h"
-#endif
+namespace __llvm_libc {
 
-#endif // LLVM_LIBC_SRC_SUPPORT_OSUTIL_IO_H
+void write_to_stderr(const char *msg);
+
+} // namespace __llvm_libc
+
+#endif // LLVM_LIBC_SRC_SUPPORT_OSUTIL_LINUX_IO_H

libc/src/__support/OSUtil/gpu/io.cpp

+//===-------------- GPU implementation of IO utils --------------*- 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 "io.h"
+
+#include "src/__support/RPC/rpc_client.h"
+#include "src/string/string_utils.h"
+
+namespace __llvm_libc {
+
+void write_to_stderr(const char *msg) {
+  uint64_t length = internal::string_length(msg) + 1;
+  uint64_t buffer_len = sizeof(rpc::Buffer) - sizeof(uint64_t);
+  for (uint64_t i = 0; i < length; i += buffer_len)
+    rpc::client.run(
+        [&](rpc::Buffer *buffer) {
+          buffer->data[0] = rpc::Opcode::print;
+          inline_memcpy(reinterpret_cast<char *>(&buffer->data[1]), &msg[i],
+                        (length > buffer_len ? buffer_len : length));
+        },
+        [](rpc::Buffer *) {});
+}
+
+} // namespace __llvm_libc

libc/src/__support/OSUtil/gpu/quick_exit.cpp

 //===---------- GPU implementation of a quick exit function -----*- 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_SRC_SUPPORT_OSUTIL_GPU_QUICK_EXIT_H
 #define LLVM_LIBC_SRC_SUPPORT_OSUTIL_GPU_QUICK_EXIT_H
 
 #include "quick_exit.h"
 
+#include "src/__support/RPC/rpc_client.h"
 #include "src/__support/macros/properties/architectures.h"
 
 namespace __llvm_libc {
 
 void quick_exit(int status) {
+  // TODO: Support asynchronous calls so we don't wait and exit from the GPU
+  // immediately.
+  rpc::client.run(
+      [&](rpc::Buffer *buffer) {
+        buffer->data[0] = rpc::Opcode::exit;
+        buffer->data[1] = status;
+      },
+      [](rpc::Buffer *) {});
+
 #if defined(LIBC_TARGET_ARCH_IS_NVPTX)
   asm("exit" ::: "memory");
 #elif defined(LIBC_TARGET_ARCH_IS_AMDGPU)
   // This will terminate the entire wavefront, may not be valid with divergent
   // work items.
   asm("s_endpgm" ::: "memory");
 #endif
   __builtin_unreachable();
 }
 
 } // namespace __llvm_libc
 
 #endif // LLVM_LIBC_SRC_SUPPORT_OSUTIL_GPU_QUICK_EXIT_H

libc/src/__support/OSUtil/io.h

 //===---------------- Implementation of IO utils ----------------*- 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_SRC_SUPPORT_OSUTIL_IO_H
 #define LLVM_LIBC_SRC_SUPPORT_OSUTIL_IO_H
 
-#ifdef __unix__
+#include "src/__support/macros/properties/architectures.h"
+
+#if defined(LIBC_TARGET_ARCH_IS_GPU)
+#include "gpu/io.h"
+#elif defined(__unix__)
 #include "linux/io.h"
 #endif
 
 #endif // LLVM_LIBC_SRC_SUPPORT_OSUTIL_IO_H

libc/src/__support/RPC/CMakeLists.txt

+add_header_library(
+  rpc
+  HDRS
+    rpc.h
+  DEPENDS
+    libc.src.__support.common
+    libc.src.__support.CPP.atomic
+)
+
+add_object_library(
+  rpc_client
+  SRCS
+    rpc_client.cpp
+  HDRS
+    rpc_client.h
+  DEPENDS
+    .rpc
+)

libc/src/__support/RPC/rpc.h

+//===-- Shared memory RPC client / server interface -------------*- 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_SRC_SUPPORT_RPC_RPC_H
+#define LLVM_LIBC_SRC_SUPPORT_RPC_RPC_H
+
+#include "src/__support/CPP/atomic.h"
+
+#include <stdint.h>
+
+namespace __llvm_libc {
+namespace rpc {
+
+enum Opcode : uint64_t {
+  noop = 0,
+  print = 1,
+  exit = 2,
+};
+
+struct Buffer {
+  uint64_t data[8];
+};
+
+struct Process {
+  cpp::Atomic<uint32_t> *inbox;
+  cpp::Atomic<uint32_t> *outbox;
+  Buffer *buffer;
+
+  void reset(void *inbox, void *outbox, void *buffer) {
+    *this = {
+        reinterpret_cast<cpp::Atomic<uint32_t> *>(inbox),
+        reinterpret_cast<cpp::Atomic<uint32_t> *>(outbox),
+        reinterpret_cast<Buffer *>(buffer),
+    };
+  }
+};
+
+struct Client : public Process {
+  template <typename F, typename U> void run(F fill, U use);
+};
+
+struct Server : public Process {
+  template <typename W, typename C> bool run(W work, C clean);
+};
+
+template <typename F, typename U> void Client::run(F fill, U use) {
+  bool in = inbox->load(cpp::MemoryOrder::RELAXED);
+  bool out = outbox->load(cpp::MemoryOrder::RELAXED);
+  atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
+  // ready write to buffer then to outbox
+  if (!in & !out) {
+    fill(buffer);
+    atomic_thread_fence(cpp::MemoryOrder::RELEASE);
+    outbox->store(1, cpp::MemoryOrder::RELEASE);
+    out = 1;
+  }
+  // wait for result
+  if (!in & out) {
+    while (!in)
+      in = inbox->load(cpp::MemoryOrder::RELAXED);
+    atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
+  }
+  // read from buffer then write to outbox
+  if (in & out) {
+    use(buffer);
+    atomic_thread_fence(cpp::MemoryOrder::RELEASE);
+    outbox->store(0, cpp::MemoryOrder::RELEASE);
+    out = 0;
+  }
+  // wait for server to garbage collect
+  if (in & !out) {
+    while (in)
+      in = inbox->load(cpp::MemoryOrder::RELAXED);
+    atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
+  }
+}
+
+template <typename W, typename C> bool Server::run(W work, C clean) {
+  bool in = inbox->load(cpp::MemoryOrder::RELAXED);
+  bool out = outbox->load(cpp::MemoryOrder::RELAXED);
+  atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
+  // Nothing to do, return early
+  if (!in & !out)
+    return false;
+
+  // Do work then signal client.
+  if (in & !out) {
+    work(buffer);
+    atomic_thread_fence(cpp::MemoryOrder::RELEASE);
+    outbox->store(1, cpp::MemoryOrder::RELEASE);
+    out = 1;
+  }
+  // Work done, wait for client.
+  if (in & out) {
+    while (in)
+      in = inbox->load(cpp::MemoryOrder::RELAXED);
+    atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
+  }
+  // All done, clean up buffer.
+  if (!in & out) {
+    clean(buffer);
+    atomic_thread_fence(cpp::MemoryOrder::RELEASE);
+    outbox->store(0, cpp::MemoryOrder::RELEASE);
+    out = 0;
+  }
+
+  return true;
+}
+
+} // namespace rpc
+} // namespace __llvm_libc
+
+#endif

libc/src/__support/RPC/rpc_client.h

+//===-- Shared memory RPC client 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_SRC_SUPPORT_RPC_RPC_CLIENT_H
+#define LLVM_LIBC_SRC_SUPPORT_RPC_RPC_CLIENT_H
+
+#include "rpc.h"
+
+namespace __llvm_libc {
+namespace rpc {
+
+/// The libc client instance used to communicate with the server.
+extern Client client;
+
+} // namespace rpc
+} // namespace __llvm_libc
+
+#endif

libc/src/__support/RPC/rpc_client.cpp

+//===-- Shared memory RPC client 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_SRC_SUPPORT_RPC_RPC_CLIENT_H
+#define LLVM_LIBC_SRC_SUPPORT_RPC_RPC_CLIENT_H
+
+#include "rpc.h"
+
+namespace __llvm_libc {
+namespace rpc {
+
+/// The libc client instance used to communicate with the server.
+Client client;
+
+/// Externally visible symbol to signify the usage of an RPC client to
+/// whomever needs to run the server.
+extern "C" [[gnu::visibility("protected")]] bool __llvm_libc_rpc = true;
+
+} // namespace rpc
+} // namespace __llvm_libc
+
+#endif

libc/startup/gpu/amdgpu/CMakeLists.txt

 add_startup_object(
   crt1
   SRC
     start.cpp
+  DEPENDS
+    libc.src.__support.OSUtil.osutil
   COMPILE_OPTIONS
     -ffreestanding # To avoid compiler warnings about calling the main function.
     -fno-builtin
     -nogpulib # Do not include any GPU vendor libraries.
-    -nostdinc
     -mcpu=${LIBC_GPU_TARGET_ARCHITECTURE}
     -emit-llvm # AMDGPU's intermediate object file format is bitcode.
     --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 
                       "--target=${LIBC_GPU_TARGET_TRIPLE}" "-flto")

libc/startup/gpu/amdgpu/start.cpp

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

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
   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 <hsa/hsa.h>
 #include <hsa/hsa_ext_amd.h>
 
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
 #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 *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 run_server() {
+  while (server.run(
+      [&](__llvm_libc::rpc::Buffer *buffer) {
+        switch (static_cast<__llvm_libc::rpc::Opcode>(buffer->data[0])) {
+        case __llvm_libc::rpc::Opcode::print: {
+          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)
int load(int argc, char **argv, void *image, size_t size) {
   // Allocate space for the return pointer and initialize it to zero.
   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);
+
   // 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->ret = dev_ret;
+  kernel_args->inbox = server_outbox;
+  kernel_args->outbox = server_inbox;
+  kernel_args->buffer = 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, 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.
-  while (hsa_signal_wait_scacquire(packet->completion_signal,
-                                   HSA_SIGNAL_CONDITION_EQ, 0, UINT64_MAX,
-                                   HSA_WAIT_STATE_ACTIVE) != 0)
-    ;
+  // 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)
+    run_server();
 
   // 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;