Actual source code: curand.c

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

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

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

     13:   PetscFunctionBegin;
     14:   PetscCallCURAND(curandSetPseudoRandomGeneratorSeed(curand->gen, r->seed));
     15:   PetscFunctionReturn(PETSC_SUCCESS);
     16: }

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

     20: static 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) : (size_t)n; /* handle complex case */

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

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

     47: static PetscErrorCode PetscRandomDestroy_CURAND(PetscRandom r)
     48: {
     49:   PetscRandom_CURAND *curand = (PetscRandom_CURAND *)r->data;

     51:   PetscFunctionBegin;
     52:   PetscCallCURAND(curandDestroyGenerator(curand->gen));
     53:   PetscCall(PetscFree(r->data));
     54:   PetscFunctionReturn(PETSC_SUCCESS);
     55: }

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

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

     69:   Level: beginner

     71:   Note:
     72:   This random number generator is available when PETSc is configured with ``./configure --with-cuda=1``

     74: .seealso: `PetscRandomCreate()`, `PetscRandomSetType()`, `PetscRandomType`
     75: M*/

     77: PETSC_EXTERN PetscErrorCode PetscRandomCreate_CURAND(PetscRandom r)
     78: {
     79:   PetscRandom_CURAND *curand;
     80:   PetscDeviceContext  dctx;
     81:   cudaStream_t       *stream;

     83:   PetscFunctionBegin;
     84:   PetscCall(PetscDeviceInitialize(PETSC_DEVICE_CUDA));
     85:   PetscCall(PetscDeviceContextGetCurrentContextAssertType_Internal(&dctx, PETSC_DEVICE_CUDA));
     86:   PetscCall(PetscDeviceContextGetStreamHandle(dctx, (void **)&stream));
     87:   PetscCall(PetscNew(&curand));
     88:   PetscCallCURAND(curandCreateGenerator(&curand->gen, CURAND_RNG_PSEUDO_DEFAULT));
     89:   PetscCallCURAND(curandSetStream(curand->gen, *stream));
     90:   /* https://docs.nvidia.com/cuda/curand/host-api-overview.html#performance-notes2 */
     91:   PetscCallCURAND(curandSetGeneratorOrdering(curand->gen, CURAND_ORDERING_PSEUDO_SEEDED));
     92:   r->ops[0] = PetscRandomOps_Values;
     93:   PetscCall(PetscObjectChangeTypeName((PetscObject)r, PETSCCURAND));
     94:   r->data = curand;
     95:   PetscCall(PetscRandomSeed_CURAND(r));
     96:   PetscFunctionReturn(PETSC_SUCCESS);
     97: }