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"

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);
  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_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