[libc] Fix RPC interface when sending and recieving aribtrary packets

[jhuber6]

Summary

The interface exported by the RPC library allows users to simply send
and recieve fixed sized packets without worrying about the data motion
underneath. However, this was broken in the current implementation. We
can think of the send and recieve implementations in terms of waiting
for ownership of the buffer, using the buffer, and posting ownership to
the other side. Our implementation of recv was incorrect in the
following scenarios.

recv -> send we still own the buffer and should give away ownership
recv -> close The other side is not waiting for data, this will

result in multiple openings of the same port

This patch attempts to fix this with an admittedly hacky fix where we
track if the previous implementation was a recv and post conditionally.

[JonChesterfield]

s/recieve/receive/g but otherwise looks reasonable. Seems a good sign that send and recv now do different things, and I'm glad that you've changed the behaviour without modifying process (as that means I don't have to try to think through the concurrency behaviour)

libc/src/__support/RPC/rpc.h

 namespace rpc {
 
 /// A list of opcodes that we use to invoke certain actions on the server.
 enum Opcode : uint16_t {
   NOOP = 0,
   PRINT_TO_STDERR = 1,
   EXIT = 2,
   TEST_INCREMENT = 3,
+  TEST_INTERFACE = 4,
 };
 
 /// A fixed size channel used to communicate between the RPC client and server.
 struct Buffer {
   uint64_t data[8];
 };
 static_assert(sizeof(Buffer) == 64, "Buffer size mismatch");
 

 };
 
 /// The port provides the interface to communicate between the multiple
 /// processes. A port is conceptually an index into the memory provided by the
 /// underlying process that is guarded by a lock bit.
 template <bool T> struct Port {
   LIBC_INLINE Port(Process<T> &process, uint64_t lane_mask, uint64_t index,
                    uint32_t out)
-      : process(process), lane_mask(lane_mask), index(index), out(out) {}
+      : process(process), lane_mask(lane_mask), index(index), out(out),
+        receive(false) {}
   LIBC_INLINE ~Port() = default;
 
 private:
   LIBC_INLINE Port(const Port &) = delete;
   LIBC_INLINE Port &operator=(const Port &) = delete;
   LIBC_INLINE Port(Port &&) = default;
   LIBC_INLINE Port &operator=(Port &&) = default;
 
public:
   template <typename W> LIBC_INLINE void recv_and_send(W work);
   LIBC_INLINE void send_n(const void *src, uint64_t size);
   template <typename A> LIBC_INLINE void recv_n(A alloc);
 
   LIBC_INLINE uint16_t get_opcode() const {
     return process.get_packet(index).header.opcode;
   }
 
-  LIBC_INLINE void close() { process.unlock(lane_mask, index); }
+  LIBC_INLINE void close() {
+    // If the server last did a receive it needs to exchange ownership before
+    // closing the port.
+    if (receive && T)
+      out = process.invert_outbox(index, out);
+    process.unlock(lane_mask, index);
+  }
 
 private:
   Process<T> &process;
   uint64_t lane_mask;
   uint64_t index;
   uint32_t out;
+  bool receive;
 };
 
 /// The RPC client used to make requests to the server.
 struct Client : public Process<false> {
   LIBC_INLINE Client() = default;
   LIBC_INLINE Client(const Client &) = delete;
   LIBC_INLINE Client &operator=(const Client &) = delete;
   LIBC_INLINE ~Client() = default;
while (Process<T>::buffer_unavailable(in, out)) {
     sleep_briefly();
     in = process.load_inbox(index);
   }
 
   // Apply the \p fill function to initialize the buffer and release the memory.
   process.invoke_rpc(fill, process.get_packet(index));
   atomic_thread_fence(cpp::MemoryOrder::RELEASE);
   out = process.invert_outbox(index, out);
+  receive = false;
 }
 
 /// Applies \p use to the shared buffer and acknowledges the send.
 template <bool T> template <typename U> LIBC_INLINE void Port<T>::recv(U use) {
+  // We only exchange ownership of the buffer during a receive if we are waiting
+  // for a previous receive to finish.
+  if (receive)
+    out = process.invert_outbox(index, out);
+
   uint32_t in = process.load_inbox(index);
 
   // We need to wait until we own the buffer before receiving.
   while (Process<T>::buffer_unavailable(in, out)) {
     sleep_briefly();
     in = process.load_inbox(index);
   }
   atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
 
   // Apply the \p use function to read the memory out of the buffer.
   process.invoke_rpc(use, process.get_packet(index));
-  out = process.invert_outbox(index, out);
+  receive = true;
 }
 
 /// Combines a send and receive into a single function.
 template <bool T>
 template <typename F, typename U>
 LIBC_INLINE void Port<T>::send_and_recv(F fill, U use) {
   send(fill);
   recv(use);

libc/test/integration/startup/gpu/CMakeLists.txt

 )
 
 add_integration_test(
   init_fini_array_test
   SUITE libc-startup-tests
   SRCS
     init_fini_array_test.cpp
 )
+
+add_integration_test(
+  startup_rpc_interface_test
+  SUITE libc-startup-tests
+  SRCS
+   rpc_interface_test.cpp
+)

libc/test/integration/startup/gpu/rpc_interface_test.cpp

+//===-- Loader test to check the RPC interface with the loader ------------===//
+//
+// 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/__support/GPU/utils.h"
+#include "src/__support/RPC/rpc_client.h"
+#include "test/IntegrationTest/test.h"
+
+using namespace __llvm_libc;
+
+// Test to ensure that we can use aribtrary combinations of sends and recieves
+// as long as they are mirrored.
+static void test_interface(bool end_with_send) {
+  uint64_t cnt = 0;
+  rpc::Client::Port port = rpc::client.open<rpc::TEST_INTERFACE>();
+  port.send([&](rpc::Buffer *buffer) { buffer->data[0] = end_with_send; });
+  port.send([&](rpc::Buffer *buffer) { buffer->data[0] = cnt = cnt + 1; });
+  port.recv([&](rpc::Buffer *buffer) { cnt = buffer->data[0]; });
+  port.send([&](rpc::Buffer *buffer) { buffer->data[0] = cnt = cnt + 1; });
+  port.recv([&](rpc::Buffer *buffer) { cnt = buffer->data[0]; });
+  port.send([&](rpc::Buffer *buffer) { buffer->data[0] = cnt = cnt + 1; });
+  port.send([&](rpc::Buffer *buffer) { buffer->data[0] = cnt = cnt + 1; });
+  port.recv([&](rpc::Buffer *buffer) { cnt = buffer->data[0]; });
+  port.recv([&](rpc::Buffer *buffer) { cnt = buffer->data[0]; });
+  if (end_with_send)
+    port.send([&](rpc::Buffer *buffer) { buffer->data[0] = cnt = cnt + 1; });
+  else
+    port.recv([&](rpc::Buffer *buffer) { cnt = buffer->data[0]; });
+  port.close();
+
+  ASSERT_TRUE(cnt == 9 && "Invalid number of increments");
+}
+
+TEST_MAIN(int argc, char **argv, char **envp) {
+  test_interface(true);
+  test_interface(false);
+
+  return 0;
+}

libc/utils/gpu/loader/Server.h

case __llvm_libc::rpc::Opcode::EXIT: {
       break;
     }
     case __llvm_libc::rpc::Opcode::TEST_INCREMENT: {
       port->recv_and_send([](__llvm_libc::rpc::Buffer *buffer) {
         reinterpret_cast<uint64_t *>(buffer->data)[0] += 1;
       });
       break;
     }
+    case __llvm_libc::rpc::Opcode::TEST_INTERFACE: {
+      uint64_t cnt = 0;
+      bool end_with_recv;
+      port->recv([&](__llvm_libc::rpc::Buffer *buffer) {
+        end_with_recv = buffer->data[0];
+      });
+      port->recv(
+          [&](__llvm_libc::rpc::Buffer *buffer) { cnt = buffer->data[0]; });
+      port->send([&](__llvm_libc::rpc::Buffer *buffer) {
+        buffer->data[0] = cnt = cnt + 1;
+      });
+      port->recv(
+          [&](__llvm_libc::rpc::Buffer *buffer) { cnt = buffer->data[0]; });
+      port->send([&](__llvm_libc::rpc::Buffer *buffer) {
+        buffer->data[0] = cnt = cnt + 1;
+      });
+      port->recv(
+          [&](__llvm_libc::rpc::Buffer *buffer) { cnt = buffer->data[0]; });
+      port->recv(
+          [&](__llvm_libc::rpc::Buffer *buffer) { cnt = buffer->data[0]; });
+      port->send([&](__llvm_libc::rpc::Buffer *buffer) {
+        buffer->data[0] = cnt = cnt + 1;
+      });
+      port->send([&](__llvm_libc::rpc::Buffer *buffer) {
+        buffer->data[0] = cnt = cnt + 1;
+      });
+      if (end_with_recv)
+        port->recv(
+            [&](__llvm_libc::rpc::Buffer *buffer) { cnt = buffer->data[0]; });
+      else
+        port->send([&](__llvm_libc::rpc::Buffer *buffer) {
+          buffer->data[0] = cnt = cnt + 1;
+        });
+      break;
+    }
     default:
       port->recv([](__llvm_libc::rpc::Buffer *buffer) {});
     }
     port->close();
   }
 }
 #endif