Actual source code: petscvec_kokkos.hpp

      1: #if !defined(PETSCVEC_KOKKOS_HPP)
      2: #define PETSCVEC_KOKKOS_HPP

      4: #include <petscconf.h>

      6: #if defined(PETSC_HAVE_KOKKOS)
      7:   #if defined(petsccomplexlib)
      8:     #error "Error: You must include petscvec_kokkos.hpp before other petsc headers in this C++ file to use petsc complex with Kokkos"
      9:   #endif

     11:   #define PETSC_DESIRE_KOKKOS_COMPLEX   1   /* To control the definition of petsccomplexlib in petscsystypes.h */
     12: #endif

     14: #include <petscvec.h>

     16: #if defined(PETSC_HAVE_KOKKOS)
     17:   #include <Kokkos_Core.hpp>

     19: /*@C
     20:      VecGetKokkosView - Returns a constant Kokkos View that contains up-to-date data of a vector in the specified memory space.

     22:    Synopsis:
     23: #include <petscvec_kokkos.hpp>
     24:    PetscErrorCode VecGetKokkosView  (Vec v,Kokkos::View<const PetscScalar*,MemorySpace>* kv);
     25:    PetscErrorCode VecGetKokkosView  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);

     27:    Logically Collective on Vec

     29:    Input Parameter:
     30: .  v - the vector in type of VECKOKKOS

     32:    Output Parameter:
     33: .  kv - the Kokkos View with a user-specified template parameter MemorySpace

     35:    Notes:
     36:    If the vector is not of type VECKOKKOS, an error will be raised.
     37:    The functions are similar to VecGetArrayRead() and VecGetArray() respectively. One can read-only or read/write the returned Kokkos View.
     38:    Note that passing in a const View enables read-only access.
     39:    One must return the View by a matching VecRestoreKokkosView() after finishing using the View. Currently, only two memory
     40:    spaces are supported: Kokkos::HostSpace and Kokkos::DefaultExecutionSpace::memory_space.
     41:    If needed, a memory copy will be internally called to copy the latest vector data to the specified memory space.

     43:    Level: beginner

     45: .seealso: VecRestoreKokkosView(), VecRestoreArray(), VecGetKokkosViewWrite(), VecGetArrayRead(), VecGetArrays(), VecGetArrayF90(), VecGetArrayReadF90(), VecPlaceArray(), VecGetArray2d(),
     46:           VecGetArrayPair(), VecRestoreArrayPair(), VecGetArrayWrite(), VecRestoreArrayWrite()
     47: @*/
     48:   template<class MemorySpace> PetscErrorCode VecGetKokkosView (Vec,Kokkos::View<const PetscScalar*,MemorySpace>*);
     49:   template<class MemorySpace> PetscErrorCode VecGetKokkosView (Vec,Kokkos::View<PetscScalar*,MemorySpace>*);

     51: /*@C
     52:    VecRestoreKokkosView - Returns a Kokkos View gotten by VecGetKokkosView().

     54:    Synopsis:
     55: #include <petscvec_kokkos.hpp>
     56:    PetscErrorCode VecRestoreKokkosView  (Vec v,Kokkos::View<const PetscScalar*,MemorySpace>* kv);
     57:    PetscErrorCode VecRestoreKokkosView  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);

     59:    Logically Collective on Vec

     61:    Input Parameters:
     62: +  v  - the vector in type of VECKOKKOS
     63: -  kv - the Kokkos View with a user-specified template parameter MemorySpace

     65:    Notes:
     66:    If the vector is not of type VECKOKKOS, an error will be raised.
     67:    The functions are similar to VecRestoreArrayRead() and VecRestoreArray() respectively. They are the counterpart of VecGetKokkosView().

     69:    Level: beginner

     71: .seealso: VecGetKokkosView(), VecRestoreKokkosViewWrite(), VecRestoreArray(), VecGetArrayRead(), VecGetArrays(), VecGetArrayF90(), VecGetArrayReadF90(), VecPlaceArray(), VecGetArray2d(),
     72:           VecGetArrayPair(), VecRestoreArrayPair(), VecGetArrayWrite(), VecRestoreArrayWrite()
     73: @*/
     74:   template<class MemorySpace> PetscErrorCode VecRestoreKokkosView(Vec,Kokkos::View<const PetscScalar*,MemorySpace>*){return 0;}
     75:   template<class MemorySpace> PetscErrorCode VecRestoreKokkosView(Vec,Kokkos::View<PetscScalar*,MemorySpace>*);

     78: /*@C
     79:    VecGetKokkosViewWrite - Returns a Kokkos View that contains the array of a vector in the specified memory space.

     81:    Synopsis:
     82: #include <petscvec_kokkos.hpp>
     83:    PetscErrorCode VecGetKokkosViewWrite  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);

     85:    Logically Collective on Vec

     87:    Input Parameter:
     88: .  v - the vector in type of VECKOKKOS

     90:    Output Parameter:
     91: .  kv - the Kokkos View with a user-specified template parameter MemorySpace

     93:    Notes:
     94:    If the vector is not of type VECKOKKOS, an error will be raised.
     95:    The functions is similar to VecGetArrayWrite(). The returned view might contain garbage data or stale data and one is not
     96:    expected to read data from the View. Instead, one is expected to overwrite all data in the View.
     97:    One must return the View by a matching VecRestoreKokkosViewWrite() after finishing using the View.
     98:   Currently, only two memory spaces are supported: Kokkos::HostSpace and Kokkos::DefaultExecutionSpace::memory_space.

    100:    Level: beginner

    102: .seealso: VecRestoreKokkosViewWrite(), VecRestoreKokkosView(), VecGetKokkosView(), VecRestoreArray(), VecGetArrayRead(), VecGetArrays(), VecGetArrayF90(), VecGetArrayReadF90(), VecPlaceArray(), VecGetArray2d(),
    103:           VecGetArrayPair(), VecRestoreArrayPair(), VecGetArrayWrite(), VecRestoreArrayWrite()
    104: @*/
    105:   template<class MemorySpace> PetscErrorCode VecGetKokkosViewWrite    (Vec,Kokkos::View<PetscScalar*,MemorySpace>*);

    107: /*@C
    108:    VecRestoreKokkosViewWrite - Returns a Kokkos View gotten by VecGetKokkosViewWrite().

    110:    Synopsis:
    111: #include <petscvec_kokkos.hpp>
    112:    PetscErrorCode VecRestoreKokkosViewWrite  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);

    114:    Logically Collective on Vec

    116:    Input Parameters:
    117: +  v  - the vector in type of VECKOKKOS
    118: -  kv - the Kokkos View with a user-specified template parameter MemorySpace

    120:    Notes:
    121:    If the vector is not of type VECKOKKOS, an error will be raised.
    122:    The function is similar to VecRestoreArrayWrite(). It is the counterpart of VecGetKokkosViewWrite().

    124:    Level: beginner

    126: .seealso: VecGetKokkosViewWrite(), VecGetKokkosView(), VecGetKokkosView(), VecRestoreArray(), VecGetArrayRead(), VecGetArrays(), VecGetArrayF90(), VecGetArrayReadF90(), VecPlaceArray(), VecGetArray2d(),
    127:           VecGetArrayPair(), VecRestoreArrayPair(), VecGetArrayWrite(), VecRestoreArrayWrite()
    128: @*/
    129:   template<class MemorySpace> PetscErrorCode VecRestoreKokkosViewWrite(Vec,Kokkos::View<PetscScalar*,MemorySpace>*);

    131:  #if defined(PETSC_HAVE_COMPLEX) && defined(PETSC_USE_COMPLEX)
    132:   static_assert(std::alignment_of<Kokkos::complex<PetscReal>>::value == std::alignment_of<std::complex<PetscReal>>::value,
    133:    "Alignment of Kokkos::complex<PetscReal> and std::complex<PetscReal> mismatch. Reconfigure your Kokkos with -DKOKKOS_ENABLE_COMPLEX_ALIGN=OFF, or let PETSc install Kokkos for you with --download-kokkos --download-kokkos-kernels");
    134:  #endif

    136: #endif

    138: #endif