1 #include <petsc/private/randomimpl.h> 2 #include <thrust/transform.h> 3 #include <thrust/device_ptr.h> 4 #include <thrust/iterator/counting_iterator.h> 5 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; 11 12 complexscalelw(PetscScalar low, PetscScalar width) { 13 rl = PetscRealPart(low); 14 il = PetscImaginaryPart(low); 15 rw = PetscRealPart(width); 16 iw = PetscImaginaryPart(width); 17 } 18 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 25 26 struct realscalelw : public thrust::unary_function<PetscReal,PetscReal> 27 { 28 PetscReal l,w; 29 30 realscalelw(PetscReal low, PetscReal width) : l(low), w(width) {} 31 32 __host__ __device__ 33 PetscReal operator()(PetscReal x) { 34 return x*w + l; 35 } 36 }; 37 38 PETSC_INTERN PetscErrorCode PetscRandomCurandScale_Private(PetscRandom r, size_t n, PetscReal *val, PetscBool isneg) 39 { 40 PetscFunctionBegin; 41 if (!r->iset) PetscFunctionReturn(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 PetscFunctionReturn(0); 57 } 58