1 #include <petsc/private/matimpl.h>
2 #include <petscsf.h>
3
4 /* this function maps rows to locally owned rows */
MatZeroRowsMapLocal_Private(Mat A,PetscInt N,const PetscInt * rows,PetscInt * nr,PetscInt ** olrows)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;
14
15 PetscFunctionBegin;
16 /* Create SF where leaves are input rows and roots are owned rows */
17 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)A), &rank));
18 PetscCall(PetscMalloc1(n, &lrows));
19 for (r = 0; r < n; ++r) lrows[r] = -1;
20 if (!A->nooffproczerorows) PetscCall(PetscMalloc1(N, &rrows));
21 for (r = 0; r < N; ++r) {
22 const PetscInt idx = rows[r];
23 PetscCheck(idx >= 0 && A->rmap->N > idx, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Row %" PetscInt_FMT " out of range [0,%" PetscInt_FMT ")", idx, A->rmap->N);
24 if (idx < owners[p] || owners[p + 1] <= idx) { /* short-circuit the search if the last p owns this row too */
25 PetscCall(PetscLayoutFindOwner(A->rmap, idx, &p));
26 }
27 if (A->nooffproczerorows) {
28 PetscCheck(p == rank, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "MAT_NO_OFF_PROC_ZERO_ROWS set, but row %" PetscInt_FMT " is not owned by rank %d", idx, rank);
29 lrows[len++] = idx - owners[p];
30 } else {
31 rrows[r].rank = p;
32 rrows[r].index = rows[r] - owners[p];
33 }
34 }
35 if (!A->nooffproczerorows) {
36 PetscCall(PetscSFCreate(PetscObjectComm((PetscObject)A), &sf));
37 PetscCall(PetscSFSetGraph(sf, n, N, NULL, PETSC_OWN_POINTER, rrows, PETSC_OWN_POINTER));
38 /* Collect flags for rows to be zeroed */
39 PetscCall(PetscSFReduceBegin(sf, MPIU_INT, (PetscInt *)rows, lrows, MPI_LOR));
40 PetscCall(PetscSFReduceEnd(sf, MPIU_INT, (PetscInt *)rows, lrows, MPI_LOR));
41 PetscCall(PetscSFDestroy(&sf));
42 /* Compress and put in row numbers */
43 for (r = 0; r < n; ++r)
44 if (lrows[r] >= 0) lrows[len++] = r;
45 }
46 if (nr) *nr = len;
47 if (olrows) *olrows = lrows;
48 PetscFunctionReturn(PETSC_SUCCESS);
49 }
50