diff --git a/WORKSPACE b/WORKSPACE index 57c014e00d23..6fd5f3c28b3a 100644 --- a/WORKSPACE +++ b/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", diff --git a/buckbuild.bzl b/buckbuild.bzl index 72d003d5879b..135325dc39ae 100644 --- a/buckbuild.bzl +++ b/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"], diff --git a/c10/BUCK.oss b/c10/BUCK.oss index 4453622f8853..4b2cbd049b85 100644 --- a/c10/BUCK.oss +++ b/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 = [ diff --git a/c10/CMakeLists.txt b/c10/CMakeLists.txt index 47bc6dfe565e..f00c662e70e0 100644 --- a/c10/CMakeLists.txt +++ b/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:") diff --git a/c10/test/util/Semaphore_test.cpp b/c10/test/util/Semaphore_test.cpp new file mode 100644 index 000000000000..9647270940d7 --- /dev/null +++ b/c10/test/util/Semaphore_test.cpp @@ -0,0 +1,35 @@ +#include +#include +#include + +#include + +using namespace ::testing; + +TEST(SemaphoreTest, TestConcurrency) { + auto num_threads = std::thread::hardware_concurrency(); + auto num_incr = 10000; + + c10::Semaphore sem; + + std::vector 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()); +} diff --git a/c10/util/Semaphore.h b/c10/util/Semaphore.h new file mode 100644 index 000000000000..041d9abecf51 --- /dev/null +++ b/c10/util/Semaphore.h @@ -0,0 +1,71 @@ +#pragma once + +#include + +/* + a simple semaphore interface. +*/ + +// note: __cpp_lib_semaphore will not be defined in some apple platforms +// even if >= C++20. +#if __has_include() && defined(__cpp_lib_semaphore) && __cpp_lib_semaphore >= 201907L +#define C10_SEMAPHORE_USE_STL +#endif + +#ifdef C10_SEMAPHORE_USE_STL +#include +#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 // @manual +#pragma pop_macro("BLOCK_SIZE") +#else +#include // @manual +#endif + +#include // @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 diff --git a/c10/util/build.bzl b/c10/util/build.bzl index 5e1dc6fbfbf5..e56864f4eb6b 100644 --- a/c10/util/build.bzl +++ b/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": [], diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index b035b33aa078..15bbfaa82ddd 100644 --- a/cmake/Dependencies.cmake +++ b/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) diff --git a/third_party/BUCK.oss b/third_party/BUCK.oss index 0c6d2cd096a2..91862eaeb835 100644 --- a/third_party/BUCK.oss +++ b/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 = "", diff --git a/third_party/moodycamel.BUILD b/third_party/moodycamel.BUILD new file mode 100644 index 000000000000..d3028205016f --- /dev/null +++ b/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"], +) diff --git a/torch/CMakeLists.txt b/torch/CMakeLists.txt index a8ba3e0e529f..0da483c09743 100644 --- a/torch/CMakeLists.txt +++ b/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)