KokkosVector.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"
 
#include "libmesh/petsc_vector.h"
#include "libmesh/dof_map.h"
 
namespace Moose::Kokkos
{
 
class System;
 
class Vector
{
public:
  Vector() = default;
  ~Vector() { destroy(); }
  void destroy();
 
#ifdef MOOSE_KOKKOS_SCOPE
  bool isAlloc() const { return _is_alloc; }
  void create(libMesh::NumericVector<PetscScalar> & vector, const System & system, bool assemble);
  void copyToDevice();
  void copyToHost();
  void restore();
  void close();
 
  KOKKOS_FUNCTION PetscScalar & operator()(dof_id_type i) const
  {
    return i < _local.size() ? _local[i] : _ghost(i);
  }
  KOKKOS_FUNCTION PetscScalar & operator[](dof_id_type i) const
  {
    return i < _local.size() ? _local[i] : _ghost(i);
  }
  auto & operator=(PetscScalar scalar)
  {
    _local = scalar;
    _ghost = scalar;
 
    return *this;
  }
 
  struct PackBuffer
  {
  };
  struct UnpackBuffer
  {
  };
 
  KOKKOS_FUNCTION void operator()(PackBuffer, const PetscCount tid) const;
  KOKKOS_FUNCTION void operator()(UnpackBuffer, const PetscCount tid) const;
#endif
 
private:
  struct DeviceAssembly
  {
    Array<Array<libMesh::dof_id_type>> list;
    Array<int> count;
    Array<int> offset;
    Array<PetscScalar> buffer;
    void create(const Array<Array<libMesh::dof_id_type>> & list);
    void destroy();
  };
 
  DeviceAssembly _send;
  DeviceAssembly _recv;
 
  unsigned int _current_proc;
 
  Vec _global_vector = PETSC_NULLPTR;
  Vec _local_vector = PETSC_NULLPTR;
 
  PetscScalar * _array = PETSC_NULLPTR;
  const System * _system;
  const libMesh::Parallel::Communicator * _comm = nullptr;
  Array<PetscScalar> _local;
  Array<PetscScalar> _ghost;
 
  bool _assemble = false;
  bool _is_ghosted = false;
  bool _is_host = false;
  bool _is_alloc = false;
};
 
} // namespace Moose::Kokkos