Execution.h

Go to the documentation of this file.

#pragma once

 
#include <algorithm>
#include <concepts>
#include <cstddef>
#include <iterator>
#include <span>
#include <type_traits>
#include <utility>
#include <vector>
 
#include <EmrysFilter/Core/Allocation.h>
#include <EmrysFilter/Core/Macros.h>
 
namespace Emrys
{


namespace execution
{

struct Sequential
{};

inline constexpr Sequential sequential{};
 
}  // namespace execution
 
namespace internal
{

class CountingIterator
{
public:
 
    using iterator_concept = std::random_access_iterator_tag;   
    using iterator_category = std::random_access_iterator_tag;  
    using value_type = Eigen::Index;                            
    using difference_type = Eigen::Index;                       
    using pointer = void;                                       
    using reference = Eigen::Index;                             
 
    CountingIterator() = default;

    explicit CountingIterator(Eigen::Index index)
        : m_index{index}
    {}

    [[nodiscard]] Eigen::Index
    operator*() const
    {
        return m_index;
    }

    [[nodiscard]] Eigen::Index
    operator[](difference_type n) const
    {
        return m_index + n;
    }

    CountingIterator &
    operator++()
    {
        ++m_index;
        return *this;
    }

    CountingIterator
    operator++(int)
    {
        CountingIterator const before = *this;
        ++m_index;
        return before;
    }

    CountingIterator &
    operator--()
    {
        --m_index;
        return *this;
    }

    CountingIterator
    operator--(int)
    {
        CountingIterator const before = *this;
        --m_index;
        return before;
    }

    CountingIterator &
    operator+=(difference_type n)
    {
        m_index += n;
        return *this;
    }

    CountingIterator &
    operator-=(difference_type n)
    {
        m_index -= n;
        return *this;
    }

    [[nodiscard]] friend CountingIterator
    operator+(CountingIterator it, difference_type n)
    {
        return CountingIterator{it.m_index + n};
    }

    [[nodiscard]] friend CountingIterator
    operator+(difference_type n, CountingIterator it)
    {
        return CountingIterator{it.m_index + n};
    }

    [[nodiscard]] friend CountingIterator
    operator-(CountingIterator it, difference_type n)
    {
        return CountingIterator{it.m_index - n};
    }

    [[nodiscard]] friend difference_type
    operator-(CountingIterator a, CountingIterator b)
    {
        return a.m_index - b.m_index;
    }

    [[nodiscard]] friend auto
    operator<=>(CountingIterator const &, CountingIterator const &) = default;
 
private:
 
    Eigen::Index m_index = 0;  
};
 
static_assert(std::random_access_iterator<CountingIterator>);

template <typename Range>
class ElementIterator
{
public:
 
    using iterator_concept = std::random_access_iterator_tag;   
    using iterator_category = std::random_access_iterator_tag;  
    using value_type = typename Range::value_type;              
    using difference_type = Eigen::Index;                       
    using pointer = void;                                       
    using reference = value_type;                               
 
    ElementIterator() = default;

    ElementIterator(Range const & range, Eigen::Index index)
        : m_range{&range}
        , m_index{index}
    {}

    [[nodiscard]] value_type
    operator*() const
    {
        return (*m_range)[m_index];
    }

    [[nodiscard]] value_type
    operator[](difference_type n) const
    {
        return (*m_range)[m_index + n];
    }

    ElementIterator &
    operator++()
    {
        ++m_index;
        return *this;
    }

    ElementIterator
    operator++(int)
    {
        ElementIterator const before = *this;
        ++m_index;
        return before;
    }

    ElementIterator &
    operator--()
    {
        --m_index;
        return *this;
    }

    ElementIterator
    operator--(int)
    {
        ElementIterator const before = *this;
        --m_index;
        return before;
    }

    ElementIterator &
    operator+=(difference_type n)
    {
        m_index += n;
        return *this;
    }

    ElementIterator &
    operator-=(difference_type n)
    {
        m_index -= n;
        return *this;
    }

    [[nodiscard]] friend ElementIterator
    operator+(ElementIterator it, difference_type n)
    {
        return it += n;
    }

    [[nodiscard]] friend ElementIterator
    operator+(difference_type n, ElementIterator it)
    {
        return it += n;
    }

    [[nodiscard]] friend ElementIterator
    operator-(ElementIterator it, difference_type n)
    {
        return it -= n;
    }

    [[nodiscard]] friend difference_type
    operator-(ElementIterator const & a, ElementIterator const & b)
    {
        return a.m_index - b.m_index;
    }

    [[nodiscard]] friend bool
    operator==(ElementIterator const & a, ElementIterator const & b)
    {
        return a.m_index == b.m_index;
    }

    [[nodiscard]] friend auto
    operator<=>(ElementIterator const & a, ElementIterator const & b)
    {
        return a.m_index <=> b.m_index;
    }
 
private:
 
    Range const * m_range = nullptr;  
    Eigen::Index m_index = 0;         
};

template <typename P>
concept StandardExecutionPolicy = requires(P && policy, CountingIterator i, void (*f)(Eigen::Index)) {
    std::for_each(std::forward<P>(policy), i, i, f);
};

struct IndexFunction
{
    void
    operator()(Eigen::Index i) const;
};

template <typename P>
concept UserExecutionPolicy = requires(std::remove_cvref_t<P> const & policy, IndexFunction const & f) {
    policy.forEachIndex(Eigen::Index{}, f);
};

template <typename P>
concept UserPolicyWithAllocator =
    UserExecutionPolicy<P> && requires(std::remove_cvref_t<P> const & policy) { policy.get_allocator(); };
 
}  // namespace internal

template <typename P>
concept ExecutionPolicy = std::same_as<std::remove_cvref_t<P>, execution::Sequential>
                       || internal::UserExecutionPolicy<P> || internal::StandardExecutionPolicy<P>;
 
namespace internal
{

template <typename P>
inline constexpr bool isSequential = std::same_as<std::remove_cvref_t<P>, execution::Sequential>;

template <typename... Args>
inline constexpr bool leadsWithPolicy = false;

template <typename First, typename... Rest>
inline constexpr bool leadsWithPolicy<First, Rest...> = ExecutionPolicy<First>;

template <typename T, typename P>
[[nodiscard]] auto
scratchVector(P const & policy, std::size_t size)
{
    if constexpr (UserPolicyWithAllocator<P>)
    {
        using Allocator = ReBound<std::remove_cvref_t<decltype(policy.get_allocator())>, T>;
        return std::vector<T, Allocator>(size, Allocator(policy.get_allocator()));
    }
    else
    {
        return std::vector<T>(size);
    }
}

inline constexpr Eigen::Index chunkSize = 1024;

[[nodiscard]] constexpr Eigen::Index
chunkCount(Eigen::Index count)
{
    return (count + chunkSize - 1) / chunkSize;
}

template <ExecutionPolicy P, typename F>
void
forEachIndex([[maybe_unused]] P && policy, Eigen::Index count, F && f)
{
    EMRYS_FILTER_ASSERT(count >= 0);
    if constexpr (isSequential<P>)
    {
        for (Eigen::Index i = 0; i < count; ++i)
        {
            f(i);
        }
    }
    else if constexpr (UserExecutionPolicy<P>)
    {
        std::as_const(policy).forEachIndex(count, [&f](Eigen::Index i) {
            f(i);
        });
    }
    else
    {
        std::for_each(std::forward<P>(policy), CountingIterator{0}, CountingIterator{count}, [&f](Eigen::Index i) {
            f(i);
        });
    }
}

template <ExecutionPolicy P, typename F>
void
forEachChunk(P && policy, Eigen::Index count, F && f)
{
    EMRYS_FILTER_ASSERT(count >= 0);
    if constexpr (isSequential<P>)
    {
        if (count > 0)
        {
            f(Eigen::Index{0}, count);
        }
    }
    else
    {
        forEachIndex(std::forward<P>(policy), chunkCount(count), [&](Eigen::Index chunk) {
            Eigen::Index const begin = chunk * chunkSize;
            f(begin, std::min(chunkSize, count - begin));
        });
    }
}

template <ExecutionPolicy P, typename Partial, typename Combine>
[[nodiscard]] auto
reduceRanges(P && policy, Eigen::Index count, Partial const & partial, Combine const & combine)
{
    using T = std::remove_cvref_t<std::invoke_result_t<Partial const &, Eigen::Index, Eigen::Index>>;
    EMRYS_FILTER_ASSERT(count >= 0);
    if constexpr (isSequential<P>)
    {
        return T{partial(Eigen::Index{0}, count)};
    }
    else
    {
        Eigen::Index const chunks = chunkCount(count);
        // An empty range, as one chunk, is one call to partial(), as under Sequential.
        if (chunks <= 1)
        {
            return T{partial(Eigen::Index{0}, count)};
        }
        auto partials = scratchVector<T>(policy, static_cast<std::size_t>(chunks));
        forEachChunk(std::forward<P>(policy), count, [&](Eigen::Index begin, Eigen::Index length) {
            partials[static_cast<std::size_t>(begin / chunkSize)] = partial(begin, length);
        });
        T result = std::move(partials[0]);
        for (std::size_t c = 1; c < partials.size(); ++c)
        {
            result = combine(std::move(result), partials[c]);
        }
        return result;
    }
}

template <ExecutionPolicy P, typename T, typename Combine, typename F>
[[nodiscard]] T
transformReduce(P && policy, Eigen::Index count, T const & identity, Combine const & combine, F const & f)
{
    return reduceRanges(
        std::forward<P>(policy),
        count,
        [&](Eigen::Index begin, Eigen::Index length) {
            T result = identity;
            for (Eigen::Index i = begin; i < begin + length; ++i)
            {
                result = combine(std::move(result), T(f(i)));
            }
            return result;
        },
        combine
    );
}

template <ExecutionPolicy P, typename T>
void
cumulativeSum(P && policy, std::span<T const> in, std::span<T> out)
{
    EMRYS_FILTER_ASSERT(in.size() == out.size() && "the scan writes one sum per element");
    Eigen::Index const count = Eigen::Index(in.size());
    // Read each element before writing its sum, so that in and out may be the same.
    auto const scan = [&](Eigen::Index begin, Eigen::Index length, T const * offset) {
        auto const first = std::size_t(begin);
        T sum = offset ? T(*offset + in[first]) : in[first];
        out[first] = sum;
        for (std::size_t i = first + 1; i < std::size_t(begin + length); ++i)
        {
            sum += in[i];
            out[i] = sum;
        }
    };
    if (count == 0)
    {
        return;
    }
    if constexpr (isSequential<P>)
    {
        scan(0, count, nullptr);
    }
    else
    {
        Eigen::Index const chunks = chunkCount(count);
        if (chunks == 1)
        {
            scan(0, count, nullptr);
            return;
        }
        auto offsets = scratchVector<T>(policy, static_cast<std::size_t>(chunks));
        forEachChunk(policy, count, [&](Eigen::Index begin, Eigen::Index length) {
            auto const first = std::size_t(begin);
            T total = in[first];
            for (std::size_t i = first + 1; i < std::size_t(begin + length); ++i)
            {
                total += in[i];
            }
            offsets[std::size_t(begin / chunkSize)] = total;
        });
        // Chunk c starts from the sum of the chunks before it; chunk 0 starts from nothing.
        T carry = offsets[0];
        for (std::size_t c = 1; c < offsets.size(); ++c)
        {
            T const total = offsets[c];
            offsets[c] = carry;
            carry += total;
        }
        forEachChunk(std::forward<P>(policy), count, [&](Eigen::Index begin, Eigen::Index length) {
            std::size_t const chunk = std::size_t(begin / chunkSize);
            scan(begin, length, chunk == 0 ? nullptr : &offsets[chunk]);
        });
    }
}

template <ExecutionPolicy P, typename FromDerived, typename To>
void
gather(P && policy, std::span<Eigen::Index const> indices, Eigen::MatrixBase<FromDerived> const & from, To && to)
{
    EMRYS_FILTER_ASSERT(to.cols() == Eigen::Index(indices.size()) && "gather one column per index");
    forEachIndex(std::forward<P>(policy), Eigen::Index(indices.size()), [&](Eigen::Index j) {
        Eigen::Index const source = indices[std::size_t(j)];
        EMRYS_FILTER_ASSERT(source >= 0 && source < from.cols() && "a gathered index is out of range");
        to.col(j) = from.col(source);
    });
}

template <ExecutionPolicy P, typename Keep, typename Move>
Eigen::Index
copyIfIndex(P && policy, Eigen::Index count, Keep const & keep, Move const & move)
{
    EMRYS_FILTER_ASSERT(count >= 0);
    if constexpr (isSequential<P>)
    {
        Eigen::Index out = 0;
        for (Eigen::Index i = 0; i < count; ++i)
        {
            if (keep(i))
            {
                move(i, out++);
            }
        }
        return out;
    }
    else
    {
        if (count == 0)
        {
            return 0;
        }
        auto positions = scratchVector<Eigen::Index>(policy, static_cast<std::size_t>(count));
        forEachIndex(policy, count, [&](Eigen::Index i) {
            positions[std::size_t(i)] = keep(i) ? 1 : 0;
        });
        cumulativeSum(policy, std::span<Eigen::Index const>{positions}, std::span<Eigen::Index>{positions});
        forEachIndex(std::forward<P>(policy), count, [&](Eigen::Index i) {
            Eigen::Index const before = i == 0 ? 0 : positions[std::size_t(i - 1)];
            if (positions[std::size_t(i)] != before)
            {
                move(i, before);
            }
        });
        return positions.back();
    }
}

template <ExecutionPolicy P, typename T>
void
fill(P && policy, std::span<T> out, T const & value)
{
    forEachChunk(std::forward<P>(policy), Eigen::Index(out.size()), [&](Eigen::Index begin, Eigen::Index length) {
        std::fill_n(out.begin() + begin, length, value);
    });
}
 
}  // namespace internal
 
}  // namespace Emrys