[libc] Make the opcode parameter a compile time constant

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.

Reviewed By: JonChesterfield

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

4 files changed, 10 insertions, 11 deletions

diff --git a/libc/src/__support/OSUtil/gpu/io.cpp b/libc/src/__support/OSUtil/gpu/io.cpp
index 62eedd5113f6..995a97389184 100644
--- a/libc/src/__support/OSUtil/gpu/io.cpp
+++ b/libc/src/__support/OSUtil/gpu/io.cpp
@@ -15,7 +15,7 @@
 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();
 }
diff --git a/libc/src/__support/OSUtil/gpu/quick_exit.cpp b/libc/src/__support/OSUtil/gpu/quick_exit.cpp
index 70f8bfd46647..5cee9010bbf6 100644
--- a/libc/src/__support/OSUtil/gpu/quick_exit.cpp
+++ b/libc/src/__support/OSUtil/gpu/quick_exit.cpp
@@ -17,7 +17,7 @@
 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;
   });
diff --git a/libc/src/__support/RPC/rpc.h b/libc/src/__support/RPC/rpc.h
index 19be921047b2..ef5f704f8c8c 100644
--- a/libc/src/__support/RPC/rpc.h
+++ b/libc/src/__support/RPC/rpc.h
@@ -273,8 +273,8 @@ struct Client : public Process<false> {
   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.
@@ -411,10 +411,9 @@ LIBC_INLINE void Port<T>::recv_n(A alloc) {
 /// 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.
@@ -445,9 +444,9 @@ Client::try_open(uint16_t opcode) {
   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();
   }
diff --git a/libc/test/integration/startup/gpu/rpc_test.cpp b/libc/test/integration/startup/gpu/rpc_test.cpp
index 9dc2214fde41..0f051bd72c97 100644
--- a/libc/test/integration/startup/gpu/rpc_test.cpp
+++ b/libc/test/integration/startup/gpu/rpc_test.cpp
@@ -17,7 +17,7 @@ static void test_add_simple() {
       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;
@@ -32,7 +32,7 @@ static void test_add_simple() {
 
 // 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();
 }