Add async_node implementation and tests

uxlfoundation · Jan 15, 2025 · d9ebc3b · d9ebc3b
1 parent b8149b0
commit d9ebc3b
Show file tree

Hide file tree

Showing 9 changed files with 579 additions and 493 deletions.
diff --git a/include/oneapi/tbb/detail/_flow_graph_body_impl.h b/include/oneapi/tbb/detail/_flow_graph_body_impl.h
@@ -168,47 +168,63 @@ class function_body_leaf< continue_msg, Output, B > : public function_body< cont
     B body;
 };
 
-class multifunction_node_tag;
+#if __TBB_PREVIEW_FLOW_GRAPH_TRY_PUT_AND_WAIT
+class metainfo_tag_type;
+#endif
+
+// TODO: add description
+struct invoke_body_with_tag_helper {
+    using first_priority = int;
+    using second_priority = double;
+
+#if __TBB_PREVIEW_FLOW_GRAPH_TRY_PUT_AND_WAIT
+    template <typename Body, typename... Args>
+    static auto invoke(first_priority, Body&& body, metainfo_tag_type&& tag, Args&&... args)
+        noexcept(noexcept(tbb::detail::invoke(body, std::forward<Args>(args)..., std::move(tag))))
+        -> decltype(tbb::detail::invoke(body, std::forward<Args>(args)..., std::move(tag)))
+    {
+        tbb::detail::invoke(body, std::forward<Args>(args)..., std::move(tag));
+    }
+#endif
+    template <typename Body, typename... Args>
+    static void invoke(second_priority, Body&& body __TBB_FLOW_GRAPH_METAINFO_ARG(metainfo_tag_type&&),
+                       Args&&... args)
+        noexcept(noexcept(tbb::detail::invoke(body, std::forward<Args>(args)...)))
+    {
+        tbb::detail::invoke(body, std::forward<Args>(args)...);
+    }
+};
+
+// TODO: add comment
+template <typename Body, typename... Args>
+void invoke_body_with_tag(Body&& body __TBB_FLOW_GRAPH_METAINFO_ARG(metainfo_tag_type&& tag), Args&&... args)
+    noexcept(noexcept(invoke_body_with_tag_helper::invoke(1, std::forward<Body>(body) __TBB_FLOW_GRAPH_METAINFO_ARG(std::move(tag)),
+                                                          std::forward<Args>(args)...)))
+{
+    invoke_body_with_tag_helper::invoke(/*overload priority helper*/1,
+                                        std::forward<Body>(body) __TBB_FLOW_GRAPH_METAINFO_ARG(std::move(tag)),
+                                        std::forward<Args>(args)...);
+}
+
 
 //! function_body that takes an Input and a set of output ports
 template<typename Input, typename OutputSet>
 class multifunction_body : no_assign {
 public:
     virtual ~multifunction_body () {}
-    virtual void operator()(const Input &/* input*/, OutputSet &/*oset*/ __TBB_FLOW_GRAPH_METAINFO_ARG(multifunction_node_tag&& /*tag*/)) = 0;
+    virtual void operator()(const Input &/* input*/, OutputSet &/*oset*/ __TBB_FLOW_GRAPH_METAINFO_ARG(metainfo_tag_type&& /*tag*/)) = 0;
     virtual multifunction_body* clone() = 0;
     virtual void* get_body_ptr() = 0;
 };
 
 //! leaf for multifunction.  OutputSet can be a std::tuple or a vector.
 template<typename Input, typename OutputSet, typename B>
 class multifunction_body_leaf : public multifunction_body<Input, OutputSet> {
-    using first_priority = int;
-    using second_priority = double;
-
-    // body may explicitly put() to one or more of oset.
-    void invoke_body_impl(const Input& input, OutputSet& oset __TBB_FLOW_GRAPH_METAINFO_ARG(multifunction_node_tag&&), second_priority)
-    {
-        tbb::detail::invoke(body, input, oset);
-    }
-
-#if __TBB_PREVIEW_FLOW_GRAPH_TRY_PUT_AND_WAIT
-    template <typename InputT, typename OutputSetT>
-    auto invoke_body_impl(const InputT& input, OutputSetT& oset, multifunction_node_tag&& tag, first_priority)
-    -> decltype(tbb::detail::invoke(std::declval<B>(), input, oset, std::move(tag)), void())
-    {
-        tbb::detail::invoke(body, input, oset, std::move(tag));
-    }
-#endif
-
-    void invoke_body(const Input& input, OutputSet& oset __TBB_FLOW_GRAPH_METAINFO_ARG(multifunction_node_tag&& tag)) {
-        invoke_body_impl(input, oset __TBB_FLOW_GRAPH_METAINFO_ARG(std::move(tag)), 1);
-    }
-
 public:
     multifunction_body_leaf(const B &_body) : body(_body) { }
-    void operator()(const Input &input, OutputSet &oset __TBB_FLOW_GRAPH_METAINFO_ARG(multifunction_node_tag&& tag)) override {
-        invoke_body(input, oset __TBB_FLOW_GRAPH_METAINFO_ARG(std::move(tag)));
+    // body may explicitly put() to one or more of oset.
+    void operator()(const Input &input, OutputSet &oset __TBB_FLOW_GRAPH_METAINFO_ARG(metainfo_tag_type&& tag)) override {
+        invoke_body_with_tag(body __TBB_FLOW_GRAPH_METAINFO_ARG(std::move(tag)), input, oset);
     }
     void* get_body_ptr() override { return &body; }
     multifunction_body_leaf* clone() override {

diff --git a/include/oneapi/tbb/detail/_flow_graph_cache_impl.h b/include/oneapi/tbb/detail/_flow_graph_cache_impl.h
@@ -419,12 +419,13 @@ class broadcast_cache : public successor_cache<T, M> {
 #endif
 
     // call try_put_task and return list of received tasks
-    bool gather_successful_try_puts( const T &t, graph_task_list& tasks ) {
+    bool gather_successful_try_puts( const T &t, graph_task_list& tasks
+                                     __TBB_FLOW_GRAPH_METAINFO_ARG(const message_metainfo& metainfo) ) {
         bool is_at_least_one_put_successful = false;
         typename mutex_type::scoped_lock l(this->my_mutex, /*write=*/true);
         typename successors_type::iterator i = this->my_successors.begin();
         while ( i != this->my_successors.end() ) {
-            graph_task * new_task = (*i)->try_put_task(t);
+            graph_task * new_task = (*i)->try_put_task(t, metainfo);
             if(new_task) {
                 ++i;
                 if(new_task != SUCCESSFULLY_ENQUEUED) {

diff --git a/include/oneapi/tbb/detail/_flow_graph_node_impl.h b/include/oneapi/tbb/detail/_flow_graph_node_impl.h
@@ -556,23 +556,23 @@ struct init_output_ports {
 
 #if __TBB_PREVIEW_FLOW_GRAPH_TRY_PUT_AND_WAIT
 
-class multifunction_node_tag {
+class metainfo_tag_type {
 public:
-    multifunction_node_tag() = default;
+    metainfo_tag_type() = default;
 
-    multifunction_node_tag(const multifunction_node_tag&) = delete;
+    metainfo_tag_type(const metainfo_tag_type&) = delete;
 
-    multifunction_node_tag(multifunction_node_tag&& other)
+    metainfo_tag_type(metainfo_tag_type&& other)
         : my_metainfo(std::move(other.my_metainfo)) {}
 
-    multifunction_node_tag(const message_metainfo& metainfo) : my_metainfo(metainfo) {
+    metainfo_tag_type(const message_metainfo& metainfo) : my_metainfo(metainfo) {
         for (auto waiter : my_metainfo.waiters()) {
             waiter->reserve();
         }
     }
 
-    multifunction_node_tag& operator=(const multifunction_node_tag&) = delete;
-    multifunction_node_tag& operator=(multifunction_node_tag&& other) {
+    metainfo_tag_type& operator=(const metainfo_tag_type&) = delete;
+    metainfo_tag_type& operator=(metainfo_tag_type&& other) {
         // TODO: should this method be thread-safe?
         if (this != &other) {
             reset();
@@ -581,11 +581,11 @@ class multifunction_node_tag {
         return *this;
     }
 
-    ~multifunction_node_tag() {
+    ~metainfo_tag_type() {
         reset();
     }
 
-    void merge(const multifunction_node_tag& other_tag) {
+    void merge(const metainfo_tag_type& other_tag) {
         tbb::spin_mutex::scoped_lock lock(my_mutex);
 
         // TODO: add comment
@@ -604,13 +604,18 @@ class multifunction_node_tag {
         my_metainfo = message_metainfo{};
     }
 private:
-    template <typename Output>
-    friend class multifunction_output;
+    friend class metainfo_tag_accessor;
 
     message_metainfo my_metainfo;
     tbb::spin_mutex  my_mutex;
 };
 
+struct metainfo_tag_accessor {
+    static const message_metainfo& get_metainfo(const metainfo_tag_type& tag) {
+        return tag.my_metainfo;
+    }
+};
+
 #endif
 
 //! Implements methods for a function node that takes a type Input as input
@@ -622,7 +627,7 @@ class multifunction_input : public function_input_base<Input, Policy, A, multifu
     typedef Input input_type;
     typedef OutputPortSet output_ports_type;
 #if __TBB_PREVIEW_FLOW_GRAPH_TRY_PUT_AND_WAIT
-    typedef multifunction_node_tag tag_type;
+    typedef metainfo_tag_type tag_type;
 #endif
     typedef multifunction_body<input_type, output_ports_type> multifunction_body_type;
     typedef multifunction_input<Input, OutputPortSet, Policy, A> my_class;
@@ -664,7 +669,7 @@ class multifunction_input : public function_input_base<Input, Policy, A, multifu
                                         __TBB_FLOW_GRAPH_METAINFO_ARG(const message_metainfo& metainfo) )
     {
 #if __TBB_PREVIEW_FLOW_GRAPH_TRY_PUT_AND_WAIT
-        multifunction_node_tag tag(metainfo);
+        metainfo_tag_type tag(metainfo);
 #endif
         fgt_begin_body( my_body );
         (*my_body)(i, my_output_ports __TBB_FLOW_GRAPH_METAINFO_ARG(std::move(tag)));
@@ -920,13 +925,13 @@ class multifunction_output : public function_output<Output> {
     }
 
 #if __TBB_PREVIEW_FLOW_GRAPH_TRY_PUT_AND_WAIT
-    bool try_put(const output_type& i, const multifunction_node_tag& tag) {
-        return try_put_impl(i, tag.my_metainfo);
+    bool try_put(const output_type& i, const metainfo_tag_type& tag) {
+        return try_put_impl(i, metainfo_tag_accessor::get_metainfo(tag));
     }
 
-    bool try_put(const output_type& i, multifunction_node_tag&& tag) {
-        multifunction_node_tag local_tag = std::move(tag);
-        return try_put_impl(i, local_tag.my_metainfo);
+    bool try_put(const output_type& i, metainfo_tag_type&& tag) {
+        metainfo_tag_type local_tag = std::move(tag);
+        return try_put_impl(i, metainfo_tag_accessor::get_metainfo(local_tag));
     }
 #endif
 

diff --git a/include/oneapi/tbb/flow_graph.h b/include/oneapi/tbb/flow_graph.h
@@ -3130,8 +3130,11 @@ class async_body: public async_body_base<Gateway> {
     async_body(const Body &body, gateway_type *gateway)
         : base_type(gateway), my_body(body) { }
 
-    void operator()( const Input &v, Ports & ) noexcept(noexcept(tbb::detail::invoke(my_body, v, std::declval<gateway_type&>()))) {
-        tbb::detail::invoke(my_body, v, *this->my_gateway);
+    void operator()( const Input &v, Ports & __TBB_FLOW_GRAPH_METAINFO_ARG(metainfo_tag_type&& tag) )
+        noexcept(noexcept(invoke_body_with_tag(my_body __TBB_FLOW_GRAPH_METAINFO_ARG(std::move(tag)),
+                                               v, *this->my_gateway)))
+    {
+        invoke_body_with_tag(my_body __TBB_FLOW_GRAPH_METAINFO_ARG(std::move(tag)), v, *this->my_gateway);
     }
 
     Body get_body() { return my_body; }
@@ -3176,9 +3179,20 @@ class async_node
 
         //! Implements gateway_type::try_put for an external activity to submit a message to FG
         bool try_put(const Output &i) override {
-            return my_node->try_put_impl(i);
+            return my_node->try_put_impl(i __TBB_FLOW_GRAPH_METAINFO_ARG(message_metainfo{}));
+        }
+
+#if __TBB_PREVIEW_FLOW_GRAPH_TRY_PUT_AND_WAIT
+        bool try_put(const Output &i, const metainfo_tag_type& tag) override {
+            return my_node->try_put_impl(i, metainfo_tag_accessor::get_metainfo(tag));
         }
 
+        bool try_put(const Output &i, metainfo_tag_type&& tag) override {
+            metainfo_tag_type local_tag = std::move(tag);
+            return my_node->try_put_impl(i, metainfo_tag_accessor::get_metainfo(local_tag));
+        }
+#endif
+
     private:
         async_node* my_node;
     } my_gateway;
@@ -3187,13 +3201,14 @@ class async_node
     async_node* self() { return this; }
 
     //! Implements gateway_type::try_put for an external activity to submit a message to FG
-    bool try_put_impl(const Output &i) {
+    bool try_put_impl(const Output &i __TBB_FLOW_GRAPH_METAINFO_ARG(const message_metainfo& metainfo)) {
         multifunction_output<Output> &port_0 = output_port<0>(*this);
         broadcast_cache<output_type>& port_successors = port_0.successors();
         fgt_async_try_put_begin(this, &port_0);
         // TODO revamp: change to std::list<graph_task*>
         graph_task_list tasks;
-        bool is_at_least_one_put_successful = port_successors.gather_successful_try_puts(i, tasks);
+        bool is_at_least_one_put_successful =
+            port_successors.gather_successful_try_puts(i, tasks __TBB_FLOW_GRAPH_METAINFO_ARG(metainfo));
         __TBB_ASSERT( is_at_least_one_put_successful || tasks.empty(),
                       "Return status is inconsistent with the method operation." );
 

diff --git a/include/oneapi/tbb/flow_graph_abstractions.h b/include/oneapi/tbb/flow_graph_abstractions.h
@@ -33,6 +33,9 @@ class graph_proxy {
     virtual ~graph_proxy() {}
 };
 
+#if __TBB_PREVIEW_FLOW_GRAPH_TRY_PUT_AND_WAIT
+class metainfo_tag_type;
+#endif
 template <typename Input>
 class receiver_gateway : public graph_proxy {
 public:
@@ -41,6 +44,10 @@ class receiver_gateway : public graph_proxy {
 
     //! Submit signal from an asynchronous activity to FG.
     virtual bool try_put(const input_type&) = 0;
+#if __TBB_PREVIEW_FLOW_GRAPH_TRY_PUT_AND_WAIT
+    virtual bool try_put(const input_type&, const metainfo_tag_type&) = 0;
+    virtual bool try_put(const input_type&, metainfo_tag_type&&) = 0;
+#endif
 };
 
 } // d2

diff --git a/test/tbb/test_async_node.cpp b/test/tbb/test_async_node.cpp
@@ -32,6 +32,7 @@
 #include "common/spin_barrier.h"
 #include "common/test_follows_and_precedes_api.h"
 #include "common/concepts_common.h"
+#include "test_try_put_and_wait.h"
 
 #include <string>
 #include <thread>
@@ -802,9 +803,9 @@ TEST_CASE("Basic tests"){
 
 //! NativeParallelFor test with various concurrency settings
 //! \brief \ref requirement \ref error_guessing
-TEST_CASE("Lightweight tests"){
-    lightweight_testing::test<tbb::flow::async_node>(NUMBER_OF_MSGS);
-}
+// TEST_CASE("Lightweight tests"){
+//     lightweight_testing::test<tbb::flow::async_node>(NUMBER_OF_MSGS);
+// }
 
 //! Test reset and cancellation
 //! \brief \ref error_guessing
@@ -878,3 +879,11 @@ TEST_CASE("constraints for async_node body") {
 }
 
 #endif // __TBB_CPP20_CONCEPTS_PRESENT
+
+#if __TBB_PREVIEW_FLOW_GRAPH_TRY_PUT_AND_WAIT
+//! \brief \ref error_guessing
+TEST_CASE("test async_node try_put_and_wait") {
+    using node_type = oneapi::tbb::flow::async_node<int, int, tbb::flow::queueing>;
+    test_try_put_and_wait::test_multioutput<node_type>();
+}
+#endif