[libc] Make the opcode parameter a compile time constant

[jhuber6]

Summary

Currently the opcode is only valid if it is the same between all of the
ports. This is possible to violate if the opcode is places into a memory
location and then read in a non-uniform manner by the warp / wavefront.
Moving this to a compile time constant makes it impossible to break this
invariant.

[JonChesterfield]

This opcode thing should probably be an enum class

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

 
 #include "src/__support/CPP/string_view.h"
 #include "src/__support/RPC/rpc_client.h"
 #include "src/string/memory_utils/memcpy_implementations.h"
 
 namespace __llvm_libc {
 
 void write_to_stderr(cpp::string_view msg) {
-  rpc::Client::Port port = rpc::client.open(rpc::PRINT_TO_STDERR);
+  rpc::Client::Port port = rpc::client.open<rpc::PRINT_TO_STDERR>();
   port.send_n(msg.data(), msg.size());
   port.close();
 }
 
 } // namespace __llvm_libc

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

 #include "quick_exit.h"
 
 #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);
+  rpc::Client::Port port = rpc::client.open<rpc::EXIT>();
   port.send([&](rpc::Buffer *buffer) {
     reinterpret_cast<uint32_t *>(buffer->data)[0] = status;
   });
   port.close();
 
 #if defined(LIBC_TARGET_ARCH_IS_NVPTX)
   asm("exit;" ::: "memory");
 #elif defined(LIBC_TARGET_ARCH_IS_AMDGPU)

libc/src/__support/RPC/rpc.h

 /// 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;
 
   using Port = rpc::Port<false>;
-  LIBC_INLINE cpp::optional<Port> try_open(uint16_t opcode);
-  LIBC_INLINE Port open(uint16_t opcode);
+  template <uint16_t opcode> LIBC_INLINE cpp::optional<Port> try_open();
+  template <uint16_t opcode> LIBC_INLINE Port open();
 };
 
 /// The RPC server used to respond to the client.
 struct Server : public Process<true> {
   LIBC_INLINE Server() = default;
   LIBC_INLINE Server(const Server &) = delete;
   LIBC_INLINE Server &operator=(const Server &) = delete;
   LIBC_INLINE ~Server() = default;
if constexpr (is_process_gpu()) {
     return;
   }
 }
 
 /// Attempts to open a port to use as the client. The client can only open a
 /// port if we find an index that is in a valid sending state. That is, there
 /// are send operations pending that haven't been serviced on this port. Each
 /// port instance uses an associated \p opcode to tell the server what to do.
-/// Opening a port is only valid if the `opcode` is the sam accross every
-/// participating thread.
+template <uint16_t opcode>
 [[clang::convergent]] LIBC_INLINE cpp::optional<Client::Port>
-Client::try_open(uint16_t opcode) {
+Client::try_open() {
   // Perform a naive linear scan for a port that can be opened to send data.
   for (uint64_t index = 0; index < port_count; ++index) {
     // Attempt to acquire the lock on this index.
     uint64_t lane_mask = gpu::get_lane_mask();
     if (!try_lock(lane_mask, index))
       continue;
 
     // The mailbox state must be read with the lock held.
if (is_first_lane(lane_mask)) {
       get_packet(index).header.mask = lane_mask;
     }
     gpu::sync_lane(lane_mask);
     return Port(*this, lane_mask, index, out);
   }
   return cpp::nullopt;
 }
 
-LIBC_INLINE Client::Port Client::open(uint16_t opcode) {
+template <uint16_t opcode> LIBC_INLINE Client::Port Client::open() {
   for (;;) {
-    if (cpp::optional<Client::Port> p = try_open(opcode))
+    if (cpp::optional<Client::Port> p = try_open<opcode>())
       return cpp::move(p.value());
     sleep_briefly();
   }
 }
 
 /// Attempts to open a port to use as the server. The server can only open a
 /// port if it has a pending receive operation
 [[clang::convergent]] LIBC_INLINE cpp::optional<Server::Port>

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);
+    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.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);
+  rpc::Client::Port port = rpc::client.open<rpc::NOOP>();
   port.send([=](rpc::Buffer *buffer) { buffer->data[0] = 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;
 }