llvm-project

1
//===-- GPU implementation of fprintf -------------------------------------===//
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 "rpc_fprintf.h"
10

11
#include "src/__support/CPP/string_view.h"
12
#include "src/__support/GPU/utils.h"
13
#include "src/__support/RPC/rpc_client.h"
14
#include "src/__support/common.h"
15
#include "src/stdio/gpu/file.h"
16

17
namespace LIBC_NAMESPACE {
18

19
template <uint16_t opcode>
20
int fprintf_impl(::FILE *__restrict file, const char *__restrict format,
21
                 size_t format_size, void *args, size_t args_size) {
22
  uint64_t mask = gpu::get_lane_mask();
23
  rpc::Client::Port port = rpc::client.open<opcode>();
24

25
  if constexpr (opcode == RPC_PRINTF_TO_STREAM) {
26
    port.send([&](rpc::Buffer *buffer) {
27
      buffer->data[0] = reinterpret_cast<uintptr_t>(file);
28
    });
29
  }
30

31
  port.send_n(format, format_size);
32
  port.send_n(args, args_size);
33

34
  uint32_t ret = 0;
35
  for (;;) {
36
    const char *str = nullptr;
37
    port.recv([&](rpc::Buffer *buffer) {
38
      ret = static_cast<uint32_t>(buffer->data[0]);
39
      str = reinterpret_cast<const char *>(buffer->data[1]);
40
    });
41
    // If any lanes have a string argument it needs to be copied back.
42
    if (!gpu::ballot(mask, str))
43
      break;
44

45
    uint64_t size = str ? internal::string_length(str) + 1 : 0;
46
    port.send_n(str, size);
47
  }
48

49
  port.close();
50
  return ret;
51
}
52

53
// TODO: This is a stand-in function that uses a struct pointer and size in
54
// place of varargs. Once varargs support is added we will use that to
55
// implement the real version.
56
LLVM_LIBC_FUNCTION(int, rpc_fprintf,
57
                   (::FILE *__restrict stream, const char *__restrict format,
58
                    void *args, size_t size)) {
59
  cpp::string_view str(format);
60
  if (stream == stdout)
61
    return fprintf_impl<RPC_PRINTF_TO_STDOUT>(stream, format, str.size() + 1,
62
                                              args, size);
63
  else if (stream == stderr)
64
    return fprintf_impl<RPC_PRINTF_TO_STDERR>(stream, format, str.size() + 1,
65
                                              args, size);
66
  else
67
    return fprintf_impl<RPC_PRINTF_TO_STREAM>(stream, format, str.size() + 1,
68
                                              args, size);
69
}
70

71
} // namespace LIBC_NAMESPACE
72