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: }