[libc] Implement temporary printf on the GPU (#85331)

Summary:
This patch adds a temporary implementation that uses a struct-based
interface in lieu of varargs support. Once varargs support exists we
will move this implementation to the "real" printf implementation.

Conceptually, this patch has the client copy over its format string and
arguments to the server. The server will then scan the format string
searching for any specifiers that are actually a string. If it is a
string then we will send the pointer back to the server to tell it to
copy it back. This copied value will then replace the pointer when the
final formatting is done.

This will require a built-in extension to the varargs support to get
access to the underlying struct. The varargs used on the GPU will simply
be a struct wrapped in a varargs ABI.

Author: jhuber6

libc/config/gpu/entrypoints.txt

@@ -211,6 +211,7 @@ set(TARGET_LIBC_ENTRYPOINTS
 
     # gpu/rpc.h entrypoints
     libc.src.gpu.rpc_host_call
+    libc.src.gpu.rpc_fprintf
 )
 
 set(TARGET_LIBM_ENTRYPOINTS

libc/include/llvm-libc-types/rpc_opcodes_t.h

@@ -31,6 +31,9 @@ typedef enum {
   RPC_FTELL,
   RPC_FFLUSH,
   RPC_UNGETC,
+  RPC_PRINTF_TO_STDOUT,
+  RPC_PRINTF_TO_STDERR,
+  RPC_PRINTF_TO_STREAM,
   RPC_LAST = 0xFFFF,
 } rpc_opcode_t;
 

libc/spec/gpu_ext.td

@@ -10,6 +10,14 @@ def GPUExtensions : StandardSpec<"GPUExtensions"> {
             RetValSpec<VoidType>,
             [ArgSpec<VoidPtr>, ArgSpec<VoidPtr>, ArgSpec<SizeTType>]
         >,
+        FunctionSpec<
+            "rpc_fprintf",
+            RetValSpec<IntType>,
+            [ArgSpec<FILERestrictedPtr>,
+             ArgSpec<ConstCharRestrictedPtr>,
+             ArgSpec<VoidPtr>,
+             ArgSpec<SizeTType>]
+        >,
     ]
   >;
   let Headers = [

libc/src/__support/arg_list.h

@@ -13,6 +13,7 @@
 
 #include <stdarg.h>
 #include <stddef.h>
+#include <stdint.h>
 
 namespace LIBC_NAMESPACE {
 namespace internal {
@@ -60,6 +61,43 @@ class MockArgList {
   size_t read_count() const { return arg_counter; }
 };
 
+// Used for the GPU implementation of `printf`. This models a variadic list as a
+// simple array of pointers that are built manually by the implementation.
+class StructArgList {
+  void *ptr;
+  void *end;
+
+public:
+  LIBC_INLINE StructArgList(void *ptr, size_t size)
+      : ptr(ptr), end(reinterpret_cast<unsigned char *>(ptr) + size) {}
+  LIBC_INLINE StructArgList(const StructArgList &other) {
+    ptr = other.ptr;
+    end = other.end;
+  }
+  LIBC_INLINE StructArgList() = default;
+  LIBC_INLINE ~StructArgList() = default;
+
+  LIBC_INLINE StructArgList &operator=(const StructArgList &rhs) {
+    ptr = rhs.ptr;
+    return *this;
+  }
+
+  LIBC_INLINE void *get_ptr() const { return ptr; }
+
+  template <class T> LIBC_INLINE T next_var() {
+    ptr = reinterpret_cast<void *>(
+        ((reinterpret_cast<uintptr_t>(ptr) + alignof(T) - 1) / alignof(T)) *
+        alignof(T));
+
+    if (ptr >= end)
+      return T(-1);
+
+    T val = *reinterpret_cast<T *>(ptr);
+    ptr = reinterpret_cast<unsigned char *>(ptr) + sizeof(T);
+    return val;
+  }
+};
+
 } // namespace internal
 } // namespace LIBC_NAMESPACE
 

libc/src/gpu/CMakeLists.txt

@@ -8,3 +8,15 @@ add_entrypoint_object(
     libc.src.__support.RPC.rpc_client
     libc.src.__support.GPU.utils
 )
+
+add_entrypoint_object(
+  rpc_fprintf
+  SRCS
+    rpc_fprintf.cpp
+  HDRS
+    rpc_fprintf.h
+  DEPENDS
+    libc.src.stdio.gpu.gpu_file
+    libc.src.__support.RPC.rpc_client
+    libc.src.__support.GPU.utils
+)

libc/src/gpu/rpc_fprintf.cpp

@@ -0,0 +1,71 @@
+//===-- GPU implementation of fprintf -------------------------------------===//
+//
+// 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_fprintf.h"
+
+#include "src/__support/CPP/string_view.h"
+#include "src/__support/GPU/utils.h"
+#include "src/__support/RPC/rpc_client.h"
+#include "src/__support/common.h"
+#include "src/stdio/gpu/file.h"
+
+namespace LIBC_NAMESPACE {
+
+template <uint16_t opcode>
+int fprintf_impl(::FILE *__restrict file, const char *__restrict format,
+                 size_t format_size, void *args, size_t args_size) {
+  uint64_t mask = gpu::get_lane_mask();
+  rpc::Client::Port port = rpc::client.open<opcode>();
+
+  if constexpr (opcode == RPC_PRINTF_TO_STREAM) {
+    port.send([&](rpc::Buffer *buffer) {
+      buffer->data[0] = reinterpret_cast<uintptr_t>(file);
+    });
+  }
+
+  port.send_n(format, format_size);
+  port.send_n(args, args_size);
+
+  uint32_t ret = 0;
+  for (;;) {
+    const char *str = nullptr;
+    port.recv([&](rpc::Buffer *buffer) {
+      ret = static_cast<uint32_t>(buffer->data[0]);
+      str = reinterpret_cast<const char *>(buffer->data[1]);
+    });
+    // If any lanes have a string argument it needs to be copied back.
+    if (!gpu::ballot(mask, str))
+      break;
+
+    uint64_t size = str ? internal::string_length(str) + 1 : 0;
+    port.send_n(str, size);
+  }
+
+  port.close();
+  return ret;
+}
+
+// TODO: This is a stand-in function that uses a struct pointer and size in
+// place of varargs. Once varargs support is added we will use that to
+// implement the real version.
+LLVM_LIBC_FUNCTION(int, rpc_fprintf,
+                   (::FILE *__restrict stream, const char *__restrict format,
+                    void *args, size_t size)) {
+  cpp::string_view str(format);
+  if (stream == stdout)
+    return fprintf_impl<RPC_PRINTF_TO_STDOUT>(stream, format, str.size() + 1,
+                                              args, size);
+  else if (stream == stderr)
+    return fprintf_impl<RPC_PRINTF_TO_STDERR>(stream, format, str.size() + 1,
+                                              args, size);
+  else
+    return fprintf_impl<RPC_PRINTF_TO_STREAM>(stream, format, str.size() + 1,
+                                              args, size);
+}
+
+} // namespace LIBC_NAMESPACE

libc/src/gpu/rpc_fprintf.h

@@ -0,0 +1,22 @@
+//===-- 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_HOST_CALL_H
+#define LLVM_LIBC_SRC_GPU_RPC_HOST_CALL_H
+
+#include <stddef.h>
+#include <stdio.h>
+
+namespace LIBC_NAMESPACE {
+
+int rpc_fprintf(::FILE *__restrict stream, const char *__restrict format,
+                void *argc, size_t size);
+
+} // namespace LIBC_NAMESPACE
+
+#endif // LLVM_LIBC_SRC_GPU_RPC_HOST_CALL_H

libc/test/integration/src/stdio/CMakeLists.txt

@@ -1,3 +1,6 @@
+if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${LIBC_TARGET_OS})
+  add_subdirectory(${LIBC_TARGET_OS})
+endif()
 add_custom_target(stdio-integration-tests)
 add_dependencies(libc-integration-tests stdio-integration-tests)
 

libc/test/integration/src/stdio/gpu/CMakeLists.txt

@@ -0,0 +1,21 @@
+add_custom_target(stdio-gpu-integration-tests)
+add_dependencies(libc-integration-tests stdio-gpu-integration-tests)
+
+# Create an output directory for any temporary test files.
+file(MAKE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/testdata)
+
+# These tests are not for correctness testing, but are instead a convenient way
+# to generate hermetic binaries for comparitive binary size testing.
+add_integration_test(
+  printf_test
+  SUITE
+    stdio-gpu-integration-tests
+  SRCS
+    printf.cpp
+  DEPENDS
+    libc.src.gpu.rpc_fprintf
+    libc.src.stdio.fopen
+  LOADER_ARGS
+    --threads 32
+    --blocks 4
+)

libc/test/integration/src/stdio/gpu/printf.cpp

@@ -0,0 +1,88 @@
+//===-- RPC test to check args to printf ----------------------------------===//
+//
+// 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 "test/IntegrationTest/test.h"
+
+#include "src/__support/GPU/utils.h"
+#include "src/gpu/rpc_fprintf.h"
+#include "src/stdio/fopen.h"
+
+using namespace LIBC_NAMESPACE;
+
+FILE *file = LIBC_NAMESPACE::fopen("testdata/test_data.txt", "w");
+
+TEST_MAIN(int argc, char **argv, char **envp) {
+  ASSERT_TRUE(file && "failed to open file");
+  // Check basic printing.
+  int written = 0;
+  written = LIBC_NAMESPACE::rpc_fprintf(file, "A simple string\n", nullptr, 0);
+  ASSERT_EQ(written, 16);
+
+  const char *str = "A simple string\n";
+  written = LIBC_NAMESPACE::rpc_fprintf(file, "%s", &str, sizeof(void *));
+  ASSERT_EQ(written, 16);
+
+  // Check printing a different value with each thread.
+  uint64_t thread_id = gpu::get_thread_id();
+  written = LIBC_NAMESPACE::rpc_fprintf(file, "%8ld\n", &thread_id,
+                                        sizeof(thread_id));
+  ASSERT_EQ(written, 9);
+
+  struct {
+    uint32_t x = 1;
+    char c = 'c';
+    double f = 1.0;
+  } args1;
+  written =
+      LIBC_NAMESPACE::rpc_fprintf(file, "%d%c%.1f\n", &args1, sizeof(args1));
+  ASSERT_EQ(written, 6);
+
+  struct {
+    uint32_t x = 1;
+    const char *str = "A simple string\n";
+  } args2;
+  written =
+      LIBC_NAMESPACE::rpc_fprintf(file, "%032b%s\n", &args2, sizeof(args2));
+  ASSERT_EQ(written, 49);
+
+  // Check that the server correctly handles divergent numbers of arguments.
+  const char *format = gpu::get_thread_id() % 2 ? "%s" : "%20ld\n";
+  written = LIBC_NAMESPACE::rpc_fprintf(file, format, &str, sizeof(void *));
+  ASSERT_EQ(written, gpu::get_thread_id() % 2 ? 16 : 21);
+
+  format = gpu::get_thread_id() % 2 ? "%s" : str;
+  written = LIBC_NAMESPACE::rpc_fprintf(file, format, &str, sizeof(void *));
+  ASSERT_EQ(written, 16);
+
+  // Check that we handle null arguments correctly.
+  struct {
+    void *null = nullptr;
+  } args3;
+  written = LIBC_NAMESPACE::rpc_fprintf(file, "%p", &args3, sizeof(args3));
+  ASSERT_EQ(written, 9);
+
+#ifndef LIBC_COPT_PRINTF_NO_NULLPTR_CHECKS
+  written = LIBC_NAMESPACE::rpc_fprintf(file, "%s", &args3, sizeof(args3));
+  ASSERT_EQ(written, 6);
+#endif // LIBC_COPT_PRINTF_NO_NULLPTR_CHECKS
+
+  // Check for extremely abused variable width arguments
+  struct {
+    uint32_t x = 1;
+    uint32_t y = 2;
+    double f = 1.0;
+  } args4;
+  written = LIBC_NAMESPACE::rpc_fprintf(file, "%**d", &args4, sizeof(args4));
+  ASSERT_EQ(written, 4);
+  written = LIBC_NAMESPACE::rpc_fprintf(file, "%**d%6d", &args4, sizeof(args4));
+  ASSERT_EQ(written, 10);
+  written = LIBC_NAMESPACE::rpc_fprintf(file, "%**.**f", &args4, sizeof(args4));
+  ASSERT_EQ(written, 7);
+
+  return 0;
+}

libc/utils/gpu/server/CMakeLists.txt

@@ -1,4 +1,8 @@
-add_library(llvmlibc_rpc_server STATIC rpc_server.cpp)
+add_library(llvmlibc_rpc_server STATIC
+  ${LIBC_SOURCE_DIR}/src/stdio/printf_core/writer.cpp
+  ${LIBC_SOURCE_DIR}/src/stdio/printf_core/converter.cpp
+  rpc_server.cpp
+)
 
 # Include the RPC implemenation from libc.
 target_include_directories(llvmlibc_rpc_server PRIVATE ${LIBC_SOURCE_DIR})
@@ -9,6 +13,10 @@ target_include_directories(llvmlibc_rpc_server PUBLIC ${CMAKE_CURRENT_SOURCE_DIR
 target_compile_options(llvmlibc_rpc_server PUBLIC
                        $<$<CXX_COMPILER_ID:GNU>:-Wno-attributes>)
 target_compile_definitions(llvmlibc_rpc_server PUBLIC
+                           LIBC_COPT_USE_C_ASSERT
+                           LIBC_COPT_ARRAY_ARG_LIST
+                           LIBC_COPT_PRINTF_DISABLE_WRITE_INT
+                           LIBC_COPT_PRINTF_DISABLE_INDEX_MODE
                            LIBC_NAMESPACE=${LIBC_NAMESPACE})
 
 # Install the server and associated header.