parallel.hpp

Go to the documentation of this file.

#ifndef IRLBA_PARALLEL_HPP
#define IRLBA_PARALLEL_HPP
 
#include "utils.hpp"
#include <vector>
#include "Eigen/Dense"
 
#ifndef IRLBA_CUSTOM_PARALLEL
#include "subpar/subpar.hpp"
#endif
 
namespace irlba {
 
template<typename Task_, class Run_>
void parallelize(Task_ num_tasks, Run_ run_task) {
#ifndef IRLBA_CUSTOM_PARALLEL
    // Use cases here don't allocate or throw, so nothrow_ = true is fine.
    subpar::parallelize_simple<true>(num_tasks, std::move(run_task));
#else
    IRLBA_CUSTOM_PARALLEL(num_tasks, run_task);
#endif
}
 
template<
    class ValueArray_ = std::vector<double>, 
    class IndexArray_ = std::vector<int>, 
    class PointerArray_ = std::vector<size_t>,
    class EigenVector_ = Eigen::VectorXd
>
class ParallelSparseMatrix {
public:
    ParallelSparseMatrix() {}
 
    ParallelSparseMatrix(Eigen::Index nrow, Eigen::Index ncol, ValueArray_ x, IndexArray_ i, PointerArray_ p, bool column_major, int nthreads) : 
        my_primary_dim(column_major ? ncol : nrow), 
        my_secondary_dim(column_major ? nrow : ncol), 
        my_nthreads(nthreads), 
        my_values(std::move(x)), 
        my_indices(std::move(i)), 
        my_ptrs(std::move(p)),
        my_column_major(column_major)
    {
        if (nthreads > 1) {
            fragment_threads();
        }
    }
 
public:
    Eigen::Index rows() const { 
        if (my_column_major) {
            return my_secondary_dim;
        } else {
            return my_primary_dim;
        }
    }
 
    Eigen::Index cols() const { 
        if (my_column_major) {
            return my_primary_dim;
        } else {
            return my_secondary_dim;
        }
    }
 
    const ValueArray_& get_values() const {
        return my_values;
    }
 
    const IndexArray_& get_indices() const {
        return my_indices;
    }
 
    const PointerArray_& get_pointers() const {
        return my_ptrs;
    }
 
public:
    typedef typename std::remove_const<typename std::remove_reference<decltype(std::declval<PointerArray_>()[0])>::type>::type PointerType;
 
private:
    Eigen::Index my_primary_dim, my_secondary_dim;
    int my_nthreads;
    ValueArray_ my_values;
    IndexArray_ my_indices;
    PointerArray_ my_ptrs;
    bool my_column_major;
 
    typedef typename std::remove_const<typename std::remove_reference<decltype(std::declval<IndexArray_>()[0])>::type>::type IndexType;
 
    std::vector<size_t> my_primary_starts, my_primary_ends;
    std::vector<IndexType> my_secondary_ends;
    std::vector<std::vector<PointerType> > my_secondary_nonzero_starts;
 
public:
    const std::vector<size_t>& get_primary_starts() const {
        return my_primary_starts;
    }
 
    const std::vector<size_t>& get_primary_ends() const {
        return my_primary_ends;
    }
 
    const std::vector<std::vector<PointerType> >& get_secondary_nonzero_starts() const {
        return my_secondary_nonzero_starts;
    }
 
private:
 
    void fragment_threads() {
        auto total_nzeros = my_ptrs[my_primary_dim]; // last element - not using back() to avoid an extra requirement on PointerArray.
        PointerType per_thread = (total_nzeros / my_nthreads) + (total_nzeros % my_nthreads > 0); // i.e., ceiling.
        
        // Splitting columns across threads so each thread processes the same number of nonzero elements.
        my_primary_starts.resize(my_nthreads);
        my_primary_ends.resize(my_nthreads);
        {
            Eigen::Index primary_counter = 0;
            PointerType sofar = per_thread;
            for (int t = 0; t < my_nthreads; ++t) {
                my_primary_starts[t] = primary_counter;
                while (primary_counter < my_primary_dim && my_ptrs[primary_counter + 1] <= sofar) {
                    ++primary_counter;
                }
                my_primary_ends[t] = primary_counter;
                sofar += per_thread;
            }
        }
 
        // Splitting rows across threads so each thread processes the same number of nonzero elements.
        my_secondary_ends.resize(my_nthreads + 1);
        my_secondary_nonzero_starts.resize(my_nthreads + 1, std::vector<PointerType>(my_primary_dim));
        {
            std::vector<PointerType> secondary_nonzeros(my_secondary_dim);
            for (PointerType i = 0; i < total_nzeros; ++i) { // don't using range for loop to avoid an extra requirement on IndexArray.
                ++(secondary_nonzeros[my_indices[i]]);
            }
 
            IndexType secondary_counter = 0;
            PointerType sofar = per_thread;
            PointerType cum_rows = 0;
 
            for (int t = 0; t < my_nthreads; ++t) {
                while (secondary_counter < my_secondary_dim && cum_rows <= sofar) {
                    cum_rows += secondary_nonzeros[secondary_counter];
                    ++secondary_counter;
                }
                my_secondary_ends[t + 1] = secondary_counter;
                sofar += per_thread;
            }
 
            for (Eigen::Index c = 0; c < my_primary_dim; ++c) {
                auto primary_start = my_ptrs[c], primary_end = my_ptrs[c + 1];
                my_secondary_nonzero_starts[0][c] = primary_start;
 
                auto s = primary_start;
                for (int thread = 0; thread < my_nthreads; ++thread) {
                    auto limit = my_secondary_ends[thread + 1];
                    while (s < primary_end && my_indices[s] < limit) {
                        ++s; 
                    }
                    my_secondary_nonzero_starts[thread + 1][c] = s;
                }
            }
        }
    }
 
private:
    void indirect_multiply(const EigenVector_& rhs, std::vector<std::vector<typename EigenVector_::Scalar> >& thread_buffers, EigenVector_& output) const {
        if (my_nthreads == 1) {
            output.setZero();
            for (Eigen::Index c = 0; c < my_primary_dim; ++c) {
                auto start = my_ptrs[c];
                auto end = my_ptrs[c + 1];
                auto val = rhs.coeff(c);
                for (PointerType s = start; s < end; ++s) {
                    output.coeffRef(my_indices[s]) += my_values[s] * val;
                }
            }
            return;
        }
 
        parallelize(my_nthreads, [&](int t) -> void {
            auto secondary_start = my_secondary_ends[t];
            auto secondary_end = my_secondary_ends[t + 1];
            size_t secondary_len = secondary_end - secondary_start;
 
            // Using a separate buffer for the other threads to avoid false
            // sharing. On first use, each buffer is allocated within each
            // thread to give malloc a chance of using thread-specific arenas.
            typename EigenVector_::Scalar* optr;
            if (t != 0) {
                auto& curbuffer = thread_buffers[t - 1];
                curbuffer.resize(secondary_len);
                optr = curbuffer.data();
            } else {
                optr = output.data() + secondary_start;
            }
            std::fill_n(optr, secondary_len, static_cast<typename EigenVector_::Scalar>(0));
 
            const auto& nz_starts = my_secondary_nonzero_starts[t];
            const auto& nz_ends = my_secondary_nonzero_starts[t + 1];
            for (Eigen::Index c = 0; c < my_primary_dim; ++c) {
                auto nz_start = nz_starts[c];
                auto nz_end = nz_ends[c];
                auto val = rhs.coeff(c);
                for (PointerType s = nz_start; s < nz_end; ++s) {
                    optr[my_indices[s] - secondary_start] += my_values[s] * val;
                }
            }
 
            if (t != 0) {
                std::copy_n(optr, secondary_len, output.data() + secondary_start);
            }
        });
 
        return;
    }
 
    void direct_multiply(const EigenVector_& rhs, EigenVector_& output) const {
        if (my_nthreads == 1) {
            for (Eigen::Index c = 0; c < my_primary_dim; ++c) {
                output.coeffRef(c) = column_dot_product<typename EigenVector_::Scalar>(c, rhs);
            }
            return;
        }
 
        parallelize(my_nthreads, [&](int t) -> void {
            auto curstart = my_primary_starts[t];
            auto curend = my_primary_ends[t];
            for (size_t c = curstart; c < curend; ++c) {
                output.coeffRef(c) = column_dot_product<typename EigenVector_::Scalar>(c, rhs);
            }
        });
 
        return;
    }
 
    template<typename Scalar_>
    Scalar_ column_dot_product(size_t c, const EigenVector_& rhs) const {
        PointerType primary_start = my_ptrs[c], primary_end = my_ptrs[c + 1];
        Scalar_ dot = 0;
        for (PointerType s = primary_start; s < primary_end; ++s) {
            dot += my_values[s] * rhs.coeff(my_indices[s]);
        }
        return dot;
    }
 
public:
    struct Workspace {
        EigenVector_ buffer;
        std::vector<std::vector<typename EigenVector_::Scalar> > thread_buffers;
    };
 
    Workspace workspace() const {
        Workspace output;
        if (my_nthreads > 1 && my_column_major) {
            output.thread_buffers.resize(my_nthreads - 1);
        }
        return output;
    }
 
    struct AdjointWorkspace {
        EigenVector_ buffer;
        std::vector<std::vector<typename EigenVector_::Scalar> > thread_buffers;
    };
 
    AdjointWorkspace adjoint_workspace() const {
        AdjointWorkspace output;
        if (my_nthreads > 1 && !my_column_major) {
            output.thread_buffers.resize(my_nthreads - 1);
        }
        return output;
    }
 
public:
    template<class Right_>
    void multiply(const Right_& rhs, Workspace& work, EigenVector_& output) const {
        const auto& realized_rhs = internal::realize_rhs(rhs, work.buffer);
        if (my_column_major) {
            indirect_multiply(realized_rhs, work.thread_buffers, output);
        } else {
            direct_multiply(realized_rhs, output);
        }
    }
 
    template<class Right_>
    void adjoint_multiply(const Right_& rhs, AdjointWorkspace& work, EigenVector_& output) const {
        const auto& realized_rhs = internal::realize_rhs(rhs, work.buffer);
        if (my_column_major) {
            direct_multiply(realized_rhs, output);
        } else {
            indirect_multiply(realized_rhs, work.thread_buffers, output);
        }
    }
 
public:
    template<class EigenMatrix_>
    EigenMatrix_ realize() const {
        auto nr = rows(), nc = cols();
        EigenMatrix_ output(nr, nc);
        output.setZero();
 
        if (my_column_major) {
            for (Eigen::Index c = 0; c < nc; ++c) {
                PointerType col_start = my_ptrs[c], col_end = my_ptrs[c + 1];
                for (PointerType s = col_start; s < col_end; ++s) {
                    output.coeffRef(my_indices[s], c) = my_values[s];
                }
            }
        } else {
            for (Eigen::Index r = 0; r < nr; ++r) {
                PointerType row_start = my_ptrs[r], row_end = my_ptrs[r + 1];
                for (PointerType s = row_start; s < row_end; ++s) {
                    output.coeffRef(r, my_indices[s]) = my_values[s];
                }
            }
        }
 
        return output;
    }
};
 
class EigenThreadScope {
#ifndef _OPENMP
public:
    EigenThreadScope([[maybe_unused]] int num_threads) {}
 
#else
public:
    EigenThreadScope([[maybe_unused]] int num_threads) : my_previous(Eigen::nbThreads()) {
#ifdef IRLBA_CUSTOM_PARALLEL
#ifdef IRLBA_CUSTOM_PARALLEL_USES_OPENMP
        Eigen::setNbThreads(num_threads);
#else
        Eigen::setNbThreads(1);
#endif
#else
#ifdef SUBPAR_USES_OPENMP_SIMPLE
        Eigen::setNbThreads(num_threads);
#else
        Eigen::setNbThreads(1);
#endif
#endif
    }
 
private:
    int my_previous;
 
public:
    ~EigenThreadScope() { 
        Eigen::setNbThreads(my_previous);
    }
#endif
 
public:
    EigenThreadScope(const EigenThreadScope&) = delete;
    EigenThreadScope(EigenThreadScope&&) = delete;
    EigenThreadScope& operator=(const EigenThreadScope&) = delete;
    EigenThreadScope& operator=(EigenThreadScope&&) = delete;
};
 
}
 
#endif