matrix_csr_impl.h

// Copyright (C) 2021-2023 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 <iostream>
#include <numeric>
#include <span>
#include <utility>
#include <vector>
 
namespace dolfinx::la
{
namespace impl
{
template <int BS0, int BS1, 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,
                typename Y::value_type num_rows);
 
template <int BS0, int BS1, 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,
                        typename Y::value_type num_rows);
 
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,
                           typename Y::value_type num_rows, int bs0, int bs1);
 
} // namespace impl
 
//-----------------------------------------------------------------------------
template <int BS0, int BS1, typename OP, typename U, typename V, typename W,
          typename X, typename Y>
void impl::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)
{
  const std::size_t nc = xcols.size();
  assert(x.size() == xrows.size() * xcols.size() * BS0 * BS1);
  for (std::size_t r = 0; r < xrows.size(); ++r)
  {
    // Row index and current data row
    auto row = xrows[r];
    using T = typename X::value_type;
    const T* xr = x.data() + r * nc * BS0 * BS1;
 
#ifndef NDEBUG
    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[row]);
    auto cit1 = std::next(cols.begin(), 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])
        throw std::runtime_error("Entry not in sparsity");
 
      std::size_t d = std::distance(cols.begin(), it);
      int di = d * BS0 * BS1;
      int xi = c * BS1;
      assert(di < data.size());
      for (int i = 0; i < BS0; ++i)
      {
        for (int j = 0; j < BS1; ++j)
          op(data[di + j], xr[xi + j]);
        di += BS1;
        xi += nc * BS1;
      }
    }
  }
}
//-----------------------------------------------------------------------------
// Insert with block insertion into a regular CSR (block size 1)
template <int BS0, int BS1, typename OP, typename U, typename V, typename W,
          typename X, typename Y>
void impl::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)
{
  const std::size_t nc = xcols.size();
  assert(x.size() == xrows.size() * xcols.size() * BS0 * BS1);
  for (std::size_t r = 0; r < xrows.size(); ++r)
  {
    // Row index and current data row
    auto row = xrows[r] * BS0;
 
#ifndef NDEBUG
    if (row >= num_rows)
      throw std::runtime_error("Local row out of range");
#endif
 
    for (int 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]);
      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)
          throw std::runtime_error("Entry not in sparsity");
 
        std::size_t d = std::distance(cols.begin(), it);
        assert(d < data.size());
        int xi = c * BS1;
        for (int j = 0; j < BS1; ++j)
          op(data[d + j], xr[xi + j]);
      }
    }
  }
}
//-----------------------------------------------------------------------------
// Add individual entries in block-CSR storage
template <typename OP, typename U, typename V, typename W, typename X,
          typename Y>
void impl::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)
{
  const std::size_t nc = xcols.size();
  const int 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);
    using T = typename X::value_type;
    const T* xr = x.data() + r * nc;
 
#ifndef NDEBUG
    if (rdiv.quot >= 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]);
    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)
        throw std::runtime_error("Entry not in sparsity");
 
      std::size_t d = std::distance(cols.begin(), it);
      const int di = d * nbs + rdiv.rem * bs1 + cdiv.rem;
      assert(di < data.size());
      op(data[di], xr[c]);
    }
  }
}
//-----------------------------------------------------------------------------
} // namespace dolfinx::la