Actual source code: zerorows.c
1: #include <petsc/private/matimpl.h>
2: #include <petscsf.h>
4: /* this function maps rows to locally owned rows */
5: PETSC_INTERN PetscErrorCode MatZeroRowsMapLocal_Private(Mat A, PetscInt N, const PetscInt *rows, PetscInt *nr, PetscInt **olrows)
6: {
7: PetscInt *owners = A->rmap->range;
8: PetscInt n = A->rmap->n;
9: PetscSF sf;
10: PetscInt *lrows;
11: PetscSFNode *rrows;
12: PetscMPIInt rank, p = 0;
13: PetscInt r, len = 0;
15: /* Create SF where leaves are input rows and roots are owned rows */
16: MPI_Comm_rank(PetscObjectComm((PetscObject)A), &rank);
17: PetscMalloc1(n, &lrows);
18: for (r = 0; r < n; ++r) lrows[r] = -1;
19: if (!A->nooffproczerorows) PetscMalloc1(N, &rrows);
20: for (r = 0; r < N; ++r) {
21: const PetscInt idx = rows[r];
23: if (idx < owners[p] || owners[p + 1] <= idx) { /* short-circuit the search if the last p owns this row too */
24: PetscLayoutFindOwner(A->rmap, idx, &p);
25: }
26: if (A->nooffproczerorows) {
28: lrows[len++] = idx - owners[p];
29: } else {
30: rrows[r].rank = p;
31: rrows[r].index = rows[r] - owners[p];
32: }
33: }
34: if (!A->nooffproczerorows) {
35: PetscSFCreate(PetscObjectComm((PetscObject)A), &sf);
36: PetscSFSetGraph(sf, n, N, NULL, PETSC_OWN_POINTER, rrows, PETSC_OWN_POINTER);
37: /* Collect flags for rows to be zeroed */
38: PetscSFReduceBegin(sf, MPIU_INT, (PetscInt *)rows, lrows, MPI_LOR);
39: PetscSFReduceEnd(sf, MPIU_INT, (PetscInt *)rows, lrows, MPI_LOR);
40: PetscSFDestroy(&sf);
41: /* Compress and put in row numbers */
42: for (r = 0; r < n; ++r)
43: if (lrows[r] >= 0) lrows[len++] = r;
44: }
45: if (nr) *nr = len;
46: if (olrows) *olrows = lrows;
47: return 0;
48: }