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, cudaStream_t);

 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 */
 24:   PetscDeviceContext  dctx;
 25:   cudaStream_t       *stream;

 27:   PetscFunctionBegin;
 28:   PetscCall(PetscDeviceContextGetCurrentContextAssertType_Internal(&dctx, PETSC_DEVICE_CUDA));
 29:   PetscCall(PetscDeviceContextGetStreamHandle(dctx, (void **)&stream));
 30:   PetscCallCURAND(curandSetStream(curand->gen, *stream));
 31: #if PetscDefined(USE_REAL_SINGLE)
 32:   PetscCallCURAND(curandGenerateUniform(curand->gen, val, nn));
 33: #else
 34:   PetscCallCURAND(curandGenerateUniformDouble(curand->gen, val, nn));
 35: #endif
 36:   if (r->iset) PetscCall(PetscRandomCurandScale_Private(r, nn, val, (PetscBool)(n < 0), *stream));
 37:   PetscCall(PetscDeviceContextSynchronizeIfWithBarrier_Internal(dctx));
 38:   PetscFunctionReturn(PETSC_SUCCESS);
 39: }

 41: static PetscErrorCode PetscRandomGetValues_CURAND(PetscRandom r, PetscInt n, PetscScalar *val)
 42: {
 43:   PetscFunctionBegin;
 44: #if PetscDefined(USE_COMPLEX)
 45:   /* pass negative size to flag complex scaling (if needed) */
 46:   PetscCall(PetscRandomGetValuesReal_CURAND(r, -n, (PetscReal *)val));
 47: #else
 48:   PetscCall(PetscRandomGetValuesReal_CURAND(r, n, val));
 49: #endif
 50:   PetscFunctionReturn(PETSC_SUCCESS);
 51: }

 53: static PetscErrorCode PetscRandomDestroy_CURAND(PetscRandom r)
 54: {
 55:   PetscRandom_CURAND *curand = (PetscRandom_CURAND *)r->data;

 57:   PetscFunctionBegin;
 58:   PetscCallCURAND(curandDestroyGenerator(curand->gen));
 59:   PetscCall(PetscFree(r->data));
 60:   PetscFunctionReturn(PETSC_SUCCESS);
 61: }

 63: static struct _PetscRandomOps PetscRandomOps_Values = {
 64:   PetscDesignatedInitializer(seed, PetscRandomSeed_CURAND),
 65:   PetscDesignatedInitializer(getvalue, NULL),
 66:   PetscDesignatedInitializer(getvaluereal, NULL),
 67:   PetscDesignatedInitializer(getvalues, PetscRandomGetValues_CURAND),
 68:   PetscDesignatedInitializer(getvaluesreal, PetscRandomGetValuesReal_CURAND),
 69:   PetscDesignatedInitializer(destroy, PetscRandomDestroy_CURAND),
 70: };

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

 75:   Level: beginner

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

 80: .seealso: `PetscRandomCreate()`, `PetscRandomSetType()`, `PetscRandomType`
 81: M*/

 83: PETSC_EXTERN PetscErrorCode PetscRandomCreate_CURAND(PetscRandom r)
 84: {
 85:   PetscRandom_CURAND *curand;
 86:   PetscDeviceContext  dctx;
 87:   cudaStream_t       *stream;

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