xref: /llvm-project/libc/src/stdio/gpu/vfprintf_utils.h (revision be0c67c90e045b03b0ffecc06ca6f93e440f48d8)
1 //===--- GPU helper functions for printf using RPC ------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "hdr/types/FILE.h"
10 #include "src/__support/RPC/rpc_client.h"
11 #include "src/__support/arg_list.h"
12 #include "src/stdio/gpu/file.h"
13 #include "src/string/string_utils.h"
14 
15 namespace LIBC_NAMESPACE {
16 
17 template <uint16_t opcode>
18 LIBC_INLINE int vfprintf_impl(::FILE *__restrict file,
19                               const char *__restrict format, size_t format_size,
20                               va_list vlist) {
21   uint64_t mask = gpu::get_lane_mask();
22   rpc::Client::Port port = rpc::client.open<opcode>();
23 
24   if constexpr (opcode == RPC_PRINTF_TO_STREAM ||
25                 opcode == RPC_PRINTF_TO_STREAM_PACKED) {
26     port.send([&](rpc::Buffer *buffer, uint32_t) {
27       buffer->data[0] = reinterpret_cast<uintptr_t>(file);
28     });
29   }
30 
31   size_t args_size = 0;
32   port.send_n(format, format_size);
33   port.recv([&](rpc::Buffer *buffer, uint32_t) {
34     args_size = static_cast<size_t>(buffer->data[0]);
35   });
36   port.send_n(vlist, args_size);
37 
38   uint32_t ret = 0;
39   for (;;) {
40     const char *str = nullptr;
41     port.recv([&](rpc::Buffer *buffer, uint32_t) {
42       ret = static_cast<uint32_t>(buffer->data[0]);
43       str = reinterpret_cast<const char *>(buffer->data[1]);
44     });
45     // If any lanes have a string argument it needs to be copied back.
46     if (!gpu::ballot(mask, str))
47       break;
48 
49     uint64_t size = str ? internal::string_length(str) + 1 : 0;
50     port.send_n(str, size);
51   }
52 
53   port.close();
54   return ret;
55 }
56 
57 LIBC_INLINE int vfprintf_internal(::FILE *__restrict stream,
58                                   const char *__restrict format,
59                                   size_t format_size, va_list vlist) {
60   // The AMDPGU backend uses a packed struct for its varargs. We pass it as a
61   // separate opcode so the server knows how much to advance the pointers.
62 #if defined(LIBC_TARGET_ARCH_IS_AMDGPU)
63   if (stream == stdout)
64     return vfprintf_impl<RPC_PRINTF_TO_STDOUT_PACKED>(stream, format,
65                                                       format_size, vlist);
66   else if (stream == stderr)
67     return vfprintf_impl<RPC_PRINTF_TO_STDERR_PACKED>(stream, format,
68                                                       format_size, vlist);
69   else
70     return vfprintf_impl<RPC_PRINTF_TO_STREAM_PACKED>(stream, format,
71                                                       format_size, vlist);
72 #else
73   if (stream == stdout)
74     return vfprintf_impl<RPC_PRINTF_TO_STDOUT>(stream, format, format_size,
75                                                vlist);
76   else if (stream == stderr)
77     return vfprintf_impl<RPC_PRINTF_TO_STDERR>(stream, format, format_size,
78                                                vlist);
79   else
80     return vfprintf_impl<RPC_PRINTF_TO_STREAM>(stream, format, format_size,
81                                                vlist);
82 #endif
83 }
84 
85 } // namespace LIBC_NAMESPACE
86