pairwise_sum.hpp

Go to the documentation of this file.

#ifndef QUICKSTATS_PAIRWISE_SUM_HPP
#define QUICKSTATS_PAIRWISE_SUM_HPP
 
#include <vector>
#include <cstddef>
#include <optional>
#include <cassert>
#include <array>
#include <stdexcept>
 
namespace quickstats {
 
template<typename Output_ = double>
struct PairwiseSumWorkspace {
    struct State {
        State(const std::size_t left_end) : left_end(left_end) {}
        std::size_t left_end;
        std::optional<Output_> left_sum;
    };
    std::vector<State> states;
};
 
struct PairwiseSumOptions {
    std::size_t max_sum_length = 128;
};
 
// Mathematically equivalent to std::accumulate but reorders summations for greater instruction-level parallelism.
// See performance tests in https://github.com/tatami-inc/test-multiplication/tree/master/other/accumulators.
template<typename Output_, std::size_t width_>
Output_ recursive_sum(std::array<Output_, width_>& dots) {
    if constexpr(width_ == 1) {
        return dots[0];
    } else if constexpr(width_ == 2) {
        return dots[0] + dots[1];
    } else {
        constexpr auto half_width = width_ / 2;
        std::array<Output_, half_width> tmp;
        for (std::size_t s = 0; s < half_width; ++s) { // Increase potential for vectorization.
            tmp[s] = dots[s] + dots[s + half_width];
        }
        if constexpr(width_ % 2 == 1) {
            return recursive_sum(tmp) + dots[width_ - 1];
        } else {
            return recursive_sum(tmp);
        }
    }
}
template<std::size_t accumulators_ = 4, class Input_, typename Output_>
Output_ pairwise_sum_abstract(const std::size_t num_total, Input_ input, PairwiseSumWorkspace<Output_>& work, const PairwiseSumOptions& options) {
    static_assert(accumulators_ > 0);
 
    work.states.clear();
    if (num_total < accumulators_) {
        if constexpr(accumulators_ == 1) {
            return 0;
        } else {
            Output_ out = 0;
            for (std::size_t i = 0; i < num_total; ++i) {
                out += input(i);
            }
            return out;
        }
    }
 
    const std::size_t limit = options.max_sum_length;
    if (accumulators_ > limit / 2) {
        throw std::runtime_error("'max_sum_length' must be greater than '2 * accumulators_'");
    }
 
    std::size_t start = 0, end = num_total; 
    Output_ out = 0;
    while (1) {
        const std::size_t len = end - start;
        if (len > limit) {
            work.states.emplace_back(end);
            end = start + len / 2;
            continue;
        }
 
        // 'start + accumulators_ <= end' should always be true. 
        //
        // Let's start by considering the initial case where 'num_total <= limit'.
        // As we already know that 'num_total >= accumulators_', we get 'start + accumulators_ == accumulators_ <= num_total == end'.
        //
        // Alright, what about the left-hand-side of the recursion?
        // We defined 'end = start + len / 2', and we already know that 'len > limit' and 'limit / 2 >= accumulators_';
        // hence, we know that that 'start + len / 2 >= start + accumulators_'.
        //
        // The right-hand-side of the recursion is easier as we know its length is greater than or equal to 'len / 2' (as integer division is truncating).
        // So, if the LHS fulfills the requirement, then the RHS must definitely fulfill it.
        assert(start + accumulators_ <= end);
 
        Output_ tmp;
        if constexpr(accumulators_ == 1) {
            tmp = input(start); // We know that start < end, so we can skip one addition. 
            for (std::size_t i = start + 1; i < end; ++i) {
                tmp += input(i);
            }
 
        } else {
            // This accumulator logic was originally implemented in https://github.com/tatami-inc/tatami_mult.
            // We added peeling as we can guarantee that we have enough observations and thus can omit the conditional.
            std::array<Output_, accumulators_> partials; 
            for (std::size_t a = 0; a < accumulators_; ++a) { // peeling the first loop as we know that start + accumulators_ <= end.
                partials[a] = input(start + a);
            }
 
            const std::size_t num_cycles = len / accumulators_;
            const std::size_t remainder = len % accumulators_;
            for (std::size_t c = 1; c < num_cycles; ++c) {
                for (std::size_t a = 0; a < accumulators_; ++a) {
                    const std::size_t idx = start + c * accumulators_ + a;
                    partials[a] += input(idx);
                }
            }
 
            // Technically, we could structure the splits to reduce the number of calls to the epilogue loops.
            // However, this would weaken the symmetry of the splitting and compromise the precision improvements.
            tmp = 0;
            for (std::size_t i = 0; i < remainder; ++i) {
                const std::size_t idx = start + num_cycles * accumulators_ + i;
                tmp += input(idx);
            }
 
            tmp += recursive_sum(partials);
        }
 
        start = end;
        while (work.states.size() && work.states.back().left_sum.has_value()) {
            tmp += *(work.states.back().left_sum);
            start = work.states.back().left_end;
            work.states.pop_back();
        }
 
        if (work.states.empty()) {
            out = tmp;
            break;
        }
 
        work.states.back().left_sum = tmp;
        end = work.states.back().left_end;
    }
 
    return out;
}
 
template<std::size_t accumulators_ = 4, typename Input_, typename Output_>
Output_ pairwise_sum(const std::size_t num_total, const Input_* const ptr, PairwiseSumWorkspace<Output_>& work, const PairwiseSumOptions& options) {
    return pairwise_sum_abstract<accumulators_>(
        num_total,
        [&](const std::size_t i) -> auto {
            return ptr[i];
        },
        work,
        options
    );
}
 
/* COMMENTS:
 * I tried to write a multi-threaded version of this where each direct summation was submitted to a separate worker until all workers were occupied,
 * and then added the results once they became available from each worker. 
 * This worksharing is fine-grained but imposes a high cost for inter-thread communication relative to the summation for small `limit_`. 
 * As a consequence, the performance of this multi-threaded version is worse than its serial counterpart.
 *
 * I could have implemented alternative approaches that involve less communication but require more memory.
 * For example, we could split elements into the subarrays ahead of time, distribute the summations to threads once, and then sum the results once all workers are done.
 * This greatly reduces the cross-talk between threads but requires an extra allocation to store the results.
 *
 * TBH, the easiest and most performant approach to parallelization is to just split your input array into one subarray per worker,
 * perform the sum within each worker for that subarray, and then add the sums afterwards.
 * This won't give exactly the same result as serial execution but we've crossed that bridge already.
 * (If exact results are required, we can split it into ceil(log2(num_workers)) subarrays,
 * which allows us to follow the same halving as pairwise_sum() to get the exact same result at the cost of suboptimal worksharing.)
 *
 * In any case, summation is already so fast that I don't think we need to spend a lot of effort in thinking about parallelization.
 * Especially given that, in real applications, the other threads will typically be occupied elsewhere.
 * Indeed, we don't deal with parallelization in other parts of this library, so it would be odd to implement it here. 
 *
 * I also have a sneaking suspicion that the serial code is already pseudo-parallelized via out-of-order execution,
 * where the next summation starts before the first one has ended.
 * I say this because pairwise_sum() somehow manages to be slightly faster than std::accumulate() in our R bindings.
 */
 
// For back-compatibility.
template<std::size_t limit_ = 128, std::size_t accumulators_ = 4, class Input_, class Modifier_, typename Output_>
Output_ pairwise_sum(const std::size_t num_total, const Input_* const ptr, Modifier_ mod, PairwiseSumWorkspace<Output_>& work) {
    return pairwise_sum_abstract<accumulators_>(
        num_total,
        [&](const std::size_t i) -> auto { 
            return mod(i, ptr[i]);
        },
        work,
        [&]{
            PairwiseSumOptions opt;
            opt.max_sum_length = limit_;
            return opt;
        }()
    );
}
 
template<std::size_t limit_ = 128, std::size_t accumulators_ = 4, typename Input_, typename Output_>
Output_ pairwise_sum(const std::size_t num_total, const Input_* const ptr, PairwiseSumWorkspace<Output_>& work) {
    return pairwise_sum<accumulators_>(
        num_total,
        ptr,
        work,
        [&]{
            PairwiseSumOptions opt;
            opt.max_sum_length = limit_;
            return opt;
        }()
    );
}
}
 
#endif