[libc-commits] [libc] dff3909 - [libc] Support suspending threads during RPC spin loops

[Joseph Huber via libc-commits]

Author: Joseph Huber
Date: 2023-03-30T11:40:24-05:00
New Revision: dff3909c3ed935138608a02fcb940589a95f1a14

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

LOG: [libc] Support suspending threads during RPC spin loops

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.

Reviewed By: JonChesterfield, tianshilei1992

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

Added: 
    libc/src/__support/RPC/rpc_util.h

Modified: 
    libc/src/__support/RPC/CMakeLists.txt
    libc/src/__support/RPC/rpc.h

Removed: 
    


################################################################################
diff  --git a/libc/src/__support/RPC/CMakeLists.txt b/libc/src/__support/RPC/CMakeLists.txt
index f5837628c2050..c1a971b54b365 100644
--- a/libc/src/__support/RPC/CMakeLists.txt
+++ b/libc/src/__support/RPC/CMakeLists.txt
@@ -2,6 +2,7 @@ add_header_library(
   rpc
   HDRS
     rpc.h
+    rpc_util.h
   DEPENDS
     libc.src.__support.common
     libc.src.__support.CPP.atomic

diff  --git a/libc/src/__support/RPC/rpc.h b/libc/src/__support/RPC/rpc.h
index 561e90c156034..43660fd8e1c9c 100644
--- a/libc/src/__support/RPC/rpc.h
+++ b/libc/src/__support/RPC/rpc.h
@@ -18,6 +18,7 @@
 #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>
@@ -101,8 +102,10 @@ template <typename F, typename U> LIBC_INLINE void Client::run(F fill, U use) {
   }
   // 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.
@@ -114,8 +117,10 @@ template <typename F, typename U> LIBC_INLINE void Client::run(F fill, U use) {
   }
   // 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);
   }
 }

diff  --git a/libc/src/__support/RPC/rpc_util.h b/libc/src/__support/RPC/rpc_util.h
new file mode 100644
index 0000000000000..53a993292fb35
--- /dev/null
+++ b/libc/src/__support/RPC/rpc_util.h
@@ -0,0 +1,32 @@
+//===-- 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