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

Commit 3b86f92

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 3b86f92

File tree

11 files changed

+138
-0
lines changed

11 files changed

+138
-0
lines changed

‎WORKSPACE

Lines changed: 6 additions & 0 deletions
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

Lines changed: 1 addition & 0 deletions
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

Lines changed: 1 addition & 0 deletions
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

Lines changed: 1 addition & 0 deletions
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

Lines changed: 35 additions & 0 deletions
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

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

‎c10/util/build.bzl

Lines changed: 1 addition & 0 deletions
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

Lines changed: 5 additions & 0 deletions
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

Lines changed: 9 additions & 0 deletions
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

Lines changed: 7 additions & 0 deletions
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

Lines changed: 1 addition & 0 deletions
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