xref: /petsc/src/sys/classes/random/impls/curand/curand.c (revision f6b722a57d07b8c3a218465fb7fc5d7800a7778a)
1 #include <../src/sys/classes/random/randomimpl.h>
2 #include <curand.h>
3 
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)
11 
12 typedef struct {
13   curandGenerator_t gen;
14 } PetscRandom_CURAND;
15 
16 PetscErrorCode PetscRandomSeed_CURAND(PetscRandom r)
17 {
18   curandStatus_t     cerr;
19   PetscRandom_CURAND *curand = (PetscRandom_CURAND*)r->data;
20 
21   PetscFunctionBegin;
22   cerr = curandSetPseudoRandomGeneratorSeed(curand->gen,r->seed);CHKERRCURAND(cerr);
23   PetscFunctionReturn(0);
24 }
25 
26 PETSC_INTERN PetscErrorCode PetscRandomCurandScale_Private(PetscRandom,size_t,PetscReal*,PetscBool);
27 
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 ? -2*(size_t)n : n; /* handle complex case */
33 
34   PetscFunctionBegin;
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 ierr = PetscRandomCurandScale_Private(r,nn,val,(PetscBool)(n<0));CHKERRQ(ierr);
42   }
43   PetscFunctionReturn(0);
44 }
45 
46 PetscErrorCode PetscRandomGetValues_CURAND(PetscRandom r, PetscInt n, PetscScalar *val)
47 {
48   PetscErrorCode ierr;
49 
50   PetscFunctionBegin;
51 #if defined(PETSC_USE_COMPLEX)
52   /* pass negative size to flag complex scaling (if needed) */
53   ierr = PetscRandomGetValuesReal_CURAND(r,-n,(PetscReal*)val);CHKERRQ(ierr);
54 #else
55   ierr = PetscRandomGetValuesReal_CURAND(r,n,val);CHKERRQ(ierr);
56 #endif
57   PetscFunctionReturn(0);
58 }
59 
60 PetscErrorCode PetscRandomDestroy_CURAND(PetscRandom r)
61 {
62   PetscErrorCode     ierr;
63   curandStatus_t     cerr;
64   PetscRandom_CURAND *curand = (PetscRandom_CURAND*)r->data;
65 
66   PetscFunctionBegin;
67   cerr = curandDestroyGenerator(curand->gen);CHKERRCURAND(cerr);
68   ierr = PetscFree(r->data);CHKERRQ(ierr);
69   PetscFunctionReturn(0);
70 }
71 
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 };
81 
82 /*MC
83    PETSCCURAND - access to the CUDA random number generator
84 
85   Level: beginner
86 
87 .seealso: PetscRandomCreate(), PetscRandomSetType()
88 M*/
89 
90 PETSC_EXTERN PetscErrorCode PetscRandomCreate_CURAND(PetscRandom r)
91 {
92   PetscErrorCode     ierr;
93   curandStatus_t     cerr;
94   PetscRandom_CURAND *curand;
95 
96   PetscFunctionBegin;
97   ierr = PetscCUDAInitializeCheck();CHKERRQ(ierr);
98   ierr = PetscNewLog(r,&curand);CHKERRQ(ierr);
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   ierr = PetscMemcpy(r->ops,&PetscRandomOps_Values,sizeof(PetscRandomOps_Values));CHKERRQ(ierr);
103   ierr = PetscObjectChangeTypeName((PetscObject)r,PETSCCURAND);CHKERRQ(ierr);
104   r->data = curand;
105   r->seed = 1234ULL; /* taken from example */
106   ierr = PetscRandomSeed_CURAND(r);CHKERRQ(ierr);
107   PetscFunctionReturn(0);
108 }
109