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 */ 5 6 #include <petscdm.h> 7 #include <petscdmda.h> 8 #include <petscsnes.h> 9 10 #include <thrust/device_ptr.h> 11 #include <thrust/for_each.h> 12 #include <thrust/tuple.h> 13 #include <thrust/iterator/constant_iterator.h> 14 #include <thrust/iterator/counting_iterator.h> 15 #include <thrust/iterator/zip_iterator.h> 16 17 extern PetscErrorCode ComputeFunction(SNES,Vec,Vec,void*), ComputeJacobian(SNES,Vec,Mat,Mat,void*); 18 PetscBool useCUDA = PETSC_FALSE; 19 20 int main(int argc,char **argv) 21 { 22 SNES snes; 23 Vec x,f; 24 Mat J; 25 DM da; 26 PetscErrorCode ierr; 27 char *tmp,typeName[256]; 28 PetscBool flg; 29 30 ierr = PetscInitialize(&argc,&argv,(char*)0,help);if (ierr) return ierr; 31 ierr = PetscOptionsGetString(NULL,NULL,"-dm_vec_type",typeName,sizeof(typeName),&flg);CHKERRQ(ierr); 32 if (flg) { 33 ierr = PetscStrstr(typeName,"cuda",&tmp);CHKERRQ(ierr); 34 if (tmp) useCUDA = PETSC_TRUE; 35 } 36 37 ierr = DMDACreate1d(PETSC_COMM_WORLD,DM_BOUNDARY_NONE,8,1,1,NULL,&da);CHKERRQ(ierr); 38 ierr = DMSetFromOptions(da);CHKERRQ(ierr); 39 ierr = DMSetUp(da);CHKERRQ(ierr); 40 ierr = DMCreateGlobalVector(da,&x); VecDuplicate(x,&f);CHKERRQ(ierr); 41 ierr = DMCreateMatrix(da,&J);CHKERRQ(ierr); 42 43 ierr = SNESCreate(PETSC_COMM_WORLD,&snes);CHKERRQ(ierr); 44 ierr = SNESSetFunction(snes,f,ComputeFunction,da);CHKERRQ(ierr); 45 ierr = SNESSetJacobian(snes,J,J,ComputeJacobian,da);CHKERRQ(ierr); 46 ierr = SNESSetFromOptions(snes);CHKERRQ(ierr); 47 ierr = SNESSolve(snes,NULL,x);CHKERRQ(ierr); 48 49 ierr = MatDestroy(&J);CHKERRQ(ierr); 50 ierr = VecDestroy(&x);CHKERRQ(ierr); 51 ierr = VecDestroy(&f);CHKERRQ(ierr); 52 ierr = SNESDestroy(&snes);CHKERRQ(ierr); 53 ierr = DMDestroy(&da);CHKERRQ(ierr); 54 55 ierr = PetscFinalize(); 56 return ierr; 57 } 58 59 struct ApplyStencil 60 { 61 template <typename Tuple> 62 __host__ __device__ 63 void operator()(Tuple t) 64 { 65 /* f = (2*x_i - x_(i+1) - x_(i-1))/h - h*exp(x_i) */ 66 thrust::get<0>(t) = 1; 67 if ((thrust::get<4>(t) > 0) && (thrust::get<4>(t) < thrust::get<5>(t)-1)) { 68 thrust::get<0>(t) = (((PetscScalar)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)); 69 } else if (thrust::get<4>(t) == 0) { 70 thrust::get<0>(t) = thrust::get<1>(t) / (thrust::get<6>(t)); 71 } else if (thrust::get<4>(t) == thrust::get<5>(t)-1) { 72 thrust::get<0>(t) = thrust::get<1>(t) / (thrust::get<6>(t)); 73 } 74 } 75 }; 76 77 PetscErrorCode ComputeFunction(SNES snes,Vec x,Vec f,void *ctx) 78 { 79 PetscInt i,Mx,xs,xm,xstartshift,xendshift,fstart,lsize; 80 PetscScalar *xx,*ff,hx; 81 DM da = (DM) ctx; 82 Vec xlocal; 83 PetscErrorCode ierr; 84 PetscMPIInt rank,size; 85 MPI_Comm comm; 86 PetscScalar const *xarray; 87 PetscScalar *farray; 88 89 ierr = 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);CHKERRQ(ierr); 90 hx = 1.0/(PetscReal)(Mx-1); 91 ierr = DMGetLocalVector(da,&xlocal);CHKERRQ(ierr); 92 ierr = DMGlobalToLocalBegin(da,x,INSERT_VALUES,xlocal);CHKERRQ(ierr); 93 ierr = DMGlobalToLocalEnd(da,x,INSERT_VALUES,xlocal);CHKERRQ(ierr); 94 95 if (useCUDA) { 96 ierr = VecCUDAGetArrayRead(xlocal,&xarray);CHKERRQ(ierr); 97 ierr = VecCUDAGetArrayWrite(f,&farray);CHKERRQ(ierr); 98 ierr = PetscObjectGetComm((PetscObject)da,&comm);CHKERRQ(ierr); 99 ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr); 100 ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr); 101 if (rank) xstartshift = 1; 102 else xstartshift = 0; 103 if (rank != size-1) xendshift = 1; 104 else xendshift = 0; 105 ierr = VecGetOwnershipRange(f,&fstart,NULL);CHKERRQ(ierr); 106 ierr = VecGetLocalSize(x,&lsize);CHKERRQ(ierr); 107 try { 108 thrust::for_each( 109 thrust::make_zip_iterator( 110 thrust::make_tuple( 111 thrust::device_ptr<PetscScalar>(farray), 112 thrust::device_ptr<const PetscScalar>(xarray + xstartshift), 113 thrust::device_ptr<const PetscScalar>(xarray + xstartshift + 1), 114 thrust::device_ptr<const PetscScalar>(xarray + xstartshift - 1), 115 thrust::counting_iterator<int>(fstart), 116 thrust::constant_iterator<int>(Mx), 117 thrust::constant_iterator<PetscScalar>(hx))), 118 thrust::make_zip_iterator( 119 thrust::make_tuple( 120 thrust::device_ptr<PetscScalar>(farray + lsize), 121 thrust::device_ptr<const PetscScalar>(xarray + lsize - xendshift), 122 thrust::device_ptr<const PetscScalar>(xarray + lsize - xendshift + 1), 123 thrust::device_ptr<const PetscScalar>(xarray + lsize - xendshift - 1), 124 thrust::counting_iterator<int>(fstart) + lsize, 125 thrust::constant_iterator<int>(Mx), 126 thrust::constant_iterator<PetscScalar>(hx))), 127 ApplyStencil()); 128 } 129 catch (char *all) { 130 ierr = PetscPrintf(PETSC_COMM_WORLD, "Thrust is not working\n");CHKERRQ(ierr); 131 } 132 ierr = VecCUDARestoreArrayRead(xlocal,&xarray);CHKERRQ(ierr); 133 ierr = VecCUDARestoreArrayWrite(f,&farray);CHKERRQ(ierr); 134 } else { 135 ierr = DMDAVecGetArray(da,xlocal,&xx);CHKERRQ(ierr); 136 ierr = DMDAVecGetArray(da,f,&ff);CHKERRQ(ierr); 137 ierr = DMDAGetCorners(da,&xs,NULL,NULL,&xm,NULL,NULL);CHKERRQ(ierr); 138 139 for (i=xs; i<xs+xm; i++) { 140 if (i == 0 || i == Mx-1) ff[i] = xx[i]/hx; 141 else ff[i] = (2.0*xx[i] - xx[i-1] - xx[i+1])/hx - hx*PetscExpScalar(xx[i]); 142 } 143 ierr = DMDAVecRestoreArray(da,xlocal,&xx);CHKERRQ(ierr); 144 ierr = DMDAVecRestoreArray(da,f,&ff);CHKERRQ(ierr); 145 } 146 ierr = DMRestoreLocalVector(da,&xlocal);CHKERRQ(ierr); 147 return 0; 148 149 } 150 PetscErrorCode ComputeJacobian(SNES snes,Vec x,Mat J,Mat B,void *ctx) 151 { 152 DM da = (DM) ctx; 153 PetscInt i,Mx,xm,xs; 154 PetscScalar hx,*xx; 155 Vec xlocal; 156 PetscErrorCode ierr; 157 158 ierr = 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);CHKERRQ(ierr); 159 hx = 1.0/(PetscReal)(Mx-1); 160 ierr = DMGetLocalVector(da,&xlocal);CHKERRQ(ierr); 161 ierr = DMGlobalToLocalBegin(da,x,INSERT_VALUES,xlocal);CHKERRQ(ierr); 162 ierr = DMGlobalToLocalEnd(da,x,INSERT_VALUES,xlocal);CHKERRQ(ierr); 163 ierr = DMDAVecGetArray(da,xlocal,&xx);CHKERRQ(ierr); 164 ierr = DMDAGetCorners(da,&xs,NULL,NULL,&xm,NULL,NULL);CHKERRQ(ierr); 165 166 for (i=xs; i<xs+xm; i++) { 167 if (i == 0 || i == Mx-1) { 168 ierr = MatSetValue(J,i,i,1.0/hx,INSERT_VALUES);CHKERRQ(ierr); 169 } else { 170 ierr = MatSetValue(J,i,i-1,-1.0/hx,INSERT_VALUES);CHKERRQ(ierr); 171 ierr = MatSetValue(J,i,i,2.0/hx - hx*PetscExpScalar(xx[i]),INSERT_VALUES);CHKERRQ(ierr); 172 ierr = MatSetValue(J,i,i+1,-1.0/hx,INSERT_VALUES);CHKERRQ(ierr); 173 } 174 } 175 ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 176 ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 177 ierr = DMDAVecRestoreArray(da,xlocal,&xx);CHKERRQ(ierr); 178 ierr = DMRestoreLocalVector(da,&xlocal);CHKERRQ(ierr); 179 return 0; 180 } 181 182 183 184 /*TEST 185 186 build: 187 requires: cuda 188 189 testset: 190 args: -snes_monitor_short -dm_mat_type aijcusparse -dm_vec_type cuda 191 output_file: output/ex47cu_1.out 192 test: 193 suffix: 1 194 nsize: 1 195 test: 196 suffix: 2 197 nsize: 2 198 199 TEST*/ 200