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