Line data    Source code

       1             : //* This file is part of the MOOSE framework
       2             : //* https://www.mooseframework.org
       3             : //*
       4             : //* All rights reserved, see COPYRIGHT for full restrictions
       5             : //* https://github.com/idaholab/moose/blob/master/COPYRIGHT
       6             : //*
       7             : //* Licensed under LGPL 2.1, please see LICENSE for details
       8             : //* https://www.gnu.org/licenses/lgpl-2.1.html
       9             : 
      10             : #pragma once
      11             : 
      12             : #include "KokkosArray.h"
      13             : 
      14             : #include "libmesh/petsc_vector.h"
      15             : #include "libmesh/dof_map.h"
      16             : 
      17             : namespace Moose::Kokkos
      18             : {
      19             : 
      20             : class System;
      21             : 
      22             : /**
      23             :  * The Kokkos wrapper class for PETSc vector
      24             :  */
      25             : class Vector
      26             : {
      27             : public:
      28             :   /**
      29             :    * Default constructor
      30             :    */
      31      135900 :   Vector() = default;
      32             :   /**
      33             :    * Destructor
      34             :    */
      35      134940 :   ~Vector() { destroy(); }
      36             :   /**
      37             :    * Free all data and reset
      38             :    */
      39             :   void destroy();
      40             : 
      41             : #ifdef MOOSE_KOKKOS_SCOPE
      42             :   /**
      43             :    * Get whether the vector was allocated
      44             :    * @returns Whether the vector was allocated
      45             :    */
      46             :   bool isAlloc() const { return _is_alloc; }
      47             :   /**
      48             :    * Create the vector from a libMesh PetscVector
      49             :    * @param vector The libMesh PetscVector
      50             :    * @param system The Kokkos system
      51             :    * @param assemble Whether the vector will be assembled
      52             :    */
      53             :   void create(libMesh::NumericVector<PetscScalar> & vector, const System & system, bool assemble);
      54             :   /**
      55             :    * Copy from the host libMesh PetscVector
      56             :    */
      57             :   void copyToDevice();
      58             :   /**
      59             :    * Copy to the host libMesh PetscVector
      60             :    */
      61             :   void copyToHost();
      62             :   /**
      63             :    * Restore the underlying PETSc vector
      64             :    */
      65             :   void restore();
      66             :   /**
      67             :    * Assemble the underlying PETSc vector
      68             :    */
      69             :   void close();
      70             : 
      71             :   /**
      72             :    * Get an entry with a given index
      73             :    * @param i The entry index local to this process
      74             :    * @returns The reference of the entry
      75             :    */
      76             :   KOKKOS_FUNCTION PetscScalar & operator()(dof_id_type i) const
      77             :   {
      78             :     return i < _local.size() ? _local[i] : _ghost(i);
      79             :   }
      80             :   /**
      81             :    * Get an entry with a given index
      82             :    * @param i The entry index local to this process
      83             :    * @returns The reference of the entry
      84             :    */
      85   346988879 :   KOKKOS_FUNCTION PetscScalar & operator[](dof_id_type i) const
      86             :   {
      87   346988879 :     return i < _local.size() ? _local[i] : _ghost(i);
      88             :   }
      89             :   /**
      90             :    * Assign a scalar value uniformly
      91             :    * @param scalar The scalar value to be assigned
      92             :    */
      93             :   auto & operator=(PetscScalar scalar)
      94             :   {
      95             :     _local = scalar;
      96             :     _ghost = scalar;
      97             : 
      98             :     return *this;
      99             :   }
     100             : 
     101             :   /**
     102             :    * Kokkos functions for direct assembly on device
     103             :    */
     104             :   ///@{
     105             :   struct PackBuffer
     106             :   {
     107             :   };
     108             :   struct UnpackBuffer
     109             :   {
     110             :   };
     111             : 
     112             :   KOKKOS_FUNCTION void operator()(PackBuffer, const PetscCount tid) const;
     113             :   KOKKOS_FUNCTION void operator()(UnpackBuffer, const PetscCount tid) const;
     114             :   ///@}
     115             : #endif
     116             : 
     117             : private:
     118             :   /**
     119             :    * Data for direct assembly on device
     120             :    */
     121             :   ///@{
     122             :   struct DeviceAssembly
     123             :   {
     124             :     /**
     125             :      * List of DOFs to send/receive for each process
     126             :      */
     127             :     Array<Array<libMesh::dof_id_type>> list;
     128             :     /**
     129             :      * Number of DOFs to send/receive for each process
     130             :      */
     131             :     Array<int> count;
     132             :     /**
     133             :      * Starting offset of each process into the communication buffer
     134             :      */
     135             :     Array<int> offset;
     136             :     /**
     137             :      * Communication buffer
     138             :      */
     139             :     Array<PetscScalar> buffer;
     140             :     /**
     141             :      * Allocate data
     142             :      */
     143             :     void create(const Array<Array<libMesh::dof_id_type>> & list);
     144             :     /**
     145             :      * Free data
     146             :      */
     147             :     void destroy();
     148             :   };
     149             : 
     150             :   DeviceAssembly _send;
     151             :   DeviceAssembly _recv;
     152             : 
     153             :   unsigned int _current_proc;
     154             :   ///@}
     155             : 
     156             :   /**
     157             :    * PETSc vectors
     158             :    */
     159             :   ///@{
     160             :   Vec _global_vector = PETSC_NULLPTR;
     161             :   Vec _local_vector = PETSC_NULLPTR;
     162             :   ///@}
     163             :   /**
     164             :    * Raw data of local PETSc vector
     165             :    */
     166             :   PetscScalar * _array = PETSC_NULLPTR;
     167             :   /**
     168             :    * Pointer to the Kokkos system
     169             :    */
     170             :   const System * _system;
     171             :   /**
     172             :    * Pointer to the libMesh communicator
     173             :    */
     174             :   const libMesh::Parallel::Communicator * _comm = nullptr;
     175             :   /**
     176             :    * Data vectors on device
     177             :    */
     178             :   ///@{
     179             :   Array<PetscScalar> _local;
     180             :   Array<PetscScalar> _ghost;
     181             :   ///@}
     182             :   /**
     183             :    * Flag whether the vector will be assembled
     184             :    */
     185             :   bool _assemble = false;
     186             :   /**
     187             :    * Flag whether the PETSc vector is ghosted
     188             :    */
     189             :   bool _is_ghosted = false;
     190             :   /**
     191             :    * Flag whether the PETSc vector is a host vector
     192             :    */
     193             :   bool _is_host = false;
     194             :   /**
     195             :    * Flag whether the vector was allocated
     196             :    */
     197             :   bool _is_alloc = false;
     198             : };
     199             : 
     200             : } // namespace Moose::Kokkos