1: #include <petsc/private/randomimpl.h>
  2: #include <curand.h>

  4: #define CHKERRCURAND(stat) \
  5: do { \
  6:    if (PetscUnlikely(stat != CURAND_STATUS_SUCCESS)) { \
  7:      if (((stat == CURAND_STATUS_INITIALIZATION_FAILED) || (stat == CURAND_STATUS_ALLOCATION_FAILED)) && PetscCUDAInitialized) SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_GPU_RESOURCE,"cuRAND error %d. Reports not initialized or alloc failed; this indicates the GPU has run out resources",(int)stat); \
  8:      else SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_GPU,"cuRand error %d",(int)stat); \
  9:    } \
 10: } while (0)

 12: typedef struct {
 13:   curandGenerator_t gen;
 14: } PetscRandom_CURAND;

 16: PetscErrorCode PetscRandomSeed_CURAND(PetscRandom r)
 17: {
 18:   curandStatus_t     cerr;
 19:   PetscRandom_CURAND *curand = (PetscRandom_CURAND*)r->data;

 22:   cerr = curandSetPseudoRandomGeneratorSeed(curand->gen,r->seed);CHKERRCURAND(cerr);
 23:   return(0);
 24: }

 26: PETSC_INTERN PetscErrorCode PetscRandomCurandScale_Private(PetscRandom,size_t,PetscReal*,PetscBool);

 28: PetscErrorCode  PetscRandomGetValuesReal_CURAND(PetscRandom r, PetscInt n, PetscReal *val)
 29: {
 30:   curandStatus_t     cerr;
 31:   PetscRandom_CURAND *curand = (PetscRandom_CURAND*)r->data;
 32:   size_t             nn = n < 0 ? (size_t)(-2*n) : n; /* handle complex case */

 35: #if defined(PETSC_USE_REAL_SINGLE)
 36:   cerr = curandGenerateUniform(curand->gen,val,nn);CHKERRCURAND(cerr);
 37: #else
 38:   cerr = curandGenerateUniformDouble(curand->gen,val,nn);CHKERRCURAND(cerr);
 39: #endif
 40:   if (r->iset) {
 41:     PetscErrorCode PetscRandomCurandScale_Private(r,nn,val,(PetscBool)(n<0));
 42:   }
 43:   return(0);
 44: }

 46: PetscErrorCode PetscRandomGetValues_CURAND(PetscRandom r, PetscInt n, PetscScalar *val)
 47: {

 51: #if defined(PETSC_USE_COMPLEX)
 52:   /* pass negative size to flag complex scaling (if needed) */
 53:   PetscRandomGetValuesReal_CURAND(r,-n,(PetscReal*)val);
 54: #else
 55:   PetscRandomGetValuesReal_CURAND(r,n,val);
 56: #endif
 57:   return(0);
 58: }

 60: PetscErrorCode PetscRandomDestroy_CURAND(PetscRandom r)
 61: {
 62:   PetscErrorCode     ierr;
 63:   curandStatus_t     cerr;
 64:   PetscRandom_CURAND *curand = (PetscRandom_CURAND*)r->data;

 67:   cerr = curandDestroyGenerator(curand->gen);CHKERRCURAND(cerr);
 68:   PetscFree(r->data);
 69:   return(0);
 70: }

 72: static struct _PetscRandomOps PetscRandomOps_Values = {
 73:   PetscRandomSeed_CURAND,
 74:   NULL,
 75:   NULL,
 76:   PetscRandomGetValues_CURAND,
 77:   PetscRandomGetValuesReal_CURAND,
 78:   PetscRandomDestroy_CURAND,
 79:   NULL
 80: };

 82: /*MC
 83:    PETSCCURAND - access to the CUDA random number generator

 85:   Level: beginner

 87: .seealso: PetscRandomCreate(), PetscRandomSetType()
 88: M*/

 90: PETSC_EXTERN PetscErrorCode PetscRandomCreate_CURAND(PetscRandom r)
 91: {
 92:   PetscErrorCode     ierr;
 93:   curandStatus_t     cerr;
 94:   PetscRandom_CURAND *curand;

 97:   PetscCUDAInitializeCheck();
 98:   PetscNewLog(r,&curand);
 99:   cerr = curandCreateGenerator(&curand->gen,CURAND_RNG_PSEUDO_DEFAULT);CHKERRCURAND(cerr);
100:   /* https://docs.nvidia.com/cuda/curand/host-api-overview.html#performance-notes2 */
101:   cerr = curandSetGeneratorOrdering(curand->gen,CURAND_ORDERING_PSEUDO_SEEDED);CHKERRCURAND(cerr);
102:   PetscMemcpy(r->ops,&PetscRandomOps_Values,sizeof(PetscRandomOps_Values));
103:   PetscObjectChangeTypeName((PetscObject)r,PETSCCURAND);
104:   r->data = curand;
105:   r->seed = 1234ULL; /* taken from example */
106:   PetscRandomSeed_CURAND(r);
107:   return(0);
108: }