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