Actual source code: ex47cu.cu

petsc-dev 2014-02-02
Report Typos and Errors
  1: static char help[] = "Solves -Laplacian u - exp(u) = 0,  0 < x < 1 using GPU\n\n";
  2: /*
  3:    Same as ex47.c except it also uses the GPU to evaluate the function
  4: */

  6: #include <petscdmda.h>
  7: #include <petscsnes.h>
  8: #include <petsccusp.h>

 10: #include <thrust/for_each.h>
 11: #include <thrust/tuple.h>
 12: #include <thrust/iterator/constant_iterator.h>
 13: #include <thrust/iterator/counting_iterator.h>
 14: #include <thrust/iterator/zip_iterator.h>

 16: extern PetscErrorCode ComputeFunction(SNES,Vec,Vec,void*), ComputeJacobian(SNES,Vec,Mat*,Mat*,MatStructure*,void*);
 17: PetscBool useCUSP = PETSC_FALSE;

 19: int main(int argc,char **argv)
 20: {
 21:   SNES           snes;
 22:   Vec            x,f;
 23:   Mat            J;
 24:   DM             da;
 26:   char           *tmp,typeName[256];
 27:   PetscBool      flg;

 29:   PetscInitialize(&argc,&argv,(char*)0,help);
 30:   PetscOptionsGetString(NULL,"-dm_vec_type",typeName,256,&flg);
 31:   if (flg) {
 32:     PetscStrstr(typeName,"cusp",&tmp);
 33:     if (tmp) useCUSP = PETSC_TRUE;
 34:   }

 36:   DMDACreate1d(PETSC_COMM_WORLD,DMDA_BOUNDARY_NONE,-8,1,1,NULL,&da);
 37:   DMCreateGlobalVector(da,&x); VecDuplicate(x,&f);
 38:   DMSetMatType(da,MATAIJ);
 39:   DMCreateMatrix(da,&J);

 41:   SNESCreate(PETSC_COMM_WORLD,&snes);
 42:   SNESSetFunction(snes,f,ComputeFunction,da);
 43:   SNESSetJacobian(snes,J,J,ComputeJacobian,da);
 44:   SNESSetFromOptions(snes);
 45:   SNESSolve(snes,NULL,x);

 47:   MatDestroy(&J);
 48:   VecDestroy(&x);
 49:   VecDestroy(&f);
 50:   SNESDestroy(&snes);
 51:   DMDestroy(&da);

 53:   PetscFinalize();
 54:   return 0;
 55: }

 57: struct ApplyStencil
 58: {
 59:   template <typename Tuple>
 60:   __host__ __device__
 61:   void operator()(Tuple t)
 62:   {
 63:     /* f = (2*x_i - x_(i+1) - x_(i-1))/h - h*exp(x_i) */
 64:     thrust::get<0>(t) = 1;
 65:     if ((thrust::get<4>(t) > 0) && (thrust::get<4>(t) < thrust::get<5>(t)-1)) {
 66:       thrust::get<0>(t) = (2.0*thrust::get<1>(t) - thrust::get<2>(t) - thrust::get<3>(t)) / (thrust::get<6>(t)) - (thrust::get<6>(t))*exp(thrust::get<1>(t));
 67:     } else if (thrust::get<4>(t) == 0) {
 68:       thrust::get<0>(t) = thrust::get<1>(t) / (thrust::get<6>(t));
 69:     } else if (thrust::get<4>(t) == thrust::get<5>(t)-1) {
 70:       thrust::get<0>(t) = thrust::get<1>(t) / (thrust::get<6>(t));
 71:     }
 72:   }
 73: };

 75: PetscErrorCode ComputeFunction(SNES snes,Vec x,Vec f,void *ctx)
 76: {
 77:   PetscInt       i,Mx,xs,xm,xstartshift,xendshift,fstart,lsize;
 78:   PetscScalar    *xx,*ff,hx;
 79:   DM             da = (DM) ctx;
 80:   Vec            xlocal;
 82:   PetscMPIInt    rank,size;
 83:   MPI_Comm       comm;
 84:   cusp::array1d<PetscScalar,cusp::device_memory> *xarray,*farray;

 86:   DMDAGetInfo(da,PETSC_IGNORE,&Mx,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE);
 87:   hx   = 1.0/(PetscReal)(Mx-1);
 88:   DMGetLocalVector(da,&xlocal);
 89:   DMGlobalToLocalBegin(da,x,INSERT_VALUES,xlocal);
 90:   DMGlobalToLocalEnd(da,x,INSERT_VALUES,xlocal);

 92:   if (useCUSP) {
 93:     VecCUSPGetArrayRead(xlocal,&xarray);
 94:     VecCUSPGetArrayWrite(f,&farray);
 95:     PetscObjectGetComm((PetscObject)da,&comm);
 96:     MPI_Comm_size(comm,&size);
 97:     MPI_Comm_rank(comm,&rank);
 98:     if (rank) xstartshift = 1;
 99:     else xstartshift = 0;
100:     if (rank != size-1) xendshift = 1;
101:     else xendshift = 0;
102:     VecGetOwnershipRange(f,&fstart,NULL);
103:     VecGetLocalSize(x,&lsize);
104:     try {
105:       thrust::for_each(
106:         thrust::make_zip_iterator(
107:           thrust::make_tuple(
108:             farray->begin(),
109:             xarray->begin()+xstartshift,
110:             xarray->begin()+xstartshift + 1,
111:             xarray->begin()+xstartshift - 1,
112:             thrust::counting_iterator<int>(fstart),
113:             thrust::constant_iterator<int>(Mx),
114:             thrust::constant_iterator<PetscScalar>(hx))),
115:         thrust::make_zip_iterator(
116:           thrust::make_tuple(
117:             farray->end(),
118:             xarray->end()-xendshift,
119:             xarray->end()-xendshift + 1,
120:             xarray->end()-xendshift - 1,
121:             thrust::counting_iterator<int>(fstart) + lsize,
122:             thrust::constant_iterator<int>(Mx),
123:             thrust::constant_iterator<PetscScalar>(hx))),
124:         ApplyStencil());
125:     }
126:     catch (char *all) {
127:       PetscPrintf(PETSC_COMM_WORLD, "Thrust is not working\n");
128:     }
129:     VecCUSPRestoreArrayRead(xlocal,&xarray);
130:     VecCUSPRestoreArrayWrite(f,&farray);
131:   } else {
132:     DMDAVecGetArray(da,xlocal,&xx);
133:     DMDAVecGetArray(da,f,&ff);
134:     DMDAGetCorners(da,&xs,NULL,NULL,&xm,NULL,NULL);

136:     for (i=xs; i<xs+xm; i++) {
137:       if (i == 0 || i == Mx-1) ff[i] = xx[i]/hx;
138:       else ff[i] =  (2.0*xx[i] - xx[i-1] - xx[i+1])/hx - hx*PetscExpScalar(xx[i]);
139:     }
140:     DMDAVecRestoreArray(da,xlocal,&xx);
141:     DMDAVecRestoreArray(da,f,&ff);
142:   }
143:   DMRestoreLocalVector(da,&xlocal);
144:   //  VecView(x,0);printf("f\n");
145:   //  VecView(f,0);
146:   return 0;

148: }
149: PetscErrorCode ComputeJacobian(SNES snes,Vec x,Mat *J,Mat *B,MatStructure *flag,void *ctx)
150: {
151:   DM             da = (DM) ctx;
152:   PetscInt       i,Mx,xm,xs;
153:   PetscScalar    hx,*xx;
154:   Vec            xlocal;

157:   DMDAGetInfo(da,PETSC_IGNORE,&Mx,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE);
158:   hx   = 1.0/(PetscReal)(Mx-1);
159:   DMGetLocalVector(da,&xlocal);DMGlobalToLocalBegin(da,x,INSERT_VALUES,xlocal);
160:   DMGlobalToLocalEnd(da,x,INSERT_VALUES,xlocal);
161:   DMDAVecGetArray(da,xlocal,&xx);
162:   DMDAGetCorners(da,&xs,NULL,NULL,&xm,NULL,NULL);

164:   for (i=xs; i<xs+xm; i++) {
165:     if (i == 0 || i == Mx-1) {
166:       MatSetValue(*J,i,i,1.0/hx,INSERT_VALUES);
167:     } else {
168:       MatSetValue(*J,i,i-1,-1.0/hx,INSERT_VALUES);
169:       MatSetValue(*J,i,i,2.0/hx - hx*PetscExpScalar(xx[i]),INSERT_VALUES);
170:       MatSetValue(*J,i,i+1,-1.0/hx,INSERT_VALUES);
171:     }
172:   }
173:   MatAssemblyBegin(*J,MAT_FINAL_ASSEMBLY);
174:   MatAssemblyEnd(*J,MAT_FINAL_ASSEMBLY);
175:   *flag = SAME_NONZERO_PATTERN;
176:   DMDAVecRestoreArray(da,xlocal,&xx);
177:   DMRestoreLocalVector(da,&xlocal);
178:   return 0;
179: }