Date: 2022-08-02T03:13:20+00:00

Actual source code: curand2.cu

      1: #include <petsc/private/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: {
     40:   if (!r->iset) return 0;
     41:   if (isneg) { /* complex case, need to scale differently */
     42: #if defined(PETSC_USE_COMPLEX)
     43:     thrust::device_ptr<PetscReal> pval = thrust::device_pointer_cast(val);
     44:     auto zibit = thrust::make_zip_iterator(thrust::make_tuple(pval,thrust::counting_iterator<size_t>(0)));
     45:     thrust::transform(zibit,zibit+n,pval,complexscalelw(r->low,r->width));
     46: #else
     47:     SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Negative array size %" PetscInt_FMT,(PetscInt)n);
     48: #endif
     49:   } else {
     50:     PetscReal rl = PetscRealPart(r->low);
     51:     PetscReal rw = PetscRealPart(r->width);
     52:     thrust::device_ptr<PetscReal> pval = thrust::device_pointer_cast(val);
     53:     thrust::transform(pval,pval+n,pval,realscalelw(rl,rw));
     54:   }
     55:   return 0;
     56: }