sparse_column.hpp

Go to the documentation of this file.

#ifndef TATAMI_MULT_MULTIPLE_VECTORS_SPARSE_COLUMN_HPP
#define TATAMI_MULT_MULTIPLE_VECTORS_SPARSE_COLUMN_HPP
 
#include <cstddef>
#include <vector>
#include <optional>
#include <algorithm>
 
#include "tatami/tatami.hpp"
#include "sanisizer/sanisizer.hpp"
#include "jiwoo/jiwoo.hpp"
 
#include "../utils.hpp"
#include "../sparse_dot_product.hpp"
 
namespace tatami_mult {
 
/* See https://github.com/tatami-inc/test-multiplication/tree/master/sparse_column/multiple_vectors
 * for an explanation of the choice of algorithm.
 */
 
struct MultiplySparseColumnWithMultipleVectorsOptions {
    int num_threads = 1;
 
    int block_size = 16;
};
 
template<typename LeftValue_, typename LeftIndex_, typename RightVectors_, typename GetRightVector_, typename GetOutputVector_>
void multiply_sparse_column_with_multiple_vectors_internal(
    const tatami::Matrix<LeftValue_, LeftIndex_>& left,
    const LeftIndex_ start,
    const LeftIndex_ length,
    const LeftIndex_ left_NR,
    const RightVectors_ right_vectors,
    GetRightVector_ get_right_vector,
    GetOutputVector_ get_output_vector,
    const MultiplySparseColumnWithMultipleVectorsOptions& options
) {
    auto ext = tatami::consecutive_extractor<true>(left, false, start, length);
    typedef I<decltype(get_output_vector(0)[0])> Output;
 
    if (options.block_size == 1) {
        auto vbuffer = tatami::create_container_of_Index_size<std::vector<LeftValue_> >(left_NR);
        auto ibuffer = tatami::create_container_of_Index_size<std::vector<LeftIndex_> >(left_NR);
        for (LeftIndex_ cd = 0; cd < length; ++cd) {
            const auto range = ext->fetch(vbuffer.data(), ibuffer.data());
            if (range.number == 0) {
                continue;
            }
 
            for (RightVectors_ rv = 0; rv < right_vectors; ++rv) {
                const auto optr = get_output_vector(rv);
                const Output mult = get_right_vector(rv)[start + cd];
                for (LeftIndex_ x = 0; x < range.number; ++x) {
                    optr[range.index[x]] += mult * static_cast<Output>(range.value[x]); 
                }
            }
        }
 
    } else {
        std::vector<std::vector<LeftValue_> > left_vbuffers;
        std::vector<std::vector<LeftIndex_> > left_ibuffers;
        std::vector<tatami::SparseRange<LeftValue_, LeftIndex_> > left_ranges;
        std::vector<LeftIndex_> left_non_empty;
        {
            const LeftIndex_ max_block_cols = sanisizer::min(length, options.block_size);
            left_vbuffers.reserve(max_block_cols);
            left_ibuffers.reserve(max_block_cols);
            for (LeftIndex_ cd = 0; cd < max_block_cols; ++cd) {
                left_vbuffers.emplace_back(tatami::cast_Index_to_container_size<std::vector<LeftValue_> >(left_NR));
                left_ibuffers.emplace_back(tatami::cast_Index_to_container_size<std::vector<LeftIndex_> >(left_NR));
            }
            sanisizer::resize(left_ranges, max_block_cols);
            left_non_empty.reserve(max_block_cols);
        }
 
        LeftIndex_ cd = 0;
        while (cd < length) {
            // Only considering the LHS columns with at least one structural non-zero.
            // Thus, our block consists of 'options.block_size' non-empty LHS columns, rather than fixed column-wise chunks of the LHS matrix.
            // This ensures that we don't waste iterations on LHS columns that don't contribute anything to the product.
            const auto left_block_info = fetch_non_empty_sparse_block(
                *ext,
                left_vbuffers,
                left_ibuffers,
                left_ranges,
                left_non_empty,
                cd,
                length,
                options.block_size,
                /* zero = */ [](LeftIndex_) -> void {} // No need to worry about zeroing as the buffers have already been zeroed.
            );
            const auto cd_num = left_block_info.num_non_empty;
 
            // If the LHS columns are all non-empty, we can speed up the loops by just using a simple counter to get the column indices.
            // Otherwise, we'll have to access the 'left_non_empty' vector to figure out the indices of each non-empty column.
            if (left_block_info.all_non_empty) {
                const LeftIndex_ cd_base = start + cd;
                for (RightVectors_ rv = 0; rv < right_vectors; ++rv) {
                    const auto outvec = get_output_vector(rv);
                    const auto rightvec = get_right_vector(rv);
                    for (LeftIndex_ cd_counter = 0; cd_counter < cd_num; ++cd_counter) {
                        const auto& currange = left_ranges[cd_counter];
                        const Output mult = rightvec[cd_base + cd_counter];
                        for (LeftIndex_ x = 0; x < currange.number; ++x) {
                            outvec[currange.index[x]] += mult * static_cast<Output>(currange.value[x]);
                        }
                    }
                }
 
            } else {
                for (auto& cdne : left_non_empty) {
                    cdne += start;
                }
                for (RightVectors_ rv = 0; rv < right_vectors; ++rv) {
                    const auto outvec = get_output_vector(rv);
                    const auto rightvec = get_right_vector(rv);
                    for (LeftIndex_ cd_counter = 0; cd_counter < cd_num; ++cd_counter) {
                        const auto& currange = left_ranges[cd_counter];
                        const Output mult = rightvec[left_non_empty[cd_counter]];
                        for (LeftIndex_ x = 0; x < currange.number; ++x) {
                            outvec[currange.index[x]] += mult * static_cast<Output>(currange.value[x]);
                        }
                    }
                }
            }
 
            cd = left_block_info.position;
        }
    }
}
template<typename LeftValue_, typename LeftIndex_, typename RightVectors_, typename GetRightVector_, typename GetOutputVector_>
void multiply_sparse_column_with_multiple_vectors(
    const tatami::Matrix<LeftValue_, LeftIndex_>& left,
    const RightVectors_ right_vectors,
    GetRightVector_ get_right_vector,
    GetOutputVector_ get_output_vector,
    const MultiplySparseColumnWithMultipleVectorsOptions& options
) {
    const auto left_NR = left.nrow();
    const auto common_dim = left.ncol();
    for (RightVectors_ rv = 0; rv < right_vectors; ++rv) {
        std::fill_n(get_output_vector(rv), left_NR, 0);
    }
 
    const bool do_parallel = options.num_threads > 1;
    typedef I<decltype(get_output_vector(0)[0])> Output;
    std::optional<std::vector<std::optional<std::vector<std::vector<Output> > > > > tmp_results;
    if (do_parallel) {
        tmp_results.emplace(sanisizer::cast<I<decltype(tmp_results->size())> >(options.num_threads - 1));
    }
 
    const auto num_used = tatami::parallelize([&](int t, LeftIndex_ start, LeftIndex_ length) -> void {
        if (!do_parallel || t == 0) {
            multiply_sparse_column_with_multiple_vectors_internal(
                left,
                start,
                length,
                left_NR,
                right_vectors,
                get_right_vector,
                get_output_vector,
                options
            );
 
        } else {
            std::vector<std::vector<Output> > tmp_output;
            tmp_output.reserve(right_vectors);
            for (RightVectors_ rv = 0; rv < right_vectors; ++rv) {
                tmp_output.emplace_back(tatami::cast_Index_to_container_size<std::vector<Output> >(left_NR));
            }
            multiply_sparse_column_with_multiple_vectors_internal(
                left,
                start,
                length,
                left_NR,
                right_vectors,
                get_right_vector,
                [&](const RightVectors_ rv) -> Output* {
                    return tmp_output[rv].data();
                },
                options
            );
            (*tmp_results)[t - 1] = std::move(tmp_output);
        }
    }, common_dim, options.num_threads);
 
    if (do_parallel) {
        for (int u = 1; u < num_used; ++u) {
            const auto& tmp = *((*tmp_results)[u - 1]);
            for (RightVectors_ rv = 0; rv < right_vectors; ++rv) {
                const auto& tmpvec = tmp[rv];
                const auto outptr = get_output_vector(rv);
                for (LeftIndex_ lr = 0; lr < left_NR; ++lr) {
                    outptr[lr] += tmpvec[lr];
                }
            }
        }
    }
}
 
template<typename LeftValue_, typename LeftIndex_, typename RightValue_, typename Output_>
void multiply_sparse_column_with_multiple_vectors(
    const tatami::Matrix<LeftValue_, LeftIndex_>& left,
    const std::vector<RightValue_*>& right,
    const std::vector<Output_*>& output,
    const MultiplySparseColumnWithMultipleVectorsOptions& options
) {
    const auto left_NR = left.nrow();
    const auto common_dim = left.ncol();
    const auto right_vectors = right.size();
    typedef I<decltype(right_vectors)> RightVectors;
    for (RightVectors rv = 0; rv < right_vectors; ++rv) {
        std::fill_n(output[rv], left_NR, 0);
    }
 
    const bool do_parallel = options.num_threads > 1;
    std::optional<std::vector<std::optional<jiwoo::EquilengthArrays<Output_> > > > tmp_results;
    if (do_parallel) {
        tmp_results.emplace(sanisizer::cast<I<decltype(tmp_results->size())> >(options.num_threads - 1));
    }
 
    const auto num_used = tatami::parallelize([&](int t, LeftIndex_ start, LeftIndex_ length) -> void {
        std::optional<jiwoo::EquilengthArrays<Output_> > tmp_output;
 
        Output_* const * output_ptrs;
        if (!do_parallel || t == 0) {
            output_ptrs = output.data();
        } else {
            tmp_output.emplace(
                sanisizer::cast<I<decltype(tmp_output->size())> >(right_vectors),
                static_cast<std::size_t>(left_NR), // cast to size_t is safe due to the tatami contract.
                0
            );
            output_ptrs = tmp_output->get();
        }
 
        multiply_sparse_column_with_multiple_vectors_internal(
            left,
            start,
            length,
            left_NR,
            right_vectors,
            [&](const RightVectors rv) -> const RightValue_* {
                return right[rv];
            },
            [&](const RightVectors rv) -> Output_* {
                return output_ptrs[rv];
            },
            options
        );
 
        if (do_parallel && t > 0) {
            (*tmp_results)[t - 1] = std::move(tmp_output);
        }
    }, common_dim, options.num_threads);
 
    if (do_parallel) {
        for (int u = 1; u < num_used; ++u) {
            const auto& tmp = *((*tmp_results)[u - 1]);
            for (RightVectors rv = 0; rv < right_vectors; ++rv) {
                const auto tmpvec = tmp[rv];
                const auto outptr = output[rv];
                for (LeftIndex_ lr = 0; lr < left_NR; ++lr) {
                    outptr[lr] += tmpvec[lr];
                }
            }
        }
    }
}
 
}
 
#endif