1 static char help[] = "Test PC redistribute on matrix with load imbalance. \n\
2 Modified from src/ksp/ksp/tutorials/ex2.c.\n\
3 Input parameters include:\n\
4 -random_exact_sol : use a random exact solution vector\n\
5 -view_exact_sol : write exact solution vector to stdout\n\
6 -n <mesh_y> : number of mesh points\n\n";
7 /*
8 Example:
9 mpiexec -n 8 ./ex3 -n 10000 -ksp_type cg -pc_type bjacobi -sub_pc_type icc -ksp_rtol 1.e-8 -log_view
10 mpiexec -n 8 ./ex3 -n 10000 -ksp_type preonly -pc_type redistribute -redistribute_ksp_type cg -redistribute_pc_type bjacobi -redistribute_sub_pc_type icc -redistribute_ksp_rtol 1.e-8 -log_view
11 */
12
13 #include <petscksp.h>
14
main(int argc,char ** args)15 int main(int argc, char **args)
16 {
17 Vec x, b, u; /* approx solution, RHS, exact solution */
18 Mat A; /* linear system matrix */
19 KSP ksp; /* linear solver context */
20 PetscRandom rctx; /* random number generator context */
21 PetscReal norm; /* norm of solution error */
22 PetscInt i, j, Ii, J, Istart, Iend, m, n = 7, its, nloc, matdistribute = 0;
23 PetscBool flg = PETSC_FALSE;
24 PetscScalar v;
25 PetscMPIInt rank, size;
26 PetscLogStage stage;
27
28 PetscFunctionBeginUser;
29 PetscCall(PetscInitialize(&argc, &args, NULL, help));
30 PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
31 PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
32 PetscCheck(size > 1, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "This example requires at least 2 MPI processes!");
33
34 PetscCall(PetscOptionsGetInt(NULL, NULL, "-n", &n, NULL));
35 PetscCall(PetscOptionsGetInt(NULL, NULL, "-matdistribute", &matdistribute, NULL));
36 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
37 Compute the matrix and right-hand-side vector that define
38 the linear system, Ax = b.
39 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
40 switch (matdistribute) {
41 case 1: /* very imbalanced process load for matrix A */
42 m = (1 + size) * size;
43 nloc = (rank + 1) * n;
44 if (rank == size - 1) { /* proc[size-1] stores all remaining rows */
45 nloc = m * n;
46 for (i = 0; i < size - 1; i++) nloc -= (i + 1) * n;
47 }
48 break;
49 default: /* proc[0] and proc[1] load much smaller row blocks, the rest processes have same loads */
50 if (rank == 0 || rank == 1) {
51 nloc = n;
52 } else {
53 nloc = 10 * n; /* 10x larger load */
54 }
55 m = 2 + (size - 2) * 10;
56 break;
57 }
58 PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
59 PetscCall(MatSetSizes(A, nloc, nloc, PETSC_DECIDE, PETSC_DECIDE));
60 PetscCall(MatSetFromOptions(A));
61 PetscCall(MatMPIAIJSetPreallocation(A, 5, NULL, 5, NULL));
62 PetscCall(MatSeqAIJSetPreallocation(A, 5, NULL));
63 PetscCall(MatSetUp(A));
64
65 PetscCall(MatGetOwnershipRange(A, &Istart, &Iend));
66 nloc = Iend - Istart;
67 PetscCall(PetscSynchronizedPrintf(PETSC_COMM_WORLD, "[%d] A Istart,Iend: %" PetscInt_FMT " %" PetscInt_FMT "; nloc %" PetscInt_FMT "\n", rank, Istart, Iend, nloc));
68 PetscCall(PetscSynchronizedFlush(PETSC_COMM_WORLD, PETSC_STDOUT));
69
70 PetscCall(PetscLogStageRegister("Assembly", &stage));
71 PetscCall(PetscLogStagePush(stage));
72 for (Ii = Istart; Ii < Iend; Ii++) {
73 v = -1.0;
74 i = Ii / n;
75 j = Ii - i * n;
76 if (i > 0) {
77 J = Ii - n;
78 PetscCall(MatSetValues(A, 1, &Ii, 1, &J, &v, INSERT_VALUES));
79 }
80 if (i < m - 1) {
81 J = Ii + n;
82 PetscCall(MatSetValues(A, 1, &Ii, 1, &J, &v, INSERT_VALUES));
83 }
84 if (j > 0) {
85 J = Ii - 1;
86 PetscCall(MatSetValues(A, 1, &Ii, 1, &J, &v, INSERT_VALUES));
87 }
88 if (j < n - 1) {
89 J = Ii + 1;
90 PetscCall(MatSetValues(A, 1, &Ii, 1, &J, &v, INSERT_VALUES));
91 }
92 v = 4.0;
93 PetscCall(MatSetValues(A, 1, &Ii, 1, &Ii, &v, INSERT_VALUES));
94 }
95 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
96 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
97 PetscCall(PetscLogStagePop());
98
99 /* A is symmetric. Set symmetric flag to enable ICC/Cholesky preconditioner */
100 PetscCall(MatSetOption(A, MAT_SYMMETRIC, PETSC_TRUE));
101
102 /* Create parallel vectors. */
103 PetscCall(VecCreate(PETSC_COMM_WORLD, &u));
104 PetscCall(VecSetSizes(u, nloc, PETSC_DECIDE));
105 PetscCall(VecSetFromOptions(u));
106 PetscCall(VecDuplicate(u, &b));
107 PetscCall(VecDuplicate(b, &x));
108
109 /* Set exact solution; then compute right-hand-side vector. */
110 PetscCall(PetscOptionsGetBool(NULL, NULL, "-random_exact_sol", &flg, NULL));
111 if (flg) {
112 PetscCall(PetscRandomCreate(PETSC_COMM_WORLD, &rctx));
113 PetscCall(PetscRandomSetFromOptions(rctx));
114 PetscCall(VecSetRandom(u, rctx));
115 PetscCall(PetscRandomDestroy(&rctx));
116 } else {
117 PetscCall(VecSet(u, 1.0));
118 }
119 PetscCall(MatMult(A, u, b));
120
121 /* View the exact solution vector if desired */
122 flg = PETSC_FALSE;
123 PetscCall(PetscOptionsGetBool(NULL, NULL, "-view_exact_sol", &flg, NULL));
124 if (flg) PetscCall(VecView(u, PETSC_VIEWER_STDOUT_WORLD));
125
126 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
127 Create the linear solver and set various options
128 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
129 PetscCall(KSPCreate(PETSC_COMM_WORLD, &ksp));
130 PetscCall(KSPSetOperators(ksp, A, A));
131 PetscCall(KSPSetTolerances(ksp, 1.e-2 / ((m + 1) * (n + 1)), PETSC_CURRENT, PETSC_CURRENT, PETSC_CURRENT));
132 PetscCall(KSPSetFromOptions(ksp));
133
134 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
135 Solve the linear system
136 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
137 PetscCall(KSPSolve(ksp, b, x));
138
139 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
140 Check solution and clean up
141 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
142 PetscCall(VecAXPY(x, -1.0, u));
143 PetscCall(VecNorm(x, NORM_2, &norm));
144 PetscCall(KSPGetIterationNumber(ksp, &its));
145 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Norm of error %g iterations %" PetscInt_FMT "\n", (double)norm, its));
146
147 /* Free work space. */
148 PetscCall(KSPDestroy(&ksp));
149 PetscCall(VecDestroy(&u));
150 PetscCall(VecDestroy(&x));
151 PetscCall(VecDestroy(&b));
152 PetscCall(MatDestroy(&A));
153 PetscCall(PetscFinalize());
154 return 0;
155 }
156
157 /*TEST
158
159 test:
160 nsize: 8
161 args: -n 100 -ksp_type cg -pc_type bjacobi -sub_pc_type icc -ksp_rtol 1.e-8
162
163 test:
164 suffix: 2
165 nsize: 8
166 args: -n 100 -ksp_type preonly -pc_type redistribute -redistribute_ksp_type cg -redistribute_pc_type bjacobi -redistribute_sub_pc_type icc -redistribute_ksp_rtol 1.e-8
167
168 TEST*/
169