[libc] Fix and simplify the implementation of 'fread' on the GPU (#66948)

Summary:
Previously, the `fread` operation was wrong in cases when we read less
data than was requested. That is, if we tried to read N bytes while the
file was in EOF, it would still copy N bytes of garbage. This is fixed
by only copying over the sizes we got from locally opening it rather
than just using the provided size.

Additionally, this patch simplifies the interface. The output functions
have special variants for writing to stdout / stderr. This is primarily
an optimization for these common cases so we can avoid sending the
stream as an argument which has a high delay. Because for input, we
already need to start with a `send` to tell the server how much data to
read, it costs us nothing to send the file along with it so this is
redundant. Re-use the file encoding scheme from the other
implementations, the one that stores the stream type in the LSBs of the
FILE pointer.

Author: jhuber6

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

@@ -15,17 +15,16 @@ typedef enum : unsigned short {
   RPC_WRITE_TO_STDOUT = 2,
   RPC_WRITE_TO_STDERR = 3,
   RPC_WRITE_TO_STREAM = 4,
-  RPC_READ_FROM_STDIN = 5,
-  RPC_READ_FROM_STREAM = 6,
-  RPC_OPEN_FILE = 7,
-  RPC_CLOSE_FILE = 8,
-  RPC_MALLOC = 9,
-  RPC_FREE = 10,
-  RPC_HOST_CALL = 11,
-  RPC_ABORT = 12,
-  RPC_FEOF = 13,
-  RPC_FERROR = 14,
-  RPC_CLEARERR = 15,
+  RPC_READ_FROM_STREAM = 5,
+  RPC_OPEN_FILE = 6,
+  RPC_CLOSE_FILE = 7,
+  RPC_MALLOC = 8,
+  RPC_FREE = 9,
+  RPC_HOST_CALL = 10,
+  RPC_ABORT = 11,
+  RPC_FEOF = 12,
+  RPC_FERROR = 13,
+  RPC_CLEARERR = 14,
 } rpc_opcode_t;
 
 #endif // __LLVM_LIBC_TYPES_RPC_OPCODE_H__

libc/src/stdio/gpu/file.h

@@ -14,6 +14,39 @@
 namespace __llvm_libc {
 namespace file {
 
+enum Stream {
+  File = 0,
+  Stdin = 1,
+  Stdout = 2,
+  Stderr = 3,
+};
+
+// When copying between the client and server we need to indicate if this is one
+// of the special streams. We do this by enocding the low order bits of the
+// pointer to indicate if we need to use the host's standard stream.
+LIBC_INLINE uintptr_t from_stream(::FILE *f) {
+  if (f == stdin)
+    return reinterpret_cast<uintptr_t>(f) | Stdin;
+  if (f == stdout)
+    return reinterpret_cast<uintptr_t>(f) | Stdout;
+  if (f == stderr)
+    return reinterpret_cast<uintptr_t>(f) | Stderr;
+  return reinterpret_cast<uintptr_t>(f);
+}
+
+// Get the associated stream out of an encoded number.
+LIBC_INLINE ::FILE *to_stream(uintptr_t f) {
+  ::FILE *stream = reinterpret_cast<FILE *>(f & ~0x3ull);
+  Stream type = static_cast<Stream>(f & 0x3ull);
+  if (type == Stdin)
+    return stdin;
+  if (type == Stdout)
+    return stdout;
+  if (type == Stderr)
+    return stderr;
+  return stream;
+}
+
 template <uint16_t opcode>
 LIBC_INLINE uint64_t write_impl(::FILE *file, const void *data, size_t size) {
   uint64_t ret = 0;
@@ -42,15 +75,13 @@ LIBC_INLINE uint64_t write(::FILE *f, const void *data, size_t size) {
     return write_impl<RPC_WRITE_TO_STREAM>(f, data, size);
 }
 
-template <uint16_t opcode>
 LIBC_INLINE uint64_t read_from_stream(::FILE *file, void *buf, size_t size) {
   uint64_t ret = 0;
   uint64_t recv_size;
-  rpc::Client::Port port = rpc::client.open<opcode>();
+  rpc::Client::Port port = rpc::client.open<RPC_READ_FROM_STREAM>();
   port.send([=](rpc::Buffer *buffer) {
     buffer->data[0] = size;
-    if constexpr (opcode == RPC_READ_FROM_STREAM)
-      buffer->data[1] = reinterpret_cast<uintptr_t>(file);
+    buffer->data[1] = from_stream(file);
   });
   port.recv_n(&buf, &recv_size, [&](uint64_t) { return buf; });
   port.recv([&](rpc::Buffer *buffer) { ret = buffer->data[0]; });
@@ -59,43 +90,7 @@ LIBC_INLINE uint64_t read_from_stream(::FILE *file, void *buf, size_t size) {
 }
 
 LIBC_INLINE uint64_t read(::FILE *f, void *data, size_t size) {
-  if (f == stdin)
-    return read_from_stream<RPC_READ_FROM_STDIN>(f, data, size);
-  else
-    return read_from_stream<RPC_READ_FROM_STREAM>(f, data, size);
-}
-
-enum Stream {
-  File = 0,
-  Stdin = 1,
-  Stdout = 2,
-  Stderr = 3,
-};
-
-// When copying between the client and server we need to indicate if this is one
-// of the special streams. We do this by enocding the low order bits of the
-// pointer to indicate if we need to use the host's standard stream.
-LIBC_INLINE uintptr_t from_stream(::FILE *f) {
-  if (f == stdin)
-    return reinterpret_cast<uintptr_t>(f) | Stdin;
-  if (f == stdout)
-    return reinterpret_cast<uintptr_t>(f) | Stdout;
-  if (f == stderr)
-    return reinterpret_cast<uintptr_t>(f) | Stderr;
-  return reinterpret_cast<uintptr_t>(f);
-}
-
-// Get the associated stream out of an encoded number.
-LIBC_INLINE ::FILE *to_stream(uintptr_t f) {
-  ::FILE *stream = reinterpret_cast<FILE *>(f & ~0x3ull);
-  Stream type = static_cast<Stream>(f & 0x3ull);
-  if (type == Stdin)
-    return stdin;
-  if (type == Stdout)
-    return stdout;
-  if (type == Stderr)
-    return stderr;
-  return stream;
+  return read_from_stream(f, data, size);
 }
 
 } // namespace file

libc/utils/gpu/server/rpc_server.cpp

@@ -80,23 +80,18 @@ struct Server {
       });
       break;
     }
-    case RPC_READ_FROM_STREAM:
-    case RPC_READ_FROM_STDIN: {
+    case RPC_READ_FROM_STREAM: {
       uint64_t sizes[lane_size] = {0};
       void *data[lane_size] = {nullptr};
-      uint64_t rets[lane_size] = {0};
       port->recv([&](rpc::Buffer *buffer, uint32_t id) {
-        sizes[id] = buffer->data[0];
-        data[id] = new char[sizes[id]];
-        FILE *file = port->get_opcode() == RPC_READ_FROM_STREAM
-                         ? reinterpret_cast<FILE *>(buffer->data[1])
-                         : stdin;
-        rets[id] = fread(data[id], 1, sizes[id], file);
+        data[id] = new char[buffer->data[0]];
+        sizes[id] = fread(data[id], 1, buffer->data[0],
+                          file::to_stream(buffer->data[1]));
       });
       port->send_n(data, sizes);
       port->send([&](rpc::Buffer *buffer, uint32_t id) {
         delete[] reinterpret_cast<uint8_t *>(data[id]);
-        std::memcpy(buffer->data, &rets[id], sizeof(uint64_t));
+        std::memcpy(buffer->data, &sizes[id], sizeof(uint64_t));
       });
       break;
     }