Actual source code: cudavecimpl.h

      4: #include <petscvec.h>
      5: #include <petsccublas.h>
      6: #include <petsc/private/vecimpl.h>

      8: typedef struct {
      9:   PetscScalar  *GPUarray;           /* this always holds the GPU data */
     10:   PetscScalar  *GPUarray_allocated; /* if the array was allocated by PETSc this is its pointer */
     11:   cudaStream_t stream;              /* A stream for doing asynchronous data transfers */
     12:   PetscBool    nvshmem;             /* Is GPUarray_allocated allocated in nvshmem? It is used to allocate Mvctx->lvec in nvshmem */
     13: } Vec_CUDA;

     15: PETSC_INTERN PetscErrorCode VecCUDAGetArrays_Private(Vec,const PetscScalar**,const PetscScalar**,PetscOffloadMask*);
     16: PETSC_INTERN PetscErrorCode VecDotNorm2_SeqCUDA(Vec,Vec,PetscScalar*, PetscScalar*);
     17: PETSC_INTERN PetscErrorCode VecPointwiseDivide_SeqCUDA(Vec,Vec,Vec);
     18: PETSC_INTERN PetscErrorCode VecWAXPY_SeqCUDA(Vec,PetscScalar,Vec,Vec);
     19: PETSC_INTERN PetscErrorCode VecMDot_SeqCUDA(Vec,PetscInt,const Vec[],PetscScalar*);
     20: PETSC_EXTERN PetscErrorCode VecSet_SeqCUDA(Vec,PetscScalar);
     21: PETSC_INTERN PetscErrorCode VecMAXPY_SeqCUDA(Vec,PetscInt,const PetscScalar*,Vec*);
     22: PETSC_INTERN PetscErrorCode VecAXPBYPCZ_SeqCUDA(Vec,PetscScalar,PetscScalar,PetscScalar,Vec,Vec);
     23: PETSC_INTERN PetscErrorCode VecPointwiseMult_SeqCUDA(Vec,Vec,Vec);
     24: PETSC_INTERN PetscErrorCode VecPlaceArray_SeqCUDA(Vec,const PetscScalar*);
     25: PETSC_INTERN PetscErrorCode VecResetArray_SeqCUDA(Vec);
     26: PETSC_INTERN PetscErrorCode VecReplaceArray_SeqCUDA(Vec,const PetscScalar*);
     27: PETSC_INTERN PetscErrorCode VecDot_SeqCUDA(Vec,Vec,PetscScalar*);
     28: PETSC_INTERN PetscErrorCode VecTDot_SeqCUDA(Vec,Vec,PetscScalar*);
     29: PETSC_INTERN PetscErrorCode VecScale_SeqCUDA(Vec,PetscScalar);
     30: PETSC_EXTERN PetscErrorCode VecCopy_SeqCUDA(Vec,Vec);
     31: PETSC_INTERN PetscErrorCode VecSwap_SeqCUDA(Vec,Vec);
     32: PETSC_EXTERN PetscErrorCode VecAXPY_SeqCUDA(Vec,PetscScalar,Vec);
     33: PETSC_INTERN PetscErrorCode VecAXPBY_SeqCUDA(Vec,PetscScalar,PetscScalar,Vec);
     34: PETSC_INTERN PetscErrorCode VecDuplicate_SeqCUDA(Vec,Vec*);
     35: PETSC_INTERN PetscErrorCode VecConjugate_SeqCUDA(Vec xin);
     36: PETSC_INTERN PetscErrorCode VecNorm_SeqCUDA(Vec,NormType,PetscReal*);
     37: PETSC_INTERN PetscErrorCode VecCUDACopyToGPU(Vec);
     38: PETSC_INTERN PetscErrorCode VecCUDAAllocateCheck(Vec);
     39: PETSC_EXTERN PetscErrorCode VecCreate_SeqCUDA(Vec);
     40: PETSC_INTERN PetscErrorCode VecCreate_SeqCUDA_Private(Vec,const PetscScalar*);
     41: PETSC_INTERN PetscErrorCode VecCreate_MPICUDA(Vec);
     42: PETSC_INTERN PetscErrorCode VecCreate_MPICUDA_Private(Vec,PetscBool,PetscInt,const PetscScalar*);
     43: PETSC_INTERN PetscErrorCode VecCreate_CUDA(Vec);
     44: PETSC_INTERN PetscErrorCode VecDestroy_SeqCUDA(Vec);
     45: PETSC_INTERN PetscErrorCode VecDestroy_MPICUDA(Vec);
     46: PETSC_INTERN PetscErrorCode VecAYPX_SeqCUDA(Vec,PetscScalar,Vec);
     47: PETSC_INTERN PetscErrorCode VecSetRandom_SeqCUDA(Vec,PetscRandom);
     48: PETSC_INTERN PetscErrorCode VecGetLocalVector_SeqCUDA(Vec,Vec);
     49: PETSC_INTERN PetscErrorCode VecRestoreLocalVector_SeqCUDA(Vec,Vec);
     50: PETSC_INTERN PetscErrorCode VecGetArrayWrite_SeqCUDA(Vec,PetscScalar**);
     51: PETSC_INTERN PetscErrorCode VecGetArray_SeqCUDA(Vec,PetscScalar**);
     52: PETSC_INTERN PetscErrorCode VecRestoreArray_SeqCUDA(Vec,PetscScalar**);
     53: PETSC_INTERN PetscErrorCode VecGetArrayAndMemType_SeqCUDA(Vec,PetscScalar**,PetscMemType*);
     54: PETSC_INTERN PetscErrorCode VecRestoreArrayAndMemType_SeqCUDA(Vec,PetscScalar**);
     55: PETSC_INTERN PetscErrorCode VecCopy_SeqCUDA_Private(Vec,Vec);
     56: PETSC_INTERN PetscErrorCode VecDestroy_SeqCUDA_Private(Vec);
     57: PETSC_INTERN PetscErrorCode VecResetArray_SeqCUDA_Private(Vec);
     58: PETSC_INTERN PetscErrorCode VecMax_SeqCUDA(Vec,PetscInt*,PetscReal*);
     59: PETSC_INTERN PetscErrorCode VecMin_SeqCUDA(Vec,PetscInt*,PetscReal*);

     61: #if defined(PETSC_HAVE_NVSHMEM)
     62: PETSC_EXTERN PetscErrorCode PetscNvshmemInitializeCheck(void);
     63: PETSC_EXTERN PetscErrorCode PetscNvshmemMalloc(size_t,void**);
     64: PETSC_EXTERN PetscErrorCode PetscNvshmemCalloc(size_t,void**);
     65: PETSC_EXTERN PetscErrorCode PetscNvshmemFree_Private(void*);
     66: #define      PetscNvshmemFree(ptr)      ((ptr) && (PetscNvshmemFree_Private(ptr),(ptr)=NULL,0))
     67: PETSC_INTERN PetscErrorCode PetscNvshmemSum(PetscInt,PetscScalar*,const PetscScalar*);
     68: PETSC_INTERN PetscErrorCode PetscNvshmemMax(PetscInt,PetscReal*,const PetscReal*);
     69: PETSC_INTERN PetscErrorCode VecNormAsync_NVSHMEM(Vec,NormType,PetscReal*);
     70: PETSC_INTERN PetscErrorCode VecAllocateNVSHMEM_SeqCUDA(Vec);
     71: #endif

     73: /* complex single */
     74: #if defined(PETSC_USE_COMPLEX)
     75: #if defined(PETSC_USE_REAL_SINGLE)
     76: #define cublasXaxpy(a,b,c,d,e,f,g)               cublasCaxpy((a),(b),(cuComplex*)(c),(cuComplex*)(d),(e),(cuComplex*)(f),(g))
     77: #define cublasXscal(a,b,c,d,e)                   cublasCscal((a),(b),(cuComplex*)(c),(cuComplex*)(d),(e))
     78: #define cublasXdotu(a,b,c,d,e,f,g)               cublasCdotu((a),(b),(cuComplex*)(c),(d),(cuComplex*)(e),(f),(cuComplex*)(g))
     79: #define cublasXdot(a,b,c,d,e,f,g)                cublasCdotc((a),(b),(cuComplex*)(c),(d),(cuComplex*)(e),(f),(cuComplex*)(g))
     80: #define cublasXswap(a,b,c,d,e,f)                 cublasCswap((a),(b),(cuComplex*)(c),(d),(cuComplex*)(e),(f))
     81: #define cublasXnrm2(a,b,c,d,e)                   cublasScnrm2((a),(b),(cuComplex*)(c),(d),(e))
     82: #define cublasIXamax(a,b,c,d,e)                  cublasIcamax((a),(b),(cuComplex*)(c),(d),(e))
     83: #define cublasXasum(a,b,c,d,e)                   cublasScasum((a),(b),(cuComplex*)(c),(d),(e))
     84: #define cublasXgemv(a,b,c,d,e,f,g,h,i,j,k,l)     cublasCgemv((a),(b),(c),(d),(cuComplex*)(e),(cuComplex*)(f),(g),(cuComplex*)(h),(i),(cuComplex*)(j),(cuComplex*)(k),(l))
     85: #define cublasXgemm(a,b,c,d,e,f,g,h,i,j,k,l,m,n) cublasCgemm((a),(b),(c),(d),(e),(f),(cuComplex*)(g),(cuComplex*)(h),(i),(cuComplex*)(j),(k),(cuComplex*)(l),(cuComplex*)(m),(n))
     86: #define cublasXgeam(a,b,c,d,e,f,g,h,i,j,k,l,m)   cublasCgeam((a),(b),(c),(d),(e),(cuComplex*)(f),(cuComplex*)(g),(h),(cuComplex*)(i),(cuComplex*)(j),(k),(cuComplex*)(l),(m))
     87: #else /* complex double */
     88: #define cublasXaxpy(a,b,c,d,e,f,g)               cublasZaxpy((a),(b),(cuDoubleComplex*)(c),(cuDoubleComplex*)(d),(e),(cuDoubleComplex*)(f),(g))
     89: #define cublasXscal(a,b,c,d,e)                   cublasZscal((a),(b),(cuDoubleComplex*)(c),(cuDoubleComplex*)(d),(e))
     90: #define cublasXdotu(a,b,c,d,e,f,g)               cublasZdotu((a),(b),(cuDoubleComplex*)(c),(d),(cuDoubleComplex*)(e),(f),(cuDoubleComplex*)(g))
     91: #define cublasXdot(a,b,c,d,e,f,g)                cublasZdotc((a),(b),(cuDoubleComplex*)(c),(d),(cuDoubleComplex*)(e),(f),(cuDoubleComplex*)(g))
     92: #define cublasXswap(a,b,c,d,e,f)                 cublasZswap((a),(b),(cuDoubleComplex*)(c),(d),(cuDoubleComplex*)(e),(f))
     93: #define cublasXnrm2(a,b,c,d,e)                   cublasDznrm2((a),(b),(cuDoubleComplex*)(c),(d),(e))
     94: #define cublasIXamax(a,b,c,d,e)                  cublasIzamax((a),(b),(cuDoubleComplex*)(c),(d),(e))
     95: #define cublasXasum(a,b,c,d,e)                   cublasDzasum((a),(b),(cuDoubleComplex*)(c),(d),(e))
     96: #define cublasXgemv(a,b,c,d,e,f,g,h,i,j,k,l)     cublasZgemv((a),(b),(c),(d),(cuDoubleComplex*)(e),(cuDoubleComplex*)(f),(g),(cuDoubleComplex*)(h),(i),(cuDoubleComplex*)(j),(cuDoubleComplex*)(k),(l))
     97: #define cublasXgemm(a,b,c,d,e,f,g,h,i,j,k,l,m,n) cublasZgemm((a),(b),(c),(d),(e),(f),(cuDoubleComplex*)(g),(cuDoubleComplex*)(h),(i),(cuDoubleComplex*)(j),(k),(cuDoubleComplex*)(l),(cuDoubleComplex*)(m),(n))
     98: #define cublasXgeam(a,b,c,d,e,f,g,h,i,j,k,l,m)   cublasZgeam((a),(b),(c),(d),(e),(cuDoubleComplex*)(f),(cuDoubleComplex*)(g),(h),(cuDoubleComplex*)(i),(cuDoubleComplex*)(j),(k),(cuDoubleComplex*)(l),(m))
     99: #endif
    100: #else /* real single */
    101: #if defined(PETSC_USE_REAL_SINGLE)
    102: #define cublasXaxpy  cublasSaxpy
    103: #define cublasXscal  cublasSscal
    104: #define cublasXdotu  cublasSdot
    105: #define cublasXdot   cublasSdot
    106: #define cublasXswap  cublasSswap
    107: #define cublasXnrm2  cublasSnrm2
    108: #define cublasIXamax cublasIsamax
    109: #define cublasXasum  cublasSasum
    110: #define cublasXgemv  cublasSgemv
    111: #define cublasXgemm  cublasSgemm
    112: #define cublasXgeam  cublasSgeam
    113: #else /* real double */
    114: #define cublasXaxpy  cublasDaxpy
    115: #define cublasXscal  cublasDscal
    116: #define cublasXdotu  cublasDdot
    117: #define cublasXdot   cublasDdot
    118: #define cublasXswap  cublasDswap
    119: #define cublasXnrm2  cublasDnrm2
    120: #define cublasIXamax cublasIdamax
    121: #define cublasXasum  cublasDasum
    122: #define cublasXgemv  cublasDgemv
    123: #define cublasXgemm  cublasDgemm
    124: #define cublasXgeam  cublasDgeam
    125: #endif
    126: #endif

    128: #endif