1 #pragma once
2
3 #include "../sycl_flow.hpp"
4
5 namespace tf {
6
7 // command group function object of for_each
8 template <typename I, typename C>
_for_each_cgh(I first,I last,C && op)9 auto syclFlow::_for_each_cgh(I first, I last, C&& op) {
10
11 // TODO: special case N == 0?
12 size_t N = std::distance(first, last);
13 size_t B = _default_group_size(N);
14
15 return [=, op=std::forward<C>(op)] (sycl::handler& handler) mutable {
16 size_t _N = (N % B == 0) ? N : (N + B - N % B);
17 handler.parallel_for(
18 sycl::nd_range<1>{sycl::range<1>(_N), sycl::range<1>(B)},
19 [=] (sycl::nd_item<1> item) {
20 size_t i = item.get_global_id(0);
21 if(i < N) {
22 op(*(first + i));
23 }
24 }
25 );
26 };
27 }
28
29 // command group function object of for_each_index
30 template <typename I, typename C>
_for_each_index_cgh(I first,I last,I step,C && op)31 auto syclFlow::_for_each_index_cgh(I first, I last, I step, C&& op) {
32
33 if(is_range_invalid(first, last, step)) {
34 TF_THROW("invalid range [", first, ", ", last, ") with step size ", step);
35 }
36
37 // TODO: special case when N is 0?
38 size_t N = distance(first, last, step);
39 size_t B = _default_group_size(N);
40
41 return [=, op=std::forward<C>(op)] (sycl::handler& handler) mutable {
42 size_t _N = (N % B == 0) ? N : (N + B - N % B);
43 handler.parallel_for(
44 sycl::nd_range<1>{sycl::range<1>(_N), sycl::range<1>(B)},
45 [=] (sycl::nd_item<1> item) {
46 size_t i = item.get_global_id(0);
47 if(i < N) {
48 op(static_cast<I>(i)*step + first);
49 }
50 }
51 );
52 };
53 }
54
55 // ----------------------------------------------------------------------------
56 // for_each and for_each_index algorithms
57 // ----------------------------------------------------------------------------
58
59 // Function: for_each
60 template <typename I, typename C>
for_each(I first,I last,C && op)61 syclTask syclFlow::for_each(I first, I last, C&& op) {
62 return on(_for_each_cgh(first, last, std::forward<C>(op)));
63 }
64
65 // Function: for_each_index
66 template <typename I, typename C>
for_each_index(I beg,I end,I inc,C && op)67 syclTask syclFlow::for_each_index(I beg, I end, I inc, C&& op) {
68 return on(_for_each_index_cgh(beg, end, inc, std::forward<C>(op)));
69 }
70
71 // ----------------------------------------------------------------------------
72 // rebind
73 // ----------------------------------------------------------------------------
74
75 // Function: for_each
76 template <typename I, typename C>
for_each(syclTask task,I first,I last,C && op)77 void syclFlow::for_each(syclTask task, I first, I last, C&& op) {
78 on(task, _for_each_cgh(first, last, std::forward<C>(op)));
79 }
80
81 // Function: for_each_index
82 template <typename I, typename C>
for_each_index(syclTask task,I beg,I end,I inc,C && op)83 void syclFlow::for_each_index(syclTask task, I beg, I end, I inc, C&& op) {
84 on(task, _for_each_index_cgh(beg, end, inc, std::forward<C>(op)));
85 }
86
87
88 } // end of namespace tf -----------------------------------------------------
89