KokkosSparse::BsrMatrix¶
Defined in header KokkosSparse_BsrMatrix.hpp
template <class ScalarType, class OrdinalType, class Device, class MemoryTraits = void,
class SizeType = KokkosKernels::default_size_type>
class BsrMatrix;
KokkosSparse::BsrMatrix
provides a block compressed sparse row implementation of a sparse matrix, as described, for example, in Saad (2nd ed.).
Template Parameters¶
- ScalarType:
type of the values stored in the matrix.
- OrdinalType:
type of the indices storied in the matrix (column indices).
- Device:
type of the
Kokkos::Device
of the matrix- MemoryTraits:
type of memory traits used for the views of the matrix.
- SizeType:
type used to store row offsets.
Member Types¶
Memeber type |
Definition |
---|---|
execution_space |
Alias for Device::execution_space. |
memory_space |
Alias for Device::memory_space. |
device_type |
Device associated with the matrix. |
value_type |
Type of the values stored in the matrix, alias for ScalarType. |
ordinal_type |
Type of the indices stored in the matrix, alias for OrdinalType. |
size_type |
Type of the offsets stored in the matrix, alias for SizeType. |
memory_traits |
Alias for MemoryTraits. |
HostMirror |
CrsMatrix type templated on ScalarType, OrdinalType, host_mirror_space, MemoryTraits and SizeType. |
StaticCrsGraph |
Type of the underlying static crs graph. |
staticcsrgraph |
Alias for StaticCrsGraph. |
index_type |
Type of the view storing the column indices in the underlying static crs graph. |
const_ordinal_type |
Const variant for the type of the column indices stored in index_type. |
non_const_ordinal_type |
Non-const variant for the type of the column indices stored in index_type. |
row_map_type |
Type of the view storing the row offsets in the underlying static crs graph. |
const_size_type |
Const variant for the type of the offsets stored in row_map_type. |
non_const_size_type |
Non-const variant for the type of the offsets stored in row_map_type. |
values_type |
Type of the view storing the values stored in the matrix. |
const_value_type |
Const variant for the type of the values stored in values_type. |
non_const_value_type |
Non-const variant for the type of the values stored in values_type. |
block_layout_type |
The Kokkos layout used when extracting blocks from the matrix. It is an alias for |
block_type |
Type used to access of a view of a block as Kokkos::View or rank 2. |
const_block_type |
Type used to access of a const view of a block as Kokkos::View or rank 2. |
Data Members¶
Data Member |
Definition |
---|---|
staticcrsgraph_type graph |
The underlying static crs graph of the matrix storing its sparsity pattern. |
values_type values |
The values stored in the matrix. |
const ordinal blockDim_ |
The size of the blocks stored by the matrix. |
Member Functions¶
Name |
Definition |
---|---|
Construct a CrsMatrix from inputs, this can perform shallow or deep copies of input parameters depending on inputs and semantic. |
|
Assignment operator, will attempt to assign the data (graph and values) of the right hand side to the left hand side. If the assignment operator of graph or values fail, matrix assignment fails. |
|
Returns the number of rows in the matrix. |
|
Returns the number of columns in the matrix. |
|
Returns the dimension of the blocks stored in the matrix. |
|
Returns the number of structural non-zero values in the matrix (some of these might actually store zero). |
|
Returns a SparseRowView object from a row of the matrix. |
|
Returns a SparseRowViewConst object from a row of the matrix. |
|
Return a view of the i-th block in the matrix. |
|
Return a const view of the i-th block in the matrix. |
operator=¶
template <typename aScalarType, typename aOrdinalType, class aDevice, class aMemoryTraits, typename aSizeType>
BsrMatrix& operator=(const BsrMatrix<aScalarType, aOrdinalType, aDevice, aMemoryTraits, aSizeType>& mtx);
Attempts to assign the underlying graph
and values
of the input matrix mtx
to the matrix.
numRows¶
KOKKOS_INLINE_FUNCTION ordinal_type numRows() const;
Returns the number of rows in the matrix.
numCols¶
KOKKOS_INLINE_FUNCTION ordinal_type numCols() const;
Returns the number of columns in the matrix.
blockDim¶
KOKKOS_INLINE_FUNCTION ordinal_type blockDim() const;
Returns the dimension of the blocks stored by the matrix.
nnz¶
KOKKOS_INLINE_FUNCTION size_type nnz() const;
Returns the number of non-zero entries in the matrix.
block_row¶
KOKKOS_INLINE_FUNCTION
BsrRowView<BsrMatrix> block_row(const ordinal_type i) const;
Returns a view of the i-th block row of the matrix as a BsrRowView.
block_row_Const¶
KOKKOS_INLINE_FUNCTION
BsrRowViewConst<BsrMatrix> block_row_Const(const ordinal_type i) const;
Returns a view of the i-th block row of the matrix as a BsrRowViewConst.
unmanaged_block¶
KOKKOS_INLINE_FUNCTION
block_type unmanaged_block(const size_type i) const;
Return a view of the i-th block in the matrix.
unmanaged_block_const¶
KOKKOS_INLINE_FUNCTION
const_block_type unmanaged_block_const(const size_type i) const;
Return a const view of the i-th block in the matrix.
Example¶
#include <sstream>
#include <iostream>
#include <iomanip>
#include "Kokkos_Core.hpp"
#include "KokkosKernels_default_types.hpp"
#include "KokkosSparse_BsrMatrix.hpp"
using Scalar = KokkosKernels::default_scalar;
using Ordinal = KokkosKernels::default_lno_t;
using Offset = KokkosKernels::default_size_type;
using Layout = KokkosKernels::default_layout;
template <class bsrmatrix_type>
struct bsr_fill {
bsrmatrix_type bsr_mat;
bsr_fill(bsrmatrix_type bsr_mat_) : bsr_mat(bsr_mat_) {}
KOKKOS_INLINE_FUNCTION
void operator()(const int& rowIdx) const {
if (rowIdx == 0) { // Left boundary condition
auto block_tmp = bsr_mat.unmanaged_block(0);
block_tmp(0, 0) = 1.0;
block_tmp(0, 1) = 0.0;
block_tmp(1, 0) = 0.0;
block_tmp(1, 1) = 1.0;
} else if (rowIdx == bsr_mat.numRows() - 1) { // Right boundary condition
auto block_tmp = bsr_mat.unmanaged_block(bsr_mat.graph.row_map(rowIdx) + 1);
block_tmp(0, 0) = 1.0;
block_tmp(1, 1) = 1.0;
} else {
auto block_tmp = bsr_mat.unmanaged_block(bsr_mat.graph.row_map(rowIdx));
block_tmp(0, 0) = -1.0;
block_tmp(0, 1) = -1.0 / 2.0;
block_tmp(1, 0) = 0.0;
block_tmp(1, 1) = -1.0;
block_tmp = bsr_mat.unmanaged_block(bsr_mat.graph.row_map(rowIdx) + 1);
block_tmp(0, 0) = 2.0;
block_tmp(0, 1) = 0.0;
block_tmp(1, 0) = 0.0;
block_tmp(1, 1) = 2.0;
block_tmp = bsr_mat.unmanaged_block(bsr_mat.graph.row_map(rowIdx) + 2);
block_tmp(0, 0) = -1.0;
block_tmp(0, 1) = 1.0 / 2.0;
block_tmp(1, 0) = 0.0;
block_tmp(1, 1) = -1.0;
}
}
};
template <class bsrmatrix_type, class diag_blocks_type>
struct diagonal_extractor {
using graph_type = typename bsrmatrix_type::staticcrsgraph_type;
using row_map_type = typename graph_type::row_map_type;
using entries_type = typename graph_type::entries_type;
using bsr_block_type = typename bsrmatrix_type::block_type;
bsrmatrix_type bsr_mat;
row_map_type row_map;
entries_type entries;
diag_blocks_type diag_blocks;
diagonal_extractor(bsrmatrix_type bsr_mat_, diag_blocks_type diag_blocks_)
: bsr_mat(bsr_mat_),
row_map(bsr_mat_.graph.row_map),
entries(bsr_mat_.graph.entries),
diag_blocks(diag_blocks_) {}
KOKKOS_INLINE_FUNCTION
void operator()(const int& rowIdx) const {
for (Offset entryIdx = row_map(rowIdx); entryIdx < row_map(rowIdx + 1); ++entryIdx) {
if (entries(entryIdx) == rowIdx) {
bsr_block_type bsr_diag_block = bsr_mat.unmanaged_block(entryIdx);
for (int i = 0; i < bsr_mat.blockDim(); ++i) {
for (int j = 0; j < bsr_mat.blockDim(); ++j) {
diag_blocks(rowIdx, i, j) = bsr_diag_block(i, j);
}
}
}
}
}
};
int main(int argc, char* argv[]) {
using device_type =
typename Kokkos::Device<Kokkos::DefaultExecutionSpace, typename Kokkos::DefaultExecutionSpace::memory_space>;
using bsrmatrix_type = typename KokkosSparse::Experimental::BsrMatrix<Scalar, Ordinal, device_type, void, Offset>;
using graph_type = typename bsrmatrix_type::staticcrsgraph_type;
using row_map_type = typename graph_type::row_map_type;
using entries_type = typename graph_type::entries_type;
Kokkos::initialize(argc, argv);
{
//
// We will create a 1D discretization for the coupled thermo-elastic
// diffusion
//
// -\div(EA \grad_s(u) - \alpha(T-T0)I) = f_u
// -\kappa\Delta(T) = f_T
//
// The problem is discretized using finite differences as follows:
// \frac{d^2 u}{dx^2}\approx \frac{u_{i+1}-2u_i+u_{i-1}}{h_x^2}
// \frac{dT}{dx}\approx\frac{T_{i+1}-T_{i-1}}{2h_x}
// \frac{d^2T}{dx^2}\approx\frac{T_{i+1}-2T_i+T_{i-1}}{h_x^2}
//
// This leads to the combined stencil (assuming all unit coefficients):
//
// [-1 1/2] [2 0] [-1 -1/2]
// [ 0 -1] [0 2] [ 0 -1]
//
// First the graph for the mesh will be constructed.
// Second a BsrMatrix will be constructed from the graph
// Third the values of the BsrMatrix will be filled.
constexpr Ordinal blockSize = 2;
constexpr Ordinal numRows = 10;
constexpr Offset numNNZ = 3 * numRows - 2;
bsrmatrix_type bsr_mat;
{
typename row_map_type::non_const_type row_map(Kokkos::view_alloc(Kokkos::WithoutInitializing, "row pointers"),
numRows + 1);
typename entries_type::non_const_type entries(Kokkos::view_alloc(Kokkos::WithoutInitializing, "column indices"),
numNNZ);
typename row_map_type::HostMirror row_map_h = Kokkos::create_mirror_view(row_map);
typename entries_type::HostMirror entries_h = Kokkos::create_mirror_view(entries);
// First Step: build the CrsGraph
{
// Build the row pointers and store numNNZ
row_map_h(0) = 0;
for (Ordinal rowIdx = 0; rowIdx < numRows; ++rowIdx) {
if (rowIdx == 0) {
row_map_h(rowIdx + 1) = row_map_h(rowIdx) + 2;
entries_h(row_map_h(rowIdx)) = rowIdx;
entries_h(row_map_h(rowIdx) + 1) = rowIdx + 1;
} else if (rowIdx == numRows - 1) {
row_map_h(rowIdx + 1) = row_map_h(rowIdx) + 2;
entries_h(row_map_h(rowIdx)) = rowIdx - 1;
entries_h(row_map_h(rowIdx) + 1) = rowIdx;
} else {
row_map_h(rowIdx + 1) = row_map_h(rowIdx) + 3;
entries_h(row_map_h(rowIdx)) = rowIdx - 1;
entries_h(row_map_h(rowIdx) + 1) = rowIdx;
entries_h(row_map_h(rowIdx) + 2) = rowIdx + 1;
}
}
if (row_map_h(numRows) != numNNZ) {
std::ostringstream error_msg;
error_msg << "error: row_map(numRows) != numNNZ, row_map_h(numRows)=" << row_map_h(numRows)
<< ", numNNZ=" << numNNZ;
throw std::runtime_error(error_msg.str());
}
Kokkos::deep_copy(row_map, row_map_h);
Kokkos::deep_copy(entries, entries_h);
}
graph_type myGraph(entries, row_map);
// Second Step: build the BsrMatrix from graph and block size
bsr_mat = bsrmatrix_type("block matrix", myGraph, blockSize);
bsr_fill fillFunctor(bsr_mat);
Kokkos::parallel_for(Kokkos::RangePolicy<int>(0, numRows), fillFunctor);
std::cout << "BsrMatrix graph: " << std::endl;
for (int rowIdx = 0; rowIdx < numRows; ++rowIdx) {
std::cout << " [";
for (int colIdx = 0; colIdx < entries_h(row_map_h(rowIdx)); ++colIdx) {
std::cout << " ";
}
std::cout << "*";
for (Offset entryIdx = row_map_h(rowIdx); entryIdx < row_map_h(rowIdx + 1) - 1; ++entryIdx) {
for (int colIdx = entries_h(entryIdx) + 1; colIdx < entries_h(entryIdx + 1); ++colIdx) {
std::cout << " ";
}
std::cout << "*";
}
for (int colIdx = entries_h(row_map_h(rowIdx + 1) - 1) + 1; colIdx < numRows; ++colIdx) {
std::cout << " ";
}
std::cout << "]" << std::endl;
}
}
// Extract diagonal block and store them in a rank-3 view
using diag_blocks_type = Kokkos::View<Scalar***, typename bsrmatrix_type::block_layout_type, device_type>;
diag_blocks_type diag_blocks("diagonal blocks", numRows, blockSize, blockSize);
diagonal_extractor myFunc(bsr_mat, diag_blocks);
Kokkos::parallel_for(Kokkos::RangePolicy<int>(0, numRows), myFunc);
auto diag_blocks_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, diag_blocks);
std::cout << "\nBsrMatrix diagonal blocks: " << std::endl;
for (int blockId = 0; blockId < diag_blocks_h.extent_int(0); ++blockId) {
std::cout << " [" << diag_blocks_h(blockId, 0, 0) << ", " << diag_blocks_h(blockId, 0, 1) << "]" << std::endl;
std::cout << " [" << diag_blocks_h(blockId, 1, 0) << ", " << diag_blocks_h(blockId, 1, 1) << "]\n" << std::endl;
}
}
Kokkos::finalize();
return 0;
}