Skip to content

Commit

Permalink
Merge pull request #277 from paulsengroup/perf/to-xxx-matrices
Browse files Browse the repository at this point in the history
Improve performance of ToSparseMatrix by reducing memory allocations
  • Loading branch information
robomics authored Oct 7, 2024
2 parents 5e298f6 + e5cd665 commit de5a6e0
Show file tree
Hide file tree
Showing 7 changed files with 511 additions and 177 deletions.
4 changes: 4 additions & 0 deletions src/libhictk/transformers/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@ if(HICTK_WITH_ARROW)
find_package(Arrow REQUIRED QUIET)
endif()

find_package(spdlog REQUIRED QUIET)

add_library(transformers INTERFACE)
add_library(hictk::transformers ALIAS transformers)

Expand All @@ -19,6 +21,8 @@ target_include_directories(transformers INTERFACE "$<BUILD_INTERFACE:${CMAKE_CUR
"$<INSTALL_INTERFACE:include>")
target_link_libraries(transformers INTERFACE hictk::cooler hictk::file hictk::hic)

target_link_system_libraries(transformers INTERFACE spdlog::spdlog_header_only)

if(HICTK_WITH_ARROW_SHARED)
target_link_system_libraries(transformers INTERFACE "$<$<BOOL:${HICTK_WITH_ARROW}>:Arrow::arrow_shared>")
else()
Expand Down
59 changes: 34 additions & 25 deletions src/libhictk/transformers/include/hictk/transformers/common.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
#include <cassert>
#include <cstdint>
#include <type_traits>
#include <utility>

#include "hictk/pixel.hpp"
#include "hictk/suppress_warnings.hpp"
Expand All @@ -31,14 +32,24 @@ template <typename T>
inline constexpr bool has_weights_member_fx<T, std::void_t<decltype(std::declval<T>().weights())>> =
true;

template <typename N, typename PixelSelector, typename MatrixT, typename SetterOp>
inline void fill_matrix_square(const PixelSelector& sel, MatrixT& buffer, std::int64_t num_rows,
std::int64_t num_cols, std::int64_t offset1, std::int64_t offset2,
template <typename PixelSelector>
[[nodiscard]] inline bool selector_is_symmetric_upper(const PixelSelector& sel) noexcept {
if constexpr (std::is_same_v<PixelSelector, cooler::PixelSelector>) {
return sel.is_symmetric_upper();
}
return true;
}

template <typename PixelIt, typename MatrixT1, typename MatrixT2, typename SetterOp>
inline void fill_matrix_square(PixelIt first_pixel, PixelIt last_pixel, MatrixT1& matrix_ut,
MatrixT2& matrix_lt, std::int64_t num_rows, std::int64_t num_cols,
std::int64_t offset1, std::int64_t offset2,
bool populate_lower_triangle, bool populate_upper_triangle,
SetterOp matrix_setter) {
using N = decltype(std::declval<PixelIt>()->count);
assert(populate_lower_triangle || populate_upper_triangle);

std::for_each(sel.template begin<N>(), sel.template end<N>(), [&](const ThinPixel<N>& p) {
std::for_each(std::move(first_pixel), std::move(last_pixel), [&](const ThinPixel<N>& p) {
const auto i1 = static_cast<std::int64_t>(p.bin1_id) - offset1;
const auto i2 = static_cast<std::int64_t>(p.bin2_id) - offset2;

Expand All @@ -50,36 +61,38 @@ inline void fill_matrix_square(const PixelSelector& sel, MatrixT& buffer, std::i
if (populate_upper_triangle && i1 <= i2) {
HICTK_DISABLE_WARNING_PUSH
HICTK_DISABLE_WARNING_CONVERSION
matrix_setter(buffer, i1, i2, p.count);
matrix_setter(matrix_ut, i1, i2, p.count);
HICTK_DISABLE_WARNING_POP
inserted = true;
}
if (populate_lower_triangle && i1 >= i2 && !inserted) {
HICTK_DISABLE_WARNING_PUSH
HICTK_DISABLE_WARNING_CONVERSION
matrix_setter(buffer, i1, i2, p.count);
matrix_setter(matrix_lt, i1, i2, p.count);
HICTK_DISABLE_WARNING_POP
}
});
}

template <typename N, typename PixelSelector, typename MatrixT, typename SetterOp>
inline void fill_matrix_symmetric_upper(const PixelSelector& sel, MatrixT& buffer,
template <typename PixelIt, typename MatrixT1, typename MatrixT2, typename SetterOp>
inline void fill_matrix_symmetric_upper(PixelIt first_pixel, PixelIt last_pixel,
MatrixT1& matrix_ut, MatrixT2& matrix_lt,
std::int64_t num_rows, std::int64_t num_cols,
std::int64_t offset1, std::int64_t offset2,
bool populate_lower_triangle, bool populate_upper_triangle,
SetterOp matrix_setter) {
using N = decltype(std::declval<PixelIt>()->count);
assert(populate_lower_triangle || populate_upper_triangle);

std::for_each(sel.template begin<N>(), sel.template end<N>(), [&](const ThinPixel<N>& p) {
std::for_each(std::move(first_pixel), std::move(last_pixel), [&](const ThinPixel<N>& p) {
const auto i1 = static_cast<std::int64_t>(p.bin1_id) - offset1;
const auto i2 = static_cast<std::int64_t>(p.bin2_id) - offset2;
bool inserted = false;
if (populate_upper_triangle) {
if (i1 >= 0 && i1 < num_rows && i2 >= 0 && i2 < num_cols) {
HICTK_DISABLE_WARNING_PUSH
HICTK_DISABLE_WARNING_CONVERSION
matrix_setter(buffer, i1, i2, p.count);
matrix_setter(matrix_ut, i1, i2, p.count);
HICTK_DISABLE_WARNING_POP
inserted = true;
}
Expand All @@ -96,34 +109,30 @@ inline void fill_matrix_symmetric_upper(const PixelSelector& sel, MatrixT& buffe
if (i3 >= 0 && i3 < num_rows && i4 >= 0 && i4 < num_cols) {
HICTK_DISABLE_WARNING_PUSH
HICTK_DISABLE_WARNING_CONVERSION
matrix_setter(buffer, i3, i4, p.count);
matrix_setter(matrix_lt, i3, i4, p.count);
HICTK_DISABLE_WARNING_POP
}
}
});
}

template <typename N, typename PixelSelector, typename MatrixT, typename SetterOp>
inline void fill_matrix(const PixelSelector& sel, MatrixT& buffer, std::int64_t num_rows,
template <typename PixelIt, typename MatrixT1, typename MatrixT2, typename SetterOp>
inline void fill_matrix(PixelIt first_pixel, PixelIt last_pixel, bool symmetric_upper,
MatrixT1& matrix_ut, MatrixT2& matrix_lt, std::int64_t num_rows,
std::int64_t num_cols, std::int64_t offset1, std::int64_t offset2,
bool populate_lower_triangle, bool populate_upper_triangle,
SetterOp matrix_setter) {
assert(populate_lower_triangle || populate_upper_triangle);

const auto selector_is_symmetric_upper = [&]() constexpr {
if constexpr (std::is_same_v<PixelSelector, cooler::PixelSelector>) {
return sel.is_symmetric_upper();
}
return true;
}();

if (HICTK_UNLIKELY(!selector_is_symmetric_upper)) {
fill_matrix_square<N>(sel, buffer, num_rows, num_cols, offset1, offset2,
populate_lower_triangle, populate_upper_triangle, matrix_setter);
if (HICTK_UNLIKELY(!symmetric_upper)) {
fill_matrix_square(std::move(first_pixel), std::move(last_pixel), matrix_ut, matrix_lt,
num_rows, num_cols, offset1, offset2, populate_lower_triangle,
populate_upper_triangle, matrix_setter);
return;
}
fill_matrix_symmetric_upper<N>(sel, buffer, num_rows, num_cols, offset1, offset2,
populate_lower_triangle, populate_upper_triangle, matrix_setter);
fill_matrix_symmetric_upper(std::move(first_pixel), std::move(last_pixel), matrix_ut, matrix_lt,
num_rows, num_cols, offset1, offset2, populate_lower_triangle,
populate_upper_triangle, matrix_setter);
}

} // namespace internal
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -40,13 +40,21 @@ inline ToDenseMatrix<N, PixelSelector>::ToDenseMatrix(std::shared_ptr<const Pixe

template <typename N, typename PixelSelector>
inline auto ToDenseMatrix<N, PixelSelector>::operator()() -> MatrixT {
assert(!!_sel);
if (!_sel) {
return {};
}

const auto populate_lower_triangle =
_span == QuerySpan::lower_triangle || _span == QuerySpan::full;
const auto populate_upper_triangle =
_span == QuerySpan::upper_triangle || _span == QuerySpan::full;

auto matrix_setter = [](MatrixT& matrix, std::int64_t i1, std::int64_t i2, N count) {
auto matrix_setter = [](MatrixT& matrix, std::int64_t i1, std::int64_t i2, N count) noexcept {
assert(i1 >= 0);
assert(i1 < matrix.rows());
assert(i2 >= 0);
assert(i2 < matrix.cols());

matrix(i1, i2) = count;
};

Expand All @@ -60,15 +68,18 @@ inline auto ToDenseMatrix<N, PixelSelector>::operator()() -> MatrixT {
coord4 = coord3;

const auto new_sel = _sel->fetch(coord3, coord4);
internal::fill_matrix<N>(new_sel, matrix, matrix.rows(), matrix.cols(), row_offset(),
col_offset(), populate_lower_triangle, populate_upper_triangle,
matrix_setter);
internal::fill_matrix(new_sel.template begin<N>(), new_sel.template end<N>(),
internal::selector_is_symmetric_upper(new_sel), matrix, matrix,
matrix.rows(), matrix.cols(), row_offset(), col_offset(),
populate_lower_triangle, populate_upper_triangle, matrix_setter);
return matrix;
}
}

internal::fill_matrix<N>(*_sel, matrix, matrix.rows(), matrix.cols(), row_offset(), col_offset(),
populate_lower_triangle, populate_upper_triangle, matrix_setter);
internal::fill_matrix(_sel->template begin<N>(), _sel->template end<N>(),
internal::selector_is_symmetric_upper(*_sel), matrix, matrix, matrix.rows(),
matrix.cols(), row_offset(), col_offset(), populate_lower_triangle,
populate_upper_triangle, matrix_setter);
return matrix;
}

Expand Down
Loading

0 comments on commit de5a6e0

Please sign in to comment.