[libc] Support suspending threads during RPC spin loops

[jhuber6]

Summary

The RPC interface relies on waiting on atomic signals to coordinate
which side of the protocol is in control of the shared buffer. The GPU client
supports briefly suspending the executing thread group. This is used by the
thread scheduler to identify which thread groups can be switched out so that
others may execute. This allows us to ensure that other threads get a chance
to make forward progress while these threads wait on the atomic signal.

This is currently only relevant on the client-side. We could use an
alternative implementation on the server that uses the standard
nanosleep on supported hosts.

[tianshilei1992]

It looks like the two platforms are using different sleep duration? Other than that, LGTM.

[jhuber6]

In D147238#4233699, @tianshilei1992 wrote:

It looks like the two platforms are using different sleep duration? Other than that, LGTM.

The documentation is a little different. The AMDGCN implementation states that a value of 2 will sleep between 64 and 128 cycles so assuming a 2 GHz clock it'll sleep ~60 ns in the worst case. These numbers are just guesses, we could probably refine them. I think PCI(e) atomic access is in the order of microseconds.

[JonChesterfield]

Thanks. Couple of nits above.

libc/src/__support/RPC/rpc.h

105–106

this probably gets rotated, might be better written do_while

sleep before or after load? the fence above probably takes time

libc/src/__support/RPC/rpc_util.h

21

probably worth leaving a comment here about how the magic numbers were chosen

libc/src/__support/RPC/CMakeLists.txt

 add_header_library(
   rpc
   HDRS
     rpc.h
+    rpc_util.h
   DEPENDS
     libc.src.__support.common
     libc.src.__support.CPP.atomic
 )
 
 add_object_library(
   rpc_client
   SRCS
     rpc_client.cpp
   HDRS
     rpc_client.h
   DEPENDS
     .rpc
 )

libc/src/__support/RPC/rpc.h

 // signals to indicate which side, the client or the server is in ownership of
 // the buffer.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIBC_SRC_SUPPORT_RPC_RPC_H
 #define LLVM_LIBC_SRC_SUPPORT_RPC_RPC_H
 
+#include "rpc_util.h"
 #include "src/__support/CPP/atomic.h"
 
 #include <stdint.h>
 
 namespace __llvm_libc {
 namespace rpc {
 
 /// A list of opcodes that we use to invoke certain actions on the server. We
template <typename F, typename U> LIBC_INLINE void Client::run(F fill, U use) {
   if (!in & !out) {
     fill(buffer);
     atomic_thread_fence(cpp::MemoryOrder::RELEASE);
     outbox->store(1, cpp::MemoryOrder::RELAXED);
     out = 1;
   }
   // Wait for the server to work on the buffer and respond.
   if (!in & out) {
-    while (!in)
+    while (!in) {
+      sleep_briefly();
       in = inbox->load(cpp::MemoryOrder::RELAXED);
+    }
     atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
   }
   // Apply \p use to the buffer and signal the server.
   if (in & out) {
     use(buffer);
     atomic_thread_fence(cpp::MemoryOrder::RELEASE);
     outbox->store(0, cpp::MemoryOrder::RELAXED);
     out = 0;
   }
   // Wait for the server to signal the end of the protocol.
   if (in & !out) {
-    while (in)
+    while (in) {
+      sleep_briefly();
       in = inbox->load(cpp::MemoryOrder::RELAXED);
+    }
     atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
   }
 }
 
 /// Run the RPC server protocol to communicate with the client. This is
 /// non-blocking and only checks the server a single time. We perform the
 /// following high level actions to complete a communication:
 ///   - Query if the inbox is 1 and exit if there is no work to do.

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/macros/attributes.h"
+#include "src/__support/macros/properties/architectures.h"
+
+namespace __llvm_libc {
+namespace rpc {
+
+/// Suspend the thread briefly to assist the thread scheduler during busy loops.
+LIBC_INLINE void sleep_briefly() {
+#if defined(LIBC_TARGET_ARCH_IS_NVPTX) && __CUDA_ARCH__ >= 700
+  asm("nanosleep.u32 64;" ::: "memory");
+#elif defined(LIBC_TARGET_ARCH_IS_AMDGPU)
+  __builtin_amdgcn_s_sleep(2);
+#else
+  // Simply do nothing if sleeping isn't supported on this platform.
+#endif
+}
+
+} // namespace rpc
+} // namespace __llvm_libc
+
+#endif