Actual source code: rosenbrock4.h
1: #pragma once
3: #include <petsctao.h>
4: #include <petscsf.h>
5: #include <petscdevice_cupm.h>
7: /*
8: User-defined application context - contains data needed by the
9: application-provided call-back routines that evaluate the function,
10: gradient, and hessian.
11: */
13: typedef struct _Rosenbrock {
14: PetscInt bs; // each block of bs variables is one chained multidimensional rosenbrock problem
15: PetscInt i_start, i_end;
16: PetscInt c_start, c_end;
17: PetscReal alpha; // condition parameter
18: } Rosenbrock;
20: typedef struct _AppCtx *AppCtx;
21: struct _AppCtx {
22: MPI_Comm comm;
23: PetscInt n; /* dimension */
24: PetscInt n_local;
25: PetscInt n_local_comp;
26: Rosenbrock problem;
27: Vec Hvalues; /* vector for writing COO values of this MPI process */
28: Vec gvalues; /* vector for writing gradient values of this mpi process */
29: Vec fvector;
30: PetscSF off_process_scatter;
31: PetscSF gscatter;
32: Vec off_process_values; /* buffer for off-process values if chained */
33: PetscBool test_lmvm;
34: PetscLogEvent event_f, event_g, event_fg;
35: };
37: /* -------------- User-defined routines ---------- */
39: static PETSC_HOSTDEVICE_INLINE_DECL PetscReal RosenbrockObjective(PetscScalar alpha, PetscScalar x_1, PetscScalar x_2)
40: {
41: PetscScalar d = x_2 - x_1 * x_1;
42: PetscScalar e = 1.0 - x_1;
43: return alpha * d * d + e * e;
44: }
46: static const PetscLogDouble RosenbrockObjectiveFlops = 7.0;
48: static PETSC_HOSTDEVICE_INLINE_DECL void RosenbrockGradient(PetscScalar alpha, PetscScalar x_1, PetscScalar x_2, PetscScalar g[2])
49: {
50: PetscScalar d = x_2 - x_1 * x_1;
51: PetscScalar e = 1.0 - x_1;
52: PetscScalar g2 = alpha * d * 2.0;
54: g[0] = -2.0 * x_1 * g2 - 2.0 * e;
55: g[1] = g2;
56: }
58: static const PetscInt RosenbrockGradientFlops = 9.0;
60: static PETSC_HOSTDEVICE_INLINE_DECL PetscReal RosenbrockObjectiveGradient(PetscScalar alpha, PetscScalar x_1, PetscScalar x_2, PetscScalar g[2])
61: {
62: PetscScalar d = x_2 - x_1 * x_1;
63: PetscScalar e = 1.0 - x_1;
64: PetscScalar ad = alpha * d;
65: PetscScalar g2 = ad * 2.0;
67: g[0] = -2.0 * x_1 * g2 - 2.0 * e;
68: g[1] = g2;
69: return ad * d + e * e;
70: }
72: static const PetscLogDouble RosenbrockObjectiveGradientFlops = 12.0;
74: static PETSC_HOSTDEVICE_INLINE_DECL void RosenbrockHessian(PetscScalar alpha, PetscScalar x_1, PetscScalar x_2, PetscScalar h[4])
75: {
76: PetscScalar d = x_2 - x_1 * x_1;
77: PetscScalar g2 = alpha * d * 2.0;
78: PetscScalar h2 = -4.0 * alpha * x_1;
80: h[0] = -2.0 * (g2 + x_1 * h2) + 2.0;
81: h[1] = h[2] = h2;
82: h[3] = 2.0 * alpha;
83: }
85: static const PetscLogDouble RosenbrockHessianFlops = 11.0;
87: PetscErrorCode AppCtxCreate(MPI_Comm comm, AppCtx *ctx)
88: {
89: AppCtx user;
90: PetscDeviceContext dctx;
92: PetscFunctionBegin;
93: PetscCall(PetscNew(ctx));
94: user = *ctx;
95: user->comm = PETSC_COMM_WORLD;
97: /* Initialize problem parameters */
98: user->n = 2;
99: user->problem.alpha = 99.0;
100: user->problem.bs = 2; // bs = 2 is block Rosenbrock, bs = n is chained Rosenbrock
101: user->test_lmvm = PETSC_FALSE;
102: /* Check for command line arguments to override defaults */
103: PetscOptionsBegin(user->comm, NULL, "Rosenbrock example", NULL);
104: PetscCall(PetscOptionsInt("-n", "Rosenbrock problem size", NULL, user->n, &user->n, NULL));
105: PetscCall(PetscOptionsInt("-bs", "Rosenbrock block size (2 <= bs <= n)", NULL, user->problem.bs, &user->problem.bs, NULL));
106: PetscCall(PetscOptionsReal("-alpha", "Rosenbrock off-diagonal coefficient", NULL, user->problem.alpha, &user->problem.alpha, NULL));
107: PetscCall(PetscOptionsBool("-test_lmvm", "Test LMVM solve against LMVM mult", NULL, user->test_lmvm, &user->test_lmvm, NULL));
108: PetscOptionsEnd();
109: PetscCheck(user->problem.bs >= 1, comm, PETSC_ERR_ARG_INCOMP, "Block size %" PetscInt_FMT " is not bigger than 1", user->problem.bs);
110: PetscCheck((user->n % user->problem.bs) == 0, comm, PETSC_ERR_ARG_INCOMP, "Block size %" PetscInt_FMT " does not divide problem size % " PetscInt_FMT, user->problem.bs, user->n);
111: PetscCall(PetscLogEventRegister("Rbock_Obj", TAO_CLASSID, &user->event_f));
112: PetscCall(PetscLogEventRegister("Rbock_Grad", TAO_CLASSID, &user->event_g));
113: PetscCall(PetscLogEventRegister("Rbock_ObjGrad", TAO_CLASSID, &user->event_fg));
114: PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
115: PetscCall(PetscDeviceContextSetUp(dctx));
116: PetscFunctionReturn(PETSC_SUCCESS);
117: }
119: PetscErrorCode AppCtxDestroy(AppCtx *ctx)
120: {
121: AppCtx user;
123: PetscFunctionBegin;
124: user = *ctx;
125: *ctx = NULL;
126: PetscCall(VecDestroy(&user->Hvalues));
127: PetscCall(VecDestroy(&user->gvalues));
128: PetscCall(VecDestroy(&user->fvector));
129: PetscCall(VecDestroy(&user->off_process_values));
130: PetscCall(PetscSFDestroy(&user->off_process_scatter));
131: PetscCall(PetscSFDestroy(&user->gscatter));
132: PetscCall(PetscFree(user));
133: PetscFunctionReturn(PETSC_SUCCESS);
134: }
136: PetscErrorCode CreateHessian(AppCtx user, Mat *Hessian)
137: {
138: Mat H;
139: PetscLayout layout;
140: PetscInt i_start, i_end, n_local_comp, nnz_local;
141: PetscInt c_start, c_end;
142: PetscInt *coo_i;
143: PetscInt *coo_j;
144: PetscInt bs = user->problem.bs;
145: VecType vec_type;
147: PetscFunctionBegin;
148: /* Partition the optimization variables and the computations.
149: There are (bs - 1) contributions to the objective function for every (bs)
150: degrees of freedom. */
151: PetscCall(PetscLayoutCreateFromSizes(user->comm, PETSC_DECIDE, user->n, 1, &layout));
152: PetscCall(PetscLayoutSetUp(layout));
153: PetscCall(PetscLayoutGetRange(layout, &i_start, &i_end));
154: user->problem.i_start = i_start;
155: user->problem.i_end = i_end;
156: user->n_local = i_end - i_start;
157: user->problem.c_start = c_start = (i_start / bs) * (bs - 1) + (i_start % bs);
158: user->problem.c_end = c_end = (i_end / bs) * (bs - 1) + (i_end % bs);
159: user->n_local_comp = n_local_comp = c_end - c_start;
161: PetscCall(MatCreate(user->comm, Hessian));
162: H = *Hessian;
163: PetscCall(MatSetLayouts(H, layout, layout));
164: PetscCall(PetscLayoutDestroy(&layout));
165: PetscCall(MatSetType(H, MATAIJ));
166: PetscCall(MatSetOption(H, MAT_HERMITIAN, PETSC_TRUE));
167: PetscCall(MatSetOption(H, MAT_SYMMETRIC, PETSC_TRUE));
168: PetscCall(MatSetOption(H, MAT_SYMMETRY_ETERNAL, PETSC_TRUE));
169: PetscCall(MatSetOption(H, MAT_STRUCTURALLY_SYMMETRIC, PETSC_TRUE));
170: PetscCall(MatSetOption(H, MAT_STRUCTURAL_SYMMETRY_ETERNAL, PETSC_TRUE));
171: PetscCall(MatSetFromOptions(H)); /* set from options so that we can change the underlying matrix type */
173: nnz_local = n_local_comp * 4;
174: PetscCall(PetscMalloc2(nnz_local, &coo_i, nnz_local, &coo_j));
175: /* Instead of having one computation thread per row of the matrix,
176: this example uses one thread per contribution to the objective
177: function. Each contribution to the objective function relates
178: two adjacent degrees of freedom, so each contribution to
179: the objective function adds a 2x2 block into the matrix.
180: We describe these 2x2 blocks in COO format. */
181: for (PetscInt c = c_start, k = 0; c < c_end; c++, k += 4) {
182: PetscInt i = (c / (bs - 1)) * bs + c % (bs - 1);
184: coo_i[k + 0] = i;
185: coo_i[k + 1] = i;
186: coo_i[k + 2] = i + 1;
187: coo_i[k + 3] = i + 1;
189: coo_j[k + 0] = i;
190: coo_j[k + 1] = i + 1;
191: coo_j[k + 2] = i;
192: coo_j[k + 3] = i + 1;
193: }
194: PetscCall(MatSetPreallocationCOO(H, nnz_local, coo_i, coo_j));
195: PetscCall(PetscFree2(coo_i, coo_j));
197: PetscCall(MatGetVecType(H, &vec_type));
198: PetscCall(VecCreate(user->comm, &user->Hvalues));
199: PetscCall(VecSetSizes(user->Hvalues, nnz_local, PETSC_DETERMINE));
200: PetscCall(VecSetType(user->Hvalues, vec_type));
202: // vector to collect contributions to the objective
203: PetscCall(VecCreate(user->comm, &user->fvector));
204: PetscCall(VecSetSizes(user->fvector, user->n_local_comp, PETSC_DETERMINE));
205: PetscCall(VecSetType(user->fvector, vec_type));
207: { /* If we are using a device (such as a GPU), run some computations that will
208: warm up its linear algebra runtime before the problem we actually want
209: to profile */
211: PetscMemType memtype;
212: const PetscScalar *a;
214: PetscCall(VecGetArrayReadAndMemType(user->fvector, &a, &memtype));
215: PetscCall(VecRestoreArrayReadAndMemType(user->fvector, &a));
217: if (memtype == PETSC_MEMTYPE_DEVICE) {
218: PetscLogStage warmup;
219: Mat A, AtA;
220: Vec x, b;
221: PetscInt warmup_size = 1000;
222: PetscDeviceContext dctx;
224: PetscCall(PetscLogStageRegister("Device Warmup", &warmup));
225: PetscCall(PetscLogStageSetActive(warmup, PETSC_FALSE));
227: PetscCall(PetscLogStagePush(warmup));
228: PetscCall(MatCreateDenseFromVecType(PETSC_COMM_SELF, vec_type, warmup_size, warmup_size, warmup_size, warmup_size, PETSC_DEFAULT, NULL, &A));
229: PetscCall(MatSetRandom(A, NULL));
230: PetscCall(MatCreateVecs(A, &x, &b));
231: PetscCall(VecSetRandom(x, NULL));
233: PetscCall(MatMult(A, x, b));
234: PetscCall(MatTransposeMatMult(A, A, MAT_INITIAL_MATRIX, PETSC_DETERMINE, &AtA));
235: PetscCall(MatShift(AtA, (PetscScalar)warmup_size));
236: PetscCall(MatSetOption(AtA, MAT_SPD, PETSC_TRUE));
237: PetscCall(MatCholeskyFactor(AtA, NULL, NULL));
238: PetscCall(MatDestroy(&AtA));
239: PetscCall(VecDestroy(&b));
240: PetscCall(VecDestroy(&x));
241: PetscCall(MatDestroy(&A));
242: PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
243: PetscCall(PetscDeviceContextSynchronize(dctx));
244: PetscCall(PetscLogStagePop());
245: }
246: }
247: PetscFunctionReturn(PETSC_SUCCESS);
248: }
250: PetscErrorCode CreateVectors(AppCtx user, Mat H, Vec *solution, Vec *gradient)
251: {
252: VecType vec_type;
253: PetscInt n_coo, *coo_i, i_start, i_end;
254: Vec x;
255: PetscInt n_recv;
256: PetscSFNode recv;
257: PetscLayout layout;
258: PetscInt c_start = user->problem.c_start, c_end = user->problem.c_end, bs = user->problem.bs;
260: PetscFunctionBegin;
261: PetscCall(MatCreateVecs(H, solution, gradient));
262: x = *solution;
263: PetscCall(VecGetOwnershipRange(x, &i_start, &i_end));
264: PetscCall(VecGetType(x, &vec_type));
265: // create scatter for communicating values
266: PetscCall(VecGetLayout(x, &layout));
267: n_recv = 0;
268: if (user->n_local_comp && i_end < user->n) {
269: PetscMPIInt rank;
270: PetscInt index;
272: n_recv = 1;
273: PetscCall(PetscLayoutFindOwnerIndex(layout, i_end, &rank, &index));
274: recv.rank = rank;
275: recv.index = index;
276: }
277: PetscCall(PetscSFCreate(user->comm, &user->off_process_scatter));
278: PetscCall(PetscSFSetGraph(user->off_process_scatter, user->n_local, n_recv, NULL, PETSC_USE_POINTER, &recv, PETSC_COPY_VALUES));
279: PetscCall(VecCreate(user->comm, &user->off_process_values));
280: PetscCall(VecSetSizes(user->off_process_values, 1, PETSC_DETERMINE));
281: PetscCall(VecSetType(user->off_process_values, vec_type));
282: PetscCall(VecZeroEntries(user->off_process_values));
284: // create COO data for writing the gradient
285: n_coo = user->n_local_comp * 2;
286: PetscCall(PetscMalloc1(n_coo, &coo_i));
287: for (PetscInt c = c_start, k = 0; c < c_end; c++, k += 2) {
288: PetscInt i = (c / (bs - 1)) * bs + (c % (bs - 1));
290: coo_i[k + 0] = i;
291: coo_i[k + 1] = i + 1;
292: }
293: PetscCall(PetscSFCreate(user->comm, &user->gscatter));
294: PetscCall(PetscSFSetGraphLayout(user->gscatter, layout, n_coo, NULL, PETSC_USE_POINTER, coo_i));
295: PetscCall(PetscSFSetUp(user->gscatter));
296: PetscCall(PetscFree(coo_i));
297: PetscCall(VecCreate(user->comm, &user->gvalues));
298: PetscCall(VecSetSizes(user->gvalues, n_coo, PETSC_DETERMINE));
299: PetscCall(VecSetType(user->gvalues, vec_type));
300: PetscFunctionReturn(PETSC_SUCCESS);
301: }
303: #if PetscDefined(USING_CUPMCC)
305: #if PetscDefined(USING_NVCC)
306: typedef cudaStream_t cupmStream_t;
307: #define PetscCUPMLaunch(...) \
308: do { \
309: __VA_ARGS__; \
310: PetscCallCUDA(cudaGetLastError()); \
311: } while (0)
312: #elif PetscDefined(USING_HCC)
313: #define PetscCUPMLaunch(...) \
314: do { \
315: __VA_ARGS__; \
316: PetscCallHIP(hipGetLastError()); \
317: } while (0)
318: typedef hipStream_t cupmStream_t;
319: #endif
321: // x: on-process optimization variables
322: // o: buffer that contains the next optimization variable after the variables on this process
323: template <typename T>
324: PETSC_DEVICE_INLINE_DECL static void rosenbrock_for_loop(Rosenbrock r, const PetscScalar x[], const PetscScalar o[], T &&func) noexcept
325: {
326: PetscInt idx = blockIdx.x * blockDim.x + threadIdx.x; // 1D grid
327: PetscInt num_threads = gridDim.x * blockDim.x;
329: for (PetscInt c = r.c_start + idx, k = idx; c < r.c_end; c += num_threads, k += num_threads) {
330: PetscInt i = (c / (r.bs - 1)) * r.bs + (c % (r.bs - 1));
331: PetscScalar x_a = x[i - r.i_start];
332: PetscScalar x_b = ((i + 1) < r.i_end) ? x[i + 1 - r.i_start] : o[0];
334: func(k, x_a, x_b);
335: }
336: return;
337: }
339: PETSC_KERNEL_DECL void RosenbrockObjective_Kernel(Rosenbrock r, const PetscScalar x[], const PetscScalar o[], PetscScalar f_vec[])
340: {
341: rosenbrock_for_loop(r, x, o, [&](PetscInt k, PetscScalar x_a, PetscScalar x_b) { f_vec[k] = RosenbrockObjective(r.alpha, x_a, x_b); });
342: }
344: PETSC_KERNEL_DECL void RosenbrockGradient_Kernel(Rosenbrock r, const PetscScalar x[], const PetscScalar o[], PetscScalar g[])
345: {
346: rosenbrock_for_loop(r, x, o, [&](PetscInt k, PetscScalar x_a, PetscScalar x_b) { RosenbrockGradient(r.alpha, x_a, x_b, &g[2 * k]); });
347: }
349: PETSC_KERNEL_DECL void RosenbrockObjectiveGradient_Kernel(Rosenbrock r, const PetscScalar x[], const PetscScalar o[], PetscScalar f_vec[], PetscScalar g[])
350: {
351: rosenbrock_for_loop(r, x, o, [&](PetscInt k, PetscScalar x_a, PetscScalar x_b) { f_vec[k] = RosenbrockObjectiveGradient(r.alpha, x_a, x_b, &g[2 * k]); });
352: }
354: PETSC_KERNEL_DECL void RosenbrockHessian_Kernel(Rosenbrock r, const PetscScalar x[], const PetscScalar o[], PetscScalar h[])
355: {
356: rosenbrock_for_loop(r, x, o, [&](PetscInt k, PetscScalar x_a, PetscScalar x_b) { RosenbrockHessian(r.alpha, x_a, x_b, &h[4 * k]); });
357: }
359: static PetscErrorCode RosenbrockObjective_Device(cupmStream_t stream, Rosenbrock r, const PetscScalar x[], const PetscScalar o[], PetscScalar f_vec[])
360: {
361: PetscInt n_comp = r.c_end - r.c_start;
363: PetscFunctionBegin;
364: if (n_comp) PetscCUPMLaunch(RosenbrockObjective_Kernel<<<(n_comp + 255) / 256, 256, 0, stream>>>(r, x, o, f_vec));
365: PetscCall(PetscLogGpuFlops(RosenbrockObjectiveFlops * n_comp));
366: PetscFunctionReturn(PETSC_SUCCESS);
367: }
369: static PetscErrorCode RosenbrockGradient_Device(cupmStream_t stream, Rosenbrock r, const PetscScalar x[], const PetscScalar o[], PetscScalar g[])
370: {
371: PetscInt n_comp = r.c_end - r.c_start;
373: PetscFunctionBegin;
374: if (n_comp) PetscCUPMLaunch(RosenbrockGradient_Kernel<<<(n_comp + 255) / 256, 256, 0, stream>>>(r, x, o, g));
375: PetscCall(PetscLogGpuFlops(RosenbrockGradientFlops * n_comp));
376: PetscFunctionReturn(PETSC_SUCCESS);
377: }
379: static PetscErrorCode RosenbrockObjectiveGradient_Device(cupmStream_t stream, Rosenbrock r, const PetscScalar x[], const PetscScalar o[], PetscScalar f_vec[], PetscScalar g[])
380: {
381: PetscInt n_comp = r.c_end - r.c_start;
383: PetscFunctionBegin;
384: if (n_comp) PetscCUPMLaunch(RosenbrockObjectiveGradient_Kernel<<<(n_comp + 255) / 256, 256, 0, stream>>>(r, x, o, f_vec, g));
385: PetscCall(PetscLogGpuFlops(RosenbrockObjectiveGradientFlops * n_comp));
386: PetscFunctionReturn(PETSC_SUCCESS);
387: }
389: static PetscErrorCode RosenbrockHessian_Device(cupmStream_t stream, Rosenbrock r, const PetscScalar x[], const PetscScalar o[], PetscScalar h[])
390: {
391: PetscInt n_comp = r.c_end - r.c_start;
393: PetscFunctionBegin;
394: if (n_comp) PetscCUPMLaunch(RosenbrockHessian_Kernel<<<(n_comp + 255) / 256, 256, 0, stream>>>(r, x, o, h));
395: PetscCall(PetscLogGpuFlops(RosenbrockHessianFlops * n_comp));
396: PetscFunctionReturn(PETSC_SUCCESS);
397: }
398: #endif
400: static PetscErrorCode RosenbrockObjective_Host(Rosenbrock r, const PetscScalar x[], const PetscScalar o[], PetscReal *f)
401: {
402: PetscReal _f = 0.0;
404: PetscFunctionBegin;
405: for (PetscInt c = r.c_start; c < r.c_end; c++) {
406: PetscInt i = (c / (r.bs - 1)) * r.bs + (c % (r.bs - 1));
407: PetscScalar x_a = x[i - r.i_start];
408: PetscScalar x_b = ((i + 1) < r.i_end) ? x[i + 1 - r.i_start] : o[0];
410: _f += RosenbrockObjective(r.alpha, x_a, x_b);
411: }
412: *f = _f;
413: PetscCall(PetscLogFlops((RosenbrockObjectiveFlops + 1.0) * (r.c_end - r.c_start)));
414: PetscFunctionReturn(PETSC_SUCCESS);
415: }
417: static PetscErrorCode RosenbrockGradient_Host(Rosenbrock r, const PetscScalar x[], const PetscScalar o[], PetscScalar g[])
418: {
419: PetscFunctionBegin;
420: for (PetscInt c = r.c_start, k = 0; c < r.c_end; c++, k++) {
421: PetscInt i = (c / (r.bs - 1)) * r.bs + (c % (r.bs - 1));
422: PetscScalar x_a = x[i - r.i_start];
423: PetscScalar x_b = ((i + 1) < r.i_end) ? x[i + 1 - r.i_start] : o[0];
425: RosenbrockGradient(r.alpha, x_a, x_b, &g[2 * k]);
426: }
427: PetscCall(PetscLogFlops(RosenbrockGradientFlops * (r.c_end - r.c_start)));
428: PetscFunctionReturn(PETSC_SUCCESS);
429: }
431: static PetscErrorCode RosenbrockObjectiveGradient_Host(Rosenbrock r, const PetscScalar x[], const PetscScalar o[], PetscReal *f, PetscScalar g[])
432: {
433: PetscReal _f = 0.0;
435: PetscFunctionBegin;
436: for (PetscInt c = r.c_start, k = 0; c < r.c_end; c++, k++) {
437: PetscInt i = (c / (r.bs - 1)) * r.bs + (c % (r.bs - 1));
438: PetscScalar x_a = x[i - r.i_start];
439: PetscScalar x_b = ((i + 1) < r.i_end) ? x[i + 1 - r.i_start] : o[0];
441: _f += RosenbrockObjectiveGradient(r.alpha, x_a, x_b, &g[2 * k]);
442: }
443: *f = _f;
444: PetscCall(PetscLogFlops(RosenbrockObjectiveGradientFlops * (r.c_end - r.c_start)));
445: PetscFunctionReturn(PETSC_SUCCESS);
446: }
448: static PetscErrorCode RosenbrockHessian_Host(Rosenbrock r, const PetscScalar x[], const PetscScalar o[], PetscScalar h[])
449: {
450: PetscFunctionBegin;
451: for (PetscInt c = r.c_start, k = 0; c < r.c_end; c++, k++) {
452: PetscInt i = (c / (r.bs - 1)) * r.bs + (c % (r.bs - 1));
453: PetscScalar x_a = x[i - r.i_start];
454: PetscScalar x_b = ((i + 1) < r.i_end) ? x[i + 1 - r.i_start] : o[0];
456: RosenbrockHessian(r.alpha, x_a, x_b, &h[4 * k]);
457: }
458: PetscCall(PetscLogFlops(RosenbrockHessianFlops * (r.c_end - r.c_start)));
459: PetscFunctionReturn(PETSC_SUCCESS);
460: }
462: /* -------------------------------------------------------------------- */
464: PetscErrorCode FormObjective(Tao tao, Vec X, PetscReal *f, void *ptr)
465: {
466: AppCtx user = (AppCtx)ptr;
467: const PetscScalar *x;
468: const PetscScalar *o = NULL;
469: PetscMemType memtype_x;
471: PetscFunctionBeginUser;
472: PetscCall(PetscLogEventBegin(user->event_f, tao, NULL, NULL, NULL));
473: PetscCall(VecScatterBegin(user->off_process_scatter, X, user->off_process_values, INSERT_VALUES, SCATTER_FORWARD));
474: PetscCall(VecScatterEnd(user->off_process_scatter, X, user->off_process_values, INSERT_VALUES, SCATTER_FORWARD));
475: PetscCall(VecGetArrayReadAndMemType(user->off_process_values, &o, NULL));
476: PetscCall(VecGetArrayReadAndMemType(X, &x, &memtype_x));
477: if (memtype_x == PETSC_MEMTYPE_HOST) {
478: PetscCall(RosenbrockObjective_Host(user->problem, x, o, f));
479: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, f, 1, MPI_DOUBLE, MPI_SUM, user->comm));
480: #if PetscDefined(USING_CUPMCC)
481: } else if (memtype_x == PETSC_MEMTYPE_DEVICE) {
482: PetscScalar *_fvec;
483: PetscScalar f_scalar;
484: cupmStream_t *stream;
485: PetscDeviceContext dctx;
487: PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
488: PetscCall(PetscDeviceContextGetStreamHandle(dctx, (void **)&stream));
489: PetscCall(VecGetArrayWriteAndMemType(user->fvector, &_fvec, NULL));
490: PetscCall(RosenbrockObjective_Device(*stream, user->problem, x, o, _fvec));
491: PetscCall(VecRestoreArrayWriteAndMemType(user->fvector, &_fvec));
492: PetscCall(VecSum(user->fvector, &f_scalar));
493: *f = PetscRealPart(f_scalar);
494: #endif
495: } else SETERRQ(user->comm, PETSC_ERR_SUP, "Unsupported memtype %d", (int)memtype_x);
496: PetscCall(VecRestoreArrayReadAndMemType(X, &x));
497: PetscCall(VecRestoreArrayReadAndMemType(user->off_process_values, &o));
498: PetscCall(PetscLogEventEnd(user->event_f, tao, NULL, NULL, NULL));
499: PetscFunctionReturn(PETSC_SUCCESS);
500: }
502: PetscErrorCode FormGradient(Tao tao, Vec X, Vec G, void *ptr)
503: {
504: AppCtx user = (AppCtx)ptr;
505: PetscScalar *g;
506: const PetscScalar *x;
507: const PetscScalar *o = NULL;
508: PetscMemType memtype_x, memtype_g;
510: PetscFunctionBeginUser;
511: PetscCall(PetscLogEventBegin(user->event_g, tao, NULL, NULL, NULL));
512: PetscCall(VecScatterBegin(user->off_process_scatter, X, user->off_process_values, INSERT_VALUES, SCATTER_FORWARD));
513: PetscCall(VecScatterEnd(user->off_process_scatter, X, user->off_process_values, INSERT_VALUES, SCATTER_FORWARD));
514: PetscCall(VecGetArrayReadAndMemType(user->off_process_values, &o, NULL));
515: PetscCall(VecGetArrayReadAndMemType(X, &x, &memtype_x));
516: PetscCall(VecGetArrayWriteAndMemType(user->gvalues, &g, &memtype_g));
517: PetscAssert(memtype_x == memtype_g, user->comm, PETSC_ERR_ARG_INCOMP, "solution vector and gradient must have save memtype");
518: if (memtype_x == PETSC_MEMTYPE_HOST) {
519: PetscCall(RosenbrockGradient_Host(user->problem, x, o, g));
520: #if PetscDefined(USING_CUPMCC)
521: } else if (memtype_x == PETSC_MEMTYPE_DEVICE) {
522: cupmStream_t *stream;
523: PetscDeviceContext dctx;
525: PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
526: PetscCall(PetscDeviceContextGetStreamHandle(dctx, (void **)&stream));
527: PetscCall(RosenbrockGradient_Device(*stream, user->problem, x, o, g));
528: #endif
529: } else SETERRQ(user->comm, PETSC_ERR_SUP, "Unsupported memtype %d", (int)memtype_x);
530: PetscCall(VecRestoreArrayWriteAndMemType(user->gvalues, &g));
531: PetscCall(VecRestoreArrayReadAndMemType(X, &x));
532: PetscCall(VecRestoreArrayReadAndMemType(user->off_process_values, &o));
533: PetscCall(VecZeroEntries(G));
534: PetscCall(VecScatterBegin(user->gscatter, user->gvalues, G, ADD_VALUES, SCATTER_REVERSE));
535: PetscCall(VecScatterEnd(user->gscatter, user->gvalues, G, ADD_VALUES, SCATTER_REVERSE));
536: PetscCall(PetscLogEventEnd(user->event_g, tao, NULL, NULL, NULL));
537: PetscFunctionReturn(PETSC_SUCCESS);
538: }
540: /*
541: FormObjectiveGradient - Evaluates the function, f(X), and gradient, G(X).
543: Input Parameters:
544: . tao - the Tao context
545: . X - input vector
546: . ptr - optional user-defined context, as set by TaoSetObjectiveGradient()
548: Output Parameters:
549: . G - vector containing the newly evaluated gradient
550: . f - function value
552: Note:
553: Some optimization methods ask for the function and the gradient evaluation
554: at the same time. Evaluating both at once may be more efficient that
555: evaluating each separately.
556: */
557: PetscErrorCode FormObjectiveGradient(Tao tao, Vec X, PetscReal *f, Vec G, void *ptr)
558: {
559: AppCtx user = (AppCtx)ptr;
560: PetscScalar *g;
561: const PetscScalar *x;
562: const PetscScalar *o = NULL;
563: PetscMemType memtype_x, memtype_g;
565: PetscFunctionBeginUser;
566: PetscCall(PetscLogEventBegin(user->event_fg, tao, NULL, NULL, NULL));
567: PetscCall(VecScatterBegin(user->off_process_scatter, X, user->off_process_values, INSERT_VALUES, SCATTER_FORWARD));
568: PetscCall(VecScatterEnd(user->off_process_scatter, X, user->off_process_values, INSERT_VALUES, SCATTER_FORWARD));
569: PetscCall(VecGetArrayReadAndMemType(user->off_process_values, &o, NULL));
570: PetscCall(VecGetArrayReadAndMemType(X, &x, &memtype_x));
571: PetscCall(VecGetArrayWriteAndMemType(user->gvalues, &g, &memtype_g));
572: PetscAssert(memtype_x == memtype_g, user->comm, PETSC_ERR_ARG_INCOMP, "solution vector and gradient must have save memtype");
573: if (memtype_x == PETSC_MEMTYPE_HOST) {
574: PetscCall(RosenbrockObjectiveGradient_Host(user->problem, x, o, f, g));
575: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, f, 1, MPI_DOUBLE, MPI_SUM, PETSC_COMM_WORLD));
576: #if PetscDefined(USING_CUPMCC)
577: } else if (memtype_x == PETSC_MEMTYPE_DEVICE) {
578: PetscScalar *_fvec;
579: PetscScalar f_scalar;
580: cupmStream_t *stream;
581: PetscDeviceContext dctx;
583: PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
584: PetscCall(PetscDeviceContextGetStreamHandle(dctx, (void **)&stream));
585: PetscCall(VecGetArrayWriteAndMemType(user->fvector, &_fvec, NULL));
586: PetscCall(RosenbrockObjectiveGradient_Device(*stream, user->problem, x, o, _fvec, g));
587: PetscCall(VecRestoreArrayWriteAndMemType(user->fvector, &_fvec));
588: PetscCall(VecSum(user->fvector, &f_scalar));
589: *f = PetscRealPart(f_scalar);
590: #endif
591: } else SETERRQ(user->comm, PETSC_ERR_SUP, "Unsupported memtype %d", (int)memtype_x);
593: PetscCall(VecRestoreArrayWriteAndMemType(user->gvalues, &g));
594: PetscCall(VecRestoreArrayReadAndMemType(X, &x));
595: PetscCall(VecRestoreArrayReadAndMemType(user->off_process_values, &o));
596: PetscCall(VecZeroEntries(G));
597: PetscCall(VecScatterBegin(user->gscatter, user->gvalues, G, ADD_VALUES, SCATTER_REVERSE));
598: PetscCall(VecScatterEnd(user->gscatter, user->gvalues, G, ADD_VALUES, SCATTER_REVERSE));
599: PetscCall(PetscLogEventEnd(user->event_fg, tao, NULL, NULL, NULL));
600: PetscFunctionReturn(PETSC_SUCCESS);
601: }
603: static PetscErrorCode FormHessianSingle(Tao tao, Vec X, Mat H, void *ptr)
604: {
605: AppCtx user = (AppCtx)ptr;
606: PetscScalar *h;
607: const PetscScalar *x;
608: const PetscScalar *o = NULL;
609: PetscMemType memtype_x, memtype_h;
611: PetscFunctionBeginUser;
612: PetscCall(VecScatterBegin(user->off_process_scatter, X, user->off_process_values, INSERT_VALUES, SCATTER_FORWARD));
613: PetscCall(VecScatterEnd(user->off_process_scatter, X, user->off_process_values, INSERT_VALUES, SCATTER_FORWARD));
614: PetscCall(VecGetArrayReadAndMemType(user->off_process_values, &o, NULL));
615: PetscCall(VecGetArrayReadAndMemType(X, &x, &memtype_x));
616: PetscCall(VecGetArrayWriteAndMemType(user->Hvalues, &h, &memtype_h));
617: PetscAssert(memtype_x == memtype_h, user->comm, PETSC_ERR_ARG_INCOMP, "solution vector and hessian must have save memtype");
618: if (memtype_x == PETSC_MEMTYPE_HOST) {
619: PetscCall(RosenbrockHessian_Host(user->problem, x, o, h));
620: #if PetscDefined(USING_CUPMCC)
621: } else if (memtype_x == PETSC_MEMTYPE_DEVICE) {
622: cupmStream_t *stream;
623: PetscDeviceContext dctx;
625: PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
626: PetscCall(PetscDeviceContextGetStreamHandle(dctx, (void **)&stream));
627: PetscCall(RosenbrockHessian_Device(*stream, user->problem, x, o, h));
628: #endif
629: } else SETERRQ(user->comm, PETSC_ERR_SUP, "Unsupported memtype %d", (int)memtype_x);
631: PetscCall(MatSetValuesCOO(H, h, INSERT_VALUES));
632: PetscCall(VecRestoreArrayWriteAndMemType(user->Hvalues, &h));
634: PetscCall(VecRestoreArrayReadAndMemType(X, &x));
635: PetscCall(VecRestoreArrayReadAndMemType(user->off_process_values, &o));
636: PetscFunctionReturn(PETSC_SUCCESS);
637: }
639: /* ------------------------------------------------------------------- */
640: /*
641: FormHessian - Evaluates Hessian matrix.
643: Input Parameters:
644: . tao - the Tao context
645: . x - input vector
646: . ptr - optional user-defined context, as set by TaoSetHessian()
648: Output Parameters:
649: . H - Hessian matrix
651: Note: Providing the Hessian may not be necessary. Only some solvers
652: require this matrix.
653: */
654: PetscErrorCode FormHessian(Tao tao, Vec X, Mat H, Mat Hpre, void *ptr)
655: {
656: PetscFunctionBeginUser;
657: if (H) PetscCall(FormHessianSingle(tao, X, H, ptr));
658: if (Hpre && Hpre != H) {
659: if (H) PetscCall(MatCopy(H, Hpre, SAME_NONZERO_PATTERN));
660: else PetscCall(FormHessianSingle(tao, X, Hpre, ptr));
661: }
662: PetscFunctionReturn(PETSC_SUCCESS);
663: }
665: PetscErrorCode TestLMVM(Tao tao)
666: {
667: KSP ksp;
668: PC pc;
669: PetscBool is_lmvm;
671: PetscFunctionBegin;
672: PetscCall(TaoGetKSP(tao, &ksp));
673: if (!ksp) PetscFunctionReturn(PETSC_SUCCESS);
674: PetscCall(KSPGetPC(ksp, &pc));
675: PetscCall(PetscObjectTypeCompare((PetscObject)pc, PCLMVM, &is_lmvm));
676: if (is_lmvm) {
677: Mat M;
678: Vec in, out, out2;
679: PetscReal mult_solve_dist;
680: Vec x;
682: PetscCall(PCLMVMGetMatLMVM(pc, &M));
683: PetscCall(TaoGetSolution(tao, &x));
684: PetscCall(VecDuplicate(x, &in));
685: PetscCall(VecDuplicate(x, &out));
686: PetscCall(VecDuplicate(x, &out2));
687: PetscCall(VecSetRandom(in, NULL));
688: PetscCall(MatMult(M, in, out));
689: PetscCall(MatSolve(M, out, out2));
691: PetscCall(VecAXPY(out2, -1.0, in));
692: PetscCall(VecNorm(out2, NORM_2, &mult_solve_dist));
693: if (mult_solve_dist < 1.e-11) {
694: PetscCall(PetscPrintf(PetscObjectComm((PetscObject)tao), "Inverse error of LMVM MatMult and MatSolve: < 1.e-11\n"));
695: } else if (mult_solve_dist < 1.e-6) {
696: PetscCall(PetscPrintf(PetscObjectComm((PetscObject)tao), "Inverse error of LMVM MatMult and MatSolve: < 1.e-6\n"));
697: } else {
698: PetscCall(PetscPrintf(PetscObjectComm((PetscObject)tao), "Inverse error of LMVM MatMult and MatSolve is not small: %e\n", (double)mult_solve_dist));
699: }
700: PetscCall(VecDestroy(&in));
701: PetscCall(VecDestroy(&out));
702: PetscCall(VecDestroy(&out2));
703: }
704: PetscFunctionReturn(PETSC_SUCCESS);
705: }
707: PetscErrorCode RosenbrockMain(void)
708: {
709: Vec x; /* solution vector */
710: Vec g; /* gradient vector */
711: Mat H; /* Hessian matrix */
712: Tao tao; /* Tao solver context */
713: AppCtx user; /* user-defined application context */
714: PetscLogStage solve;
716: PetscFunctionBegin;
717: PetscCall(PetscLogStageRegister("Rosenbrock solve", &solve));
719: PetscCall(AppCtxCreate(PETSC_COMM_WORLD, &user));
720: PetscCall(CreateHessian(user, &H));
721: PetscCall(CreateVectors(user, H, &x, &g));
723: /* The TAO code begins here */
725: PetscCall(TaoCreate(user->comm, &tao));
726: PetscCall(VecZeroEntries(x));
727: PetscCall(TaoSetSolution(tao, x));
729: /* Set routines for function, gradient, hessian evaluation */
730: PetscCall(TaoSetObjective(tao, FormObjective, user));
731: PetscCall(TaoSetObjectiveAndGradient(tao, g, FormObjectiveGradient, user));
732: PetscCall(TaoSetGradient(tao, g, FormGradient, user));
733: PetscCall(TaoSetHessian(tao, H, H, FormHessian, user));
735: PetscCall(TaoSetFromOptions(tao));
737: /* SOLVE THE APPLICATION */
738: PetscCall(PetscLogStagePush(solve));
739: PetscCall(TaoSolve(tao));
740: PetscCall(PetscLogStagePop());
742: if (user->test_lmvm) PetscCall(TestLMVM(tao));
744: PetscCall(TaoDestroy(&tao));
745: PetscCall(VecDestroy(&g));
746: PetscCall(VecDestroy(&x));
747: PetscCall(MatDestroy(&H));
748: PetscCall(AppCtxDestroy(&user));
749: PetscFunctionReturn(PETSC_SUCCESS);
750: }