Actual source code: curand.c

  1: #include <petsc/private/deviceimpl.h>
  2: #include <petsc/private/randomimpl.h>
  3: #include <petscdevice_cuda.h>
  4: #include <curand.h>

  6: typedef struct {
  7:   curandGenerator_t gen;
  8: } PetscRandom_CURAND;

 10: PetscErrorCode PetscRandomSeed_CURAND(PetscRandom r)
 11: {
 12:   PetscRandom_CURAND *curand = (PetscRandom_CURAND *)r->data;

 14:   curandSetPseudoRandomGeneratorSeed(curand->gen, r->seed);
 15:   return 0;
 16: }

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

 20: PetscErrorCode PetscRandomGetValuesReal_CURAND(PetscRandom r, PetscInt n, PetscReal *val)
 21: {
 22:   PetscRandom_CURAND *curand = (PetscRandom_CURAND *)r->data;
 23:   size_t              nn     = n < 0 ? (size_t)(-2 * n) : n; /* handle complex case */

 25: #if defined(PETSC_USE_REAL_SINGLE)
 26:   curandGenerateUniform(curand->gen, val, nn);
 27: #else
 28:   curandGenerateUniformDouble(curand->gen, val, nn);
 29: #endif
 30:   if (r->iset) PetscRandomCurandScale_Private(r, nn, val, (PetscBool)(n < 0));
 31:   return 0;
 32: }

 34: PetscErrorCode PetscRandomGetValues_CURAND(PetscRandom r, PetscInt n, PetscScalar *val)
 35: {
 36: #if defined(PETSC_USE_COMPLEX)
 37:   /* pass negative size to flag complex scaling (if needed) */
 38:   PetscRandomGetValuesReal_CURAND(r, -n, (PetscReal *)val);
 39: #else
 40:   PetscRandomGetValuesReal_CURAND(r, n, val);
 41: #endif
 42:   return 0;
 43: }

 45: PetscErrorCode PetscRandomDestroy_CURAND(PetscRandom r)
 46: {
 47:   PetscRandom_CURAND *curand = (PetscRandom_CURAND *)r->data;

 49:   curandDestroyGenerator(curand->gen);
 50:   PetscFree(r->data);
 51:   return 0;
 52: }

 54: static struct _PetscRandomOps PetscRandomOps_Values = {
 55:   PetscDesignatedInitializer(seed, PetscRandomSeed_CURAND),
 56:   PetscDesignatedInitializer(getvalue, NULL),
 57:   PetscDesignatedInitializer(getvaluereal, NULL),
 58:   PetscDesignatedInitializer(getvalues, PetscRandomGetValues_CURAND),
 59:   PetscDesignatedInitializer(getvaluesreal, PetscRandomGetValuesReal_CURAND),
 60:   PetscDesignatedInitializer(destroy, PetscRandomDestroy_CURAND),
 61: };

 63: /*MC
 64:    PETSCCURAND - access to the CUDA random number generator from a `PetscRandom` object

 66:   PETSc must be ./configure with the option --with-cuda to use this random number generator.

 68:   Level: beginner

 70: .seealso: `PetscRandomCreate()`, `PetscRandomSetType()`, `PetscRandomType`
 71: M*/

 73: PETSC_EXTERN PetscErrorCode PetscRandomCreate_CURAND(PetscRandom r)
 74: {
 75:   PetscRandom_CURAND *curand;

 77:   PetscDeviceInitialize(PETSC_DEVICE_CUDA);
 78:   PetscNew(&curand);
 79:   curandCreateGenerator(&curand->gen, CURAND_RNG_PSEUDO_DEFAULT);
 80:   /* https://docs.nvidia.com/cuda/curand/host-api-overview.html#performance-notes2 */
 81:   curandSetGeneratorOrdering(curand->gen, CURAND_ORDERING_PSEUDO_SEEDED);
 82:   PetscMemcpy(r->ops, &PetscRandomOps_Values, sizeof(PetscRandomOps_Values));
 83:   PetscObjectChangeTypeName((PetscObject)r, PETSCCURAND);
 84:   r->data = curand;
 85:   PetscRandomSeed_CURAND(r);
 86:   return 0;
 87: }