Date: 2025-03-29T05:58:18+00:00

Actual source code: petscvec_kokkos.hpp

      1: #pragma once

      3: #include <petscconf.h>

      5: /* SUBMANSEC = Vec */

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

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

     15: #include <petscvec.h>

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

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

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

     28:    Logically Collective, No Fortran Support

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

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

     36:    Level: beginner

     38:    Notes:
     39:    If the vector is not of type `VECKOKKOS`, an error will be raised.

     41:    The functions are similar to `VecGetArrayRead()` and `VecGetArray()` respectively. One can read-only or read/write the returned Kokkos View.

     43:    Passing in a const View enables read-only access.

     45:    One must return the View by a matching `VecRestoreKokkosView()` after finishing using the View. Currently, only two memory
     46:    spaces are supported: HostMirrorMemorySpace and Kokkos::DefaultExecutionSpace::memory_space.
     47:    If needed, a memory copy will be internally called to copy the latest vector data to the specified memory space.

     49: .seealso: `VecRestoreKokkosView()`, `VecRestoreArray()`, `VecGetKokkosViewWrite()`, `VecGetArrayRead()`, `VecGetArrays()`, `VecPlaceArray()`, `VecGetArray2d()`,
     50:           `VecGetArrayPair()`, `VecRestoreArrayPair()`, `VecGetArrayWrite()`, `VecRestoreArrayWrite()`
     51: @*/
     52: template <class MemorySpace>
     53: PetscErrorCode VecGetKokkosView(Vec, Kokkos::View<const PetscScalar *, MemorySpace> *);
     54: template <class MemorySpace>
     55: PetscErrorCode VecGetKokkosView(Vec, Kokkos::View<PetscScalar *, MemorySpace> *);

     57: /*@C
     58:    VecRestoreKokkosView - Returns a Kokkos View gotten by `VecGetKokkosView()`.

     60:    Synopsis:
     61: #include <petscvec_kokkos.hpp>
     62:    PetscErrorCode VecRestoreKokkosView  (Vec v,Kokkos::View<const PetscScalar*,MemorySpace>* kv);
     63:    PetscErrorCode VecRestoreKokkosView  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);

     65:    Logically Collective, No Fortran Support

     67:    Input Parameters:
     68: +  v  - the vector in type of `VECKOKKOS`
     69: -  kv - the Kokkos View with a user-specified template parameter MemorySpace

     71:    Level: beginner

     73:    Note:
     74:    If the vector is not of type `VECKOKKOS`, an error will be raised.
     75:    The functions are similar to `VecRestoreArrayRead()` and `VecRestoreArray()` respectively. They are the counterpart of `VecGetKokkosView()`.

     77: .seealso: `VecGetKokkosView()`, `VecRestoreKokkosViewWrite()`, `VecRestoreArray()`, `VecGetArrayRead()`, `VecGetArrays()`, `VecPlaceArray()`, `VecGetArray2d()`,
     78:           `VecGetArrayPair()`, `VecRestoreArrayPair()`, `VecGetArrayWrite()`, `VecRestoreArrayWrite()`
     79: @*/
     80: template <class MemorySpace>
     81: PetscErrorCode VecRestoreKokkosView(Vec, Kokkos::View<const PetscScalar *, MemorySpace> *)
     82: {
     83:   return PETSC_SUCCESS;
     84: }
     85: template <class MemorySpace>
     86: PetscErrorCode VecRestoreKokkosView(Vec, Kokkos::View<PetscScalar *, MemorySpace> *);

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

     91:    Synopsis:
     92: #include <petscvec_kokkos.hpp>
     93:    PetscErrorCode VecGetKokkosViewWrite  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);

     95:    Logically Collective, No Fortran Support

     97:    Input Parameter:
     98: .  v - the vector in type of `VECKOKKOS`

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

    103:    Level: beginner

    105:    Notes:
    106:    If the vector is not of type `VECKOKKOS`, an error will be raised.

    108:    The functions is similar to `VecGetArrayWrite()`. The returned view might contain garbage data or stale data and one is not
    109:    expected to read data from the View. Instead, one is expected to overwrite all data in the View.
    110:    One must return the View by a matching `VecRestoreKokkosViewWrite()` after finishing using the View.

    112:    Currently, only two memory spaces are supported: HostMirrorMemorySpace and Kokkos::DefaultExecutionSpace::memory_space.

    114: .seealso: `VecRestoreKokkosViewWrite()`, `VecRestoreKokkosView()`, `VecGetKokkosView()`, `VecRestoreArray()`, `VecGetArrayRead()`, `VecGetArrays()`, `VecPlaceArray()`, `VecGetArray2d()`,
    115:           `VecGetArrayPair()`, `VecRestoreArrayPair()`, `VecGetArrayWrite()`, `VecRestoreArrayWrite()`
    116: @*/
    117: template <class MemorySpace>
    118: PetscErrorCode VecGetKokkosViewWrite(Vec, Kokkos::View<PetscScalar *, MemorySpace> *);

    120: /*@C
    121:    VecRestoreKokkosViewWrite - Returns a Kokkos View gotten with `VecGetKokkosViewWrite()`.

    123:    Synopsis:
    124: #include <petscvec_kokkos.hpp>
    125:    PetscErrorCode VecRestoreKokkosViewWrite  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);

    127:    Logically Collective, No Fortran Support

    129:    Input Parameters:
    130: +  v  - the vector in type of `VECKOKKOS`
    131: -  kv - the Kokkos View with a user-specified template parameter MemorySpace

    133:    Level: beginner

    135:    Notes:
    136:    If the vector is not of type `VECKOKKOS`, an error will be raised.

    138:    The function is similar to `VecRestoreArrayWrite()`. It is the counterpart of `VecGetKokkosViewWrite()`.

    140: .seealso: `VecGetKokkosViewWrite()`, `VecGetKokkosView()`, `VecGetKokkosView()`, `VecRestoreArray()`, `VecGetArrayRead()`, `VecGetArrays()`, `VecPlaceArray()`, `VecGetArray2d()`,
    141:           `VecGetArrayPair()`, `VecRestoreArrayPair()`, `VecGetArrayWrite()`, `VecRestoreArrayWrite()`
    142: @*/
    143: template <class MemorySpace>
    144: PetscErrorCode VecRestoreKokkosViewWrite(Vec, Kokkos::View<PetscScalar *, MemorySpace> *);

    146:   #if defined(PETSC_HAVE_COMPLEX) && defined(PETSC_USE_COMPLEX)
    147: static_assert(std::alignment_of<Kokkos::complex<PetscReal>>::value == std::alignment_of<std::complex<PetscReal>>::value,
    148:               "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");
    149:   #endif

    151: #endif