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 char *tmp,typeName[256]; 27 PetscBool flg; 28 29 PetscFunctionBeginUser; 30 PetscCall(PetscInitialize(&argc,&argv,(char*)0,help)); 31 PetscCall(PetscOptionsGetString(NULL,NULL,"-dm_vec_type",typeName,sizeof(typeName),&flg)); 32 if (flg) { 33 PetscCall(PetscStrstr(typeName,"cuda",&tmp)); 34 if (tmp) useCUDA = PETSC_TRUE; 35 } 36 37 PetscCall(DMDACreate1d(PETSC_COMM_WORLD,DM_BOUNDARY_NONE,8,1,1,NULL,&da)); 38 PetscCall(DMSetFromOptions(da)); 39 PetscCall(DMSetUp(da)); 40 PetscCall(DMCreateGlobalVector(da,&x)); 41 PetscCall(VecDuplicate(x,&f)); 42 PetscCall(DMCreateMatrix(da,&J)); 43 44 PetscCall(SNESCreate(PETSC_COMM_WORLD,&snes)); 45 PetscCall(SNESSetFunction(snes,f,ComputeFunction,da)); 46 PetscCall(SNESSetJacobian(snes,J,J,ComputeJacobian,da)); 47 PetscCall(SNESSetFromOptions(snes)); 48 PetscCall(SNESSolve(snes,NULL,x)); 49 50 PetscCall(MatDestroy(&J)); 51 PetscCall(VecDestroy(&x)); 52 PetscCall(VecDestroy(&f)); 53 PetscCall(SNESDestroy(&snes)); 54 PetscCall(DMDestroy(&da)); 55 56 PetscCall(PetscFinalize()); 57 return 0; 58 } 59 60 struct ApplyStencil 61 { 62 template <typename Tuple> 63 __host__ __device__ 64 void operator()(Tuple t) 65 { 66 /* f = (2*x_i - x_(i+1) - x_(i-1))/h - h*exp(x_i) */ 67 thrust::get<0>(t) = 1; 68 if ((thrust::get<4>(t) > 0) && (thrust::get<4>(t) < thrust::get<5>(t)-1)) { 69 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)); 70 } else if (thrust::get<4>(t) == 0) { 71 thrust::get<0>(t) = thrust::get<1>(t) / (thrust::get<6>(t)); 72 } else if (thrust::get<4>(t) == thrust::get<5>(t)-1) { 73 thrust::get<0>(t) = thrust::get<1>(t) / (thrust::get<6>(t)); 74 } 75 } 76 }; 77 78 PetscErrorCode ComputeFunction(SNES snes,Vec x,Vec f,void *ctx) 79 { 80 PetscInt i,Mx,xs,xm,xstartshift,xendshift,fstart,lsize; 81 PetscScalar *xx,*ff,hx; 82 DM da = (DM) ctx; 83 Vec xlocal; 84 PetscMPIInt rank,size; 85 MPI_Comm comm; 86 PetscScalar const *xarray; 87 PetscScalar *farray; 88 89 PetscCall(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)); 90 hx = 1.0/(PetscReal)(Mx-1); 91 PetscCall(DMGetLocalVector(da,&xlocal)); 92 PetscCall(DMGlobalToLocalBegin(da,x,INSERT_VALUES,xlocal)); 93 PetscCall(DMGlobalToLocalEnd(da,x,INSERT_VALUES,xlocal)); 94 95 if (useCUDA) { 96 PetscCall(PetscObjectGetComm((PetscObject)da,&comm)); 97 PetscCallMPI(MPI_Comm_size(comm,&size)); 98 PetscCallMPI(MPI_Comm_rank(comm,&rank)); 99 PetscCall(VecCUDAGetArrayRead(xlocal,&xarray)); 100 PetscCall(VecCUDAGetArrayWrite(f,&farray)); 101 if (rank) xstartshift = 1; 102 else xstartshift = 0; 103 if (rank != size-1) xendshift = 1; 104 else xendshift = 0; 105 PetscCall(VecGetOwnershipRange(f,&fstart,NULL)); 106 PetscCall(VecGetLocalSize(x,&lsize)); 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 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Thrust is not working\n")); 131 } 132 PetscCall(VecCUDARestoreArrayRead(xlocal,&xarray)); 133 PetscCall(VecCUDARestoreArrayWrite(f,&farray)); 134 } else { 135 PetscCall(DMDAVecGetArray(da,xlocal,&xx)); 136 PetscCall(DMDAVecGetArray(da,f,&ff)); 137 PetscCall(DMDAGetCorners(da,&xs,NULL,NULL,&xm,NULL,NULL)); 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 PetscCall(DMDAVecRestoreArray(da,xlocal,&xx)); 144 PetscCall(DMDAVecRestoreArray(da,f,&ff)); 145 } 146 PetscCall(DMRestoreLocalVector(da,&xlocal)); 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 157 PetscCall(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 PetscCall(DMGetLocalVector(da,&xlocal)); 160 PetscCall(DMGlobalToLocalBegin(da,x,INSERT_VALUES,xlocal)); 161 PetscCall(DMGlobalToLocalEnd(da,x,INSERT_VALUES,xlocal)); 162 PetscCall(DMDAVecGetArray(da,xlocal,&xx)); 163 PetscCall(DMDAGetCorners(da,&xs,NULL,NULL,&xm,NULL,NULL)); 164 165 for (i=xs; i<xs+xm; i++) { 166 if (i == 0 || i == Mx-1) { 167 PetscCall(MatSetValue(J,i,i,1.0/hx,INSERT_VALUES)); 168 } else { 169 PetscCall(MatSetValue(J,i,i-1,-1.0/hx,INSERT_VALUES)); 170 PetscCall(MatSetValue(J,i,i,2.0/hx - hx*PetscExpScalar(xx[i]),INSERT_VALUES)); 171 PetscCall(MatSetValue(J,i,i+1,-1.0/hx,INSERT_VALUES)); 172 } 173 } 174 PetscCall(MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY)); 175 PetscCall(MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY)); 176 PetscCall(DMDAVecRestoreArray(da,xlocal,&xx)); 177 PetscCall(DMRestoreLocalVector(da,&xlocal)); 178 return 0; 179 } 180 181 /*TEST 182 183 build: 184 requires: cuda 185 186 testset: 187 args: -snes_monitor_short -dm_mat_type aijcusparse -dm_vec_type cuda 188 output_file: output/ex47cu_1.out 189 test: 190 suffix: 1 191 nsize: 1 192 test: 193 suffix: 2 194 nsize: 2 195 196 TEST*/ 197