forked from NVIDIA/stdexec
-
Notifications
You must be signed in to change notification settings - Fork 3
Expand file tree
/
Copy pathstatic_thread_pool.hpp
More file actions
1892 lines (1664 loc) · 63.3 KB
/
static_thread_pool.hpp
File metadata and controls
1892 lines (1664 loc) · 63.3 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
/*
* Copyright (c) 2021-2022 Facebook, Inc. and its affiliates.
* Copyright (c) 2021-2024 NVIDIA Corporation
* Copyright (c) 2023 Maikel Nadolski
*
* Licensed under the Apache License Version 2.0 with LLVM Exceptions
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* https://llvm.org/LICENSE.txt
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#pragma once
#include "../stdexec/__detail/__atomic.hpp"
#include "../stdexec/__detail/__config.hpp"
#include "../stdexec/__detail/__intrusive_queue.hpp"
#include "../stdexec/__detail/__manual_lifetime.hpp" // IWYU pragma: keep
#include "../stdexec/__detail/__meta.hpp" // IWYU pragma: keep
#include "../stdexec/execution.hpp"
#include "detail/atomic_intrusive_queue.hpp"
#include "detail/bwos_lifo_queue.hpp"
#include "detail/numa.hpp"
#include "detail/xorshift.hpp"
#include "sender_for.hpp"
#include "sequence/iterate.hpp"
#include "sequence_senders.hpp"
#include <algorithm>
#include <compare>
#include <condition_variable>
#include <cstdint>
#include <exception>
#include <mutex>
#include <span>
#include <thread>
#include <type_traits>
#include <vector>
namespace experimental::execution
{
struct bwos_params
{
std::size_t numBlocks{32};
std::size_t blockSize{8};
};
struct CANNOT_DISPATCH_THE_BULK_ALGORITHM_TO_THE_STATIC_THREAD_POOL_SCHEDULER;
struct BECAUSE_THERE_IS_NO_STATIC_THREAD_POOL_SCHEDULER_IN_THE_ENVIRONMENT;
struct ADD_A_CONTINUES_ON_TRANSITION_TO_THE_STATIC_THREAD_POOL_SCHEDULER_BEFORE_THE_BULK_ALGORITHM;
struct CANNOT_DISPATCH_THE_ITERATE_ALGORITHM_TO_THE_STATIC_THREAD_POOL_SCHEDULER;
struct BECAUSE_THERE_IS_NO_STATIC_THREAD_POOL_SCHEDULER_IN_THE_ENVIRONMENT;
struct ADD_A_CONTINUES_ON_TRANSITION_TO_THE_STATIC_THREAD_POOL_SCHEDULER_BEFORE_THE_ITERATE_ALGORITHM;
namespace _pool_
{
using namespace STDEXEC;
// Splits `n` into `size` chunks distributing `n % size` evenly between ranks.
// Returns `[begin, end)` range in `n` for a given `rank`.
// Example:
// ```cpp
// // n_items thread n_threads
// even_share( 11, 0, 3); // -> [0, 4) -> 4 items
// even_share( 11, 1, 3); // -> [4, 8) -> 4 items
// even_share( 11, 2, 3); // -> [8, 11) -> 3 items
// ```
template <class Shape>
constexpr auto even_share(Shape n, std::size_t rank, std::size_t size) noexcept //
-> std::pair<Shape, Shape>
{
STDEXEC_ASSERT(n >= 0);
using ushape_t = std::make_unsigned_t<Shape>;
auto const avg_per_thread = static_cast<ushape_t>(n) / size;
auto const n_big_share = avg_per_thread + 1;
auto const big_shares = static_cast<ushape_t>(n) % size;
auto const is_big_share = rank < big_shares;
auto const begin = is_big_share
? n_big_share * rank
: n_big_share * big_shares + (rank - big_shares) * avg_per_thread;
auto const end = begin + (is_big_share ? n_big_share : avg_per_thread);
return std::make_pair(static_cast<Shape>(begin), static_cast<Shape>(end));
}
#if !STDEXEC_NO_STDCPP_RANGES()
namespace schedule_all_
{
template <class Range>
class sequence;
} // namespace schedule_all_
#endif
struct task_base
{
task_base* next = nullptr;
void (*execute_)(task_base*, std::uint32_t tid) noexcept = nullptr;
};
struct remote_queue
{
explicit remote_queue(std::size_t nthreads) noexcept
: queues_(nthreads)
{}
explicit remote_queue(remote_queue* next, std::size_t nthreads) noexcept
: next_(next)
, queues_(nthreads)
{}
remote_queue* next_{};
std::vector<__atomic_intrusive_queue<&task_base::next>> queues_{};
std::thread::id id_{std::this_thread::get_id()};
// This marks whether the submitter is a thread in the pool or not.
std::size_t index_{(std::numeric_limits<std::size_t>::max)()};
};
struct remote_queue_list
{
private:
__std::atomic<remote_queue*> head_;
remote_queue* tail_;
std::size_t nthreads_;
remote_queue this_remotes_;
public:
explicit remote_queue_list(std::size_t nthreads) noexcept
: head_{&this_remotes_}
, tail_{&this_remotes_}
, nthreads_(nthreads)
, this_remotes_(nthreads)
{}
~remote_queue_list() noexcept
{
remote_queue* head = head_.load(__std::memory_order_acquire);
while (head != tail_)
{
remote_queue* tmp = std::exchange(head, head->next_);
delete tmp;
}
}
auto pop_all_reversed(std::size_t tid) noexcept -> __intrusive_queue<&task_base::next>
{
remote_queue* head = head_.load(__std::memory_order_acquire);
__intrusive_queue<&task_base::next> tasks{};
while (head != nullptr)
{
tasks.append(head->queues_[tid].pop_all_reversed());
head = head->next_;
}
return tasks;
}
auto get() -> remote_queue*
{
thread_local std::thread::id this_id = std::this_thread::get_id();
remote_queue* head = head_.load(__std::memory_order_acquire);
remote_queue* queue = head;
while (queue != tail_)
{
if (queue->id_ == this_id)
{
return queue;
}
queue = queue->next_;
}
auto* new_head = new remote_queue{head, nthreads_};
while (!head_.compare_exchange_weak(head, new_head, __std::memory_order_acq_rel))
{
new_head->next_ = head;
}
return new_head;
}
};
class _static_thread_pool
{
template <class Receiver>
struct _opstate;
template <bool Parallelize, std::integral Shape, class Fun, class Sender>
struct _bulk_sender;
template <class Shape, class Fun>
struct _is_nothrow_bulk_fn
{
template <class... Args>
requires __callable<Fun, Shape, Shape, __decay_t<Args>&...>
using __f = __mbool<
// If function invocation doesn't throw ...
__nothrow_callable<Fun, Shape, Shape, __decay_t<Args>&...> &&
// ... and decay-copying the arguments doesn't throw ...
__nothrow_decay_copyable<Args...>
// ... then there is no need to advertise completion with `exception_ptr`
>;
};
template <bool Parallelize,
class Shape,
class Fun,
bool MayThrow,
class CvSender,
class Receiver>
struct _bulk_shared_state;
template <bool Parallelize,
class Shape,
class Fun,
bool MayThrow,
class CvSender,
class Receiver>
struct _bulk_receiver;
template <bool Parallelize, std::integral Shape, class Fun, class CvSender, class Receiver>
struct _bulk_opstate;
struct _transform_bulk
{
template <STDEXEC::__one_of<bulk_chunked_t, bulk_unchunked_t> Tag,
class Data,
class CvSender>
auto operator()(Tag, Data&& data, CvSender&& sndr) const
{
auto [pol, shape, fun] = static_cast<Data&&>(data);
using policy_t = std::remove_cvref_t<decltype(pol.__get())>;
constexpr bool parallelize = std::same_as<policy_t, parallel_policy>
|| std::same_as<policy_t, parallel_unsequenced_policy>;
if constexpr (__same_as<Tag, bulk_unchunked_t>)
{
// Turn a bulk_unchunked into a bulk_chunked operation
using fun_t = STDEXEC::__bulk::__as_bulk_chunked_fn<decltype(fun)>;
using sender_t = _bulk_sender<parallelize, decltype(shape), fun_t, __decay_t<CvSender>>;
return sender_t{pool_, static_cast<CvSender&&>(sndr), shape, fun_t(std::move(fun))};
}
else
{
using fun_t = decltype(fun);
using sender_t = _bulk_sender<parallelize, decltype(shape), fun_t, __decay_t<CvSender>>;
return sender_t{pool_, static_cast<CvSender&&>(sndr), shape, std::move(fun)};
}
}
_static_thread_pool& pool_;
};
#if !STDEXEC_NO_STDCPP_RANGES()
struct _transform_iterate
{
template <class Range>
auto operator()(exec::iterate_t, Range&& range) -> schedule_all_::sequence<__decay_t<Range>>
{
return {static_cast<Range&&>(range), pool_};
}
_static_thread_pool& pool_;
};
#endif
static auto _hardware_concurrency() noexcept -> unsigned int
{
unsigned int const n = std::thread::hardware_concurrency();
return n == 0 ? 1 : n;
}
public:
struct domain : STDEXEC::default_domain
{
// transform the generic bulk_chunked sender into a parallel thread-pool bulk sender
template <sender_for Sender, class Env>
requires __one_of<tag_of_t<Sender>, bulk_chunked_t, bulk_unchunked_t>
constexpr auto
transform_sender(STDEXEC::set_value_t, Sender&& sndr, Env const & env) const noexcept
{
if constexpr (__completes_on<Sender, _static_thread_pool::scheduler, Env>)
{
auto sched = STDEXEC::get_completion_scheduler<STDEXEC::set_value_t>(get_env(sndr),
env);
static_assert(std::is_same_v<decltype(sched), _static_thread_pool::scheduler>);
return __apply(_transform_bulk{*sched.pool_}, static_cast<Sender&&>(sndr));
}
else
{
return STDEXEC::__not_a_sender<
STDEXEC::_WHAT_(
CANNOT_DISPATCH_THE_BULK_ALGORITHM_TO_THE_STATIC_THREAD_POOL_SCHEDULER),
STDEXEC::_WHY_(BECAUSE_THERE_IS_NO_STATIC_THREAD_POOL_SCHEDULER_IN_THE_ENVIRONMENT),
STDEXEC::_WHERE_(STDEXEC::_IN_ALGORITHM_, tag_of_t<Sender>),
STDEXEC::_TO_FIX_THIS_ERROR_(
ADD_A_CONTINUES_ON_TRANSITION_TO_THE_STATIC_THREAD_POOL_SCHEDULER_BEFORE_THE_BULK_ALGORITHM),
STDEXEC::_WITH_PRETTY_SENDER_<Sender>,
STDEXEC::_WITH_ENVIRONMENT_(Env)>();
}
}
#if !STDEXEC_NO_STDCPP_RANGES()
template <sender_for<exec::iterate_t> Sender, class Env>
constexpr auto
transform_sender(STDEXEC::set_value_t, Sender&& sndr, Env const & env) const noexcept
{
if constexpr (__completes_on<Sender, _static_thread_pool::scheduler, Env>)
{
auto sched = STDEXEC::get_completion_scheduler<STDEXEC::set_value_t>(get_env(sndr),
env);
return __apply(_transform_iterate{*sched.pool_}, static_cast<Sender&&>(sndr));
}
else
{
return STDEXEC::__not_a_sender<
STDEXEC::_WHAT_(
CANNOT_DISPATCH_THE_ITERATE_ALGORITHM_TO_THE_STATIC_THREAD_POOL_SCHEDULER),
STDEXEC::_WHY_(BECAUSE_THERE_IS_NO_STATIC_THREAD_POOL_SCHEDULER_IN_THE_ENVIRONMENT),
STDEXEC::_WHERE_(STDEXEC::_IN_ALGORITHM_, exec::iterate_t),
STDEXEC::_TO_FIX_THIS_ERROR_(
ADD_A_CONTINUES_ON_TRANSITION_TO_THE_STATIC_THREAD_POOL_SCHEDULER_BEFORE_THE_ITERATE_ALGORITHM),
STDEXEC::_WITH_PRETTY_SENDER_<Sender>,
STDEXEC::_WITH_ENVIRONMENT_(Env)>();
}
}
#endif
};
public:
_static_thread_pool();
_static_thread_pool(std::uint32_t threadCount,
bwos_params params = {},
numa_policy numa = get_numa_policy());
~_static_thread_pool();
struct scheduler
{
private:
template <class Receiver>
friend struct _opstate;
class _sender
{
struct env
{
_static_thread_pool& pool_;
remote_queue* queue_;
template <class CPO>
auto query(get_completion_scheduler_t<CPO>, __ignore = {}) const noexcept
-> _static_thread_pool::scheduler
{
return _static_thread_pool::scheduler{pool_, *queue_};
}
template <class CPO>
auto query(get_completion_domain_t<CPO>, __ignore = {}) const noexcept -> domain
{
return {};
}
};
public:
using sender_concept = sender_t;
template <class Receiver>
using _opstate_t = _opstate<Receiver>;
template <class _Self, class _Env>
static consteval auto get_completion_signatures() noexcept
{
if constexpr (unstoppable_token<stop_token_of_t<_Env>>)
{
return STDEXEC::completion_signatures<set_value_t()>();
}
else
{
return STDEXEC::completion_signatures<set_value_t(), set_stopped_t()>();
}
}
[[nodiscard]]
auto get_env() const noexcept -> env
{
return env{.pool_ = pool_, .queue_ = queue_};
}
template <receiver Receiver>
auto connect(Receiver rcvr) const -> _opstate_t<Receiver>
{
return _opstate_t<Receiver>{pool_,
queue_,
static_cast<Receiver&&>(rcvr),
threadIndex_,
constraints_};
}
private:
friend struct _static_thread_pool::scheduler;
explicit _sender(_static_thread_pool& pool,
remote_queue* queue,
std::size_t threadIndex,
nodemask const & constraints) noexcept
: pool_(pool)
, queue_(queue)
, threadIndex_(threadIndex)
, constraints_(constraints)
{}
_static_thread_pool& pool_;
remote_queue* queue_;
std::size_t threadIndex_{(std::numeric_limits<std::size_t>::max)()};
nodemask constraints_{};
};
friend class _static_thread_pool;
explicit scheduler(_static_thread_pool& pool,
nodemask const * mask = &nodemask::any()) noexcept
: pool_(&pool)
, queue_{pool.get_remote_queue()}
, nodemask_{mask}
{}
explicit scheduler(_static_thread_pool& pool,
remote_queue& queue,
nodemask const * mask = &nodemask::any()) noexcept
: pool_(&pool)
, queue_{&queue}
, nodemask_{mask}
{}
explicit scheduler(_static_thread_pool& pool,
remote_queue& queue,
std::size_t threadIndex) noexcept
: pool_(&pool)
, queue_{&queue}
, thread_idx_{threadIndex}
{}
_static_thread_pool* pool_;
remote_queue* queue_;
nodemask const * nodemask_ = &nodemask::any();
std::size_t thread_idx_{(std::numeric_limits<std::size_t>::max)()};
public:
auto operator==(scheduler const &) const -> bool = default;
[[nodiscard]]
auto schedule() const noexcept -> _sender
{
return _sender{*pool_, queue_, thread_idx_, *nodemask_};
}
[[nodiscard]]
auto query(get_forward_progress_guarantee_t) const noexcept -> forward_progress_guarantee
{
return forward_progress_guarantee::parallel;
}
[[nodiscard]]
auto query(get_completion_domain_t<set_value_t>, __ignore = {}) const noexcept -> domain
{
return {};
}
[[nodiscard]]
auto
query(get_completion_scheduler_t<set_value_t>, __ignore = {}) const noexcept -> scheduler
{
return scheduler{*this};
}
};
auto get_scheduler() noexcept -> scheduler
{
return scheduler{*this};
}
auto get_scheduler_on_thread(std::size_t threadIndex) noexcept -> scheduler
{
return scheduler{*this, *get_remote_queue(), threadIndex};
}
// The caller must ensure that the constraints object is valid for the lifetime of the scheduler.
auto get_constrained_scheduler(nodemask const * constraints) noexcept -> scheduler
{
return scheduler{*this, *get_remote_queue(), constraints};
}
auto get_remote_queue() noexcept -> remote_queue*
{
remote_queue* queue = remotes_.get();
std::size_t index = 0;
for (std::thread& t: threads_)
{
if (t.get_id() == queue->id_)
{
queue->index_ = index;
break;
}
++index;
}
return queue;
}
void request_stop() noexcept;
[[nodiscard]]
auto available_parallelism() const noexcept -> std::uint32_t
{
return thread_count_;
}
[[nodiscard]]
auto params() const noexcept -> bwos_params
{
return params_;
}
void enqueue(task_base* task, nodemask const & contraints = nodemask::any()) noexcept;
void enqueue(remote_queue& queue,
task_base* task,
nodemask const & contraints = nodemask::any()) noexcept;
void enqueue(remote_queue& queue, task_base* task, std::size_t thread_index) noexcept;
//! Enqueue a contiguous span of tasks across task queues.
//! Note: We use the concrete `Task` because we enqueue
//! tasks `task + 0`, `task + 1`, etc. so std::span<task_base>
//! wouldn't be correct.
//! This is O(n_threads) on the calling thread.
template <std::derived_from<task_base> Task>
void bulk_enqueue(std::span<Task> tasks) noexcept;
void bulk_enqueue(remote_queue& queue,
__intrusive_queue<&task_base::next> tasks,
std::size_t tasks_size,
nodemask const & constraints = nodemask::any()) noexcept;
private:
class workstealing_victim
{
public:
explicit workstealing_victim(
bwos::lifo_queue<task_base*, numa_allocator<task_base*>>* queue,
std::uint32_t index,
int numa_node) noexcept
: queue_(queue)
, index_(index)
, numa_node_(numa_node)
{}
auto try_steal() noexcept -> task_base*
{
return queue_->steal_front();
}
[[nodiscard]]
auto index() const noexcept -> std::uint32_t
{
return index_;
}
[[nodiscard]]
auto numa_node() const noexcept -> int
{
return numa_node_;
}
private:
bwos::lifo_queue<task_base*, numa_allocator<task_base*>>* queue_;
std::uint32_t index_;
int numa_node_;
};
struct thread_state_base
{
explicit thread_state_base(std::uint32_t index, numa_policy const & numa) noexcept
: index_(index)
, numa_node_(numa.thread_index_to_node(index))
{}
std::uint32_t index_;
int numa_node_;
};
class thread_state : private thread_state_base
{
public:
struct pop_result
{
task_base* task;
std::uint32_t queue_index;
};
explicit thread_state(_static_thread_pool* pool,
std::uint32_t index,
bwos_params params,
numa_policy const & numa) noexcept
: thread_state_base(index, numa)
, local_queue_(params.numBlocks,
params.blockSize,
numa_allocator<task_base*>(this->numa_node_))
, state_(state::running)
, pool_(pool)
{
std::random_device rd;
rng_.seed(rd);
}
auto pop() -> pop_result;
void push_local(task_base* task);
void push_local(__intrusive_queue<&task_base::next>&& tasks);
auto notify() -> bool;
void request_stop();
void victims(std::vector<workstealing_victim> const & victims)
{
for (workstealing_victim v: victims)
{
if (v.index() == index_)
{
// skip self
continue;
}
if (v.numa_node() == numa_node_)
{
near_victims_.push_back(v);
}
all_victims_.push_back(v);
}
}
[[nodiscard]]
auto index() const noexcept -> std::uint32_t
{
return index_;
}
[[nodiscard]]
auto numa_node() const noexcept -> int
{
return numa_node_;
}
auto as_victim() noexcept -> workstealing_victim
{
return workstealing_victim{&local_queue_, index_, numa_node_};
}
private:
enum state
{
running,
stealing,
sleeping,
notified
};
auto try_pop() -> pop_result;
auto try_remote() -> pop_result;
auto try_steal(std::span<workstealing_victim> victims) -> pop_result;
auto try_steal_near() -> pop_result;
auto try_steal_any() -> pop_result;
void notify_one_sleeping();
void set_stealing();
void clear_stealing();
void set_sleeping();
void clear_sleeping();
bwos::lifo_queue<task_base*, numa_allocator<task_base*>> local_queue_;
__intrusive_queue<&task_base::next> pending_queue_{};
std::mutex mut_{};
std::condition_variable cv_{};
bool stop_requested_{false};
std::vector<workstealing_victim> near_victims_{};
std::vector<workstealing_victim> all_victims_{};
__std::atomic<state> state_;
_static_thread_pool* pool_;
xorshift rng_{};
};
void run(std::uint32_t index) noexcept;
void join() noexcept;
alignas(64) __std::atomic<std::uint32_t> num_active_{};
alignas(64) remote_queue_list remotes_;
std::uint32_t thread_count_;
std::uint32_t max_steals_{thread_count_ + 1};
bwos_params params_;
std::vector<std::thread> threads_;
std::vector<std::optional<thread_state>> thread_states_;
numa_policy numa_;
struct thread_index_by_numa_node
{
int numa_node;
std::size_t thread_index;
friend auto operator==(thread_index_by_numa_node const & lhs,
thread_index_by_numa_node const & rhs) noexcept -> bool
{
return lhs.numa_node == rhs.numa_node;
}
friend auto
operator<=>(thread_index_by_numa_node const & lhs,
thread_index_by_numa_node const & rhs) noexcept -> std::strong_ordering
{
return lhs.numa_node <=> rhs.numa_node;
}
};
std::vector<thread_index_by_numa_node> thread_index_by_numa_node_;
[[nodiscard]]
auto num_threads(int numa) const noexcept -> std::size_t;
[[nodiscard]]
auto num_threads(nodemask constraints) const noexcept -> std::size_t;
[[nodiscard]]
auto get_thread_index(int numa, std::size_t index) const noexcept -> std::size_t;
auto
random_thread_index_with_constraints(nodemask const & contraints) noexcept -> std::size_t;
};
inline _static_thread_pool::_static_thread_pool()
: _static_thread_pool(_hardware_concurrency())
{}
inline _static_thread_pool::_static_thread_pool(std::uint32_t thread_count,
bwos_params params,
numa_policy numa)
: remotes_(thread_count)
, thread_count_(thread_count)
, params_(params)
, thread_states_(thread_count)
, numa_(std::move(numa))
{
STDEXEC_ASSERT(thread_count > 0);
for (std::uint32_t index = 0; index < thread_count; ++index)
{
thread_states_[index].emplace(this, index, params, numa_);
thread_index_by_numa_node_.push_back(
thread_index_by_numa_node{.numa_node = thread_states_[index]->numa_node(),
.thread_index = index});
}
// NOLINTNEXTLINE(modernize-use-ranges) we still support platforms without the std::ranges algorithms
std::sort(thread_index_by_numa_node_.begin(), thread_index_by_numa_node_.end());
std::vector<workstealing_victim> victims{};
for (auto& state: thread_states_)
{
victims.emplace_back(state->as_victim());
}
for (auto& state: thread_states_)
{
state->victims(victims);
}
threads_.reserve(thread_count);
STDEXEC_TRY
{
num_active_.store(thread_count << 16u, __std::memory_order_relaxed);
for (std::uint32_t i = 0; i < thread_count; ++i)
{
threads_.emplace_back([this, i] { run(i); });
}
}
STDEXEC_CATCH_ALL
{
request_stop();
join();
STDEXEC_THROW();
}
}
inline _static_thread_pool::~_static_thread_pool()
{
request_stop();
join();
}
inline void _static_thread_pool::request_stop() noexcept
{
for (auto& state: thread_states_)
{
state->request_stop();
}
}
inline void _static_thread_pool::run(std::uint32_t thread_index) noexcept
{
STDEXEC_ASSERT(thread_index < thread_count_);
// NOLINTNEXTLINE(bugprone-unused-return-value)
numa_.bind_to_node(thread_states_[thread_index]->numa_node());
while (true)
{
// Make a blocking call to de-queue a task if we don't already have one.
auto [task, queue_index] = thread_states_[thread_index]->pop();
if (!task)
{
return; // pop() only returns null when request_stop() was called.
}
task->execute_(task, queue_index);
}
}
inline void _static_thread_pool::join() noexcept
{
for (auto& t: threads_)
{
t.join();
}
threads_.clear();
}
inline void _static_thread_pool::enqueue(task_base* task, nodemask const & constraints) noexcept
{
this->enqueue(*get_remote_queue(), task, constraints);
}
inline auto _static_thread_pool::num_threads(int numa) const noexcept -> std::size_t
{
thread_index_by_numa_node key{.numa_node = numa, .thread_index = 0};
// NOLINTNEXTLINE(modernize-use-ranges) we still support platforms without the std::ranges algorithms
auto it = std::lower_bound(thread_index_by_numa_node_.begin(),
thread_index_by_numa_node_.end(),
key);
if (it == thread_index_by_numa_node_.end())
{
return 0;
}
auto it_end = std::upper_bound(it, thread_index_by_numa_node_.end(), key);
return static_cast<std::size_t>(std::distance(it, it_end));
}
inline auto _static_thread_pool::num_threads(nodemask constraints) const noexcept -> std::size_t
{
std::size_t const n_nodes = static_cast<unsigned>(thread_index_by_numa_node_.back().numa_node
+ 1);
std::size_t n_threads = 0;
for (std::size_t node_index = 0; node_index < n_nodes; ++node_index)
{
if (!constraints[node_index])
{
continue;
}
n_threads += num_threads(static_cast<int>(node_index));
}
return n_threads;
}
inline auto
_static_thread_pool::get_thread_index(int node_index, std::size_t thread_index) const noexcept
-> std::size_t
{
thread_index_by_numa_node key{.numa_node = node_index, .thread_index = 0};
// NOLINTNEXTLINE(modernize-use-ranges) we still support platforms without the std::ranges algorithms
auto it = std::lower_bound(thread_index_by_numa_node_.begin(),
thread_index_by_numa_node_.end(),
key);
STDEXEC_ASSERT(it != thread_index_by_numa_node_.end());
std::advance(it, thread_index);
return it->thread_index;
}
inline auto
_static_thread_pool::random_thread_index_with_constraints(nodemask const & constraints) noexcept
-> std::size_t
{
thread_local std::uint64_t start_index{std::uint64_t(std::random_device{}())};
start_index += 1;
std::size_t target_index = start_index % thread_count_;
std::size_t n_threads = num_threads(constraints);
if (n_threads != 0)
{
for (std::size_t node_index = 0; node_index < numa_.num_nodes(); ++node_index)
{
if (!constraints[node_index])
{
continue;
}
std::size_t n_threads = num_threads(static_cast<int>(node_index));
if (target_index < n_threads)
{
return get_thread_index(static_cast<int>(node_index), target_index);
}
target_index -= n_threads;
}
}
return target_index;
}
inline void _static_thread_pool::enqueue(remote_queue& queue,
task_base* task,
nodemask const & constraints) noexcept
{
static thread_local std::thread::id this_id = std::this_thread::get_id();
remote_queue* correct_queue = this_id == queue.id_ ? &queue : get_remote_queue();
std::size_t idx = correct_queue->index_;
if (idx < thread_states_.size())
{
auto this_node = static_cast<std::size_t>(thread_states_[idx]->numa_node());
if (constraints[this_node])
{
thread_states_[idx]->push_local(task);
return;
}
}
std::size_t const thread_index = random_thread_index_with_constraints(constraints);
queue.queues_[thread_index].push_front(task);
thread_states_[thread_index]->notify();
}
inline void _static_thread_pool::enqueue(remote_queue& queue,
task_base* task,
std::size_t thread_index) noexcept
{
thread_index %= thread_count_;
queue.queues_[thread_index].push_front(task);
thread_states_[thread_index]->notify();
}
template <std::derived_from<task_base> Task>
void _static_thread_pool::bulk_enqueue(std::span<Task> tasks) noexcept
{
auto& queue = *this->get_remote_queue();
for (std::uint32_t i = 0; i < tasks.size(); ++i)
{
std::uint32_t index = i % this->available_parallelism();
queue.queues_[index].push_front(&tasks[i]);
thread_states_[index]->notify();
}
// At this point the calling thread can exit and the pool will take over.
// Ultimately, the last completing thread passes the result forward.
// See `if (is_last_thread)` above.
}
inline void _static_thread_pool::bulk_enqueue(remote_queue& queue,
__intrusive_queue<&task_base::next> tasks,
std::size_t tasks_size,
nodemask const & constraints) noexcept
{
static thread_local std::thread::id const this_id = std::this_thread::get_id();
remote_queue* const correct_queue = this_id == queue.id_ ? &queue : get_remote_queue();
std::size_t const idx = correct_queue->index_;
if (idx < thread_states_.size())
{
auto this_node = static_cast<std::size_t>(thread_states_[idx]->numa_node());
if (constraints[this_node])
{
thread_states_[idx]->push_local(std::move(tasks));
return;
}
}
std::uint32_t const total_threads = available_parallelism();
for (std::uint32_t i = 0; i < total_threads; ++i)
{
auto [begin, end] = even_share(tasks_size, i, total_threads);
if (begin == end)
{
continue;
}
__intrusive_queue<&task_base::next> tmp{};
for (auto j = begin; j < end; ++j)
{
tmp.push_back(tasks.pop_front());
}
correct_queue->queues_[i].prepend(std::move(tmp));
thread_states_[i]->notify();
}
}
inline void
move_pending_to_local(__intrusive_queue<&task_base::next>& pending_queue,
bwos::lifo_queue<task_base*, numa_allocator<task_base*>>& local_queue)
{
auto const last = local_queue.push_back(pending_queue.begin(), pending_queue.end());
__intrusive_queue<&task_base::next> tmp{};
tmp.splice(tmp.begin(), pending_queue, pending_queue.begin(), last);
tmp.clear();
}
inline auto
_static_thread_pool::thread_state::try_remote() -> _static_thread_pool::thread_state::pop_result
{
pop_result result{.task = nullptr, .queue_index = index_};
__intrusive_queue<&task_base::next> remotes = pool_->remotes_.pop_all_reversed(index_);
pending_queue_.append(std::move(remotes));
if (!pending_queue_.empty())
{