[libc-commits] [libc] [libc] Export the RPC interface from `libc` (PR #71432)

Joseph Huber via libc-commits libc-commits at lists.llvm.org
Fri Mar 15 08:47:03 PDT 2024


https://github.com/jhuber6 updated https://github.com/llvm/llvm-project/pull/71432

>From 369b63b65a73e2f7575c32ff6f565ee12fdb7891 Mon Sep 17 00:00:00 2001
From: Joseph Huber <jhuber6 at vols.utk.edu>
Date: Mon, 6 Nov 2023 13:08:42 -0600
Subject: [PATCH] [libc] Export the RPC interface from `libc`

Summary:
This patch adds new extensions that allow us to export the RPC interface
from `libc` to other programs. This should allow external users of the
GPU `libc` to interface with the RPC client (which more or less behaves
like syscalls in this context). This is done by wrapping the interface
into a C-style function call.

Obviously, this approach is far less safe than the carefully crafted C++
interface. For example, we now expose the internal packet buffer, and it
is theoretically possible to open a single port with conflicting opcodes
and break the whole interface. So, extra care will be needed when
interacting with this. However, the usage is very similar as shown by
the new test.

This somewhat stretches the concept of `libc` just doing `libc` things,
but I think this is important enough to justify it. It is difficult to
split this out, as we do not want to have the situation where we have
multiple RPC clients running at one time, so for now it makes sense to
just leave it in `libc`.
---
 libc/config/gpu/entrypoints.txt               |  4 ++
 libc/include/CMakeLists.txt                   |  1 +
 libc/include/gpu/rpc.h.def                    |  3 +
 libc/include/llvm-libc-types/CMakeLists.txt   |  1 +
 libc/include/llvm-libc-types/rpc_port_t.h     | 16 ++++++
 .../llvm-libc-types/test_rpc_opcodes_t.h      |  1 +
 libc/spec/gpu_ext.td                          | 25 ++++++++-
 libc/src/__support/RPC/rpc.h                  | 19 +++++--
 libc/src/gpu/CMakeLists.txt                   | 55 +++++++++++++++++++
 libc/src/gpu/rpc_close_port.cpp               | 24 ++++++++
 libc/src/gpu/rpc_close_port.h                 | 20 +++++++
 libc/src/gpu/rpc_open_port.cpp                | 26 +++++++++
 libc/src/gpu/rpc_open_port.h                  | 20 +++++++
 libc/src/gpu/rpc_recv_n.cpp                   | 26 +++++++++
 libc/src/gpu/rpc_recv_n.h                     | 20 +++++++
 libc/src/gpu/rpc_send_n.cpp                   | 25 +++++++++
 libc/src/gpu/rpc_send_n.h                     | 20 +++++++
 .../integration/startup/gpu/CMakeLists.txt    | 15 +++++
 .../gpu/rpc_external_interface_test.cpp       | 45 +++++++++++++++
 libc/utils/gpu/loader/Loader.h                | 24 +++++++-
 20 files changed, 384 insertions(+), 6 deletions(-)
 create mode 100644 libc/include/llvm-libc-types/rpc_port_t.h
 create mode 100644 libc/src/gpu/rpc_close_port.cpp
 create mode 100644 libc/src/gpu/rpc_close_port.h
 create mode 100644 libc/src/gpu/rpc_open_port.cpp
 create mode 100644 libc/src/gpu/rpc_open_port.h
 create mode 100644 libc/src/gpu/rpc_recv_n.cpp
 create mode 100644 libc/src/gpu/rpc_recv_n.h
 create mode 100644 libc/src/gpu/rpc_send_n.cpp
 create mode 100644 libc/src/gpu/rpc_send_n.h
 create mode 100644 libc/test/integration/startup/gpu/rpc_external_interface_test.cpp

diff --git a/libc/config/gpu/entrypoints.txt b/libc/config/gpu/entrypoints.txt
index 4fb87cb9f5a33e..a185a5c90757a3 100644
--- a/libc/config/gpu/entrypoints.txt
+++ b/libc/config/gpu/entrypoints.txt
@@ -211,6 +211,10 @@ set(TARGET_LIBC_ENTRYPOINTS
 
     # gpu/rpc.h entrypoints
     libc.src.gpu.rpc_host_call
+    libc.src.gpu.rpc_open_port
+    libc.src.gpu.rpc_send_n
+    libc.src.gpu.rpc_recv_n
+    libc.src.gpu.rpc_close_port
 )
 
 set(TARGET_LIBM_ENTRYPOINTS
diff --git a/libc/include/CMakeLists.txt b/libc/include/CMakeLists.txt
index b2cb10459c53e3..8ed335bb9cf188 100644
--- a/libc/include/CMakeLists.txt
+++ b/libc/include/CMakeLists.txt
@@ -614,6 +614,7 @@ if(LIBC_TARGET_OS_IS_GPU)
     DEPENDS
       .llvm_libc_common_h
       .llvm-libc-types.rpc_opcodes_t
+      .llvm-libc-types.rpc_port_t
   )
 endif()
 
diff --git a/libc/include/gpu/rpc.h.def b/libc/include/gpu/rpc.h.def
index 0438cd65e7be22..f178007ee7083f 100644
--- a/libc/include/gpu/rpc.h.def
+++ b/libc/include/gpu/rpc.h.def
@@ -11,7 +11,10 @@
 
 #include <__llvm-libc-common.h>
 
+#include <llvm-libc-types/size_t.h>
+
 #include <llvm-libc-types/rpc_opcodes_t.h>
+#include <llvm-libc-types/rpc_port_t.h>
 
 %%public_api()
 
diff --git a/libc/include/llvm-libc-types/CMakeLists.txt b/libc/include/llvm-libc-types/CMakeLists.txt
index 7fef976d7b3250..ff1354f4fe19be 100644
--- a/libc/include/llvm-libc-types/CMakeLists.txt
+++ b/libc/include/llvm-libc-types/CMakeLists.txt
@@ -94,6 +94,7 @@ add_header(socklen_t HDR socklen_t.h)
 add_header(struct_sockaddr_un HDR struct_sockaddr_un.h)
 add_header(struct_sockaddr HDR struct_sockaddr.h)
 add_header(rpc_opcodes_t HDR rpc_opcodes_t.h)
+add_header(rpc_port_t HDR rpc_port_t.h)
 add_header(ACTION HDR ACTION.h)
 add_header(ENTRY HDR ENTRY.h)
 add_header(struct_hsearch_data HDR struct_hsearch_data.h)
diff --git a/libc/include/llvm-libc-types/rpc_port_t.h b/libc/include/llvm-libc-types/rpc_port_t.h
new file mode 100644
index 00000000000000..9d3c3eaa61e661
--- /dev/null
+++ b/libc/include/llvm-libc-types/rpc_port_t.h
@@ -0,0 +1,16 @@
+//===-- Definition of type rpc_port_t -------------------------------------===//
+//
+// 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_TYPES_RPC_PORT_T_H__
+#define __LLVM_LIBC_TYPES_RPC_PORT_T_H__
+
+typedef struct {
+  __UINT8_TYPE__ reserved[32];
+} rpc_port_t;
+
+#endif // __LLVM_LIBC_TYPES_RPC_PORT_T_H__
diff --git a/libc/include/llvm-libc-types/test_rpc_opcodes_t.h b/libc/include/llvm-libc-types/test_rpc_opcodes_t.h
index 7129768dc8b988..126e36f1df87bf 100644
--- a/libc/include/llvm-libc-types/test_rpc_opcodes_t.h
+++ b/libc/include/llvm-libc-types/test_rpc_opcodes_t.h
@@ -15,6 +15,7 @@ typedef enum : unsigned short {
   RPC_TEST_NOOP = 1 << 15,
   RPC_TEST_INCREMENT,
   RPC_TEST_INTERFACE,
+  RPC_TEST_EXTERNAL,
   RPC_TEST_STREAM,
 } rpc_test_opcode_t;
 
diff --git a/libc/spec/gpu_ext.td b/libc/spec/gpu_ext.td
index dce81ff7786203..450c5135c49128 100644
--- a/libc/spec/gpu_ext.td
+++ b/libc/spec/gpu_ext.td
@@ -1,8 +1,11 @@
+def RPCPortT : NamedType<"rpc_port_t">;
+def RPCPortPtrT : PtrType<RPCPortT>;
+
 def GPUExtensions : StandardSpec<"GPUExtensions"> {
   HeaderSpec RPC = HeaderSpec<
     "gpu/rpc.h",
     [], // Macros
-    [], // Types
+    [RPCPortT, SizeTType], // Types
     [], // Enumerations
     [
         FunctionSpec<
@@ -10,6 +13,26 @@ def GPUExtensions : StandardSpec<"GPUExtensions"> {
             RetValSpec<VoidType>,
             [ArgSpec<VoidPtr>, ArgSpec<VoidPtr>, ArgSpec<SizeTType>]
         >,
+        FunctionSpec<
+            "rpc_open_port",
+            RetValSpec<RPCPortT>,
+            [ArgSpec<UnsignedIntType>]
+        >,
+        FunctionSpec<
+            "rpc_send_n",
+            RetValSpec<VoidType>,
+            [ArgSpec<RPCPortPtrT>, ArgSpec<VoidPtr>, ArgSpec<SizeTType>]
+        >,
+        FunctionSpec<
+            "rpc_recv_n",
+            RetValSpec<VoidType>,
+            [ArgSpec<RPCPortPtrT>, ArgSpec<VoidPtr>, ArgSpec<SizeTPtr>]
+        >,
+        FunctionSpec<
+            "rpc_close_port",
+            RetValSpec<VoidType>,
+            [ArgSpec<RPCPortPtrT>]
+        >,
     ]
   >;
   let Headers = [
diff --git a/libc/src/__support/RPC/rpc.h b/libc/src/__support/RPC/rpc.h
index 5dcae518bb6f8f..f6a0ee71129add 100644
--- a/libc/src/__support/RPC/rpc.h
+++ b/libc/src/__support/RPC/rpc.h
@@ -323,6 +323,7 @@ template <bool T> struct Port {
   LIBC_INLINE void send_n(const void *src, uint64_t size);
   template <typename A>
   LIBC_INLINE void recv_n(void **dst, uint64_t *size, A &&alloc);
+  LIBC_INLINE void recv_n(void *dst, uint64_t *size);
 
   LIBC_INLINE uint16_t get_opcode() const {
     return process.header[index].opcode;
@@ -359,7 +360,9 @@ struct Client {
       : process(port_count, buffer) {}
 
   using Port = rpc::Port<false>;
-  template <uint16_t opcode> LIBC_INLINE Port open();
+
+  template <uint16_t opcode> LIBC_INLINE Port open() { return open(opcode); }
+  LIBC_INLINE Port open(uint16_t opcode);
 
 private:
   Process<false> process;
@@ -484,6 +487,14 @@ LIBC_INLINE void Port<T>::send_n(const void *const *src, uint64_t *size) {
   }
 }
 
+/// Helper routine to simplify the interface when recieving from the GPU using
+/// thread private pointers to the underly value and the destination pointer
+/// contains enough data to recieve the values.
+template <bool T> LIBC_INLINE void Port<T>::recv_n(void *dst, uint64_t *size) {
+  void **dst_ptr = &dst;
+  recv_n(dst_ptr, size, [=](uint64_t) { return dst; });
+}
+
 /// Receives an arbitrarily sized data buffer across the shared channel in
 /// multiples of the packet length. The \p alloc function is called with the
 /// size of the data so that we can initialize the size of the \p dst buffer.
@@ -522,9 +533,9 @@ LIBC_INLINE void Port<T>::recv_n(void **dst, uint64_t *size, A &&alloc) {
 /// is, there are send operations pending that haven't been serviced on this
 /// port. Each port instance uses an associated \p opcode to tell the server
 /// what to do. The Client interface provides the appropriate lane size to the
-/// port using the platform's returned value.
-template <uint16_t opcode>
-[[clang::convergent]] LIBC_INLINE Client::Port Client::open() {
+/// port using the platform's returned value. It is required that \p opcode is
+/// uniform between all the lanes for this to work.
+[[clang::convergent]] LIBC_INLINE Client::Port Client::open(uint16_t opcode) {
   // Repeatedly perform a naive linear scan for a port that can be opened to
   // send data.
   for (uint32_t index = gpu::get_cluster_id();; ++index) {
diff --git a/libc/src/gpu/CMakeLists.txt b/libc/src/gpu/CMakeLists.txt
index e20228516b5112..36ea30f8e3d6e6 100644
--- a/libc/src/gpu/CMakeLists.txt
+++ b/libc/src/gpu/CMakeLists.txt
@@ -8,3 +8,58 @@ add_entrypoint_object(
     libc.src.__support.RPC.rpc_client
     libc.src.__support.GPU.utils
 )
+
+add_entrypoint_object(
+  rpc_open_port
+  SRCS
+    rpc_open_port.cpp
+  HDRS
+    rpc_open_port.h
+  DEPENDS
+    libc.src.__support.RPC.rpc_client
+    libc.src.__support.GPU.utils
+)
+
+add_entrypoint_object(
+  rpc_close_port
+  SRCS
+    rpc_close_port.cpp
+  HDRS
+    rpc_close_port.h
+  DEPENDS
+    libc.src.__support.RPC.rpc_client
+    libc.src.__support.GPU.utils
+)
+
+add_entrypoint_object(
+  rpc_get_buffer
+  SRCS
+    rpc_get_buffer.cpp
+  HDRS
+    rpc_get_buffer.h
+  DEPENDS
+    libc.src.__support.RPC.rpc_client
+    libc.src.__support.GPU.utils
+)
+
+add_entrypoint_object(
+  rpc_send_n
+  SRCS
+    rpc_send_n.cpp
+  HDRS
+    rpc_send_n.h
+  DEPENDS
+    libc.src.__support.RPC.rpc_client
+    libc.src.__support.GPU.utils
+)
+
+add_entrypoint_object(
+  rpc_recv_n
+  SRCS
+    rpc_recv_n.cpp
+  HDRS
+    rpc_recv_n.h
+  DEPENDS
+    libc.src.__support.RPC.rpc_client
+    libc.src.__support.GPU.utils
+)
diff --git a/libc/src/gpu/rpc_close_port.cpp b/libc/src/gpu/rpc_close_port.cpp
new file mode 100644
index 00000000000000..2a5265c547ca96
--- /dev/null
+++ b/libc/src/gpu/rpc_close_port.cpp
@@ -0,0 +1,24 @@
+//===---------- GPU implementation of the external RPC port interface -----===//
+//
+// 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/gpu/rpc_close_port.h"
+
+#include "src/__support/GPU/utils.h"
+#include "src/__support/RPC/rpc_client.h"
+#include "src/__support/common.h"
+
+namespace LIBC_NAMESPACE {
+
+static_assert(sizeof(rpc_port_t) == sizeof(rpc::Client::Port), "ABI mismatch");
+
+LLVM_LIBC_FUNCTION(void, rpc_close_port, (rpc_port_t * handle)) {
+  rpc::Client::Port *port = reinterpret_cast<rpc::Client::Port *>(handle);
+  port->close();
+}
+
+} // namespace LIBC_NAMESPACE
diff --git a/libc/src/gpu/rpc_close_port.h b/libc/src/gpu/rpc_close_port.h
new file mode 100644
index 00000000000000..2be39ce3698a38
--- /dev/null
+++ b/libc/src/gpu/rpc_close_port.h
@@ -0,0 +1,20 @@
+//===-- Implementation header for RPC functions -----------------*- 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_GPU_RPC_CLOSE_PORT_H
+#define LLVM_LIBC_SRC_GPU_RPC_CLOSE_PORT_H
+
+#include <gpu/rpc.h>
+
+namespace LIBC_NAMESPACE {
+
+void rpc_close_port(rpc_port_t *handle);
+
+} // namespace LIBC_NAMESPACE
+
+#endif // LLVM_LIBC_SRC_GPU_RPC_CLOSE_PORT_H
diff --git a/libc/src/gpu/rpc_open_port.cpp b/libc/src/gpu/rpc_open_port.cpp
new file mode 100644
index 00000000000000..8d38680b19e234
--- /dev/null
+++ b/libc/src/gpu/rpc_open_port.cpp
@@ -0,0 +1,26 @@
+//===---------- GPU implementation of the external RPC port interface -----===//
+//
+// 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/gpu/rpc_open_port.h"
+
+#include "src/__support/CPP/bit.h"
+#include "src/__support/GPU/utils.h"
+#include "src/__support/RPC/rpc_client.h"
+#include "src/__support/common.h"
+
+namespace LIBC_NAMESPACE {
+
+static_assert(sizeof(rpc_port_t) == sizeof(rpc::Client::Port), "ABI mismatch");
+
+LLVM_LIBC_FUNCTION(rpc_port_t, rpc_open_port, (unsigned opcode)) {
+  uint32_t uniform = gpu::broadcast_value(gpu::get_lane_mask(), opcode);
+  rpc::Client::Port port = rpc::client.open(static_cast<uint16_t>(uniform));
+  return cpp::bit_cast<rpc_port_t>(port);
+}
+
+} // namespace LIBC_NAMESPACE
diff --git a/libc/src/gpu/rpc_open_port.h b/libc/src/gpu/rpc_open_port.h
new file mode 100644
index 00000000000000..553fa59120aaf6
--- /dev/null
+++ b/libc/src/gpu/rpc_open_port.h
@@ -0,0 +1,20 @@
+//===-- Implementation header for RPC functions -----------------*- 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_GPU_RPC_OPEN_PORT_H
+#define LLVM_LIBC_SRC_GPU_RPC_OPEN_PORT_H
+
+#include <gpu/rpc.h>
+
+namespace LIBC_NAMESPACE {
+
+rpc_port_t rpc_open_port(unsigned opcode);
+
+} // namespace LIBC_NAMESPACE
+
+#endif // LLVM_LIBC_SRC_GPU_RPC_OPEN_PORT_H
diff --git a/libc/src/gpu/rpc_recv_n.cpp b/libc/src/gpu/rpc_recv_n.cpp
new file mode 100644
index 00000000000000..64a93b825ffe0b
--- /dev/null
+++ b/libc/src/gpu/rpc_recv_n.cpp
@@ -0,0 +1,26 @@
+//===---------- GPU implementation of the external RPC port interface -----===//
+//
+// 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/gpu/rpc_recv_n.h"
+
+#include "src/__support/GPU/utils.h"
+#include "src/__support/RPC/rpc_client.h"
+#include "src/__support/common.h"
+
+namespace LIBC_NAMESPACE {
+
+static_assert(sizeof(rpc_port_t) == sizeof(rpc::Client::Port), "ABI mismatch");
+static_assert(sizeof(size_t) == sizeof(uint64_t), "Size mismatch");
+
+LLVM_LIBC_FUNCTION(void, rpc_recv_n,
+                   (rpc_port_t * handle, void *dst, size_t *size)) {
+  rpc::Client::Port *port = reinterpret_cast<rpc::Client::Port *>(handle);
+  port->recv_n(dst, reinterpret_cast<uint64_t *>(size));
+}
+
+} // namespace LIBC_NAMESPACE
diff --git a/libc/src/gpu/rpc_recv_n.h b/libc/src/gpu/rpc_recv_n.h
new file mode 100644
index 00000000000000..287cc642245923
--- /dev/null
+++ b/libc/src/gpu/rpc_recv_n.h
@@ -0,0 +1,20 @@
+//===-- Implementation header for RPC functions -----------------*- 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_GPU_RPC_RECV_H
+#define LLVM_LIBC_SRC_GPU_RPC_RECV_H
+
+#include <gpu/rpc.h>
+
+namespace LIBC_NAMESPACE {
+
+void rpc_recv_n(rpc_port_t *handle, void *dst, size_t *size);
+
+} // namespace LIBC_NAMESPACE
+
+#endif // LLVM_LIBC_SRC_GPU_RPC_RECV_H
diff --git a/libc/src/gpu/rpc_send_n.cpp b/libc/src/gpu/rpc_send_n.cpp
new file mode 100644
index 00000000000000..a02749360ba4a1
--- /dev/null
+++ b/libc/src/gpu/rpc_send_n.cpp
@@ -0,0 +1,25 @@
+//===---------- GPU implementation of the external RPC port interface -----===//
+//
+// 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/gpu/rpc_send_n.h"
+
+#include "src/__support/GPU/utils.h"
+#include "src/__support/RPC/rpc_client.h"
+#include "src/__support/common.h"
+
+namespace LIBC_NAMESPACE {
+
+static_assert(sizeof(rpc_port_t) == sizeof(rpc::Client::Port), "ABI mismatch");
+
+LLVM_LIBC_FUNCTION(void, rpc_send_n,
+                   (rpc_port_t * handle, const void *src, size_t size)) {
+  rpc::Client::Port *port = reinterpret_cast<rpc::Client::Port *>(handle);
+  port->send_n(src, size);
+}
+
+} // namespace LIBC_NAMESPACE
diff --git a/libc/src/gpu/rpc_send_n.h b/libc/src/gpu/rpc_send_n.h
new file mode 100644
index 00000000000000..fcc683aab56444
--- /dev/null
+++ b/libc/src/gpu/rpc_send_n.h
@@ -0,0 +1,20 @@
+//===-- Implementation header for RPC functions -----------------*- 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_GPU_RPC_SEND_H
+#define LLVM_LIBC_SRC_GPU_RPC_SEND_H
+
+#include <gpu/rpc.h>
+
+namespace LIBC_NAMESPACE {
+
+void rpc_send_n(rpc_port_t *handle, const void *src, size_t size);
+
+} // namespace LIBC_NAMESPACE
+
+#endif // LLVM_LIBC_SRC_GPU_RPC_SEND_H
diff --git a/libc/test/integration/startup/gpu/CMakeLists.txt b/libc/test/integration/startup/gpu/CMakeLists.txt
index 7555986b16df45..d56c6b24ef7a83 100644
--- a/libc/test/integration/startup/gpu/CMakeLists.txt
+++ b/libc/test/integration/startup/gpu/CMakeLists.txt
@@ -44,6 +44,21 @@ add_integration_test(
    rpc_interface_test.cpp
 )
 
+add_integration_test(
+  startup_rpc_external_interface_test
+  SUITE libc-startup-tests
+  SRCS
+    rpc_external_interface_test.cpp
+  DEPENDS
+    libc.src.gpu.rpc_open_port
+    libc.src.gpu.rpc_send_n
+    libc.src.gpu.rpc_recv_n
+    libc.src.gpu.rpc_close_port
+  LOADER_ARGS
+    --threads 32
+    --blocks 8
+)
+
 add_integration_test(
   startup_rpc_stream_test
   SUITE libc-startup-tests
diff --git a/libc/test/integration/startup/gpu/rpc_external_interface_test.cpp b/libc/test/integration/startup/gpu/rpc_external_interface_test.cpp
new file mode 100644
index 00000000000000..77deb01133329d
--- /dev/null
+++ b/libc/test/integration/startup/gpu/rpc_external_interface_test.cpp
@@ -0,0 +1,45 @@
+//===-- Loader test to check the external RPC interface with the loader ---===//
+//
+// 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 <gpu/rpc.h>
+
+#include "src/gpu/rpc_close_port.h"
+#include "src/gpu/rpc_open_port.h"
+#include "src/gpu/rpc_recv_n.h"
+#include "src/gpu/rpc_send_n.h"
+
+#include "include/llvm-libc-types/test_rpc_opcodes_t.h"
+#include "src/__support/GPU/utils.h"
+#include "src/__support/RPC/rpc_client.h"
+#include "test/IntegrationTest/test.h"
+
+using namespace LIBC_NAMESPACE;
+
+// Test to ensure that we can use aribtrary combinations of sends and recieves
+// as long as they are mirrored using the external interface.
+static void test_interface() {
+  uint32_t num_additions =
+      10 + 10 * gpu::get_thread_id() + 10 * gpu::get_block_id();
+  uint64_t cnt = 0;
+  for (uint32_t i = 0; i < num_additions; ++i) {
+    size_t size = sizeof(uint64_t);
+    rpc_port_t port = LIBC_NAMESPACE::rpc_open_port(RPC_TEST_EXTERNAL);
+    LIBC_NAMESPACE::rpc_send_n(&port, &cnt, size);
+    LIBC_NAMESPACE::rpc_recv_n(&port, &cnt, &size);
+    LIBC_NAMESPACE::rpc_close_port(&port);
+    ASSERT_TRUE(size == sizeof(uint64_t));
+  }
+
+  ASSERT_TRUE(cnt == num_additions && "Invalid number of increments");
+}
+
+TEST_MAIN(int argc, char **argv, char **envp) {
+  test_interface();
+
+  return 0;
+}
diff --git a/libc/utils/gpu/loader/Loader.h b/libc/utils/gpu/loader/Loader.h
index 93380383701978..d55ea14799daab 100644
--- a/libc/utils/gpu/loader/Loader.h
+++ b/libc/utils/gpu/loader/Loader.h
@@ -11,8 +11,8 @@
 
 #include "utils/gpu/server/llvmlibc_rpc_server.h"
 
-#include "llvm-libc-types/rpc_opcodes_t.h"
 #include "include/llvm-libc-types/test_rpc_opcodes_t.h"
+#include "llvm-libc-types/rpc_opcodes_t.h"
 
 #include <cstddef>
 #include <cstdint>
@@ -222,6 +222,28 @@ inline void register_rpc_callbacks(uint32_t device_id) {
         }
       },
       nullptr);
+
+  // Register the stream test handler.
+  rpc_register_callback(
+      device_id, static_cast<rpc_opcode_t>(RPC_TEST_EXTERNAL),
+      [](rpc_port_t port, void *data) {
+        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);
+        for (uint64_t i = 0; i < lane_size; ++i) {
+          if (dst[i])
+            *reinterpret_cast<uint64_t *>(dst[i]) += 1;
+        }
+        rpc_send_n(port, dst, sizes);
+        for (uint64_t i = 0; i < lane_size; ++i) {
+          if (dst[i])
+            delete[] reinterpret_cast<uint8_t *>(dst[i]);
+        }
+      },
+      nullptr);
 }
 
 #endif



More information about the libc-commits mailing list