[libc] Implement (v|f)printf on the GPU

Summary:
This patch implements the `printf` family of functions on the GPU using
the new variadic support. This patch adapts the old handling in the
`rpc_fprintf` placeholder, but adds an extra RPC call to get the size of
the buffer to copy. This prevents the GPU from needing to parse the
string. While it's theoretically possible for the pass to know the size
of the struct, it's prohibitively difficult to do while maintaining ABI
compatibility with NVIDIA's varargs.

Depends on llvm#96015.

Author: jhuber6

libc/config/gpu/entrypoints.txt

@@ -1,13 +1,3 @@
-if(LIBC_TARGET_ARCHITECTURE_IS_AMDGPU)
-  set(extra_entrypoints
-      # stdio.h entrypoints
-      libc.src.stdio.snprintf
-      libc.src.stdio.sprintf
-      libc.src.stdio.vsnprintf
-      libc.src.stdio.vsprintf
-  )
-endif()
-
 set(TARGET_LIBC_ENTRYPOINTS
     # assert.h entrypoints
     libc.src.assert.__assert_fail
@@ -185,9 +175,16 @@ set(TARGET_LIBC_ENTRYPOINTS
     libc.src.errno.errno
 
     # stdio.h entrypoints
-    ${extra_entrypoints}
     libc.src.stdio.clearerr
     libc.src.stdio.fclose
+    libc.src.stdio.printf
+    libc.src.stdio.vprintf
+    libc.src.stdio.fprintf
+    libc.src.stdio.vfprintf
+    libc.src.stdio.sprintf
+    libc.src.stdio.snprintf
+    libc.src.stdio.vsprintf
+    libc.src.stdio.vsnprintf
     libc.src.stdio.feof
     libc.src.stdio.ferror
     libc.src.stdio.fflush

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

@@ -34,6 +34,9 @@ typedef enum {
   RPC_PRINTF_TO_STDOUT,
   RPC_PRINTF_TO_STDERR,
   RPC_PRINTF_TO_STREAM,
+  RPC_PRINTF_TO_STDOUT_PACKED,
+  RPC_PRINTF_TO_STDERR_PACKED,
+  RPC_PRINTF_TO_STREAM_PACKED,
   RPC_REMOVE,
   RPC_LAST = 0xFFFF,
 } rpc_opcode_t;

libc/src/__support/arg_list.h

@@ -18,6 +18,11 @@
 namespace LIBC_NAMESPACE {
 namespace internal {
 
+template <typename V, typename A>
+LIBC_INLINE constexpr V align_up(V val, A align) {
+  return ((val + V(align) - 1) / V(align)) * V(align);
+}
+
 class ArgList {
   va_list vlist;
 
@@ -54,7 +59,34 @@ class MockArgList {
   }
 
   template <class T> LIBC_INLINE T next_var() {
-    ++arg_counter;
+    arg_counter++;
+    return T(arg_counter);
+  }
+
+  size_t read_count() const { return arg_counter; }
+};
+
+// Used by the GPU implementation to parse how many bytes need to be read from
+// the variadic argument buffer.
+template <bool packed> class DummyArgList {
+  size_t arg_counter = 0;
+
+public:
+  LIBC_INLINE DummyArgList() = default;
+  LIBC_INLINE DummyArgList(va_list) { ; }
+  LIBC_INLINE DummyArgList(DummyArgList &other) {
+    arg_counter = other.arg_counter;
+  }
+  LIBC_INLINE ~DummyArgList() = default;
+
+  LIBC_INLINE DummyArgList &operator=(DummyArgList &rhs) {
+    arg_counter = rhs.arg_counter;
+    return *this;
+  }
+
+  template <class T> LIBC_INLINE T next_var() {
+    arg_counter = packed ? arg_counter + sizeof(T)
+                         : align_up(arg_counter, alignof(T)) + sizeof(T);
     return T(arg_counter);
   }
 
@@ -63,7 +95,7 @@ class MockArgList {
 
 // 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 {
+template <bool packed> class StructArgList {
   void *ptr;
   void *end;
 
@@ -85,15 +117,18 @@ class StructArgList {
   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 (!packed)
+      ptr = reinterpret_cast<void *>(
+          align_up(reinterpret_cast<uintptr_t>(ptr), alignof(T)));
     if (ptr >= end)
       return T(-1);
 
-    T val = *reinterpret_cast<T *>(ptr);
-    ptr = reinterpret_cast<unsigned char *>(ptr) + sizeof(T);
+    // Memcpy because pointer alignment may be illegal given a packed struct.
+    T val;
+    __builtin_memcpy(&val, ptr, sizeof(T));
+
+    ptr =
+        reinterpret_cast<void *>(reinterpret_cast<uintptr_t>(ptr) + sizeof(T));
     return val;
   }
 };

libc/src/gpu/rpc_fprintf.cpp

@@ -29,6 +29,9 @@ int fprintf_impl(::FILE *__restrict file, const char *__restrict format,
   }
 
   port.send_n(format, format_size);
+  port.recv([&](rpc::Buffer *buffer) {
+    args_size = static_cast<size_t>(buffer->data[0]);
+  });
   port.send_n(args, args_size);
 
   uint32_t ret = 0;
@@ -50,7 +53,7 @@ int fprintf_impl(::FILE *__restrict file, const char *__restrict format,
   return ret;
 }
 
-// TODO: This is a stand-in function that uses a struct pointer and size in
+// TODO: Delete this and port OpenMP to use `printf`.
 // place of varargs. Once varargs support is added we will use that to
 // implement the real version.
 LLVM_LIBC_FUNCTION(int, rpc_fprintf,

libc/src/stdio/CMakeLists.txt

@@ -159,17 +159,6 @@ add_entrypoint_object(
     libc.src.stdio.printf_core.writer
 )
 
-add_entrypoint_object(
-  fprintf
-  SRCS
-    fprintf.cpp
-  HDRS
-    fprintf.h
-  DEPENDS
-    libc.src.__support.arg_list
-    libc.src.stdio.printf_core.vfprintf_internal
-)
-
 add_entrypoint_object(
   vsprintf
   SRCS
@@ -192,17 +181,6 @@ add_entrypoint_object(
     libc.src.stdio.printf_core.writer
 )
 
-add_entrypoint_object(
-  vfprintf
-  SRCS
-    vfprintf.cpp
-  HDRS
-    vfprintf.h
-  DEPENDS
-    libc.src.__support.arg_list
-    libc.src.stdio.printf_core.vfprintf_internal
-)
-
 add_stdio_entrypoint_object(
   fileno
   SRCS
@@ -261,6 +239,7 @@ add_stdio_entrypoint_object(fputc)
 add_stdio_entrypoint_object(putc)
 add_stdio_entrypoint_object(putchar)
 add_stdio_entrypoint_object(printf)
+add_stdio_entrypoint_object(fprintf)
 add_stdio_entrypoint_object(fgetc)
 add_stdio_entrypoint_object(fgetc_unlocked)
 add_stdio_entrypoint_object(getc)
@@ -273,3 +252,4 @@ add_stdio_entrypoint_object(stdin)
 add_stdio_entrypoint_object(stdout)
 add_stdio_entrypoint_object(stderr)
 add_stdio_entrypoint_object(vprintf)
+add_stdio_entrypoint_object(vfprintf)

libc/src/stdio/generic/CMakeLists.txt

@@ -396,6 +396,31 @@ add_entrypoint_object(
     ${printf_deps}
 )
 
+add_entrypoint_object(
+  fprintf
+  SRCS
+    fprintf.cpp
+  HDRS
+    ../fprintf.h
+  DEPENDS
+    libc.src.__support.arg_list
+    libc.src.stdio.printf_core.vfprintf_internal
+    ${printf_deps}
+)
+
+add_entrypoint_object(
+  vfprintf
+  SRCS
+    vfprintf.cpp
+  HDRS
+    ../vfprintf.h
+  DEPENDS
+    libc.src.__support.arg_list
+    libc.src.stdio.printf_core.vfprintf_internal
+    ${printf_deps}
+)
+
+
 add_entrypoint_object(
   fgets
   SRCS

libc/src/stdio/fprintf.cpp renamed to libc/src/stdio/generic/fprintf.cpp

@@ -10,6 +10,14 @@ add_header_library(
     .stderr
 )
 
+add_header_library(
+  vfprintf_utils
+  HDRS
+    vfprintf_utils.h
+  DEPENDS
+    .gpu_file
+)
+
 add_entrypoint_object(
   feof
   SRCS
@@ -273,6 +281,46 @@ add_entrypoint_object(
     .gpu_file
 )
 
+add_entrypoint_object(
+  printf
+  SRCS
+    printf.cpp
+  HDRS
+    ../printf.h
+  DEPENDS
+    .vfprintf_utils
+)
+
+add_entrypoint_object(
+  vprintf
+  SRCS
+    vprintf.cpp
+  HDRS
+    ../vprintf.h
+  DEPENDS
+    .vfprintf_utils
+)
+
+add_entrypoint_object(
+  fprintf
+  SRCS
+    fprintf.cpp
+  HDRS
+    ../fprintf.h
+  DEPENDS
+    .vfprintf_utils
+)
+
+add_entrypoint_object(
+  vfprintf
+  SRCS
+    vfprintf.cpp
+  HDRS
+    ../vfprintf.h
+  DEPENDS
+    .vfprintf_utils
+)
+
 add_entrypoint_object(
   stdin
   SRCS

libc/src/stdio/vfprintf.cpp renamed to libc/src/stdio/generic/vfprintf.cpp

@@ -0,0 +1,31 @@
+//===-- 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 "src/stdio/fprintf.h"
+
+#include "src/__support/CPP/string_view.h"
+#include "src/__support/arg_list.h"
+#include "src/errno/libc_errno.h"
+#include "src/stdio/gpu/vfprintf_utils.h"
+
+#include <stdio.h>
+
+namespace LIBC_NAMESPACE {
+
+LLVM_LIBC_FUNCTION(int, fprintf,
+                   (::FILE *__restrict stream, const char *__restrict format,
+                    ...)) {
+  va_list vlist;
+  va_start(vlist, format);
+  cpp::string_view str_view(format);
+  int ret_val = vfprintf_internal(stream, format, str_view.size() + 1, vlist);
+  va_end(vlist);
+  return ret_val;
+}
+
+} // namespace LIBC_NAMESPACE

libc/src/stdio/gpu/CMakeLists.txt

@@ -0,0 +1,29 @@
+//===-- GPU Implementation of 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 "src/stdio/printf.h"
+
+#include "src/__support/CPP/string_view.h"
+#include "src/__support/arg_list.h"
+#include "src/errno/libc_errno.h"
+#include "src/stdio/gpu/vfprintf_utils.h"
+
+#include <stdio.h>
+
+namespace LIBC_NAMESPACE {
+
+LLVM_LIBC_FUNCTION(int, printf, (const char *__restrict format, ...)) {
+  va_list vlist;
+  va_start(vlist, format);
+  cpp::string_view str_view(format);
+  int ret_val = vfprintf_internal(stdout, format, str_view.size() + 1, vlist);
+  va_end(vlist);
+  return ret_val;
+}
+
+} // namespace LIBC_NAMESPACE

libc/src/stdio/gpu/fprintf.cpp

@@ -0,0 +1,28 @@
+//===-- GPU Implementation of vfprintf ------------------------------------===//
+//
+// 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/stdio/vfprintf.h"
+
+#include "src/__support/CPP/string_view.h"
+#include "src/__support/arg_list.h"
+#include "src/errno/libc_errno.h"
+#include "src/stdio/gpu/vfprintf_utils.h"
+
+#include <stdio.h>
+
+namespace LIBC_NAMESPACE {
+
+LLVM_LIBC_FUNCTION(int, vfprintf,
+                   (::FILE *__restrict stream, const char *__restrict format,
+                    va_list vlist)) {
+  cpp::string_view str_view(format);
+  int ret_val = vfprintf_internal(stream, format, str_view.size() + 1, vlist);
+  return ret_val;
+}
+
+} // namespace LIBC_NAMESPACE