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