xref: /petsc/src/mat/tests/ex183.c (revision cc7980154d91e4d53ff48c8dc1d2bc5deb304502)
1 static char help[] = "Example of extracting an array of MPI submatrices from a given MPI matrix.\n"
2                      "This test can only be run in parallel.\n"
3                      "\n";
4 
5 #include <petscmat.h>
6 
7 PetscErrorCode MyISView(IS *rowis, IS *colis, PetscInt gs, PetscInt ss, PetscViewer viewer)
8 {
9   PetscViewer subviewer = NULL;
10 
11   PetscFunctionBeginUser;
12   PetscCall(PetscViewerASCIIPrintf(viewer, "Row IS %" PetscInt_FMT "\n", gs));
13   if (ss > -1) {
14     PetscCall(PetscViewerGetSubViewer(viewer, PetscObjectComm((PetscObject)rowis[ss]), &subviewer));
15     PetscCall(ISView(rowis[ss], subviewer));
16     PetscCall(PetscViewerRestoreSubViewer(viewer, PetscObjectComm((PetscObject)rowis[ss]), &subviewer));
17   }
18   PetscCall(PetscViewerFlush(viewer));
19   PetscCall(PetscViewerASCIIPrintf(viewer, "Col IS %" PetscInt_FMT "\n", gs));
20   if (ss > -1) {
21     PetscCall(PetscViewerGetSubViewer(viewer, PetscObjectComm((PetscObject)rowis[ss]), &subviewer));
22     PetscCall(ISView(colis[ss], subviewer));
23     PetscCall(PetscViewerRestoreSubViewer(viewer, PetscObjectComm((PetscObject)rowis[ss]), &subviewer));
24   }
25   PetscCall(PetscViewerFlush(viewer));
26   PetscFunctionReturn(PETSC_SUCCESS);
27 }
28 
29 int main(int argc, char **args)
30 {
31   Mat          A, *submats;
32   MPI_Comm     subcomm;
33   PetscMPIInt  rank, size, subrank, subsize, color;
34   PetscInt     m, n, N, bs, rstart, rend, i, j, k, total_subdomains, hash, nsubdomains = 1;
35   PetscInt     nis, *cols, gnsubdomains, gsubdomainnums[1], gsubdomainperm[1], s, gs;
36   PetscInt    *rowindices, *colindices, idx, rep;
37   PetscScalar *vals;
38   IS           rowis[1], colis[1];
39   PetscViewer  viewer;
40   PetscBool    permute_indices, flg;
41 
42   PetscFunctionBeginUser;
43   PetscCall(PetscInitialize(&argc, &args, (char *)0, help));
44   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
45   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
46 
47   PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "ex183", "Mat");
48   m = 5;
49   PetscCall(PetscOptionsInt("-m", "Local matrix size", "MatSetSizes", m, &m, &flg));
50   total_subdomains = size - 1;
51   PetscCall(PetscOptionsRangeInt("-total_subdomains", "Number of submatrices where 0 < n < comm size", "MatCreateSubMatricesMPI", total_subdomains, &total_subdomains, &flg, 1, size));
52   permute_indices = PETSC_FALSE;
53   PetscCall(PetscOptionsBool("-permute_indices", "Whether to permute indices before breaking them into subdomains", "ISCreateGeneral", permute_indices, &permute_indices, &flg));
54   hash = 7;
55   PetscCall(PetscOptionsInt("-hash", "Permutation factor, which has to be relatively prime to M = size*m (total matrix size)", "ISCreateGeneral", hash, &hash, &flg));
56   rep = 2;
57   PetscCall(PetscOptionsRangeInt("-rep", "Number of times to carry out submatrix extractions; currently only 1 & 2 are supported", NULL, rep, &rep, &flg, 1, 2));
58   PetscOptionsEnd();
59 
60   viewer = PETSC_VIEWER_STDOUT_WORLD;
61   /* Create logically sparse, but effectively dense matrix for easy verification of submatrix extraction correctness. */
62   PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
63   PetscCall(MatSetSizes(A, m, m, PETSC_DECIDE, PETSC_DECIDE));
64   PetscCall(MatSetFromOptions(A));
65   PetscCall(MatSetUp(A));
66   PetscCall(MatGetSize(A, NULL, &N));
67   PetscCall(MatGetLocalSize(A, NULL, &n));
68   PetscCall(MatGetBlockSize(A, &bs));
69   PetscCall(MatSeqAIJSetPreallocation(A, n, NULL));
70   PetscCall(MatMPIAIJSetPreallocation(A, n, NULL, N - n, NULL));
71   PetscCall(MatSeqBAIJSetPreallocation(A, bs, n / bs, NULL));
72   PetscCall(MatMPIBAIJSetPreallocation(A, bs, n / bs, NULL, (N - n) / bs, NULL));
73   PetscCall(MatSeqSBAIJSetPreallocation(A, bs, n / bs, NULL));
74   PetscCall(MatMPISBAIJSetPreallocation(A, bs, n / bs, NULL, (N - n) / bs, NULL));
75 
76   PetscCall(PetscMalloc2(N, &cols, N, &vals));
77   PetscCall(MatGetOwnershipRange(A, &rstart, &rend));
78   for (j = 0; j < N; ++j) cols[j] = j;
79   for (i = rstart; i < rend; i++) {
80     for (j = 0; j < N; ++j) vals[j] = i * 10000 + j;
81     PetscCall(MatSetValues(A, 1, &i, N, cols, vals, INSERT_VALUES));
82   }
83   PetscCall(PetscFree2(cols, vals));
84   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
85   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
86 
87   PetscCall(PetscViewerASCIIPrintf(viewer, "Initial matrix:\n"));
88   PetscCall(MatView(A, viewer));
89 
90   /*
91      Create subcomms and ISs so that each rank participates in one IS.
92      The IS either coalesces adjacent rank indices (contiguous),
93      or selects indices by scrambling them using a hash.
94   */
95   k     = size / total_subdomains + (size % total_subdomains > 0); /* There are up to k ranks to a color */
96   color = rank / k;
97   PetscCallMPI(MPI_Comm_split(PETSC_COMM_WORLD, color, rank, &subcomm));
98   PetscCallMPI(MPI_Comm_size(subcomm, &subsize));
99   PetscCallMPI(MPI_Comm_rank(subcomm, &subrank));
100   PetscCall(MatGetOwnershipRange(A, &rstart, &rend));
101   nis = 1;
102   PetscCall(PetscMalloc2(rend - rstart, &rowindices, rend - rstart, &colindices));
103 
104   for (j = rstart; j < rend; ++j) {
105     if (permute_indices) {
106       idx = (j * hash);
107     } else {
108       idx = j;
109     }
110     rowindices[j - rstart] = idx % N;
111     colindices[j - rstart] = (idx + m) % N;
112   }
113   PetscCall(ISCreateGeneral(subcomm, rend - rstart, rowindices, PETSC_COPY_VALUES, &rowis[0]));
114   PetscCall(ISCreateGeneral(subcomm, rend - rstart, colindices, PETSC_COPY_VALUES, &colis[0]));
115   PetscCall(ISSort(rowis[0]));
116   PetscCall(ISSort(colis[0]));
117   PetscCall(PetscFree2(rowindices, colindices));
118   /*
119     Now view the ISs.  To avoid deadlock when viewing a list of objects on different subcomms,
120     we need to obtain the global numbers of our local objects and wait for the corresponding global
121     number to be viewed.
122   */
123   PetscCall(PetscViewerASCIIPrintf(viewer, "Subdomains"));
124   if (permute_indices) PetscCall(PetscViewerASCIIPrintf(viewer, " (hash=%" PetscInt_FMT ")", hash));
125   PetscCall(PetscViewerASCIIPrintf(viewer, ":\n"));
126   PetscCall(PetscViewerFlush(viewer));
127 
128   nsubdomains = 1;
129   for (s = 0; s < nsubdomains; ++s) gsubdomainperm[s] = s;
130   PetscCall(PetscObjectsListGetGlobalNumbering(PETSC_COMM_WORLD, 1, (PetscObject *)rowis, &gnsubdomains, gsubdomainnums));
131   PetscCall(PetscSortIntWithPermutation(nsubdomains, gsubdomainnums, gsubdomainperm));
132   for (gs = 0, s = 0; gs < gnsubdomains; ++gs) {
133     PetscInt ss;
134     if (s < nsubdomains) {
135       ss = gsubdomainperm[s];
136       if (gs == gsubdomainnums[ss]) { /* Global subdomain gs being viewed is my subdomain with local number ss. */
137         ++s;
138       } else ss = -1;
139     } else ss = -1;
140     PetscCall(MyISView(rowis, colis, gs, ss, viewer));
141   }
142   PetscCall(PetscViewerFlush(viewer));
143   PetscCall(ISSort(rowis[0]));
144   PetscCall(ISSort(colis[0]));
145   nsubdomains = 1;
146   PetscCall(MatCreateSubMatricesMPI(A, nsubdomains, rowis, colis, MAT_INITIAL_MATRIX, &submats));
147   /*
148     Now view the matrices.  To avoid deadlock when viewing a list of objects on different subcomms,
149     we need to obtain the global numbers of our local objects and wait for the corresponding global
150     number to be viewed. Also all MPI processes need to call PetscViewerGetSubViewer() the same number of times
151   */
152   PetscCall(PetscViewerASCIIPrintf(viewer, "Submatrices (repetition 1):\n"));
153   for (s = 0; s < nsubdomains; ++s) gsubdomainperm[s] = s;
154   PetscCall(PetscObjectsListGetGlobalNumbering(PETSC_COMM_WORLD, 1, (PetscObject *)submats, &gnsubdomains, gsubdomainnums));
155   PetscCall(PetscSortIntWithPermutation(nsubdomains, gsubdomainnums, gsubdomainperm));
156   for (gs = 0, s = 0; gs < gnsubdomains; ++gs) {
157     PetscViewer subviewer = NULL;
158     if (s < nsubdomains) {
159       PetscInt ss;
160       ss = gsubdomainperm[s];
161       if (gs == gsubdomainnums[ss]) { /* Global subdomain gs being viewed is my subdomain with local number ss. */
162         PetscCall(PetscViewerGetSubViewer(viewer, PetscObjectComm((PetscObject)submats[ss]), &subviewer));
163         PetscCall(MatView(submats[ss], subviewer));
164         PetscCall(PetscViewerRestoreSubViewer(viewer, PetscObjectComm((PetscObject)submats[ss]), &subviewer));
165         ++s;
166       } else {
167         PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &subviewer));
168         PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &subviewer));
169       }
170     } else {
171       PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &subviewer));
172       PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &subviewer));
173     }
174   }
175   if (rep == 1) goto cleanup;
176   nsubdomains = 1;
177   PetscCall(MatCreateSubMatricesMPI(A, nsubdomains, rowis, colis, MAT_REUSE_MATRIX, &submats));
178   /*
179     Now view the matrices.  To avoid deadlock when viewing a list of objects on different subcomms,
180     we need to obtain the global numbers of our local objects and wait for the corresponding global
181     number to be viewed.
182   */
183   PetscCall(PetscViewerASCIIPrintf(viewer, "Submatrices (repetition 2):\n"));
184   for (s = 0; s < nsubdomains; ++s) gsubdomainperm[s] = s;
185   PetscCall(PetscObjectsListGetGlobalNumbering(PETSC_COMM_WORLD, 1, (PetscObject *)submats, &gnsubdomains, gsubdomainnums));
186   PetscCall(PetscSortIntWithPermutation(nsubdomains, gsubdomainnums, gsubdomainperm));
187   for (gs = 0, s = 0; gs < gnsubdomains; ++gs) {
188     PetscViewer subviewer = NULL;
189     if (s < nsubdomains) {
190       PetscInt ss;
191       ss = gsubdomainperm[s];
192       if (gs == gsubdomainnums[ss]) { /* Global subdomain gs being viewed is my subdomain with local number ss. */
193         PetscCall(PetscViewerGetSubViewer(viewer, PetscObjectComm((PetscObject)submats[ss]), &subviewer));
194         PetscCall(MatView(submats[ss], subviewer));
195         PetscCall(PetscViewerRestoreSubViewer(viewer, PetscObjectComm((PetscObject)submats[ss]), &subviewer));
196         ++s;
197       } else {
198         PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &subviewer));
199         PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &subviewer));
200       }
201     } else {
202       PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &subviewer));
203       PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &subviewer));
204     }
205   }
206 cleanup:
207   for (k = 0; k < nsubdomains; ++k) PetscCall(MatDestroy(submats + k));
208   PetscCall(PetscFree(submats));
209   for (k = 0; k < nis; ++k) {
210     PetscCall(ISDestroy(rowis + k));
211     PetscCall(ISDestroy(colis + k));
212   }
213   PetscCall(MatDestroy(&A));
214   PetscCallMPI(MPI_Comm_free(&subcomm));
215   PetscCall(PetscFinalize());
216   return 0;
217 }
218 
219 /*TEST
220 
221    test:
222       nsize: 2
223       args: -total_subdomains 1
224       output_file: output/ex183_2_1.out
225 
226    test:
227       suffix: 2
228       nsize: 3
229       args: -total_subdomains 2
230       output_file: output/ex183_3_2.out
231 
232    test:
233       suffix: 3
234       nsize: 4
235       args: -total_subdomains 2
236       output_file: output/ex183_4_2.out
237 
238    test:
239       suffix: 4
240       nsize: 6
241       args: -total_subdomains 2
242       output_file: output/ex183_6_2.out
243 
244 TEST*/
245