Actual source code: curand2.cu
1: #include <../src/sys/classes/random/randomimpl.h>
2: #include <thrust/transform.h>
3: #include <thrust/device_ptr.h>
4: #include <thrust/iterator/counting_iterator.h>
6: #if defined(PETSC_USE_COMPLEX)
7: struct complexscalelw : public thrust::unary_function<thrust::tuple<PetscReal, size_t>,PetscReal>
8: {
9: PetscReal rl,rw;
10: PetscReal il,iw;
12: complexscalelw(PetscScalar low, PetscScalar width) {
13: rl = PetscRealPart(low);
14: il = PetscImaginaryPart(low);
15: rw = PetscRealPart(width);
16: iw = PetscImaginaryPart(width);
17: }
19: __host__ __device__
20: PetscReal operator()(thrust::tuple<PetscReal, size_t> x) {
21: return x.get<1>()%2 ? x.get<0>()*iw + il : x.get<0>()*rw + rl;
22: }
23: };
24: #endif
26: struct realscalelw : public thrust::unary_function<PetscReal,PetscReal>
27: {
28: PetscReal l,w;
30: realscalelw(PetscReal low, PetscReal width) : l(low), w(width) {}
32: __host__ __device__
33: PetscReal operator()(PetscReal x) {
34: return x*w + l;
35: }
36: };
38: PETSC_INTERN PetscErrorCode PetscRandomCurandScale_Private(PetscRandom r, size_t n, PetscReal *val, PetscBool isneg)
39: {
41: if (!r->iset) return(0);
42: if (isneg) { /* complex case, need to scale differently */
43: #if defined(PETSC_USE_COMPLEX)
44: thrust::device_ptr<PetscReal> pval = thrust::device_pointer_cast(val);
45: auto zibit = thrust::make_zip_iterator(thrust::make_tuple(pval,thrust::counting_iterator<size_t>(0)));
46: thrust::transform(zibit,zibit+n,pval,complexscalelw(r->low,r->width));
47: #else
48: SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Negative array size %D",(PetscInt)n);
49: #endif
50: } else {
51: PetscReal rl = PetscRealPart(r->low);
52: PetscReal rw = PetscRealPart(r->width);
53: thrust::device_ptr<PetscReal> pval = thrust::device_pointer_cast(val);
54: thrust::transform(pval,pval+n,pval,realscalelw(rl,rw));
55: }
56: return(0);
57: }