Skip to content
Snippets Groups Projects

Replace SparseMatrixDescriptor by CRSMatrixDescriptor

15 files
+ 640
484
Compare changes
  • Side-by-side
  • Inline

Files

+ 51
33
#ifndef CRS_MATRIX_HPP
#define CRS_MATRIX_HPP
#include <Kokkos_StaticCrsGraph.hpp>
#include <utils/Array.hpp>
#include <utils/PugsAssert.hpp>
#include <algebra/SparseMatrixDescriptor.hpp>
#include <algebra/Vector.hpp>
#include <iostream>
template <typename DataType = double, typename IndexType = size_t>
template <typename DataType, typename IndexType>
class CRSMatrixDescriptor;
template <typename DataType = double, typename IndexType = int>
class CRSMatrix
{
public:
using MutableDataType = std::remove_const_t<DataType>;
using index_type = IndexType;
using data_type = DataType;
private:
using HostMatrix = Kokkos::StaticCrsGraph<const IndexType, Kokkos::HostSpace>;
const IndexType m_nb_rows;
const IndexType m_nb_columns;
HostMatrix m_host_matrix;
Array<const IndexType> m_row_map;
Array<const DataType> m_values;
size_t m_nb_rows;
size_t m_nb_columns;
Array<const IndexType> m_column_indices;
public:
PUGS_INLINE
@@ -33,14 +36,14 @@ class CRSMatrix
}
PUGS_INLINE
size_t
IndexType
numberOfRows() const
{
return m_nb_rows;
}
PUGS_INLINE
size_t
IndexType
numberOfColumns() const
{
return m_nb_columns;
@@ -53,15 +56,15 @@ class CRSMatrix
}
auto
rowIndices() const
rowMap() const
{
return encapsulate(m_host_matrix.row_map);
return m_row_map;
}
auto
row(size_t i) const
columnIndices() const
{
return m_host_matrix.rowConst(i);
return m_column_indices;
}
template <typename DataType2>
@@ -71,18 +74,17 @@ class CRSMatrix
static_assert(std::is_same<std::remove_const_t<DataType>, std::remove_const_t<DataType2>>(),
"Cannot multiply matrix and vector of different type");
Assert(x.size() == m_host_matrix.numRows());
Assert(x.size() - m_nb_columns == 0);
Vector<MutableDataType> Ax{m_host_matrix.numRows()};
auto host_row_map = m_host_matrix.row_map;
Vector<MutableDataType> Ax(m_nb_rows);
parallel_for(
m_host_matrix.numRows(), PUGS_LAMBDA(const IndexType& i_row) {
const auto row_begin = host_row_map(i_row);
const auto row_end = host_row_map(i_row + 1);
m_nb_rows, PUGS_LAMBDA(const IndexType& i_row) {
const auto row_begin = m_row_map[i_row];
const auto row_end = m_row_map[i_row + 1];
MutableDataType sum{0};
for (IndexType j = row_begin; j < row_end; ++j) {
sum += m_values[j] * x[m_host_matrix.entries(j)];
sum += m_values[j] * x[m_column_indices[j]];
}
Ax[i_row] = sum;
});
@@ -90,21 +92,37 @@ class CRSMatrix
return Ax;
}
CRSMatrix(const SparseMatrixDescriptor<DataType, IndexType>& M)
: m_nb_rows{M.numberOfRows()}, m_nb_columns{M.numberOfColumns()}
friend PUGS_INLINE std::ostream&
operator<<(std::ostream& os, const CRSMatrix& A)
{
{
auto host_matrix =
Kokkos::create_staticcrsgraph<Kokkos::StaticCrsGraph<IndexType, Kokkos::HostSpace>>("connectivity_matrix",
M.graphVector());
// This is a bit crappy but it is the price to pay to avoid
m_host_matrix.entries = host_matrix.entries;
m_host_matrix.row_map = host_matrix.row_map;
m_host_matrix.row_block_offsets = host_matrix.row_block_offsets;
for (IndexType i = 0; i < A.m_nb_rows; ++i) {
const auto row_begin = A.m_row_map[i];
const auto row_end = A.m_row_map[i + 1];
os << i << "|";
for (IndexType j = row_begin; j < row_end; ++j) {
os << ' ' << A.m_column_indices[j] << ':' << A.m_values[j];
}
os << '\n';
}
m_values = M.valueArray();
return os;
}
private:
friend class CRSMatrixDescriptor<DataType, IndexType>;
CRSMatrix(const IndexType nb_rows,
const IndexType nb_columns,
const Array<const IndexType>& row_map,
const Array<const DataType>& values,
const Array<const IndexType>& column_indices)
: m_nb_rows{nb_rows},
m_nb_columns{nb_columns},
m_row_map{row_map},
m_values{values},
m_column_indices{column_indices}
{}
public:
~CRSMatrix() = default;
};
Loading