1 /*
2     Copyright (c) 2023 Intel Corporation
3 
4     Licensed under the Apache License, Version 2.0 (the "License");
5     you may not use this file except in compliance with the License.
6     You may obtain a copy of the License at
7 
8         http://www.apache.org/licenses/LICENSE-2.0
9 
10     Unless required by applicable law or agreed to in writing, software
11     distributed under the License is distributed on an "AS IS" BASIS,
12     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13     See the License for the specific language governing permissions and
14     limitations under the License.
15 */
16 
17 #ifndef __TBB_task_emulation_layer_H
18 #define __TBB_task_emulation_layer_H
19 
20 #include "tbb/task_group.h"
21 #include "tbb/task_arena.h"
22 
23 #include <atomic>
24 
25 namespace task_emulation {
26 
27 struct task_group_pool {
28     task_group_pool() : pool_size(std::thread::hardware_concurrency()), task_submitters(new tbb::task_group[pool_size]) {}
29 
30     ~task_group_pool() {
31         for (std::size_t i = 0; i < pool_size; ++i) {
32             task_submitters[i].wait();
33         }
34 
35         delete [] task_submitters;
36     }
37 
38     tbb::task_group& operator[] (std::size_t idx) { return task_submitters[idx]; }
39 
40     const std::size_t pool_size;
41     tbb::task_group* task_submitters;
42 };
43 
44 static task_group_pool tg_pool;
45 
46 class base_task {
47 public:
48     base_task() = default;
49 
50     base_task(const base_task& t) : m_parent(t.m_parent), m_child_counter(t.m_child_counter.load())
51     {}
52 
53     virtual ~base_task() = default;
54 
55     void operator() () const {
56         base_task* parent_snapshot = m_parent;
57         const_cast<base_task*>(this)->execute();
58         if (m_parent && parent_snapshot == m_parent && m_child_counter == 0) {
59             if (m_parent->remove_reference() == 0) {
60                 m_parent->operator()();
61                 delete m_parent;
62             }
63         }
64 
65         if (m_child_counter == 0 && m_type == task_type::allocated) {
66             delete this;
67         }
68     }
69 
70     virtual void execute() = 0;
71 
72     template <typename C, typename... Args>
73     C* allocate_continuation(std::uint64_t ref, Args&&... args) {
74         C* continuation = new C{std::forward<Args>(args)...};
75         continuation->m_type = task_type::continuation;
76         continuation->reset_parent(reset_parent());
77         continuation->m_child_counter = ref;
78         return continuation;
79     }
80 
81     template <typename F, typename... Args>
82     F create_child(Args&&... args) {
83         return create_child_impl<F>(std::forward<Args>(args)...);
84     }
85 
86     template <typename F, typename... Args>
87     F create_child_and_increment(Args&&... args) {
88         add_reference();
89         return create_child_impl<F>(std::forward<Args>(args)...);
90     }
91 
92     template <typename F, typename... Args>
93     F* allocate_child(Args&&... args) {
94         return allocate_child_impl<F>(std::forward<Args>(args)...);
95     }
96 
97     template <typename F, typename... Args>
98     F* allocate_child_and_increment(Args&&... args) {
99         add_reference();
100         return allocate_child_impl<F>(std::forward<Args>(args)...);
101     }
102 
103     template <typename C>
104     void recycle_as_child_of(C& c) {
105         reset_parent(&c);
106     }
107 
108     void recycle_as_continuation() {
109         m_type = task_type::continuation;
110     }
111 
112     void add_reference() {
113         ++m_child_counter;
114     }
115 
116     std::uint64_t remove_reference() {
117         return --m_child_counter;
118     }
119 
120 protected:
121     enum class task_type {
122         created,
123         allocated,
124         continuation
125     };
126 
127     task_type m_type;
128 
129 private:
130     template <typename F, typename... Args>
131     friend F create_root_task(tbb::task_group& tg, Args&&... args);
132 
133     template <typename F, typename... Args>
134     friend F* allocate_root_task(tbb::task_group& tg, Args&&... args);
135 
136     template <typename F, typename... Args>
137     F create_child_impl(Args&&... args) {
138         F obj{std::forward<Args>(args)...};
139         obj.m_type = task_type::created;
140         obj.reset_parent(this);
141         return obj;
142     }
143 
144     template <typename F, typename... Args>
145     F* allocate_child_impl(Args&&... args) {
146         F* obj = new F{std::forward<Args>(args)...};
147         obj->m_type = task_type::allocated;
148         obj->reset_parent(this);
149         return obj;
150     }
151 
152     base_task* reset_parent(base_task* ptr = nullptr) {
153         auto p = m_parent;
154         m_parent = ptr;
155         return p;
156     }
157 
158     base_task* m_parent{nullptr};
159     std::atomic<std::uint64_t> m_child_counter{0};
160 };
161 
162 class root_task : public base_task {
163 public:
164     root_task(tbb::task_group& tg) : m_tg(tg), m_callback(m_tg.defer([] { /* Create empty callback to preserve reference for wait. */})) {
165         add_reference();
166         m_type = base_task::task_type::continuation;
167     }
168 
169 private:
170     void execute() override {
171         m_tg.run(std::move(m_callback));
172     }
173 
174     tbb::task_group& m_tg;
175     tbb::task_handle m_callback;
176 };
177 
178 template <typename F, typename... Args>
179 F create_root_task(tbb::task_group& tg, Args&&... args) {
180     F obj{std::forward<Args>(args)...};
181     obj.m_type = base_task::task_type::created;
182     obj.reset_parent(new root_task{tg});
183     return obj;
184 }
185 
186 template <typename F, typename... Args>
187 F* allocate_root_task(tbb::task_group& tg, Args&&... args) {
188     F* obj = new F{std::forward<Args>(args)...};
189     obj->m_type = base_task::task_type::allocated;
190     obj->reset_parent(new root_task{tg});
191     return obj;
192 }
193 
194 template <typename F>
195 void run_task(F&& f) {
196     tg_pool[tbb::this_task_arena::current_thread_index()].run(std::forward<F>(f));
197 }
198 
199 template <typename F>
200 void run_task(F* f) {
201     tg_pool[tbb::this_task_arena::current_thread_index()].run(std::ref(*f));
202 }
203 
204 template <typename F>
205 void run_and_wait(tbb::task_group& tg, F* f) {
206    tg.run_and_wait(std::ref(*f));
207 }
208 } // namespace task_emulation
209 
210 #endif // __TBB_task_emulation_layer_H
211