Move task_proxy logic to task

Signed-off-by: Ilya Isaev <ilya.isaev@intel.com>

include/oneapi/tbb/detail/_task.h

@@ -63,6 +63,7 @@ TBB_EXPORT void __TBB_EXPORTED_FUNC resume(suspend_point_type* tag);
 TBB_EXPORT suspend_point_type* __TBB_EXPORTED_FUNC current_suspend_point();
 TBB_EXPORT void __TBB_EXPORTED_FUNC notify_waiters(std::uintptr_t wait_ctx_addr);
 
+class mail_outbox;
 class thread_data;
 class task_dispatcher;
 class external_waiter;
@@ -218,9 +219,15 @@ protected:
 public:
     virtual task* execute(execution_data&) = 0;
     virtual task* cancel(execution_data&) = 0;
+    virtual task* lock(execution_data&) { return this; };
 
+   slot_id affinity_tag;
+   slot_id slot;
+
+    std::atomic<task*> next_in_mailbox;
+    r1::mail_outbox* outbox;
 private:
-    std::uint64_t m_reserved[6]{};
+    std::uint64_t m_reserved[3]{};
     friend struct r1::task_accessor;
 };
 static_assert(sizeof(task) == task_alignment, "task size is broken");

include/oneapi/tbb/parallel_for.h

@@ -64,6 +64,9 @@ namespace d1 {
 /** @ingroup algorithms */
 template<typename Range, typename Body, typename Partitioner>
 struct start_for : public task {
+    std::atomic_int my_state;
+    int max_references;
+
     Range my_range;
     const Body my_body;
     node* my_parent;
@@ -73,24 +76,32 @@ struct start_for : public task {
 
     task* execute(execution_data&) override;
     task* cancel(execution_data&) override;
+    task* lock(execution_data&) override;
+
     void finalize(const execution_data&);
 
     //! Constructor for root task.
-    start_for( const Range& range, const Body& body, Partitioner& partitioner, small_object_allocator& alloc ) :
+    start_for( const Range& range, const Body& body, Partitioner& partitioner, small_object_allocator& alloc, int ref_size = 1 ) :
+        my_state(1),
+        max_references(ref_size),
         my_range(range),
         my_body(body),
         my_partition(partitioner),
         my_allocator(alloc) {}
     //! Splitting constructor used to generate children.
     /** parent_ becomes left child.  Newly constructed object is right child. */
-    start_for( start_for& parent_, typename Partitioner::split_type& split_obj, small_object_allocator& alloc ) :
+    start_for( start_for& parent_, typename Partitioner::split_type& split_obj, small_object_allocator& alloc, int ref_size = 1 ) :
+        my_state(ref_size),
+        max_references(ref_size),
         my_range(parent_.my_range, get_range_split_object<Range>(split_obj)),
         my_body(parent_.my_body),
         my_partition(parent_.my_partition, split_obj),
         my_allocator(alloc) {}
     //! Construct right child from the given range as response to the demand.
     /** parent_ remains left child.  Newly constructed object is right child. */
-    start_for( start_for& parent_, const Range& r, depth_t d, small_object_allocator& alloc ) :
+    start_for( start_for& parent_, const Range& r, depth_t d, small_object_allocator& alloc, int ref_size = 1 ) :
+        my_state(ref_size),
+        max_references(ref_size),
         my_range(r),
         my_body(parent_.my_body),
         my_partition(parent_.my_partition, split()),
@@ -98,15 +109,15 @@ struct start_for : public task {
     {
         my_partition.align_depth( d );
     }
-    static void run(const Range& range, const Body& body, Partitioner& partitioner) {
+    static void run(const Range& range, const Body& body, Partitioner& partitioner, int ref_size = 1) {
         task_group_context context(PARALLEL_FOR);
-        run(range, body, partitioner, context);
+        run(range, body, partitioner, context, ref_size);
     }
 
-    static void run(const Range& range, const Body& body, Partitioner& partitioner, task_group_context& context) {
+    static void run(const Range& range, const Body& body, Partitioner& partitioner, task_group_context& context, int ref_size = 1) {
         if ( !range.empty() ) {
             small_object_allocator alloc{};
-            start_for& for_task = *alloc.new_object<start_for>(range, body, partitioner, alloc);
+            start_for& for_task = *alloc.new_object<start_for>(range, body, partitioner, alloc, ref_size);
 
             // defer creation of the wait node until task allocation succeeds
             wait_node wn;
@@ -134,7 +145,7 @@ private:
     void offer_work_impl(execution_data& ed, Args&&... constructor_args) {
         // New right child
         small_object_allocator alloc{};
-        start_for& right_child = *alloc.new_object<start_for>(ed, std::forward<Args>(constructor_args)..., alloc);
+        start_for& right_child = *alloc.new_object<start_for>(ed, std::forward<Args>(constructor_args)..., alloc, this->max_references);
 
         // New root node as a continuation and ref count. Left and right child attach to the new parent.
         right_child.my_parent = my_parent = alloc.new_object<tree_node>(ed, my_parent, 2, alloc);
@@ -170,14 +181,38 @@ task* start_for<Range, Body, Partitioner>::execute(execution_data& ed) {
     }
     my_partition.check_being_stolen(*this, ed);
     my_partition.execute(*this, my_range, ed);
-    finalize(ed);
+    fold_tree<tree_node>(my_parent, ed);
+    if (--my_state == 0) {
+        this->~start_for();
+        my_allocator.deallocate(this, ed);
+    }
     return nullptr;
 }
 
 //! cancel task for parallel_for
 template<typename Range, typename Body, typename Partitioner>
 task* start_for<Range, Body, Partitioner>::cancel(execution_data& ed) {
-    finalize(ed);
+    // finalize(ed);
+    fold_tree<tree_node>(my_parent, ed);
+    if (--my_state == 0) {
+        this->~start_for();
+        my_allocator.deallocate(this, ed);
+    }
+    return nullptr;
+}
+
+//! Try lock task
+template<typename Range, typename Body, typename Partitioner>
+task* start_for<Range, Body, Partitioner>::lock(execution_data& ed) {
+    auto snapshot = my_state--;
+    if (snapshot == max_references) {
+        return this;
+    } else {
+        if (snapshot == 1) {
+            this->~start_for();
+            my_allocator.deallocate(this, ed);
+        }
+    }
     return nullptr;
 }
 
@@ -252,7 +287,7 @@ void parallel_for( const Range& range, const Body& body, const auto_partitioner&
 template<typename Range, typename Body>
     __TBB_requires(tbb_range<Range> && parallel_for_body<Body, Range>)
 void parallel_for( const Range& range, const Body& body, const static_partitioner& partitioner ) {
-    start_for<Range,Body,const static_partitioner>::run(range,body,partitioner);
+    start_for<Range,Body,const static_partitioner>::run(range,body,partitioner, 3);
 }
 
 //! Parallel iteration over range with affinity_partitioner.
@@ -260,7 +295,7 @@ void parallel_for( const Range& range, const Body& body, const static_partitione
 template<typename Range, typename Body>
     __TBB_requires(tbb_range<Range> && parallel_for_body<Body, Range>)
 void parallel_for( const Range& range, const Body& body, affinity_partitioner& partitioner ) {
-    start_for<Range,Body,affinity_partitioner>::run(range,body,partitioner);
+    start_for<Range,Body,affinity_partitioner>::run(range,body,partitioner, 3);
 }
 
 //! Parallel iteration over range with default partitioner and user-supplied context.
@@ -292,7 +327,7 @@ void parallel_for( const Range& range, const Body& body, const auto_partitioner&
 template<typename Range, typename Body>
     __TBB_requires(tbb_range<Range> && parallel_for_body<Body, Range>)
 void parallel_for( const Range& range, const Body& body, const static_partitioner& partitioner, task_group_context& context ) {
-    start_for<Range,Body,const static_partitioner>::run(range, body, partitioner, context);
+    start_for<Range,Body,const static_partitioner>::run(range, body, partitioner, context, 3);
 }
 
 //! Parallel iteration over range with affinity_partitioner and user-supplied context.
@@ -300,7 +335,7 @@ void parallel_for( const Range& range, const Body& body, const static_partitione
 template<typename Range, typename Body>
     __TBB_requires(tbb_range<Range> && parallel_for_body<Body, Range>)
 void parallel_for( const Range& range, const Body& body, affinity_partitioner& partitioner, task_group_context& context ) {
-    start_for<Range,Body,affinity_partitioner>::run(range,body,partitioner, context);
+    start_for<Range,Body,affinity_partitioner>::run(range,body,partitioner, context, 3);
 }
 
 //! Implementation of parallel iteration over stepped range of integers with explicit step and partitioner

src/tbb/arena.cpp

@@ -316,7 +316,14 @@ bool arena::is_out_of_work() {
                 if (my_pool_state.load(std::memory_order_acquire) != busy)
                     return false; // the work was published
             }
-            bool work_absent = k == n;
+            bool mailboxes_empty = true;
+            for (unsigned int i = 0; i < my_num_slots; ++i) {
+                if (!mailbox(i).empty()) {
+                    mailboxes_empty = false;
+                    break;
+                }
+            }
+            bool work_absent = k == n && mailboxes_empty;
             // Test and test-and-set.
             if (my_pool_state.load(std::memory_order_acquire) == busy) {
                 bool no_stream_tasks = !has_enqueued_tasks() && my_resume_task_stream.empty();

src/tbb/arena.h

@@ -572,12 +572,12 @@ inline d1::task* arena::steal_task(unsigned arena_index, FastRandom& frnd, execu
         return nullptr;
     }
     if (task_accessor::is_proxy_task(*t)) {
-        task_proxy &tp = *(task_proxy*)t;
-        d1::slot_id slot = tp.slot;
-        t = tp.extract_task<task_proxy::pool_bit>();
+        // task_proxy &tp = *(task_proxy*)t;
+        d1::slot_id slot = t->slot;
+        t = t->lock(ed);
         if (!t) {
             // Proxy was empty, so it's our responsibility to free it
-            tp.allocator.delete_object(&tp, ed);
+            // tp.allocator.delete_object(&tp, ed);
             return nullptr;
         }
         // Note affinity is called for any stolen task (proxy or general)

src/tbb/arena_slot.cpp

@@ -43,14 +43,14 @@ d1::task* arena_slot::get_task_impl(size_t T, execution_data_ext& ed, bool& task
         return nullptr;
     }
 
-    task_proxy& tp = static_cast<task_proxy&>(*result);
-    d1::slot_id aff_id = tp.slot;
-    if ( d1::task *t = tp.extract_task<task_proxy::pool_bit>() ) {
+    // task_proxy& tp = static_cast<task_proxy&>(*result);
+    d1::slot_id aff_id = result->slot;
+    if ( d1::task *t = result->lock(ed) ) {
         ed.affinity_slot = aff_id;
         return t;
     }
     // Proxy was empty, so it's our responsibility to free it
-    tp.allocator.delete_object(&tp, ed);
+    // tp.allocator.delete_object(&tp, ed);
 
     if ( tasks_omitted ) {
         task_pool_ptr[T] = nullptr;
@@ -173,9 +173,9 @@ d1::task* arena_slot::steal_task(arena& a, isolation_type isolation, std::size_t
                 if (!task_accessor::is_proxy_task(*result)) {
                     break;
                 }
-                task_proxy& tp = *static_cast<task_proxy*>(result);
+                // task_proxy& tp = *static_cast<task_proxy*>(result);
                 // If mailed task is likely to be grabbed by its destination thread, skip it.
-                if (!task_proxy::is_shared(tp.task_and_tag) || !tp.outbox->recipient_is_idle() || a.mailbox(slot_index).recipient_is_idle()) {
+                if (!task_proxy::is_shared(std::intptr_t(result) | result->affinity_tag) || !result->outbox->recipient_is_idle() || a.mailbox(slot_index).recipient_is_idle()) {
                     break;
                 }
             }

src/tbb/mailbox.h

@@ -95,7 +95,7 @@ struct task_proxy : public d1::task {
 //! Internal representation of mail_outbox, without padding.
 class unpadded_mail_outbox {
 protected:
-    typedef std::atomic<task_proxy*> atomic_proxy_ptr;
+    typedef std::atomic<d1::task*> atomic_proxy_ptr;
 
     //! Pointer to first task_proxy in mailbox, or nullptr if box is empty.
     atomic_proxy_ptr my_first;
@@ -112,8 +112,8 @@ protected:
 /** Padded to occupy a cache line. */
 class mail_outbox : padded<unpadded_mail_outbox> {
 
-    task_proxy* internal_pop( isolation_type isolation ) {
-        task_proxy* curr = my_first.load(std::memory_order_acquire);
+    d1::task* internal_pop( isolation_type isolation ) {
+        d1::task* curr = my_first.load(std::memory_order_acquire);
         if ( !curr )
             return nullptr;
         atomic_proxy_ptr* prev_ptr = &my_first;
@@ -128,7 +128,7 @@ class mail_outbox : padded<unpadded_mail_outbox> {
         }
         // There is a first item in the mailbox.  See if there is a second.
         // The next_in_mailbox should be read with acquire to guarantee (*second) consistency.
-        if ( task_proxy* second = curr->next_in_mailbox.load(std::memory_order_acquire) ) {
+        if ( d1::task* second = curr->next_in_mailbox.load(std::memory_order_acquire) ) {
             // There are at least two items, so first item can be popped easily.
             prev_ptr->store(second, std::memory_order_relaxed);
         } else {
@@ -156,7 +156,7 @@ public:
 
     //! Push task_proxy onto the mailbox queue of another thread.
     /** Implementation is wait-free. */
-    void push( task_proxy* t ) {
+    void push( d1::task* t ) {
         assert_pointer_valid(t);
         t->next_in_mailbox.store(nullptr, std::memory_order_relaxed);
         atomic_proxy_ptr* const link = my_last.exchange(&t->next_in_mailbox);
@@ -188,9 +188,9 @@ public:
     //! Drain the mailbox
     void drain() {
         // No fences here because other threads have already quit.
-        for( ; task_proxy* t = my_first; ) {
+        for( ; d1::task* t = my_first; ) {
             my_first.store(t->next_in_mailbox, std::memory_order_relaxed);
-            t->allocator.delete_object(t);
+            // t->allocator.delete_object(t);
         }
     }
 
@@ -218,7 +218,7 @@ public:
         my_putter = nullptr;
     }
     //! Get next piece of mail, or nullptr if mailbox is empty.
-    task_proxy* pop( isolation_type isolation ) {
+    d1::task* pop( isolation_type isolation ) {
         return my_putter->internal_pop( isolation );
     }
     //! Return true if mailbox is empty

src/tbb/task_dispatcher.cpp

@@ -53,18 +53,20 @@ void __TBB_EXPORTED_FUNC spawn(d1::task& t, d1::task_group_context& ctx, d1::slo
 
     if ( id != d1::no_slot && id != tls->my_arena_index && id < a->my_num_slots) {
         // Allocate proxy task
-        d1::small_object_allocator alloc{};
-        auto proxy = alloc.new_object<task_proxy>(static_cast<d1::execution_data&>(ed));
+        // d1::small_object_allocator alloc{};
+        // auto proxy = alloc.new_object<task_proxy>(static_cast<d1::execution_data&>(ed));
+        auto proxy = &t;
+
         // Mark as a proxy
         task_accessor::set_proxy_trait(*proxy);
         // Mark isolation for the proxy task
         task_accessor::isolation(*proxy) = ed.isolation;
         // Deallocation hint (tls) from the task allocator
-        proxy->allocator = alloc;
+        // proxy->allocator = alloc;
         proxy->slot = id;
         proxy->outbox = &a->mailbox(id);
         // Mark proxy as present in both locations (sender's task pool and destination mailbox)
-        proxy->task_and_tag = intptr_t(&t) | task_proxy::location_mask;
+        proxy->affinity_tag = task_proxy::location_mask;
         // Mail the proxy - after this point t may be destroyed by another thread at any moment.
         proxy->outbox->push(proxy);
         // Spawn proxy to the local task pool

src/tbb/task_dispatcher.h

@@ -443,14 +443,14 @@ inline d1::task* task_dispatcher::get_critical_task(d1::task* t, execution_data_
 #endif
 
 inline d1::task* task_dispatcher::get_mailbox_task(mail_inbox& my_inbox, execution_data_ext& ed, isolation_type isolation) {
-    while (task_proxy* const tp = my_inbox.pop(isolation)) {
-        if (d1::task* result = tp->extract_task<task_proxy::mailbox_bit>()) {
+    while (d1::task* tp = my_inbox.pop(isolation)) {
+        if (d1::task* result = tp->lock(ed)) {
             ed.original_slot = (unsigned short)(-2);
             ed.affinity_slot = ed.task_disp->m_thread_data->my_arena_index;
             return result;
         }
         // We have exclusive access to the proxy, and can destroy it.
-        tp->allocator.delete_object(tp, ed);
+        // tp->allocator.delete_object(tp, ed);
     }
     return NULL;
 }

test/tbb/test_parallel_for.cpp

@@ -459,6 +459,24 @@ TEST_CASE("parallel_for constraints") {
 }
 #endif // __TBB_CPP20_CONCEPTS_PRESENT
 
+TEST_CASE("dummy test 1") {
+    constexpr int N = 128;
+    printf("Concurrency: %d\n", tbb::this_task_arena::max_concurrency());
+    auto f = [](tbb::blocked_range<int>& r) {
+        printf("%d %d %lu\n", r.begin(), r.end(), r.size());
+    };
+    tbb::parallel_for(tbb::blocked_range<int>(0, N), f, tbb::static_partitioner{});
+}
+
+TEST_CASE("dummy test 2") {
+    std::atomic_int sum{};
+    tbb::parallel_for(tbb::blocked_range<int>(0, 100000), [&](tbb::blocked_range<int>& r) {
+        for (int i = r.begin(); i < r.end(); ++i)
+            ++sum;
+    }, tbb::static_partitioner{});
+    CHECK(sum == 100000);
+}
+
 #if _MSC_VER
 #pragma warning (pop)
 #endif