diff --git a/cpp/dolfinx/la/MatrixCSR.h b/cpp/dolfinx/la/MatrixCSR.h index 3b68409c14..80830dda1c 100644 --- a/cpp/dolfinx/la/MatrixCSR.h +++ b/cpp/dolfinx/la/MatrixCSR.h @@ -16,6 +16,7 @@ #include #include #include +#include #include #include @@ -275,15 +276,17 @@ class MatrixCSR assert(x.size() == rows.size() * cols.size() * BS0 * BS1); if (_bs[0] == BS0 and _bs[1] == BS1) { - impl::insert_csr(_data, _cols, _row_ptr, x, rows, cols, set_fn, - num_rows); + impl::insert_csr(_data, _cols, _row_ptr, x, rows, cols, set_fn, num_rows, + std::integral_constant{}, + std::integral_constant{}); } else if (_bs[0] == 1 and _bs[1] == 1) { // Set blocked data in a regular CSR matrix (_bs[0]=1, _bs[1]=1) // with correct sparsity - impl::insert_blocked_csr(_data, _cols, _row_ptr, x, rows, cols, - set_fn, num_rows); + impl::insert_blocked_csr(_data, _cols, _row_ptr, x, rows, cols, set_fn, + num_rows, std::integral_constant{}, + std::integral_constant{}); } else { @@ -316,24 +319,28 @@ class MatrixCSR check_not_finalized(); auto add_fn = [](value_type& y, const value_type& x) { y += x; }; + std::int32_t num_rows + = _index_maps[0]->size_local() + _index_maps[0]->num_ghosts(); assert(x.size() == rows.size() * cols.size() * BS0 * BS1); if (_bs[0] == BS0 and _bs[1] == BS1) { - impl::insert_csr(_data, _cols, _row_ptr, x, rows, cols, add_fn, - _row_ptr.size()); + impl::insert_csr(_data, _cols, _row_ptr, x, rows, cols, add_fn, num_rows, + std::integral_constant{}, + std::integral_constant{}); } else if (_bs[0] == 1 and _bs[1] == 1) { // Add blocked data to a regular CSR matrix (_bs[0]=1, _bs[1]=1) - impl::insert_blocked_csr(_data, _cols, _row_ptr, x, rows, cols, - add_fn, _row_ptr.size()); + impl::insert_blocked_csr(_data, _cols, _row_ptr, x, rows, cols, add_fn, + num_rows, std::integral_constant{}, + std::integral_constant{}); } else { assert(BS0 == 1 and BS1 == 1); // Add non-blocked data to a blocked CSR matrix (BS0=1, BS1=1) impl::insert_nonblocked_csr(_data, _cols, _row_ptr, x, rows, cols, add_fn, - _row_ptr.size(), _bs[0], _bs[1]); + num_rows, _bs[0], _bs[1]); } } @@ -929,14 +936,12 @@ MatrixCSR::MatrixCSR(const SparsityType& p, BlockMode mode) and it->first == ghost_index_array[i + 1]); local_col = it->second; } - auto cit0 = std::next(_cols.begin(), _row_ptr[local_row]); - auto cit1 = std::next(_cols.begin(), _row_ptr[local_row + 1]); - // Find position of column index and insert data - auto cit = std::lower_bound(cit0, cit1, local_col); - assert(cit != cit1); - assert(*cit == local_col); - std::size_t d = std::ranges::distance(_cols.begin(), cit); + std::size_t row_begin = _row_ptr[local_row]; + std::size_t row_end = _row_ptr[local_row + 1]; + std::size_t d = impl::csr_position(_cols, row_begin, row_end, local_col); + assert(d != row_end); + assert(_cols[d] == local_col); _unpack_pos.push_back(d); } diff --git a/cpp/dolfinx/la/matrix_csr_impl.h b/cpp/dolfinx/la/matrix_csr_impl.h index dc4d601488..43a66af3ef 100644 --- a/cpp/dolfinx/la/matrix_csr_impl.h +++ b/cpp/dolfinx/la/matrix_csr_impl.h @@ -1,4 +1,4 @@ -// Copyright (C) 2021-2023 Garth N. Wells and Chris N. Richardson +// Copyright (C) 2021-2026 Garth N. Wells and Chris N. Richardson // // This file is part of DOLFINx (https://www.fenicsproject.org) // @@ -6,6 +6,7 @@ #pragma once +#include #include #include #include @@ -35,10 +36,57 @@ concept BlockSizeArg = std::integral> or is_integral_constant>::value; /// @endcond +/// @brief Position of `col` within the sorted column-index range +/// `cols[begin, end)` of one CSR row. +/// +/// Matches `std::lower_bound(cols.begin() + begin, cols.begin() + end, +/// col)` in semantics -- returns `end` if `col` is greater than every +/// entry in the range, otherwise the position of the first entry +/// `>= col`. Callers must check `cols[pos] == col` themselves to +/// detect a genuinely missing entry; this function does not throw. +/// +/// A single shared search used by `insert_csr`, `insert_blocked_csr`, +/// `insert_nonblocked_csr` and `MatrixCSR`'s ghost-row unpacking, which +/// otherwise duplicate this lookup identically. A hand-written +/// "branchless" (conditional-move) rewrite of `std::lower_bound` was +/// tried here and measured ~13% *slower* in `csr_reasm` +/// (`scripts/bench_backend.py`) than plain `std::lower_bound` -- GCC +/// did not turn the conditional-move idiom into a `cmov` (confirmed by +/// disassembly: it kept a real conditional branch, structurally +/// identical to libstdc++'s own implementation, just without +/// libstdc++'s tuning), so there was no bet to win. Kept as a thin +/// wrapper around `std::lower_bound` for the sole purpose of removing +/// that duplication. +/// +/// @note Clang-20 mirrors the above in reverse: it if-converts this +/// same `std::lower_bound` into a branchless `cmov` chain, ~35-50% +/// *slower* here than GCC's branchy codegen -- the `cmov` serializes +/// each iteration's load behind the prior comparison, killing the +/// speculation a branch gets on this short, cache-resident row. +/// `__builtin_expect` recovers full speed but only inside a hand-rolled +/// loop (a comparator carrying the same hint is inlined away with no +/// effect; `-mllvm -x86-cmov-converter=true` and +/// `std::ranges::lower_bound` don't help either). Not worth +/// reimplementing `lower_bound` over a codegen quirk -- left as-is; if +/// this gets worse, prefer caching resolved offsets across reassembly +/// (avoids the search entirely) over a hand-rolled comparison here. +/// +/// @param[in] cols CSR column indices for the whole matrix. +/// @param[in] begin Start of the row's range within `cols`. +/// @param[in] end End (one past last) of the row's range within +/// `cols`. +/// @param[in] col Column index to locate. +/// @return Position within `cols` of `col`, or `end` if not present. +template +std::size_t csr_position(const V& cols, std::size_t begin, std::size_t end, + T col) +{ + const auto* c = cols.data(); + return std::lower_bound(c + begin, c + end, col) - c; +} + /// @brief Incorporate data into a CSR matrix. /// -/// @tparam BS0 Row block size (of both matrix and data). -/// @tparam BS1 Column block size (of both matrix and data). /// @tparam OP The operation (usually "set" or "add"). /// @param[out] data CSR matrix data. /// @param[in] cols CSR column indices. @@ -51,6 +99,12 @@ concept BlockSizeArg = std::integral> /// @param[in] num_rows Maximum row index that can be set. Used when /// debugging to check that rows beyond a permitted range are not being /// set. +/// @param[in] BS0 Row block size (of both matrix and data). Either a +/// runtime integral value or a `std::integral_constant` for some +/// integral type `U`. +/// @param[in] BS1 Column block size (of both matrix and data). Either a +/// runtime integral value or a `std::integral_constant` for some +/// integral type `U`. /// /// @note In the case of block data, where BS0 or BS1 are greater than /// one, the layout of the input data is still the same. For example, @@ -64,11 +118,12 @@ concept BlockSizeArg = std::integral> /// 8 9 | 10 11 /// 12 13 | 14 15 /// -template +template void insert_csr(U&& data, const V& cols, const W& row_ptr, const X& x, const Y& xrows, const Y& xcols, OP op, - [[maybe_unused]] typename Y::value_type num_rows) + [[maybe_unused]] typename Y::value_type num_rows, + BlockSizeArg auto BS0, BlockSizeArg auto BS1) { const std::size_t nc = xcols.size(); assert(x.size() == xrows.size() * xcols.size() * BS0 * BS1); @@ -84,22 +139,21 @@ void insert_csr(U&& data, const V& cols, const W& row_ptr, const X& x, throw std::runtime_error("Local row out of range"); #endif // Columns indices for row - auto cit0 = std::next(cols.begin(), row_ptr[row]); - auto cit1 = std::next(cols.begin(), row_ptr[row + 1]); + std::size_t row_begin = row_ptr[row]; + std::size_t row_end = row_ptr[row + 1]; for (std::size_t c = 0; c < nc; ++c) { // Find position of column index - auto it = std::lower_bound(cit0, cit1, xcols[c]); - if (it == cit1 or *it != xcols[c]) + std::size_t d = csr_position(cols, row_begin, row_end, xcols[c]); + if (d == row_end or cols[d] != xcols[c]) throw std::runtime_error("Entry not in sparsity"); - std::size_t d = std::ranges::distance(cols.begin(), it); std::size_t di = d * BS0 * BS1; std::size_t xi = c * BS1; assert(di < data.size()); - for (int i = 0; i < BS0; ++i) + for (decltype(+BS0) i = 0; i < BS0; ++i) { - for (int j = 0; j < BS1; ++j) + for (decltype(+BS1) j = 0; j < BS1; ++j) op(data[di + j], xr[xi + j]); di += BS1; xi += nc * BS1; @@ -115,8 +169,6 @@ void insert_csr(U&& data, const V& cols, const W& row_ptr, const X& x, /// accept the data. /// @note See ::insert_csr for data layout. /// -/// @tparam BS0 Row block size of data. -/// @tparam BS1 Column block size of data. /// @tparam OP The operation (usually "set" or "add"). /// @param[out] data CSR matrix data. /// @param[in] cols CSR column indices. @@ -129,11 +181,16 @@ void insert_csr(U&& data, const V& cols, const W& row_ptr, const X& x, /// @param[in] num_rows Maximum row index that can be set. Used when /// debugging to check that rows beyond a permitted range are not being /// set. -template +/// @param[in] BS0 Row block size of data. Either a runtime integral +/// value or a `std::integral_constant` for some integral type `U`. +/// @param[in] BS1 Column block size of data. Either a runtime integral +/// value or a `std::integral_constant` for some integral type `U`. +template void insert_blocked_csr(U&& data, const V& cols, const W& row_ptr, const X& x, const Y& xrows, const Y& xcols, OP op, - [[maybe_unused]] typename Y::value_type num_rows) + [[maybe_unused]] typename Y::value_type num_rows, + BlockSizeArg auto BS0, BlockSizeArg auto BS1) { const std::size_t nc = xcols.size(); assert(x.size() == xrows.size() * xcols.size() * BS0 * BS1); @@ -146,25 +203,24 @@ void insert_blocked_csr(U&& data, const V& cols, const W& row_ptr, const X& x, throw std::runtime_error("Local row out of range"); #endif - for (int i = 0; i < BS0; ++i) + for (decltype(+BS0) i = 0; i < BS0; ++i) { using T = typename X::value_type; const T* xr = x.data() + (r * BS0 + i) * nc * BS1; // Columns indices for row - auto cit0 = std::next(cols.begin(), row_ptr[row + i]); - auto cit1 = std::next(cols.begin(), row_ptr[row + i + 1]); + std::size_t row_begin = row_ptr[row + i]; + std::size_t row_end = row_ptr[row + i + 1]; for (std::size_t c = 0; c < nc; ++c) { // Find position of column index - auto it = std::lower_bound(cit0, cit1, xcols[c] * BS1); - if (it == cit1 or *it != xcols[c] * BS1) + std::size_t d = csr_position(cols, row_begin, row_end, xcols[c] * BS1); + if (d == row_end or cols[d] != xcols[c] * BS1) throw std::runtime_error("Entry not in sparsity"); - std::size_t d = std::ranges::distance(cols.begin(), it); assert(d < data.size()); std::size_t xi = c * BS1; - for (int j = 0; j < BS1; ++j) + for (decltype(+BS1) j = 0; j < BS1; ++j) op(data[d + j], xr[xi + j]); } } @@ -188,43 +244,46 @@ void insert_blocked_csr(U&& data, const V& cols, const W& row_ptr, const X& x, /// @param[in] num_rows Maximum row index that can be set. Used when /// debugging to check that rows beyond a permitted range are not being /// set. -/// @param[in] bs0 Row block size of matrix. -/// @param[in] bs1 Column block size of matrix. +/// @param[in] bs0 Row block size of matrix. Either a runtime integral +/// value or a `std::integral_constant` for some integral type `U`. +/// @param[in] bs1 Column block size of matrix. Either a runtime integral +/// value or a `std::integral_constant` for some integral type `U`. template void insert_nonblocked_csr(U&& data, const V& cols, const W& row_ptr, const X& x, const Y& xrows, const Y& xcols, OP op, [[maybe_unused]] typename Y::value_type num_rows, - int bs0, int bs1) + BlockSizeArg auto bs0, BlockSizeArg auto bs1) { const std::size_t nc = xcols.size(); - const int nbs = bs0 * bs1; + const auto nbs = bs0 * bs1; assert(x.size() == xrows.size() * xcols.size()); for (std::size_t r = 0; r < xrows.size(); ++r) { // Row index and current data row - auto rdiv = std::div(xrows[r], bs0); + const auto row = xrows[r] / bs0; + const auto row_rem = xrows[r] % bs0; using T = typename X::value_type; const T* xr = x.data() + r * nc; #ifndef NDEBUG - if (rdiv.quot >= num_rows) + if (row >= num_rows) throw std::runtime_error("Local row out of range"); #endif // Columns indices for row - auto cit0 = std::next(cols.begin(), row_ptr[rdiv.quot]); - auto cit1 = std::next(cols.begin(), row_ptr[rdiv.quot + 1]); + std::size_t row_begin = row_ptr[row]; + std::size_t row_end = row_ptr[row + 1]; for (std::size_t c = 0; c < nc; ++c) { // Find position of column index - auto cdiv = std::div(xcols[c], bs1); - auto it = std::lower_bound(cit0, cit1, cdiv.quot); - if (it == cit1 or *it != cdiv.quot) + const auto col = xcols[c] / bs1; + const auto col_rem = xcols[c] % bs1; + std::size_t d = csr_position(cols, row_begin, row_end, col); + if (d == row_end or cols[d] != col) throw std::runtime_error("Entry not in sparsity"); - std::size_t d = std::ranges::distance(cols.begin(), it); - std::size_t di = d * nbs + rdiv.rem * bs1 + cdiv.rem; + std::size_t di = d * nbs + row_rem * bs1 + col_rem; assert(di < data.size()); op(data[di], xr[c]); }