Actual source code: veckokkosimpl.hpp

      1: #pragma once

      3: #include <../src/vec/vec/impls/mpi/pvecimpl.h>
      4: #include <petsc/private/kokkosimpl.hpp>

      6: #if defined(PETSC_USE_DEBUG)
      7:   #define VecErrorIfNotKokkos(v) \
      8:     do { \
      9:       PetscBool isKokkos = PETSC_FALSE; \
     10:       PetscCall(PetscObjectTypeCompareAny((PetscObject)(v), &isKokkos, VECSEQKOKKOS, VECMPIKOKKOS, VECKOKKOS, "")); \
     11:       PetscCheck(isKokkos, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Calling VECKOKKOS methods on a non-VECKOKKOS object"); \
     12:     } while (0)
     13: #else
     14:   #define VecErrorIfNotKokkos(v) \
     15:     do { \
     16:       (void)(v); \
     17:     } while (0)
     18: #endif

     20: /* Stuff related to Vec_Kokkos */

     22: struct Vec_Kokkos {
     23:   PetscScalarKokkosDualView v_dual;

     25:   /* COO stuff */
     26:   PetscCountKokkosView jmap1_d; /* [m+1]: i-th entry of the vector has jmap1[i+1]-jmap1[i] repeats in COO arrays */
     27:   PetscCountKokkosView perm1_d; /* [tot1]: permutation array for local entries */

     29:   PetscCountKokkosView  imap2_d;              /* [nnz2]: i-th unique entry in recvbuf is imap2[i]-th entry in the vector */
     30:   PetscCountKokkosView  jmap2_d;              /* [nnz2+1] */
     31:   PetscCountKokkosView  perm2_d;              /* [recvlen] */
     32:   PetscCountKokkosView  Cperm_d;              /* [sendlen]: permutation array to fill sendbuf[]. 'C' for communication */
     33:   PetscScalarKokkosView sendbuf_d, recvbuf_d; /* Buffers for remote values in VecSetValuesCOO() */

     35:   // (internal use only) stash the pointer when we allocate a raw device array for the vector's use and that should be freed when the vector is destroyed
     36:   PetscScalar *raw_array_d_allocated;

     38:   /* Construct Vec_Kokkos with the given array(s). n is the length of the array.
     39:     If n != 0, host array (array_h) must not be NULL.
     40:     If device array (array_d) is NULL, then a proper device mirror will be allocated.
     41:     Otherwise, the mirror will be created using the given array_d.
     42:     If both arrays are given, we assume they contain the same value (i.e., sync'ed)
     43:   */
     44:   Vec_Kokkos(PetscInt n, PetscScalar *array_h, PetscScalar *array_d = NULL) : raw_array_d_allocated(nullptr)
     45:   {
     46:     PetscScalarKokkosViewHost v_h(array_h, n);
     47:     PetscScalarKokkosView     v_d;

     49:     if (array_d) {
     50:       v_d = PetscScalarKokkosView(array_d, n); /* Use the given device array */
     51:     } else {
     52:       v_d = Kokkos::create_mirror_view(Kokkos::WithoutInitializing, DefaultMemorySpace(), v_h); /* Create a mirror in DefaultMemorySpace but do not copy values */
     53:     }
     54:     v_dual = PetscScalarKokkosDualView(v_d, v_h);
     55:     if (!array_d) v_dual.modify_host();
     56:   }

     58:   ~Vec_Kokkos()
     59:   {
     60:     if (raw_array_d_allocated) Kokkos::kokkos_free(raw_array_d_allocated);
     61:   }

     63:   /* SFINAE: Update the object with an array in the given memory space,
     64:      assuming the given array contains the latest value for this vector.
     65:    */
     66:   template <typename MemorySpace, std::enable_if_t<std::is_same<MemorySpace, Kokkos::HostSpace>::value, bool> = true, std::enable_if_t<std::is_same<MemorySpace, DefaultMemorySpace>::value, bool> = true>
     67:   PetscErrorCode UpdateArray(PetscScalar *array)
     68:   {
     69:     PetscScalarKokkosViewHost v_h(array, v_dual.extent(0));

     71:     PetscFunctionBegin;
     72:     /* Kokkos said they would add error-checking so that users won't accidentally pass two different Views in this case */
     73:     PetscCallCXX(v_dual = PetscScalarKokkosDualView(v_h, v_h));
     74:     PetscFunctionReturn(PETSC_SUCCESS);
     75:   }

     77:   template <typename MemorySpace, std::enable_if_t<std::is_same<MemorySpace, Kokkos::HostSpace>::value, bool> = true, std::enable_if_t<!std::is_same<MemorySpace, DefaultMemorySpace>::value, bool> = true>
     78:   PetscErrorCode UpdateArray(PetscScalar *array)
     79:   {
     80:     PetscScalarKokkosViewHost v_h(array, v_dual.extent(0));

     82:     PetscFunctionBegin;
     83:     PetscCallCXX(v_dual = PetscScalarKokkosDualView(v_dual.view<DefaultMemorySpace>(), v_h));
     84:     PetscCallCXX(v_dual.modify_host());
     85:     PetscFunctionReturn(PETSC_SUCCESS);
     86:   }

     88:   template <typename MemorySpace, std::enable_if_t<!std::is_same<MemorySpace, Kokkos::HostSpace>::value, bool> = true, std::enable_if_t<std::is_same<MemorySpace, DefaultMemorySpace>::value, bool> = true>
     89:   PetscErrorCode UpdateArray(PetscScalar *array)
     90:   {
     91:     PetscScalarKokkosView v_d(array, v_dual.extent(0));

     93:     PetscFunctionBegin;
     94:     PetscCallCXX(v_dual = PetscScalarKokkosDualView(v_d, v_dual.view<Kokkos::HostSpace>()));
     95:     PetscCallCXX(v_dual.modify_device());
     96:     PetscFunctionReturn(PETSC_SUCCESS);
     97:   }

     99:   PetscErrorCode SetUpCOO(const Vec_Seq *vecseq, PetscInt m)
    100:   {
    101:     PetscFunctionBegin;
    102:     PetscCallCXX(jmap1_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecseq->jmap1, m + 1)));
    103:     PetscCallCXX(perm1_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecseq->perm1, vecseq->tot1)));
    104:     PetscFunctionReturn(PETSC_SUCCESS);
    105:   }

    107:   PetscErrorCode SetUpCOO(const Vec_MPI *vecmpi, PetscInt m)
    108:   {
    109:     PetscFunctionBegin;
    110:     PetscCallCXX(jmap1_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->jmap1, m + 1)));
    111:     PetscCallCXX(perm1_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->perm1, vecmpi->tot1)));
    112:     PetscCallCXX(imap2_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->imap2, vecmpi->nnz2)));
    113:     PetscCallCXX(jmap2_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->jmap2, vecmpi->nnz2 + 1)));
    114:     PetscCallCXX(perm2_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->perm2, vecmpi->recvlen)));
    115:     PetscCallCXX(Cperm_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->Cperm, vecmpi->sendlen)));
    116:     PetscCallCXX(sendbuf_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscScalarKokkosViewHost(vecmpi->sendbuf, vecmpi->sendlen)));
    117:     PetscCallCXX(recvbuf_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscScalarKokkosViewHost(vecmpi->recvbuf, vecmpi->recvlen)));
    118:     PetscFunctionReturn(PETSC_SUCCESS);
    119:   }
    120: };

    122: PETSC_INTERN PetscErrorCode VecAbs_SeqKokkos(Vec);
    123: PETSC_INTERN PetscErrorCode VecReciprocal_SeqKokkos(Vec);
    124: PETSC_INTERN PetscErrorCode VecDotNorm2_SeqKokkos(Vec, Vec, PetscScalar *, PetscScalar *);
    125: PETSC_INTERN PetscErrorCode VecPointwiseDivide_SeqKokkos(Vec, Vec, Vec);
    126: PETSC_INTERN PetscErrorCode VecWAXPY_SeqKokkos(Vec, PetscScalar, Vec, Vec);
    127: PETSC_INTERN PetscErrorCode VecMDot_SeqKokkos(Vec, PetscInt, const Vec[], PetscScalar *);
    128: PETSC_INTERN PetscErrorCode VecMTDot_SeqKokkos(Vec, PetscInt, const Vec[], PetscScalar *);
    129: PETSC_INTERN PetscErrorCode VecSet_SeqKokkos(Vec, PetscScalar);
    130: PETSC_INTERN PetscErrorCode VecMAXPY_SeqKokkos(Vec, PetscInt, const PetscScalar *, Vec *);
    131: PETSC_INTERN PetscErrorCode VecAXPBYPCZ_SeqKokkos(Vec, PetscScalar, PetscScalar, PetscScalar, Vec, Vec);
    132: PETSC_INTERN PetscErrorCode VecPointwiseMult_SeqKokkos(Vec, Vec, Vec);
    133: PETSC_INTERN PetscErrorCode VecPlaceArray_SeqKokkos(Vec, const PetscScalar *);
    134: PETSC_INTERN PetscErrorCode VecResetArray_SeqKokkos(Vec);
    135: PETSC_INTERN PetscErrorCode VecReplaceArray_SeqKokkos(Vec, const PetscScalar *);
    136: PETSC_INTERN PetscErrorCode VecDot_SeqKokkos(Vec, Vec, PetscScalar *);
    137: PETSC_INTERN PetscErrorCode VecTDot_SeqKokkos(Vec, Vec, PetscScalar *);
    138: PETSC_INTERN PetscErrorCode VecScale_SeqKokkos(Vec, PetscScalar);
    139: PETSC_INTERN PetscErrorCode VecCopy_SeqKokkos(Vec, Vec);
    140: PETSC_INTERN PetscErrorCode VecSwap_SeqKokkos(Vec, Vec);
    141: PETSC_INTERN PetscErrorCode VecAXPY_SeqKokkos(Vec, PetscScalar, Vec);
    142: PETSC_INTERN PetscErrorCode VecAXPBY_SeqKokkos(Vec, PetscScalar, PetscScalar, Vec);
    143: PETSC_INTERN PetscErrorCode VecConjugate_SeqKokkos(Vec xin);
    144: PETSC_INTERN PetscErrorCode VecNorm_SeqKokkos(Vec, NormType, PetscReal *);
    145: PETSC_INTERN PetscErrorCode VecErrorWeightedNorms_SeqKokkos(Vec, Vec, Vec, NormType, PetscReal, Vec, PetscReal, Vec, PetscReal, PetscReal *, PetscInt *, PetscReal *, PetscInt *, PetscReal *, PetscInt *);
    146: PETSC_INTERN PetscErrorCode VecCreate_SeqKokkos(Vec);
    147: PETSC_INTERN PetscErrorCode VecCreate_SeqKokkos_Private(Vec, const PetscScalar *);
    148: PETSC_INTERN PetscErrorCode VecCreate_MPIKokkos(Vec);
    149: PETSC_INTERN PetscErrorCode VecCreate_MPIKokkos_Private(Vec, PetscBool, PetscInt, const PetscScalar *);
    150: PETSC_INTERN PetscErrorCode VecCreate_Kokkos(Vec);
    151: PETSC_INTERN PetscErrorCode VecAYPX_SeqKokkos(Vec, PetscScalar, Vec);
    152: PETSC_INTERN PetscErrorCode VecSetRandom_SeqKokkos(Vec, PetscRandom);
    153: PETSC_INTERN PetscErrorCode VecGetLocalVector_SeqKokkos(Vec, Vec);
    154: PETSC_INTERN PetscErrorCode VecRestoreLocalVector_SeqKokkos(Vec, Vec);
    155: PETSC_INTERN PetscErrorCode VecGetArrayWrite_SeqKokkos(Vec, PetscScalar **);
    156: PETSC_INTERN PetscErrorCode VecCopy_SeqKokkos_Private(Vec, Vec);
    157: PETSC_INTERN PetscErrorCode VecSetRandom_SeqKokkos_Private(Vec, PetscRandom);
    158: PETSC_INTERN PetscErrorCode VecResetArray_SeqKokkos_Private(Vec);
    159: PETSC_INTERN PetscErrorCode VecMin_SeqKokkos(Vec, PetscInt *, PetscReal *);
    160: PETSC_INTERN PetscErrorCode VecMax_SeqKokkos(Vec, PetscInt *, PetscReal *);
    161: PETSC_INTERN PetscErrorCode VecSum_SeqKokkos(Vec, PetscScalar *);
    162: PETSC_INTERN PetscErrorCode VecShift_SeqKokkos(Vec, PetscScalar);
    163: PETSC_INTERN PetscErrorCode VecGetArray_SeqKokkos(Vec, PetscScalar **);
    164: PETSC_INTERN PetscErrorCode VecRestoreArray_SeqKokkos(Vec, PetscScalar **);

    166: PETSC_INTERN PetscErrorCode VecGetArrayAndMemType_SeqKokkos(Vec, PetscScalar **, PetscMemType *);
    167: PETSC_INTERN PetscErrorCode VecRestoreArrayAndMemType_SeqKokkos(Vec, PetscScalar **);
    168: PETSC_INTERN PetscErrorCode VecGetArrayWriteAndMemType_SeqKokkos(Vec, PetscScalar **, PetscMemType *);
    169: PETSC_INTERN PetscErrorCode VecGetSubVector_Kokkos_Private(Vec, PetscBool, IS, Vec *);
    170: PETSC_INTERN PetscErrorCode VecRestoreSubVector_SeqKokkos(Vec, IS, Vec *);

    172: PETSC_INTERN PetscErrorCode VecDuplicateVecs_Kokkos_GEMV(Vec, PetscInt, Vec **);
    173: PETSC_INTERN PetscErrorCode VecMDot_SeqKokkos_GEMV(Vec, PetscInt, const Vec *, PetscScalar *);
    174: PETSC_INTERN PetscErrorCode VecMTDot_SeqKokkos_GEMV(Vec, PetscInt, const Vec *, PetscScalar *);
    175: PETSC_INTERN PetscErrorCode VecMAXPY_SeqKokkos_GEMV(Vec, PetscInt, const PetscScalar *, Vec *);

    177: PETSC_INTERN PetscErrorCode VecCreateMPIKokkosWithLayoutAndArrays_Private(PetscLayout map, const PetscScalar *, const PetscScalar *, Vec *);

    179: PETSC_INTERN PetscErrorCode VecSetOps_MPIKokkos(Vec);