Deduplicate CSR column search in MatrixCSR

cpp/dolfinx/la/MatrixCSR.h

@@ -16,6 +16,7 @@
 #include <mpi.h>
 #include <numeric>
 #include <span>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
@@ -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<BS0, BS1>(_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<int, BS0>{},
+                       std::integral_constant<int, BS1>{});
     }
     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<BS0, BS1>(_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<int, BS0>{},
+                               std::integral_constant<int, BS1>{});
     }
     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<BS0, BS1>(_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<int, BS0>{},
+                       std::integral_constant<int, BS1>{});
     }
     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<BS0, BS1>(_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<int, BS0>{},
+                               std::integral_constant<int, BS1>{});
     }
     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<U, V, W, X>::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);
   }
 

cpp/dolfinx/la/matrix_csr_impl.h

@@ -1,11 +1,12 @@
-// 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)
 //
 // SPDX-License-Identifier:    LGPL-3.0-or-later
 
 #pragma once
 
+#include <algorithm>
 #include <concepts>
 #include <numeric>
 #include <span>
@@ -35,10 +36,57 @@
                        or is_integral_constant<std::remove_cvref_t<T>>::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 <typename V, std::integral T>
+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 @@
 /// @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<U, N>` 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<U, N>` 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 @@
 /// 8  9  | 10 11
 /// 12 13 | 14 15
 ///
-template <int BS0, int BS1, typename OP, typename U, typename V, typename W,
-          typename X, typename Y>
+template <typename OP, typename U, typename V, typename W, typename X,
+          typename Y>
 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 @@
       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 @@
 /// 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 @@
 /// @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 <int BS0, int BS1, typename OP, typename U, typename V, typename W,
-          typename X, typename Y>
+/// @param[in] BS0 Row block size of data. Either a runtime integral
+/// value or a `std::integral_constant<U, N>` for some integral type `U`.
+/// @param[in] BS1 Column block size of data. Either a runtime integral
+/// value or a `std::integral_constant<U, N>` for some integral type `U`.
+template <typename OP, typename U, typename V, typename W, typename X,
+          typename Y>
 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 @@
       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 @@
 /// @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<U, N>` for some integral type `U`.
+/// @param[in] bs1 Column block size of matrix. Either a runtime integral
+/// value or a `std::integral_constant<U, N>` for some integral type `U`.
 template <typename OP, typename U, typename V, typename W, typename X,
           typename Y>
 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]);
     }