[nativert] port semaphore to c10 util (#153504)

Summary:
nativert RFC: https://github.com/zhxchen17/rfcs/blob/master/RFC-0043-torch-native-runtime.md

To land the runtime into PyTorch core, we will gradually land logical parts of the code into the Github issue and get each piece properly reviewed.

This diff adds a simple semaphore interface into c10 until c++20 where we get counting_semaphore

gonna need a oss build export to take a look at this...

Test Plan: CI

Differential Revision: D73882656

Pull Request resolved: https://github.com/pytorch/pytorch/pull/153504
Approved by: https://github.com/zhxchen17

WORKSPACE

@@ -184,6 +184,12 @@ new_local_repository(
     path = "third_party/nlohmann",
 )
 
+new_local_repository(
+    name = "moodycamel",
+    build_file = "//third_party:moodycamel.BUILD",
+    path = "third_party/concurrentqueue",
+)
+
 new_local_repository(
     name = "tensorpipe",
     build_file = "//third_party:tensorpipe.BUILD",

buckbuild.bzl

@@ -178,6 +178,7 @@ THIRD_PARTY_LIBS = {
     "psimd": ["//xplat/third-party/psimd:psimd", "//third_party:psimd"],
     "pthreadpool": ["//xplat/third-party/pthreadpool:pthreadpool", "//third_party:pthreadpool"],
     "pthreadpool_header": ["//xplat/third-party/pthreadpool:pthreadpool_header", "//third_party:pthreadpool_header"],
+    "moodycamel": ["//third-party/moodycamel:moodycamel", "//third_party:moodycamel"],
     "pyyaml": ["//third-party/pypi/pyyaml:pyyaml", "//third_party:pyyaml"],
     "rt": ["//xplat/third-party/linker_lib:rt", "//third_party:rt"],
     "ruy": ["//third-party/ruy:ruy_xplat_lib", "//third_party:ruy_lib"],

c10/BUCK.oss

@@ -15,6 +15,7 @@ cxx_library(
         "//third_party:cpuinfo",
         "//third_party:fmt",
         "//third_party:glog",
+        "//third_party:moodycamel",
     ],
     exported_deps = [],
     compiler_flags = [

c10/CMakeLists.txt

@@ -96,6 +96,7 @@ if(NOT BUILD_LIBTORCHLESS)
   endif()
   target_link_libraries(c10 PRIVATE fmt::fmt-header-only)
   target_link_libraries(c10 PRIVATE nlohmann)
+  target_link_libraries(c10 PRIVATE moodycamel)
 
   if(C10_USE_NUMA)
     message(STATUS "NUMA paths:")

c10/test/util/Semaphore_test.cpp

@@ -0,0 +1,35 @@
+#include <c10/util/Semaphore.h>
+#include <c10/util/irange.h>
+#include <gtest/gtest.h>
+
+#include <thread>
+
+using namespace ::testing;
+
+TEST(SemaphoreTest, TestConcurrency) {
+  auto num_threads = std::thread::hardware_concurrency();
+  auto num_incr = 10000;
+
+  c10::Semaphore sem;
+
+  std::vector<std::thread> threads;
+  for ([[maybe_unused]] const auto _ : c10::irange(num_threads)) {
+    threads.emplace_back([num_incr = num_incr, &sem]() {
+      for ([[maybe_unused]] const auto _ : c10::irange(num_incr)) {
+        sem.release();
+      }
+      for ([[maybe_unused]] const auto _ : c10::irange(num_incr)) {
+        sem.acquire();
+      }
+      sem.release(num_incr);
+      for ([[maybe_unused]] const auto _ : c10::irange(num_incr)) {
+        sem.acquire();
+      }
+    });
+  }
+
+  std::for_each(
+      threads.begin(), threads.end(), [](std::thread& t) { t.join(); });
+
+  EXPECT_FALSE(sem.tryAcquire());
+}

c10/util/Semaphore.h

@@ -0,0 +1,71 @@
+#pragma once
+
+#include <version>
+
+/*
+  a simple semaphore interface.
+*/
+
+// note: __cpp_lib_semaphore will not be defined in some apple platforms
+// even if >= C++20.
+#if __has_include(<semaphore>) && defined(__cpp_lib_semaphore) && __cpp_lib_semaphore >= 201907L
+#define C10_SEMAPHORE_USE_STL
+#endif
+
+#ifdef C10_SEMAPHORE_USE_STL
+#include <semaphore>
+#else
+// To use moodycamel semaphore, we need to include the header file
+// for concurrentqueue first. Hiding implementation detail here.
+#ifdef BLOCK_SIZE
+#pragma push_macro("BLOCK_SIZE")
+#undef BLOCK_SIZE
+#include <moodycamel/concurrentqueue.h> // @manual
+#pragma pop_macro("BLOCK_SIZE")
+#else
+#include <moodycamel/concurrentqueue.h> // @manual
+#endif
+
+#include <moodycamel/lightweightsemaphore.h> // @manual
+#endif
+
+namespace c10 {
+
+class Semaphore {
+ public:
+  Semaphore(int32_t initial_count = 0) : impl_(initial_count) {}
+
+  void release(int32_t n = 1) {
+#ifdef C10_SEMAPHORE_USE_STL
+    impl_.release(n);
+#else
+    impl_.signal(n);
+#endif
+  }
+
+  void acquire() {
+#ifdef C10_SEMAPHORE_USE_STL
+    impl_.acquire();
+#else
+    impl_.wait();
+#endif
+  }
+
+  bool tryAcquire() {
+#ifdef C10_SEMAPHORE_USE_STL
+    return impl_.try_acquire();
+#else
+    return impl_.tryWait();
+#endif
+  }
+
+ private:
+#ifdef C10_SEMAPHORE_USE_STL
+  std::counting_semaphore<> impl_;
+#else
+  moodycamel::LightweightSemaphore impl_;
+#endif
+};
+} // namespace c10
+
+#undef C10_SEMAPHORE_USE_STL

c10/util/build.bzl

@@ -36,6 +36,7 @@ def define_targets(rules):
             ":bit_cast",
             "//c10/macros",
             "@fmt",
+            "@moodycamel//:moodycamel",
         ] + rules.select({
             "//c10:using_gflags": ["@com_github_gflags_gflags//:gflags"],
             "//conditions:default": [],

cmake/Dependencies.cmake

@@ -1154,6 +1154,7 @@ if(USE_DISTRIBUTED AND USE_TENSORPIPE)
 
     list(APPEND Caffe2_DEPENDENCY_LIBS tensorpipe)
     list(APPEND Caffe2_DEPENDENCY_LIBS nlohmann)
+    list(APPEND Caffe2_DEPENDENCY_LIBS moodycamel)
     if(USE_CUDA)
       list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS tensorpipe_cuda)
     elseif(USE_ROCM)
@@ -1713,3 +1714,7 @@ target_include_directories(httplib SYSTEM INTERFACE ${PROJECT_SOURCE_DIR}/third_
 # Include nlohmann-json
 add_library(nlohmann INTERFACE IMPORTED)
 include_directories(nlohmann SYSTEM INTERFACE ${PROJECT_SOURCE_DIR}/third_party/nlohmann/include)
+
+# Include moodycamel
+add_library(moodycamel INTERFACE IMPORTED)
+include_directories(moodycamel SYSTEM INTERFACE ${PROJECT_SOURCE_DIR}/third_party/concurrentqueue)

third_party/BUCK.oss

@@ -75,6 +75,15 @@ cxx_library(
     visibility = ["PUBLIC"],
 )
 
+cxx_library(
+    name = "moodycamel",
+    raw_headers = glob([
+        "concurrentqueue/**/*.h",
+    ]),
+    reexport_all_header_dependencies = True,
+    visibility = ["PUBLIC"],
+)
+
 cxx_library(
     name = "pocketfft_header",
     header_namespace = "",

third_party/moodycamel.BUILD

@@ -0,0 +1,7 @@
+load("@rules_cc//cc:defs.bzl", "cc_library")
+
+cc_library(
+    name = "moodycamel",
+    hdrs = glob(["**/*.h"]),
+    visibility = ["//visibility:public"],
+)

torch/CMakeLists.txt

@@ -84,6 +84,7 @@ set(TORCH_PYTHON_LINK_LIBRARIES
     opentelemetry::api
     httplib
     nlohmann
+    moodycamel
     shm
     fmt::fmt-header-only
     ATEN_CPU_FILES_GEN_LIB)