Actual source code: zerorows.c
petsc-3.13.6 2020-09-29
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;
17: /* Create SF where leaves are input rows and roots are owned rows */
18: MPI_Comm_rank(PetscObjectComm((PetscObject)A),&rank);
19: PetscMalloc1(n, &lrows);
20: for (r = 0; r < n; ++r) lrows[r] = -1;
21: if (!A->nooffproczerorows) {PetscMalloc1(N, &rrows);}
22: for (r = 0; r < N; ++r) {
23: const PetscInt idx = rows[r];
24: if (idx < 0 || A->rmap->N <= idx) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Row %D out of range [0,%D)",idx,A->rmap->N);
25: if (idx < owners[p] || owners[p+1] <= idx) { /* short-circuit the search if the last p owns this row too */
26: PetscLayoutFindOwner(A->rmap,idx,&p);
27: }
28: if (A->nooffproczerorows) {
29: if (p != rank) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"MAT_NO_OFF_PROC_ZERO_ROWS set, but row %D is not owned by rank %d",idx,rank);
30: lrows[len++] = idx - owners[p];
31: } else {
32: rrows[r].rank = p;
33: rrows[r].index = rows[r] - owners[p];
34: }
35: }
36: if (!A->nooffproczerorows) {
37: PetscSFCreate(PetscObjectComm((PetscObject) A), &sf);
38: PetscSFSetGraph(sf, n, N, NULL, PETSC_OWN_POINTER, rrows, PETSC_OWN_POINTER);
39: /* Collect flags for rows to be zeroed */
40: PetscSFReduceBegin(sf, MPIU_INT, (PetscInt*)rows, lrows, MPI_LOR);
41: PetscSFReduceEnd(sf, MPIU_INT, (PetscInt*)rows, lrows, MPI_LOR);
42: PetscSFDestroy(&sf);
43: /* Compress and put in row numbers */
44: for (r = 0; r < n; ++r) if (lrows[r] >= 0) lrows[len++] = r;
45: }
46: if (nr) *nr = len;
47: if (olrows) *olrows = lrows;
48: return(0);
49: }