[libc] Export GPU extensions to libc for external use

The GPU port of the LLVM C library needs to export a few extensions to the interface such that users can interface with it. This patch adds the necessary logic to define a GPU extension. Currently, this only exports a `rpc_reset_client` function. This allows us to use the server in D147054 to set up the RPC interface outside of `libc`. Depends on https://reviews.llvm.org/D147054 Reviewed By: sivachandra Differential Revision: https://reviews.llvm.org/D152283
llvm · Jun 15, 2023 · dcdfc96 · dcdfc96
1 parent 719d77e
commit dcdfc96
Show file tree

Hide file tree

Showing 27 changed files with 173 additions and 51 deletions.
diff --git a/libc/config/gpu/api.td b/libc/config/gpu/api.td
@@ -1,6 +1,7 @@
 include "config/public_api.td"
 
 include "spec/stdc.td"
+include "spec/gpu_ext.td"
 
 def StringAPI : PublicAPI<"string.h"> {
   let Types = ["size_t"];

diff --git a/libc/config/gpu/entrypoints.txt b/libc/config/gpu/entrypoints.txt
@@ -90,6 +90,9 @@ set(TARGET_LIBC_ENTRYPOINTS
     libc.src.stdio.stdin
     libc.src.stdio.stdout
     libc.src.stdio.stderr
+
+    # gpu/rpc.h entrypoints
+    libc.src.gpu.rpc_reset
 )
 
 set(TARGET_LIBM_ENTRYPOINTS

diff --git a/libc/config/gpu/headers.txt b/libc/config/gpu/headers.txt
@@ -6,4 +6,7 @@ set(TARGET_PUBLIC_HEADERS
     libc.include.errno
     libc.include.stdlib
     libc.include.stdio
+
+    # Header for RPC extensions
+    libc.include.gpu_rpc
 )
diff --git a/libc/include/CMakeLists.txt b/libc/include/CMakeLists.txt
@@ -488,6 +488,19 @@ add_gen_header(
     .llvm-libc-types.wchar_t
 )
 
+if(LIBC_TARGET_ARCHITECTURE_IS_GPU)
+  file(MAKE_DIRECTORY "gpu")
+
+  add_gen_header(
+    gpu_rpc
+    DEF_FILE gpu/rpc.h.def
+    GEN_HDR gpu/rpc.h
+    DEPENDS
+      .llvm_libc_common_h
+      .llvm-libc-types.rpc_opcodes_t
+  )
+endif()
+
 if(NOT LLVM_LIBC_FULL_BUILD)
   # We don't install headers in non-fullbuild mode.
   return()

diff --git a/libc/include/gpu/rpc.h.def b/libc/include/gpu/rpc.h.def
@@ -0,0 +1,18 @@
+//===-- GPU header rpc.h --------------------------------------------------===//
+//
+// 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_GPU_RPC_H
+#define LLVM_LIBC_GPU_RPC_H
+
+#include <__llvm-libc-common.h>
+
+#include <llvm-libc-types/rpc_opcodes_t.h>
+
+%%public_api()
+
+#endif // LLVM_LIBC_GPU_RPC_H
diff --git a/libc/include/llvm-libc-types/CMakeLists.txt b/libc/include/llvm-libc-types/CMakeLists.txt
@@ -90,3 +90,4 @@ add_header(wchar_t HDR wchar_t.h)
 add_header(wint_t HDR wint_t.h)
 add_header(sa_family_t HDR sa_family_t.h)
 add_header(struct_sockaddr HDR struct_sockaddr.h)
+add_header(rpc_opcodes_t HDR rpc_opcodes_t.h)
diff --git a/libc/include/llvm-libc-types/rpc_opcodes_t.h b/libc/include/llvm-libc-types/rpc_opcodes_t.h
@@ -0,0 +1,25 @@
+//===-- Definition of RPC opcodes -----------------------------------------===//
+//
+// 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_OPCODE_H__
+#define __LLVM_LIBC_TYPES_RPC_OPCODE_H__
+
+typedef enum : unsigned short {
+  RPC_NOOP = 0,
+  RPC_EXIT = 1,
+  RPC_WRITE_TO_STDOUT = 2,
+  RPC_WRITE_TO_STDERR = 3,
+  RPC_WRITE_TO_STREAM = 4,
+  RPC_MALLOC = 5,
+  RPC_FREE = 6,
+  RPC_TEST_INCREMENT = 7,
+  RPC_TEST_INTERFACE = 8,
+  RPC_TEST_STREAM = 9,
+} rpc_opcode_t;
+
+#endif // __LLVM_LIBC_TYPES_RPC_OPCODE_H__
diff --git a/libc/spec/gpu_ext.td b/libc/spec/gpu_ext.td
@@ -0,0 +1,18 @@
+def GPUExtensions : StandardSpec<"GPUExtensions"> {
+  HeaderSpec RPC = HeaderSpec<
+    "gpu/rpc.h",
+    [], // Macros
+    [], // Types
+    [], // Enumerations
+    [
+        FunctionSpec<
+            "rpc_reset",
+            RetValSpec<VoidType>,
+            [ArgSpec<UnsignedIntType>, ArgSpec<VoidPtr>]
+        >,
+    ]
+  >;
+  let Headers = [
+    RPC,
+  ];
+}
diff --git a/libc/src/CMakeLists.txt b/libc/src/CMakeLists.txt
@@ -20,6 +20,10 @@ if(${LIBC_TARGET_OS} STREQUAL "linux")
   add_subdirectory(unistd)
 endif()
 
+if(${LIBC_TARGET_OS} STREQUAL "gpu")
+  add_subdirectory(gpu)
+endif()
+
 if(NOT LLVM_LIBC_FULL_BUILD)
   return()
 endif()

diff --git a/libc/src/__support/File/gpu/file.cpp b/libc/src/__support/File/gpu/file.cpp
@@ -37,7 +37,7 @@ namespace {
 
 int write_to_stdout(const void *data, size_t size) {
   int ret = 0;
-  rpc::Client::Port port = rpc::client.open<rpc::WRITE_TO_STDOUT>();
+  rpc::Client::Port port = rpc::client.open<RPC_WRITE_TO_STDOUT>();
   port.send_n(data, size);
   port.recv([&](rpc::Buffer *buffer) {
     ret = reinterpret_cast<int *>(buffer->data)[0];
@@ -48,7 +48,7 @@ int write_to_stdout(const void *data, size_t size) {
 
 int write_to_stderr(const void *data, size_t size) {
   int ret = 0;
-  rpc::Client::Port port = rpc::client.open<rpc::WRITE_TO_STDERR>();
+  rpc::Client::Port port = rpc::client.open<RPC_WRITE_TO_STDERR>();
   port.send_n(data, size);
   port.recv([&](rpc::Buffer *buffer) {
     ret = reinterpret_cast<int *>(buffer->data)[0];
@@ -59,7 +59,7 @@ int write_to_stderr(const void *data, size_t size) {
 
 int write_to_stream(uintptr_t file, const void *data, size_t size) {
   int ret = 0;
-  rpc::Client::Port port = rpc::client.open<rpc::WRITE_TO_STREAM>();
+  rpc::Client::Port port = rpc::client.open<RPC_WRITE_TO_STREAM>();
   port.send([&](rpc::Buffer *buffer) {
     reinterpret_cast<uintptr_t *>(buffer->data)[0] = file;
   });

diff --git a/libc/src/__support/OSUtil/gpu/io.cpp b/libc/src/__support/OSUtil/gpu/io.cpp
@@ -15,7 +15,7 @@
 namespace __llvm_libc {
 
 void write_to_stderr(cpp::string_view msg) {
-  rpc::Client::Port port = rpc::client.open<rpc::WRITE_TO_STDERR>();
+  rpc::Client::Port port = rpc::client.open<RPC_WRITE_TO_STDERR>();
   port.send_n(msg.data(), msg.size());
   port.recv([](rpc::Buffer *) { /* void */ });
   port.close();

diff --git a/libc/src/__support/OSUtil/gpu/quick_exit.cpp b/libc/src/__support/OSUtil/gpu/quick_exit.cpp
@@ -17,7 +17,7 @@
 namespace __llvm_libc {
 
 void quick_exit(int status) {
-  rpc::Client::Port port = rpc::client.open<rpc::EXIT>();
+  rpc::Client::Port port = rpc::client.open<RPC_EXIT>();
   port.send([&](rpc::Buffer *buffer) {
     reinterpret_cast<uint32_t *>(buffer->data)[0] = status;
   });

diff --git a/libc/src/__support/RPC/CMakeLists.txt b/libc/src/__support/RPC/CMakeLists.txt
@@ -22,6 +22,7 @@ add_object_library(
   HDRS
     rpc_client.h
   DEPENDS
+    libc.include.gpu_rpc
     libc.src.__support.GPU.utils
     .rpc
 )
diff --git a/libc/src/__support/RPC/rpc.h b/libc/src/__support/RPC/rpc.h
@@ -30,20 +30,6 @@
 namespace __llvm_libc {
 namespace rpc {
 
-/// A list of opcodes that we use to invoke certain actions on the server.
-enum Opcode : uint16_t {
-  NOOP = 0,
-  EXIT = 1,
-  WRITE_TO_STDOUT = 2,
-  WRITE_TO_STDERR = 3,
-  WRITE_TO_STREAM = 4,
-  MALLOC = 5,
-  FREE = 6,
-  TEST_INCREMENT = 7,
-  TEST_INTERFACE = 8,
-  TEST_STREAM = 9,
-};
-
 /// A fixed size channel used to communicate between the RPC client and server.
 struct Buffer {
   uint64_t data[8];

diff --git a/libc/src/__support/RPC/rpc_client.h b/libc/src/__support/RPC/rpc_client.h
@@ -11,6 +11,8 @@
 
 #include "rpc.h"
 
+#include <gpu/rpc.h>
+
 namespace __llvm_libc {
 namespace rpc {
 

diff --git a/libc/src/gpu/CMakeLists.txt b/libc/src/gpu/CMakeLists.txt
@@ -0,0 +1,10 @@
+add_entrypoint_object(
+  rpc_reset
+  SRCS
+    rpc_reset.cpp
+  HDRS
+    rpc_reset.h
+  DEPENDS
+    libc.src.__support.RPC.rpc_client
+    libc.src.__support.GPU.utils
+)
diff --git a/libc/src/gpu/rpc_reset.cpp b/libc/src/gpu/rpc_reset.cpp
@@ -0,0 +1,25 @@
+//===---------- GPU implementation of the external RPC functionion --------===//
+//
+// 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_reset.h"
+
+#include "src/__support/GPU/utils.h"
+#include "src/__support/RPC/rpc_client.h"
+#include "src/__support/common.h"
+
+namespace __llvm_libc {
+
+// This is the external interface to initialize the RPC client with the
+// shared buffer.
+LLVM_LIBC_FUNCTION(void, rpc_reset,
+                   (unsigned int num_ports, void *rpc_shared_buffer)) {
+  __llvm_libc::rpc::client.reset(num_ports, __llvm_libc::gpu::get_lane_size(),
+                                 rpc_shared_buffer);
+}
+
+} // namespace __llvm_libc
diff --git a/libc/src/gpu/rpc_reset.h b/libc/src/gpu/rpc_reset.h
@@ -0,0 +1,18 @@
+//===-- 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_H
+#define LLVM_LIBC_SRC_GPU_RPC_H
+
+namespace __llvm_libc {
+
+void rpc_reset(unsigned int num_ports, void *buffer);
+
+} // namespace __llvm_libc
+
+#endif // LLVM_LIBC_SRC_GPU_RPC_H
diff --git a/libc/src/stdlib/gpu/free.cpp b/libc/src/stdlib/gpu/free.cpp
@@ -13,7 +13,7 @@
 namespace __llvm_libc {
 
 LLVM_LIBC_FUNCTION(void, free, (void *ptr)) {
-  rpc::Client::Port port = rpc::client.open<rpc::FREE>();
+  rpc::Client::Port port = rpc::client.open<RPC_FREE>();
   port.send([=](rpc::Buffer *buffer) {
     buffer->data[0] = reinterpret_cast<uintptr_t>(ptr);
   });

diff --git a/libc/src/stdlib/gpu/malloc.cpp b/libc/src/stdlib/gpu/malloc.cpp
@@ -14,7 +14,7 @@ namespace __llvm_libc {
 
 LLVM_LIBC_FUNCTION(void *, malloc, (size_t size)) {
   void *ptr = nullptr;
-  rpc::Client::Port port = rpc::client.open<rpc::MALLOC>();
+  rpc::Client::Port port = rpc::client.open<RPC_MALLOC>();
   port.send_and_recv([=](rpc::Buffer *buffer) { buffer->data[0] = size; },
                      [&](rpc::Buffer *buffer) {
                        ptr = reinterpret_cast<void *>(buffer->data[0]);

diff --git a/libc/test/integration/startup/gpu/rpc_interface_test.cpp b/libc/test/integration/startup/gpu/rpc_interface_test.cpp
@@ -16,7 +16,7 @@ using namespace __llvm_libc;
 // as long as they are mirrored.
 static void test_interface(bool end_with_send) {
   uint64_t cnt = 0;
-  rpc::Client::Port port = rpc::client.open<rpc::TEST_INTERFACE>();
+  rpc::Client::Port port = rpc::client.open<RPC_TEST_INTERFACE>();
   port.send([&](rpc::Buffer *buffer) { buffer->data[0] = end_with_send; });
   port.send([&](rpc::Buffer *buffer) { buffer->data[0] = cnt = cnt + 1; });
   port.recv([&](rpc::Buffer *buffer) { cnt = buffer->data[0]; });

diff --git a/libc/test/integration/startup/gpu/rpc_stream_test.cpp b/libc/test/integration/startup/gpu/rpc_stream_test.cpp
@@ -33,7 +33,7 @@ static void test_stream() {
 
   inline_memcpy(send_ptr, str, send_size);
   ASSERT_TRUE(inline_memcmp(send_ptr, str, send_size) == 0 && "Data mismatch");
-  rpc::Client::Port port = rpc::client.open<rpc::TEST_STREAM>();
+  rpc::Client::Port port = rpc::client.open<RPC_TEST_STREAM>();
   port.send_n(send_ptr, send_size);
   port.recv_n(&recv_ptr, &recv_size,
               [](uint64_t size) { return malloc(size); });
@@ -76,7 +76,7 @@ static void test_divergent() {
   inline_memcpy(buffer, &data[offset], offset);
   ASSERT_TRUE(inline_memcmp(buffer, &data[offset], offset) == 0 &&
               "Data mismatch");
-  rpc::Client::Port port = rpc::client.open<rpc::TEST_STREAM>();
+  rpc::Client::Port port = rpc::client.open<RPC_TEST_STREAM>();
   port.send_n(buffer, offset);
   inline_memset(buffer, offset, 0);
   port.recv_n(&recv_ptr, &recv_size, [&](uint64_t) { return buffer; });

diff --git a/libc/test/integration/startup/gpu/rpc_test.cpp b/libc/test/integration/startup/gpu/rpc_test.cpp
@@ -17,7 +17,7 @@ static void test_add_simple() {
       10 + 10 * gpu::get_thread_id() + 10 * gpu::get_block_id();
   uint64_t cnt = 0;
   for (uint32_t i = 0; i < num_additions; ++i) {
-    rpc::Client::Port port = rpc::client.open<rpc::TEST_INCREMENT>();
+    rpc::Client::Port port = rpc::client.open<RPC_TEST_INCREMENT>();
     port.send_and_recv(
         [=](rpc::Buffer *buffer) {
           reinterpret_cast<uint64_t *>(buffer->data)[0] = cnt;
@@ -32,7 +32,7 @@ static void test_add_simple() {
 
 // Test to ensure that the RPC mechanism doesn't hang on divergence.
 static void test_noop(uint8_t data) {
-  rpc::Client::Port port = rpc::client.open<rpc::NOOP>();
+  rpc::Client::Port port = rpc::client.open<RPC_NOOP>();
   port.send([=](rpc::Buffer *buffer) { buffer->data[0] = data; });
   port.close();
 }

diff --git a/libc/utils/gpu/loader/CMakeLists.txt b/libc/utils/gpu/loader/CMakeLists.txt
@@ -1,6 +1,7 @@
 add_library(gpu_loader OBJECT Main.cpp)
 target_include_directories(gpu_loader PUBLIC
   ${CMAKE_CURRENT_SOURCE_DIR}
+  ${LIBC_SOURCE_DIR}/include
   ${LIBC_SOURCE_DIR}
 )
 

diff --git a/libc/utils/gpu/server/CMakeLists.txt b/libc/utils/gpu/server/CMakeLists.txt
@@ -3,4 +3,6 @@ add_library(rpc_server STATIC 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})
+# TODO: This is for the opcodes, we will copy the file here when installed.
+target_include_directories(rpc_server PUBLIC ${LIBC_SOURCE_DIR}/include)
 target_include_directories(rpc_server PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})