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 <concepts>
#include <numeric>
#include <span>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace dolfinx::la
{
namespace impl
{
template <typename T>
struct is_integral_constant : std::false_type
{
};
template <std::integral U, U N>
struct is_integral_constant<std::integral_constant<U, N>> : std::true_type
{
};

template <typename T>
concept BlockSizeArg = std::integral<std::remove_cvref_t<T>>
                       or is_integral_constant<std::remove_cvref_t<T>>::value;

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,
                [[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::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 (int j = 0; j < BS1; ++j)
          op(data[di + j], xr[xi + j]);
        di += BS1;
        xi += nc * BS1;
      }
    }
  }
}

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,
                        [[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::ranges::distance(cols.begin(), it);
        assert(d < data.size());
        std::size_t xi = c * BS1;
        for (int j = 0; j < BS1; ++j)
          op(data[d + j], xr[xi + j]);
      }
    }
  }
}

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)
{
  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::ranges::distance(cols.begin(), it);
      std::size_t di = d * nbs + rdiv.rem * bs1 + cdiv.rem;
      assert(di < data.size());
      op(data[di], xr[c]);
    }
  }
}

template <typename T>
void spmv(std::span<const T> values, std::span<const std::int64_t> row_begin,
          std::span<const std::int64_t> row_end,
          std::span<const std::int32_t> indices, std::span<const T> x,
          std::span<T> y, BlockSizeArg auto bs0, BlockSizeArg auto bs1)
{
  assert(row_begin.size() == row_end.size());
  // Block layout: row-major within each block. The element at row-offset
  // k0 (∈ [0, bs0)) and column-offset k1 (∈ [0, bs1)) of block entry j
  // is stored at values[j * bs0 * bs1 + k0 * bs1 + k1].
  for (decltype(+bs0) k0 = 0; k0 < bs0; ++k0)
  {
    for (std::size_t i = 0; i < row_begin.size(); i++)
    {
      T vi{0};
      for (std::int64_t j = row_begin[i]; j < row_end[i]; j++)
      {
        for (decltype(+bs1) k1 = 0; k1 < bs1; ++k1)
        {
          vi += values[j * bs0 * bs1 + k0 * bs1 + k1]
                * x[indices[j] * bs1 + k1];
        }
      }
 
      y[i * bs0 + k0] += vi;
    }
  }
}

template <typename T>
void spmvT(std::span<const T> values, std::span<const std::int64_t> row_begin,
           std::span<const std::int64_t> row_end,
           std::span<const std::int32_t> indices, std::span<const T> x,
           std::span<T> y, BlockSizeArg auto bs0, BlockSizeArg auto bs1)
{
  assert(row_begin.size() == row_end.size());
 
  // Block layout: row-major within each block.
  for (decltype(+bs0) k0 = 0; k0 < bs0; ++k0)
  {
    for (std::size_t i = 0; i < row_begin.size(); i++)
    {
      const T xval = x[i * bs0 + k0];
      for (std::int64_t j = row_begin[i]; j < row_end[i]; j++)
      {
        for (decltype(+bs1) k1 = 0; k1 < bs1; ++k1)
        {
          y[indices[j] * bs1 + k1]
              += values[j * bs0 * bs1 + k0 * bs1 + k1] * xval;
        }
      }
    }
  }
}
 
} // namespace impl
} // namespace dolfinx::la