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

Joseph Huber via libc-commits libc-commits at lists.llvm.org
Wed May 10 16:51:53 PDT 2023


Author: Joseph Huber
Date: 2023-05-10T18:51:38-05:00
New Revision: c8c19e1c31452c93acd3b9c4db5d9ea1c25a4582

URL: https://github.com/llvm/llvm-project/commit/c8c19e1c31452c93acd3b9c4db5d9ea1c25a4582
DIFF: https://github.com/llvm/llvm-project/commit/c8c19e1c31452c93acd3b9c4db5d9ea1c25a4582.diff

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

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.

Reviewed By: JonChesterfield

Differential Revision: https://reviews.llvm.org/D150327

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

Modified: 
    libc/src/__support/RPC/rpc.h
    libc/test/integration/startup/gpu/CMakeLists.txt
    libc/utils/gpu/loader/Server.h

Removed: 
    


################################################################################
diff  --git a/libc/src/__support/RPC/rpc.h b/libc/src/__support/RPC/rpc.h
index d448e190214ff..1285f2b7cd503 100644
--- a/libc/src/__support/RPC/rpc.h
+++ b/libc/src/__support/RPC/rpc.h
@@ -36,6 +36,7 @@ enum Opcode : uint16_t {
   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.
@@ -252,7 +253,8 @@ template <bool InvertInbox> struct Process {
 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:
@@ -278,13 +280,20 @@ template <bool T> struct Port {
     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.
@@ -325,10 +334,16 @@ template <bool T> template <typename F> LIBC_INLINE void Port<T>::send(F fill) {
   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.
@@ -340,7 +355,7 @@ template <bool T> template <typename U> LIBC_INLINE void Port<T>::recv(U use) {
 
   // 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.

diff  --git a/libc/test/integration/startup/gpu/CMakeLists.txt b/libc/test/integration/startup/gpu/CMakeLists.txt
index d2028cc941f05..12ff49c46dbe4 100644
--- a/libc/test/integration/startup/gpu/CMakeLists.txt
+++ b/libc/test/integration/startup/gpu/CMakeLists.txt
@@ -36,3 +36,10 @@ add_integration_test(
   SRCS
     init_fini_array_test.cpp
 )
+
+add_integration_test(
+  startup_rpc_interface_test
+  SUITE libc-startup-tests
+  SRCS
+   rpc_interface_test.cpp
+)

diff  --git a/libc/test/integration/startup/gpu/rpc_interface_test.cpp b/libc/test/integration/startup/gpu/rpc_interface_test.cpp
new file mode 100644
index 0000000000000..b4b03ead31c10
--- /dev/null
+++ b/libc/test/integration/startup/gpu/rpc_interface_test.cpp
@@ -0,0 +1,43 @@
+//===-- 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;
+}

diff  --git a/libc/utils/gpu/loader/Server.h b/libc/utils/gpu/loader/Server.h
index 89ef712e85963..c79277dd8fdca 100644
--- a/libc/utils/gpu/loader/Server.h
+++ b/libc/utils/gpu/loader/Server.h
@@ -57,6 +57,41 @@ void handle_server() {
       });
       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) {});
     }


        


More information about the libc-commits mailing list