KokkosMatrix.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://www.mooseframework.org
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html
 
#pragma once
 
#include "KokkosArray.h"
 
#ifdef MOOSE_KOKKOS_SCOPE
#include "KokkosUtils.h"
#endif
 
#include "libmesh/petsc_matrix.h"
 
#include <vector>
 
namespace Moose::Kokkos
{
 
class System;
 
class Matrix
{
public:
  Matrix() = default;
  ~Matrix() { destroy(); }
  Mat mat() { return _matrix; }
  void destroy();
 
#ifdef MOOSE_KOKKOS_SCOPE
  bool isAlloc() const { return _is_alloc; }
  void create(libMesh::SparseMatrix<PetscScalar> & matrix,
              const System & system,
              std::vector<PetscInt> & coo_i,
              std::vector<PetscInt> & coo_j);
  void close();
 
  KOKKOS_FUNCTION void zero(PetscInt i)
  {
    for (PetscInt j = _row_ptr[i]; j < _row_ptr[i + 1]; ++j)
      _val[j] = 0;
  }
  KOKKOS_FUNCTION PetscScalar & operator()(PetscInt i, PetscInt j) const
  {
    auto idx = find(i, j);
 
    KOKKOS_ASSERT(idx != -1);
 
    return _val[idx];
  }
  auto & operator=(PetscScalar scalar)
  {
    _val = scalar;
 
    return *this;
  }
#endif
 
private:
#ifdef MOOSE_KOKKOS_SCOPE
  KOKKOS_FUNCTION PetscInt find(PetscInt i, PetscInt j) const;
#endif
 
  Mat _matrix = PETSC_NULLPTR;
  PetscCount _nr = 0;
  Array<PetscInt> _col_idx;
  Array<PetscInt> _row_ptr;
  Array<PetscScalar> _val;
 
  bool _is_host = false;
  bool _is_alloc = false;
};
 
#ifdef MOOSE_KOKKOS_SCOPE
KOKKOS_FUNCTION inline PetscInt
Matrix::find(PetscInt i, PetscInt j) const
{
  KOKKOS_ASSERT(i < _nr);
 
  auto begin = _col_idx.data() + _row_ptr[i];
  auto end = _col_idx.data() + _row_ptr[i + 1];
  auto target = Utils::find(j, begin, end);
 
  if (target == end)
    return -1;
  else
    return target - _col_idx.data();
}
#endif
 
} // namespace Moose::Kokkos