8000 [nativert] port semaphore to c10 util (#153504) · pytorch/pytorch@399d987 · GitHub
[go: up one dir, main page]
Skip to content

Commit 399d987

Browse files
dolpmdolpm
committed
[nativert] port semaphore to c10 util (#153504)
Summary: Pull Request resolved: #153504 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
1 parent 480ae2d commit 399d987

File tree

11 files changed

+143
-0
lines changed

11 files changed

+143
-0
lines changed

‎WORKSPACE

+6
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,12 @@ new_local_repository(
184184
path = "third_party/nlohmann",
185185
)
186186

187+
new_local_repository(
188+
name = "moodycamel",
189+
build_file = "//third_party:moodycamel.BUILD",
190+
path = "third_party/concurrentqueue",
191+
)
192+
187193
new_local_repository(
188194
name = "tensorpipe",
189195
build_file = "//third_party:tensorpipe.BUILD",

‎buckbuild.bzl

+1
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,7 @@ THIRD_PARTY_LIBS = {
178178
"psimd": ["//xplat/third-party/psimd:psimd", "//third_party:psimd"],
179179
"pthreadpool": ["//xplat/third-party/pthreadpool:pthreadpool", "//third_party:pthreadpool"],
180180
"pthreadpool_header": ["//xplat/third-party/pthreadpool:pthreadpool_header", "//third_party:pthreadpool_header"],
181+
"moodycamel": ["//third-party/moodycamel:moodycamel", "//third_party:moodycamel"],
181182
"pyyaml": ["//third-party/pypi/pyyaml:pyyaml", "//third_party:pyyaml"],
182183
"rt": ["//xplat/third-party/linker_lib:rt", "//third_party:rt"],
183184
"ruy": ["//third-party/ruy:ruy_xplat_lib", "//third_party:ruy_lib"],

‎c10/BUCK.oss

+1
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ cxx_library(
1515
"//third_party:cpuinfo",
1616
"//third_party:fmt",
1717
"//third_party:glog",
18+
"//third_party:moodycamel",
1819
],
1920
exported_deps = [],
2021
compiler_flags = [

‎c10/CMakeLists.txt

+1
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,7 @@ if(NOT BUILD_LIBTORCHLESS)
9696
endif()
9797
target_link_libraries(c10 PRIVATE fmt::fmt-header-only)
9898
target_link_libraries(c10 PRIVATE nlohmann)
99+
target_link_libraries(c10 PRIVATE moodycamel)
99100

100101
if(C10_USE_NUMA)
101102
message(STATUS "NUMA paths:")

‎c10/test/util/Semaphore_test.cpp

+35
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
#include <c10/util/Semaphore.h>
2+
#include <c10/util/irange.h>
3+
#include <gtest/gtest.h>
4+
5+
#include <thread>
6+
7+
using namespace ::testing;
8+
9+
TEST(SemaphoreTest, TestConcurrency) {
10+
auto num_threads = std::thread::hardware_concurrency();
11+
auto num_incr = 10000;
12+
13+
c10::Semaphore sem;
14+
15+
std::vector<std::thread> threads;
16+
for ([[maybe_unused]] const auto _ : c10::irange(num_threads)) {
17+
threads.emplace_back([num_incr = num_incr, &sem]() {
18+
for ([[maybe_unused]] const auto _ : c10::irange(num_incr)) {
19+
sem.release();
20+
}
21+
for ([[maybe_unused]] const auto _ : c10::irange(num_incr)) {
22+
sem.acquire();
23+
}
24+
sem.release(num_incr);
25+
for ([[maybe_unused]] const auto _ : c10::irange(num_incr)) {
26+
sem.acquire();
27+
}
28+
});
29+
}
30+
31+
std::for_each(
32+
threads.begin(), threads.end(), [](std::thread& t) { t.join(); });
33+
34+
EXPECT_FALSE(sem.tryAcquire());
35+
}

‎c10/util/Semaphore.h

+76
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
#pragma once
2+
3+
#ifdef __APPLE__
4+
#include <TargetConditionals.h>
5+
#include <version>
6+
#endif
7+
8+
/*
9+
a simple semaphore interface.
10+
11+
if >= C++20 and ,
12+
use std::counting_semaphore otherwise,
13+
use moodycamel::LightweightSemaphore
14+
*/
15+
16+
#if __has_include(<semaphore>) && defined(__cpp_lib_semaphore) && __cpp_lib_semaphore >= 201907L
17+
#define C10_SEMAPHORE_USE_STL
18+
#endif
19+
20+
#ifdef C10_SEMAPHORE_USE_STL
21+
#include <semaphore>
22+
#else
23+
// To use moodycamel semaphore, we need to include the header file
24+
// for concurrentqueue first. Hiding implementation detail here.
25+
#ifdef BLOCK_SIZE
26+
#pragma push_macro("BLOCK_SIZE")
27+
#undef BLOCK_SIZE
28+
#include <moodycamel/concurrentqueue.h> // @manual
29+
#pragma pop_macro("BLOCK_SIZE")
30+
#else
31+
#include <moodycamel/concurrentqueue.h> // @manual
32+
#endif
33+
34+
#include <moodycamel/lightweightsemaphore.h> // @manual
35+
#endif
36+
37+
namespace c10 {
38+
39+
class Semaphore {
40+
public:
41+
Semaphore(int32_t initial_count = 0) : impl_(initial_count) {}
42+
43+
void release(int32_t n = 1) {
44+
#ifdef C10_SEMAPHORE_USE_STL
45+
impl_.release(n);
46+
#else
47+
impl_.signal(n);
48+
#endif
49+
}
50+
51+
void acquire() {
52+
#ifdef C10_SEMAPHORE_USE_STL
53+
impl_.acquire();
54+
#else
55+
impl_.wait();
56+
#endif
57+
}
58+
59+
bool tryAcquire() {
60+
#ifdef C10_SEMAPHORE_USE_STL
61+
return impl_.try_acquire();
62+
#else
63+
return impl_.tryWait();
64+
#endif
65+
}
66+
67+
private:
68+
#ifdef C10_SEMAPHORE_USE_STL
69+
std::counting_semaphore<> impl_;
70+
#else
71+
moodycamel::LightweightSemaphore impl_;
72+
#endif
73+
};
74+
} // namespace c10
75+
76+
#undef C10_SEMAPHORE_USE_STL

‎c10/util/build.bzl

+1
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ def define_targets(rules):
3636
":bit_cast",
3737
"//c10/macros",
3838
"@fmt",
39+
"@moodycamel//:moodycamel",
3940
] + rules.select({
4041
"//c10:using_gflags": ["@com_github_gflags_gflags//:gflags"],
4142
"//conditions:default": [],

‎cmake/Dependencies.cmake

+5
Original file line numberDiff line numberDiff line change
@@ -1198,6 +1198,7 @@ if(USE_DISTRIBUTED AND USE_TENSORPIPE)
11981198

11991199
list(APPEND Caffe2_DEPENDENCY_LIBS tensorpipe)
12001200
list(APPEND Caffe2_DEPENDENCY_LIBS nlohmann)
1201+
list(APPEND Caffe2_DEPENDENCY_LIBS moodycamel)
12011202
if(USE_CUDA)
12021203
list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS tensorpipe_cuda)
12031204
elseif(USE_ROCM)
@@ -1749,3 +1750,7 @@ target_include_directories(httplib SYSTEM INTERFACE ${PROJECT_SOURCE_DIR}/third_
17491750
# Include nlohmann-json
17501751
add_library(nlohmann INTERFACE IMPORTED)
17511752
include_directories(nlohmann SYSTEM INTERFACE ${PROJECT_SOURCE_DIR}/third_party/nlohmann/include)
1753+
1754+
# Include moodycamel
1755+
add_library(moodycamel INTERFACE IMPORTED)
1756+
include_directories(moodycamel SYSTEM INTERFACE ${PROJECT_SOURCE_DIR}/third_party/concurrentqueue)

‎third_party/BUCK.oss

+9
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,15 @@ cxx_library(
7575
visibility = ["PUBLIC"],
7676
)
7777

78+
cxx_library(
79+
name = "moodycamel",
80+
raw_headers = glob([
81+
"concurrentqueue/**/*.h",
82+
]),
83+
reexport_all_header_dependencies = True,
84+
visibility = ["PUBLIC"],
85+
)
86+
7887
cxx_library(
7988
name = "pocketfft_header",
8089
header_namespace = "",

‎third_party/moodycamel.BUILD

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
load("@rules_cc//cc:defs.bzl", "cc_library")
2+
3+
cc_library(
4+
name = "moodycamel",
5+
hdrs = glob(["**/*.h"]),
6+
visibility = ["//visibility:public"],
7+
)

‎torch/CMakeLists.txt

+1
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@ set(TORCH_PYTHON_LINK_LIBRARIES
8484
opentelemetry::api
8585
httplib
8686
nlohmann
87+
moodycamel
8788
shm
8889
fmt::fmt-header-only
8990
ATEN_CPU_FILES_GEN_LIB)

0 commit comments

Comments
 (0)
0