Introduce aggregator arguments

Aggregators help pass multiple references to arguments to the
task. The app controls the upper bound and chunk size.

Signed-off-by: Joseph Schuchart <joseph.schuchart@stonybrook.edu>

Author: devreal

ttg/ttg/aggregator.h

@@ -0,0 +1,298 @@
+// SPDX-License-Identifier: BSD-3-Clause
+#ifndef TTG_AGGREGATOR_H
+#define TTG_AGGREGATOR_H
+
+#include <limits>
+#include <vector>
+#include "ttg/fwd.h"
+#include "ttg/edge.h"
+#include "ttg/util/meta.h"
+#include "ttg/terminal.h"
+
+namespace ttg {
+
+  /// @brief Collects a variable number of values sent to the same task invocation.
+  ///
+  /// An Aggregator<T> is the value type of an input generated by ttg::make_aggregator();
+  /// it holds references to (and, indirectly, keeps alive) every value collected for one
+  /// task firing. Aggregators are always const: values are exposed for reading only.
+  template <typename ValueT>
+  struct Aggregator {
+    template <typename keyT, typename output_terminalsT, typename derivedT, typename input_valueTs, ttg::ExecutionSpace Space>
+    friend class TTG_IMPL_NS::TT;
+
+    using decay_value_type = std::decay_t<ValueT>;
+
+    // try to fit the Aggregator into 2 cache lines before falling back to a heap vector
+    static constexpr size_t short_vector_size = 6;
+
+   public:
+    using value_type = std::add_const_t<decay_value_type>;
+
+   private:
+    struct vector_element_t {
+      value_type *value;  // pointer to the value
+      void *ptr;          // pointer to implementation-specific data (the owning data copy)
+      vector_element_t() = default;
+      vector_element_t(value_type *value, void *ptr) : value(value), ptr(ptr) {}
+    };
+    using vector_t = typename std::vector<vector_element_t>;
+
+    template <typename VectorElementT>
+    struct Iterator {
+     private:
+      VectorElementT *m_ptr = nullptr;
+
+     public:
+      using value_type = std::add_const_t<decay_value_type>;
+      using reference = std::add_lvalue_reference_t<value_type>;
+      using pointer = std::add_pointer_t<value_type>;
+      using difference_type = std::ptrdiff_t;
+      using iterator_category = std::forward_iterator_tag;
+
+      Iterator() = default;
+      template <typename Ptr>
+      Iterator(Ptr ptr) : m_ptr(&(*ptr)) {}
+
+      reference operator*() const { return *m_ptr->value; }
+      pointer operator->() const { return m_ptr->value; }
+
+      Iterator &operator++() {
+        ++m_ptr;
+        return *this;
+      }
+      Iterator operator++(int) {
+        Iterator tmp = *this;
+        ++(*this);
+        return tmp;
+      }
+
+      friend bool operator==(const Iterator &a, const Iterator &b) { return a.m_ptr == b.m_ptr; }
+      friend bool operator!=(const Iterator &a, const Iterator &b) { return a.m_ptr != b.m_ptr; }
+    };
+
+   public:
+    using iterator = Iterator<const vector_element_t>;
+    using const_iterator = Iterator<const vector_element_t>;
+    using size_type = typename vector_t::size_type;
+    using pointer = value_type *;
+    using reference = std::add_lvalue_reference_t<value_type>;
+    using const_reference = std::add_const_t<reference>;
+    static constexpr const size_type undef_target = std::numeric_limits<size_type>::max();
+
+    Aggregator() = default;
+
+    Aggregator(size_type target) : m_target(target) {
+      if (target > short_vector_size) {
+        m_vec.reserve(target);
+        m_is_dynamic = true;
+      } else {
+        m_is_dynamic = false;
+      }
+    }
+
+    Aggregator(const Aggregator &) = default;
+    Aggregator(Aggregator &&) = default;
+    Aggregator &operator=(const Aggregator &) = default;
+    Aggregator &operator=(Aggregator &&) = default;
+
+    ~Aggregator() = default;
+
+   private:
+    /// Add an element to the aggregator; called by the runtime only.
+    void add_value(value_type &value, void *ptr = nullptr) {
+      if (m_is_dynamic) {
+        m_vec.emplace_back(&value, ptr);
+      } else {
+        if (m_size < short_vector_size) {
+          m_arr[m_size] = vector_element_t(&value, ptr);
+        } else {
+          move_to_dynamic();
+          m_vec.emplace_back(&value, ptr);
+        }
+      }
+      ++m_size;
+    }
+
+    bool has_target() const { return (m_target != undef_target); }
+
+    size_type target() const {
+      if (m_target == undef_target) {
+        throw std::logic_error("Aggregator has no target defined!");
+      }
+      return m_target;
+    }
+
+    auto data() {
+      if (m_is_dynamic) {
+        return ttg::span(m_vec.data(), m_size);
+      } else {
+        return ttg::span(static_cast<vector_element_t *>(&m_arr[0]), m_size);
+      }
+    }
+
+    void move_to_dynamic() {
+      assert(!m_is_dynamic);
+      vector_t vec;
+      if (has_target()) {
+        vec.reserve(m_target);
+      } else {
+        vec.reserve(m_size);
+      }
+      vec.insert(vec.begin(), &m_arr[0], &m_arr[m_size]);
+      m_vec = std::move(vec);
+      m_is_dynamic = true;
+    }
+
+    const vector_element_t *get_ptr() const {
+      return (m_is_dynamic) ? m_vec.data() : static_cast<const vector_element_t *>(&m_arr[0]);
+    }
+
+   public:
+    const_reference operator[](size_type i) const { return *((m_is_dynamic) ? m_vec[i].value : m_arr[i].value); }
+
+    const_reference at(size_type i) const { return *((m_is_dynamic) ? m_vec.at(i).value : m_arr[i].value); }
+
+    size_type size() const { return m_size; }
+
+    const_iterator begin() const { return const_iterator(get_ptr()); }
+    const_iterator cbegin() const { return const_iterator(get_ptr()); }
+    const_iterator end() const { return const_iterator(get_ptr() + m_size); }
+    const_iterator cend() const { return const_iterator(get_ptr() + m_size); }
+
+   private:
+    vector_t m_vec;
+    vector_element_t m_arr[short_vector_size];
+    size_type m_size = 0;
+    size_type m_target = undef_target;
+    bool m_is_dynamic = true;
+  };
+
+  namespace detail {
+
+    /// Trait to determine if a given type is an Aggregator
+    template <typename T>
+    struct is_aggregator : std::false_type {};
+
+    template <typename ValueT>
+    struct is_aggregator<Aggregator<ValueT>> : std::true_type {};
+
+    template <typename T>
+    constexpr bool is_aggregator_v = is_aggregator<T>::value;
+
+    template <typename KeyT, typename AggregatorT, typename TargetFn>
+    struct AggregatorFactory {
+      using aggregator_type = AggregatorT;
+      using key_type = KeyT;
+
+      AggregatorFactory() : m_targetfn([]() { return aggregator_type::undef_target; }) {}
+
+      AggregatorFactory(TargetFn fn) : m_targetfn(std::forward<TargetFn>(fn)) {}
+
+      auto operator()(const key_type &key) const { return aggregator_type(m_targetfn(key)); }
+
+     private:
+      TargetFn m_targetfn;
+    };
+
+    struct AggregatorTargetProvider {
+      AggregatorTargetProvider(std::size_t target = Aggregator<int>::undef_target) : m_target(target) {}
+
+      template <typename T>
+      auto operator()(const T &) const {
+        return m_target;
+      }
+
+     private:
+      std::size_t m_target;
+    };
+
+  }  // namespace detail
+
+  /// Overload of ttg::Edge for aggregated (variable-arity) inputs; created by ttg::make_aggregator.
+  template <typename KeyT, typename ValueT>
+  class Edge<KeyT, Aggregator<ValueT>> {
+   public:
+    using edge_type = ttg::Edge<KeyT, ValueT>;
+    using aggregator_type = Aggregator<ValueT>;
+    using aggregator_factory_type = std::function<aggregator_type(const KeyT &)>;
+
+    using output_terminal_type = ttg::Out<KeyT, ValueT>;
+    using key_type = KeyT;
+    using value_type = aggregator_type;
+
+    Edge(edge_type &edge)
+        : m_edge(edge), m_aggregator_factory([](const KeyT &) { return aggregator_type(); }) {}
+
+    template <typename AggregatorFactory>
+    Edge(edge_type &edge, AggregatorFactory &&aggregator_factory)
+        : m_edge(edge), m_aggregator_factory([=](const KeyT &key) { return aggregator_factory(key); }) {}
+
+    /// Return reference to the underlying edge
+    edge_type &edge() const { return m_edge; }
+
+    auto aggregator_factory() const { return m_aggregator_factory; }
+
+    /// probes if this is already has at least one input; calls the underlying edge.live()
+    bool live() const { return m_edge.live(); }
+
+    /// call the underlying edge.set_in()
+    void set_in(Out<KeyT, ValueT> *in) const { m_edge.set_in(in); }
+
+    /// call the underlying edge.set_out()
+    void set_out(TerminalBase *out) const { m_edge.set_out(out); }
+
+    /// call the underlying edge.fire()
+    template <typename Key = KeyT, typename Value = ValueT>
+    std::enable_if_t<ttg::meta::is_all_void_v<Key, Value>> fire() const {
+      m_edge.fire();
+    }
+
+   private:
+    edge_type &m_edge;
+    aggregator_factory_type m_aggregator_factory;
+  };
+
+  /// overload for remove_wrapper to expose the underlying value type
+  namespace meta {
+    template <typename T>
+    struct remove_wrapper<Aggregator<T>> {
+      using type = T;
+    };
+    // op() parameters are typically declared as `const ttg::Aggregator<T>&`, which retains the
+    // top-level const after reference removal -- unwrap that form too.
+    template <typename T>
+    struct remove_wrapper<const Aggregator<T>> {
+      using type = T;
+    };
+  }  // namespace meta
+
+  /// @brief Turns an Edge<K,V> into an aggregating input with a per-key value count.
+  /// @param targetfn maps a key to the total number of values expected for that key.
+  template <typename EdgeT, typename TargetFn,
+            typename = std::enable_if_t<std::is_invocable_v<TargetFn, const typename std::decay_t<EdgeT>::key_type>>>
+  auto make_aggregator(EdgeT &&inedge, TargetFn &&targetfn) {
+    using value_type = typename std::decay_t<EdgeT>::value_type;
+    using key_type = typename std::decay_t<EdgeT>::key_type;
+    using fact = typename detail::AggregatorFactory<key_type, Aggregator<value_type>, TargetFn>;
+    return Edge<key_type, Aggregator<value_type>>(inedge, fact(std::forward<TargetFn>(targetfn)));
+  }
+
+  /// @brief Turns an Edge<K,V> into an aggregating input with a fixed value count.
+  template <typename EdgeT>
+  auto make_aggregator(EdgeT &&inedge, size_t target) {
+    return make_aggregator(inedge, typename detail::AggregatorTargetProvider(target));
+  }
+
+  /// @brief Turns an Edge<K,V> into an aggregating input with an unbounded value count
+  /// (the task fires once all producers for a key have completed; see ttg::TT::finalize_argstream).
+  template <typename EdgeT>
+  auto make_aggregator(EdgeT &&inedge) {
+    using value_type = typename std::decay_t<EdgeT>::value_type;
+    using key_type = typename std::decay_t<EdgeT>::key_type;
+    return Edge<key_type, Aggregator<value_type>>(inedge);
+  }
+
+}  // namespace ttg
+
+#endif  // TTG_AGGREGATOR_H