[libc] Simplify RPC send and recv callback interface

[jhuber6]

Summary

We use lambda callbacks as the main way to interface with the send and
recv functions at the core of the RPC interface. This provided the
user with a pointer to the underlying data buffer to use. However, this
was somewhat annoying to use as it required constant bit casts. This
patch uses template magic to infer the desired type and automatically
cast to what the user passes in. Some static checks ensure that this is
actually legal. This will allow for functions to more easily describe
structs that match the expected argument format.

[michaelrj]

LGTM

libc/src/__support/CPP/type_traits.h

60

nit: I think these should go right after is_pointer

libc/src/__support/CPP/type_traits.h

 template <typename T> using remove_cv_t = typename remove_cv<T>::type;
 
 template <typename T> struct remove_reference : type_identity<T> {};
 template <typename T> struct remove_reference<T &> : type_identity<T> {};
 template <typename T> struct remove_reference<T &&> : type_identity<T> {};
 template <typename T>
 using remove_reference_t = typename remove_reference<T>::type;
 
+template <typename T> struct remove_pointer : type_identity<T> {};
+template <typename T> struct remove_pointer<T *> : type_identity<T> {};
+template <typename T> struct remove_pointer<const T *> : type_identity<T> {};
+template <typename T> using remove_pointer_t = typename remove_pointer<T>::type;
+
 template <typename T> struct add_rvalue_reference : type_identity<T &&> {};
 
 template <typename T> struct remove_cvref {
   using type = remove_cv_t<remove_reference_t<T>>;
 };
 template <typename T> using remove_cvref_t = typename remove_cvref<T>::type;
 
 namespace details {

libc/src/__support/OSUtil/gpu/quick_exit.cpp

 
 #include "src/__support/RPC/rpc_client.h"
 #include "src/__support/macros/properties/architectures.h"
 
 namespace __llvm_libc {
 
 void quick_exit(int status) {
   rpc::Client::Port port = rpc::client.open<rpc::EXIT>();
-  port.send([&](rpc::Buffer *buffer) {
-    reinterpret_cast<uint32_t *>(buffer->data)[0] = status;
-  });
+  port.send([=](int *buffer) { *buffer = status; });
   port.close();
 
 #if defined(LIBC_TARGET_ARCH_IS_NVPTX)
   asm("exit;" ::: "memory");
 #elif defined(LIBC_TARGET_ARCH_IS_AMDGPU)
   // This will terminate the entire wavefront, may not be valid with divergent
   // work items.
   asm("s_endpgm" ::: "memory");
 #endif
   __builtin_unreachable();
 }
 
 } // namespace __llvm_libc
 
 #endif // LLVM_LIBC_SRC_SUPPORT_OSUTIL_GPU_QUICK_EXIT_H

libc/src/__support/RPC/rpc.h

[[clang::convergent]] LIBC_INLINE void unlock(uint64_t lane_mask,
     // then an unrelated warp claiming the lock, then a second thread in this
     // warp dropping the lock again.
     uint32_t and_mask = ~(rpc::is_first_lane(lane_mask) ? 1 : 0);
     lock[index].fetch_and(and_mask, cpp::MemoryOrder::RELAXED);
     gpu::sync_lane(lane_mask);
   }
 
   /// Invokes a function accross every active buffer across the total lane size.
-  LIBC_INLINE void invoke_rpc(cpp::function<void(Buffer *)> fn,
-                              Packet &packet) {
+  template <typename T>
+  LIBC_INLINE void invoke_rpc(cpp::function<void(T)> fn, Packet &packet) {
+    static_assert(cpp::is_pointer_v<T>, "Argument must be a pointer");
+    static_assert(sizeof(cpp::remove_pointer_t<T>) <= sizeof(Buffer::data),
+                  "Argument does not fit inside the packet");
     if constexpr (is_process_gpu()) {
-      fn(&packet.payload.slot[gpu::get_lane_id()]);
+      fn(reinterpret_cast<T>(packet.payload.slot[gpu::get_lane_id()].data));
     } else {
       for (uint32_t i = 0; i < lane_size; i += gpu::get_lane_size())
         if (packet.header.mask & 1ul << i)
-          fn(&packet.payload.slot[i]);
+          fn(reinterpret_cast<T>(packet.payload.slot[i].data));
     }
   }
 
   /// Alternate version that also provides the index of the current lane.
-  LIBC_INLINE void invoke_rpc(cpp::function<void(Buffer *, uint32_t)> fn,
+  template <typename T>
+  LIBC_INLINE void invoke_rpc(cpp::function<void(T, uint32_t)> fn,
                               Packet &packet) {
+    static_assert(cpp::is_pointer_v<T>, "Argument must be a pointer");
+    static_assert(sizeof(cpp::remove_pointer_t<T>) <= sizeof(Buffer::data),
+                  "Argument does not fit inside the packet");
     if constexpr (is_process_gpu()) {
-      fn(&packet.payload.slot[gpu::get_lane_id()], gpu::get_lane_id());
+      fn(reinterpret_cast<T>(packet.payload.slot[gpu::get_lane_id()].data),
+         gpu::get_lane_id());
     } else {
       for (uint32_t i = 0; i < lane_size; i += gpu::get_lane_size())
         if (packet.header.mask & 1ul << i)
-          fn(&packet.payload.slot[i], i);
+          fn(reinterpret_cast<T>(packet.payload.slot[i].data), i);
     }
   }
 
   /// Number of bytes allocated for mailbox or buffer
   LIBC_INLINE static uint64_t memory_allocated_mailbox(uint64_t port_count) {
     return port_count * sizeof(cpp::Atomic<uint32_t>);
   }
 
private:
   LIBC_INLINE Port(Port &&) = default;
   LIBC_INLINE Port &operator=(Port &&) = default;
 
   friend struct Client;
   friend struct Server;
   friend class cpp::optional<Port<T>>;
 
 public:
-  template <typename U> LIBC_INLINE void recv(U use);
-  template <typename F> LIBC_INLINE void send(F fill);
+  template <typename U> LIBC_INLINE void recv(U &&use);
+  template <typename F> LIBC_INLINE void send(F &&fill);
   template <typename F, typename U>
   LIBC_INLINE void send_and_recv(F fill, U use);
   template <typename W> LIBC_INLINE void recv_and_send(W work);
   LIBC_INLINE void send_n(const void *const *src, uint64_t *size);
   LIBC_INLINE void send_n(const void *src, uint64_t size);
   template <typename A>
   LIBC_INLINE void recv_n(void **dst, uint64_t *size, A &&alloc);
 
struct Server : public Process<true> {
   LIBC_INLINE ~Server() = default;
 
   using Port = rpc::Port<true>;
   LIBC_INLINE cpp::optional<Port> try_open();
   LIBC_INLINE Port open();
 };
 
 /// Applies \p fill to the shared buffer and initiates a send operation.
-template <bool T> template <typename F> LIBC_INLINE void Port<T>::send(F fill) {
+template <bool T>
+template <typename F>
+LIBC_INLINE void Port<T>::send(F &&fill) {
   uint32_t in = process.load_inbox(index);
 
   // We need to wait until we own the buffer before sending.
   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));
+  process.template invoke_rpc<first_arg_t<F>>(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) {
+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));
+  process.template invoke_rpc<first_arg_t<U>>(use, process.get_packet(index));
   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);

libc/src/__support/RPC/rpc_util.h

 //===-- Shared memory RPC client / server utilities -------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIBC_SRC_SUPPORT_RPC_RPC_UTILS_H
 #define LLVM_LIBC_SRC_SUPPORT_RPC_RPC_UTILS_H
 
+#include "src/__support/CPP/type_traits.h"
 #include "src/__support/GPU/utils.h"
 #include "src/__support/macros/attributes.h"
 #include "src/__support/macros/properties/architectures.h"
 
 namespace __llvm_libc {
 namespace rpc {
 
 /// Maximum amount of data a single lane can use.
template <typename V> LIBC_INLINE V &lane_value(V *val, uint32_t id) {
   return val[id];
 }
 
 /// Helper to get the maximum value.
 template <typename T> LIBC_INLINE const T &max(const T &x, const T &y) {
   return x < y ? y : x;
 }
 
+namespace internal {
+template <typename F, typename Ret, typename Arg, typename... Args>
+Arg first_arg_helper(Ret (F::*)(Arg, Args...));
+
+template <typename F, typename Ret, typename Arg, typename... Args>
+Arg first_arg_helper(Ret (F::*)(Arg, Args...) const);
+} // namespace internal
+
+template <typename F> struct first_arg {
+  using type = decltype(internal::first_arg_helper(
+      &cpp::remove_reference<F>::type::operator()));
+};
+template <typename F> using first_arg_t = typename first_arg<F>::type;
+
 } // namespace rpc
 } // namespace __llvm_libc
 
 #endif

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

 using namespace __llvm_libc;
 
 static void test_add_simple() {
   uint32_t num_additions =
       10 + 10 * gpu::get_thread_id() + 10 * gpu::get_block_id();
   uint64_t cnt = 0;
   for (uint32_t i = 0; i < num_additions; ++i) {
     rpc::Client::Port port = rpc::client.open<rpc::TEST_INCREMENT>();
-    port.send_and_recv(
-        [=](rpc::Buffer *buffer) {
-          reinterpret_cast<uint64_t *>(buffer->data)[0] = cnt;
-        },
-        [&](rpc::Buffer *buffer) {
-          cnt = reinterpret_cast<uint64_t *>(buffer->data)[0];
-        });
+    port.send_and_recv([=](uint64_t *buffer) { *buffer = cnt; },
+                       [&](uint64_t *buffer) { cnt = *buffer; });
     port.close();
   }
   ASSERT_TRUE(cnt == num_additions && "Incorrect sum");
 }
 
 // Test to ensure that the RPC mechanism doesn't hang on divergence.
 static void test_noop(uint8_t data) {
   rpc::Client::Port port = rpc::client.open<rpc::NOOP>();
-  port.send([=](rpc::Buffer *buffer) { buffer->data[0] = data; });
+  port.send([=](uint8_t *buffer) { *buffer = data; });
   port.close();
 }
 
 TEST_MAIN(int argc, char **argv, char **envp) {
   test_add_simple();
 
   if (gpu::get_thread_id() % 2)
     test_noop(1);
   else
     test_noop(2);
 
   return 0;
 }

libc/utils/gpu/loader/Server.h

case rpc::Opcode::PRINT_TO_STDERR: {
         if (strs[i]) {
           fwrite(strs[i], sizes[i], 1, stderr);
           delete[] reinterpret_cast<uint8_t *>(strs[i]);
         }
       }
       break;
     }
     case rpc::Opcode::EXIT: {
-      port->recv([](rpc::Buffer *buffer) {
-        exit(reinterpret_cast<uint32_t *>(buffer->data)[0]);
-      });
+      port->recv([](int *status) { exit(*status); });
       break;
     }
     case rpc::Opcode::TEST_INCREMENT: {
-      port->recv_and_send([](rpc::Buffer *buffer) {
-        reinterpret_cast<uint64_t *>(buffer->data)[0] += 1;
-      });
+      port->recv_and_send([](uint64_t *cnt) { *cnt += 1; });
       break;
     }
     case rpc::Opcode::TEST_INTERFACE: {
       uint64_t cnt = 0;
       bool end_with_recv;
-      port->recv([&](rpc::Buffer *buffer) { end_with_recv = buffer->data[0]; });
-      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->recv([&](rpc::Buffer *buffer) { cnt = buffer->data[0]; });
-      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([&](uint64_t *buffer) { end_with_recv = *buffer; });
+      port->recv([&](uint64_t *buffer) { cnt = *buffer; });
+      port->send([&](uint64_t *buffer) { *buffer = cnt = cnt + 1; });
+      port->recv([&](uint64_t *buffer) { cnt = *buffer; });
+      port->send([&](uint64_t *buffer) { *buffer = cnt = cnt + 1; });
+      port->recv([&](uint64_t *buffer) { cnt = *buffer; });
+      port->recv([&](uint64_t *buffer) { cnt = *buffer; });
+      port->send([&](uint64_t *buffer) { *buffer = cnt = cnt + 1; });
+      port->send([&](uint64_t *buffer) { *buffer = cnt = cnt + 1; });
       if (end_with_recv)
-        port->recv([&](rpc::Buffer *buffer) { cnt = buffer->data[0]; });
+        port->recv([&](uint64_t *buffer) { cnt = *buffer; });
       else
-        port->send(
-            [&](rpc::Buffer *buffer) { buffer->data[0] = cnt = cnt + 1; });
+        port->send([&](uint64_t *buffer) { *buffer = cnt = cnt + 1; });
       break;
     }
     case rpc::Opcode::TEST_STREAM: {
       uint64_t sizes[rpc::MAX_LANE_SIZE] = {0};
       void *dst[rpc::MAX_LANE_SIZE] = {nullptr};
       port->recv_n(dst, sizes, [](uint64_t size) { return new char[size]; });
       port->send_n(dst, sizes);
       for (uint64_t i = 0; i < rpc::MAX_LANE_SIZE; ++i) {