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