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