Merge pull request #368 from anton-potapov/task_group_extensions_fixup

Fixup for POC: Extend tbb::task_group to support new use cases

include/oneapi/tbb/detail/_task.h

@@ -30,6 +30,7 @@
 #include <utility>
 #include <atomic>
 #include <mutex>
+#include <functional>
 
 namespace tbb {
 namespace detail {
@@ -230,6 +231,9 @@ private:
 };
 static_assert(sizeof(task) == task_alignment, "task size is broken");
 
+
+using task_handle = std::unique_ptr<task, std::function<void(task*)>>;
+
 } // namespace d1
 } // namespace detail
 } // namespace tbb

include/oneapi/tbb/task_arena.h

@@ -233,6 +233,11 @@ class task_arena : public task_arena_base {
         r1::enqueue(*alloc.new_object<enqueue_task<typename std::decay<F>::type>>(std::forward<F>(f), alloc), this);
     }
 
+    void enqueue_impl(task_handle&& th) {
+        initialize();
+        r1::enqueue(*th.release(), this);
+    }
+
     template<typename R, typename F>
     R execute_impl(F& f) {
         initialize();
@@ -370,6 +375,13 @@ public:
         enqueue_impl(std::forward<F>(f));
     }
 
+    //! Enqueues a task into the arena to process a functor, and immediately returns.
+    //! Does not require the calling thread to join the arena
+
+    void enqueue(task_handle&& th) {
+        enqueue_impl(std::move(th));
+    }
+
     //! Joins the arena and executes a mutable functor, then returns
     //! If not possible to join, wraps the functor into a task, enqueues it and waits for task completion
     //! Can decrement the arena demand for workers, causing a worker to leave and free a slot to the calling thread
@@ -427,11 +439,20 @@ inline int current_thread_index() {
     return idx == -1 ? task_arena_base::not_initialized : idx;
 }
 
+inline bool is_inside_task() {
+    return nullptr != current_context();
+}
+
 //! Returns the maximal number of threads that can work inside the arena
 inline int max_concurrency() {
     return r1::max_concurrency(nullptr);
 }
 
+inline void enqueue(task_handle&& th) {
+    task_arena arena( task_arena::attach {});
+    return arena.enqueue(std::move(th));
+}
+
 using r1::submit;
 
 } // namespace d1
@@ -440,10 +461,13 @@ using r1::submit;
 inline namespace v1 {
 using detail::d1::task_arena;
 
+using detail::d1::is_inside_task;
+
 namespace this_task_arena {
 using detail::d1::current_thread_index;
 using detail::d1::max_concurrency;
 using detail::d1::isolate;
+using detail::d1::enqueue;
 } // namespace this_task_arena
 
 } // inline namespace v1

include/oneapi/tbb/task_group.h

@@ -27,7 +27,8 @@
 
 #include "profiling.h"
 
-#include <functional>
+#include <memory>
+#include <type_traits>
 
 #if _MSC_VER && !defined(__INTEL_COMPILER)
     // Suppress warning: structure was padded due to alignment specifier
@@ -310,6 +311,28 @@ class structured_task_group;
 class isolated_task_group;
 #endif
 
+namespace {
+    template <typename F>
+    struct is_void_return : std::is_void< typename std::result_of<F()>::type > {};
+
+    template<typename F>
+    task* return_value_or_nullptr_impl(std::false_type, F&& f){
+        task_handle th = std::forward<F>(f)();
+        return th.release();
+    }
+
+    template<typename F>
+    task* return_value_or_nullptr_impl(std::true_type, F&& f){
+        std::forward<F>(f)();
+        return nullptr;
+    }
+
+    template<typename F>
+    task* return_value_or_nullptr(F&& f){
+        return  return_value_or_nullptr_impl(is_void_return<F>{}, std::forward<F>(f));
+    }
+}
+
 template<typename F>
 class function_task : public task {
     const F m_func;
@@ -327,10 +350,12 @@ class function_task : public task {
 
         allocator.deallocate(this, ed);
     }
+private:
     task* execute(execution_data& ed) override {
-        m_func();
+        task* res = return_value_or_nullptr(m_func);
+
         finalize(ed);
-        return nullptr;
+        return res;
     }
     task* cancel(execution_data& ed) override {
         finalize(ed);
@@ -357,9 +382,9 @@ class function_stack_task : public task {
         m_wait_ctx.release();
     }
     task* execute(execution_data&) override {
-        m_func();
+        task* res = return_value_or_nullptr(m_func);
         finalize();
-        return nullptr;
+        return res;
     }
     task* cancel(execution_data&) override {
         finalize();
@@ -372,34 +397,58 @@ public:
 class task_group_base : no_copy {
 protected:
     wait_context m_wait_ctx;
-    task_group_context m_context;
+    task_group_context m_own_context;
+    task_group_context* m_context{nullptr};
 
     template<typename F>
     task_group_status internal_run_and_wait(const F& f) {
         function_stack_task<F> t{ f, m_wait_ctx };
         m_wait_ctx.reserve();
         bool cancellation_status = false;
+        __TBB_ASSERT(m_context, nullptr);
         try_call([&] {
-            execute_and_wait(t, m_context, m_wait_ctx, m_context);
+            execute_and_wait(t, *m_context, m_wait_ctx, *m_context);
         }).on_completion([&] {
             // TODO: the reset method is not thread-safe. Ensure the correct behavior.
-            cancellation_status = m_context.is_group_execution_cancelled();
-            m_context.reset();
+            cancellation_status = m_context->is_group_execution_cancelled();
+            m_context->reset();
         });
         return cancellation_status ? canceled : complete;
     }
 
     template<typename F>
-    task* prepare_task(F&& f) {
+    auto prepare_task(F&& f) -> function_task<typename std::decay<F>::type>*{
         m_wait_ctx.reserve();
         small_object_allocator alloc{};
         return alloc.new_object<function_task<typename std::decay<F>::type>>(std::forward<F>(f), m_wait_ctx, alloc);
     }
 
+    template<typename F>
+    task_handle prepare_task_handle(F&& f) {
+        m_wait_ctx.reserve();
+        small_object_allocator alloc{};
+        using function_task_t =  function_task<typename std::decay<F>::type>;
+        auto* function_task_p =  alloc.new_object<function_task_t>(std::forward<F>(f), m_wait_ctx, alloc);
+
+        return task_handle {function_task_p, [alloc, this](task* t) mutable {
+            alloc.delete_object(static_cast<function_task_t*>(t));
+            m_wait_ctx.release();
+        }};
+    }
+
+
 public:
     task_group_base(uintptr_t traits = 0)
         : m_wait_ctx(0)
-        , m_context(task_group_context::bound, task_group_context::default_traits | traits)
+        , m_own_context(task_group_context::bound, task_group_context::default_traits | traits)
+        , m_context(&m_own_context)
+    {
+    }
+
+    task_group_base(task_group_context& ctx)
+        : m_wait_ctx(0)
+        , m_own_context()
+        , m_context(&ctx)
     {
     }
 
@@ -410,40 +459,58 @@ public:
 #else
             bool stack_unwinding_in_progress = std::uncaught_exception();
 #endif
+            __TBB_ASSERT(m_context, nullptr);
             // Always attempt to do proper cleanup to avoid inevitable memory corruption
             // in case of missing wait (for the sake of better testability & debuggability)
-            if (!m_context.is_group_execution_cancelled())
+            if (!m_context->is_group_execution_cancelled())
                 cancel();
-            d1::wait(m_wait_ctx, m_context);
+            d1::wait(m_wait_ctx, *m_context);
             if (!stack_unwinding_in_progress)
                 throw_exception(exception_id::missing_wait);
         }
     }
 
     task_group_status wait() {
+        __TBB_ASSERT(m_context, nullptr);
         bool cancellation_status = false;
         try_call([&] {
-            d1::wait(m_wait_ctx, m_context);
+            d1::wait(m_wait_ctx, *m_context);
         }).on_completion([&] {
             // TODO: the reset method is not thread-safe. Ensure the correct behavior.
-            cancellation_status = m_context.is_group_execution_cancelled();
-            m_context.reset();
+            cancellation_status = m_context->is_group_execution_cancelled();
+            m_context->reset();
         });
         return cancellation_status ? canceled : complete;
     }
 
     void cancel() {
-        m_context.cancel_group_execution();
+        __TBB_ASSERT(m_context, nullptr);
+        m_context->cancel_group_execution();
     }
 }; // class task_group_base
 
 class task_group : public task_group_base {
 public:
     task_group() : task_group_base(task_group_context::concurrent_wait) {}
+    task_group(task_group_context& ctx) : task_group_base(ctx) {}
 
     template<typename F>
     void run(F&& f) {
-        spawn(*prepare_task(std::forward<F>(f)), m_context);
+        __TBB_ASSERT(m_context, nullptr);
+        spawn(*prepare_task(std::forward<F>(f)), *m_context);
+    }
+
+    void run(task_handle&& h) {
+        task* t = h.release();
+        __TBB_ASSERT(t, "Attempt running empty task_handle.");
+        __TBB_ASSERT(m_context, nullptr);
+        spawn(*t, *m_context);
+    }
+
+    template<typename F>
+    task_handle defer(F&& f) {
+        return prepare_task_handle(std::forward<F>(f));
+
     }
 
     template<typename F>
@@ -500,7 +567,8 @@ public:
 
     template<typename F>
     void run(F&& f) {
-        spawn_delegate sd(prepare_task(std::forward<F>(f)), m_context);
+        __TBB_ASSERT(m_context, nullptr);
+        spawn_delegate sd(prepare_task(std::forward<F>(f)), *m_context);
         r1::isolate_within_arena(sd, this_isolation());
     }
 
@@ -545,6 +613,8 @@ using detail::d1::canceled;
 
 using detail::d1::is_current_task_group_canceling;
 using detail::r1::missing_wait;
+
+using detail::d1::task_handle;
 }
 
 } // namespace tbb

test/common/utils.h

@@ -42,6 +42,8 @@ namespace utils {
 
 #define utils_fallthrough __TBB_fallthrough
 
+using tbb::detail::try_call;
+
 template<typename It>
 typename std::iterator_traits<It>::value_type median(It first, It last) {
     std::sort(first, last);

test/tbb/test_task_arena.cpp

@@ -31,6 +31,7 @@
 #include "tbb/concurrent_set.h"
 #include "tbb/spin_mutex.h"
 #include "tbb/spin_rw_mutex.h"
+#include "tbb/task_group.h"
 
 #include <stdexcept>
 #include <cstdlib>
@@ -1875,3 +1876,58 @@ TEST_CASE("Workers oversubscription") {
         );
     });
 }
+
+TEST_CASE("enqueue task_handle") {
+    tbb::task_arena arena;
+    tbb::task_group tg;
+
+    std::atomic<bool> run{false};
+
+    auto task_handle = tg.defer([&]{ run = true;});
+
+    arena.enqueue(std::move(task_handle));
+    tg.wait();
+
+    ASSERT(run == true, "");
+
+}
+
+TEST_CASE("this_task_arena::enqueue task_handle") {
+    tbb::task_arena arena;
+    tbb::task_group tg;
+
+    std::atomic<bool> run{false};
+
+    arena.execute([&]{
+
+        auto task_handle = tg.defer([&]{ run = true;});
+
+        tbb::this_task_arena::enqueue(std::move(task_handle));
+    });
+
+    tg.wait();
+
+    ASSERT(run == true, "");
+
+}
+
+TEST_CASE("is_inside_task in task_group"){
+    ASSERT( false == tbb::is_inside_task(), "");
+
+    tbb::task_group tg;
+    tg.run_and_wait([&]{
+        ASSERT( true == tbb::is_inside_task(), "");
+    });
+}
+
+TEST_CASE("is_inside_task in arena::execute"){
+    ASSERT( false == tbb::is_inside_task(), "");
+
+//    FIXME: seems that test
+    tbb::task_arena arena;
+
+    arena.execute([&]{
+        ASSERT( false == tbb::is_inside_task(), "");
+    });
+
+}

test/tbb/test_task_group.cpp

@@ -979,3 +979,189 @@ TEST_CASE("Async task group") {
         });
     });
 }
+
+TEST_CASE("Task handle created but not run"){
+    {
+        tbb::task_group tg;
+
+        std::atomic<bool> run {false};
+
+        auto h = tg.defer([&]{
+            run = true;
+        });
+        ASSERT(run == false, "delayed task should not be run until run(task_handle) is called");
+    }
+}
+
+TEST_CASE("Task handle run"){
+    tbb::task_handle h;
+
+    tbb::task_group tg;
+    std::atomic<bool> run {false};
+
+    h = tg.defer([&]{
+        run = true;
+    });
+    ASSERT(run == false, "delayed task should not be run until run(task_handle) is called");
+    tg.run(std::move(h));
+    tg.wait();
+    ASSERT(run == true, "Delayed task should be completed when task_group::wait exits");
+}
+
+TEST_CASE("Task handle empty check"){
+    tbb::task_group tg;
+
+    tbb::task_handle h;
+
+    bool empty = (h == nullptr);
+    ASSERT(empty, "default constructed task_handle should be empty");
+
+    h = tg.defer([]{});
+
+    ASSERT(h != nullptr, "delayed task returned by task_group::delayed should not be empty");
+}
+
+TEST_CASE("Task handle blocks wait"){
+    tbb::task_group tg;
+
+    std::atomic<bool> completed  {false};
+    std::atomic<bool> start_wait {false};
+    std::atomic<bool> thread_started{false};
+
+    tbb::task_handle h = tg.defer([&]{
+        completed = true;
+    });
+
+    std::thread wait_thread {[&]{
+        ASSERT(completed == false, "Deferred task should not be run until run(task_handle) is called");
+
+        thread_started = true;
+        utils::SpinWaitUntilEq(start_wait, true);
+        tg.wait();
+        ASSERT(completed == true, "Deferred task should be completed when task_group::wait exits");
+    }};
+
+    utils::SpinWaitUntilEq(thread_started, true);
+    ASSERT(completed == false, "Deferred task should not be run until run(task_handle) is called");
+
+    tg.run(std::move(h));
+    //TODO: more accurate test (with fixed number of threads (1 ?) to guarantee correctness of following assert)
+    //ASSERT(completed == false, "Deferred task should not be run until run(task_handle) and wait is called");
+    start_wait = true;
+    wait_thread.join();
+}
+
+TEST_CASE("Task handle for scheduler bypass"){
+    tbb::task_group tg;
+    std::atomic<bool> run {false};
+
+    tg.run([&]{
+        return tg.defer([&]{
+            run = true;
+        });
+    });
+
+    tg.wait();
+    ASSERT(run == true, "task handle returned by user lambda (bypassed) should be run");
+}
+
+TEST_CASE("Task handle for scheduler bypass via run_and_wait"){
+    tbb::task_group tg;
+    std::atomic<bool> run {false};
+
+    tg.run_and_wait([&]{
+        return tg.defer([&]{
+            run = true;
+        });
+    });
+
+    ASSERT(run == true, "task handle returned by user lambda (bypassed) should be run");
+}
+
+struct SelfRunner {
+    tbb::task_group& m_tg;
+    std::atomic<unsigned>& count;
+    void operator()() const {
+        unsigned previous_count = count.fetch_sub(1);
+        if (previous_count > 1)
+            m_tg.run( *this );
+    }
+};
+
+//! Submit work to single task_group instance from inside the work
+//! \brief \ref error_guessing
+TEST_CASE("Run self using same task_group instance") {
+    const unsigned num = 10;
+    std::atomic<unsigned> count{num};
+    tbb::task_group tg;
+    SelfRunner uf{tg, count};
+    tg.run( uf );
+    tg.wait();
+    CHECK_MESSAGE(
+        count == 0,
+        "Not all tasks were spawned from inside the functor running within task_group."
+    );
+}
+
+//! Respect task_group_context passed from outside
+//! \brief \ref interface \ref requirement
+TEST_CASE("Respect task_group_context passed from outside") {
+    std::atomic<bool> outer_cancelled{false};
+    bool implicit_cancel = true;
+    std::atomic<unsigned> count{13};
+
+    tbb::task_group_context inner_ctx(tbb::task_group_context::isolated);
+    tbb::task_group inner_tg(inner_ctx);
+
+    tbb::task_group outer_tg;
+    auto outer_tg_task = [&] {
+        inner_tg.run([&] {
+            utils::SpinWaitUntilEq(outer_cancelled, true);
+            inner_tg.run( SelfRunner{inner_tg, count} );
+        });
+
+        utils::try_call([&] {
+            if (implicit_cancel) {
+                throw int();
+            } else {
+                outer_tg.cancel();
+            }
+        }).on_completion([&] {
+            outer_cancelled = true;
+        });
+    };
+
+    auto check = [&] {
+        tbb::task_group_status outer_status = tbb::task_group_status::not_complete;
+        if (implicit_cancel) {
+            try {
+                outer_status = outer_tg.wait();
+            } catch(const int&) {
+                outer_status = tbb::task_group_status::canceled;
+            }
+        } else
+            outer_status = outer_tg.wait();
+        CHECK_MESSAGE(
+            outer_status == tbb::task_group_status::canceled,
+            "Outer task group should have been cancelled."
+        );
+
+        tbb::task_group_status inner_status = inner_tg.wait();
+        CHECK_MESSAGE(
+            inner_status == tbb::task_group_status::complete,
+            "Inner task group should have completed despite the cancellation of the outer one."
+        );
+
+        CHECK_MESSAGE(0 == count, "Some of the inner group tasks were not executed.");
+    };
+
+    outer_tg.run(outer_tg_task);
+    check();
+
+    count = 10;
+    implicit_cancel = false;
+    outer_cancelled = false;
+
+    outer_tg.run(outer_tg_task);
+    check();
+}