1 #include <petsc/private/matimpl.h> 2 #include <petscsf.h> 3 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 PetscInt *owners = A->rmap->range; 7 PetscInt n = A->rmap->n; 8 PetscSF sf; 9 PetscInt *lrows; 10 PetscSFNode *rrows; 11 PetscMPIInt rank, p = 0; 12 PetscInt r, len = 0; 13 14 PetscFunctionBegin; 15 /* Create SF where leaves are input rows and roots are owned rows */ 16 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)A), &rank)); 17 PetscCall(PetscMalloc1(n, &lrows)); 18 for (r = 0; r < n; ++r) lrows[r] = -1; 19 if (!A->nooffproczerorows) PetscCall(PetscMalloc1(N, &rrows)); 20 for (r = 0; r < N; ++r) { 21 const PetscInt idx = rows[r]; 22 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); 23 if (idx < owners[p] || owners[p + 1] <= idx) { /* short-circuit the search if the last p owns this row too */ 24 PetscCall(PetscLayoutFindOwner(A->rmap, idx, &p)); 25 } 26 if (A->nooffproczerorows) { 27 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); 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 PetscCall(PetscSFCreate(PetscObjectComm((PetscObject)A), &sf)); 36 PetscCall(PetscSFSetGraph(sf, n, N, NULL, PETSC_OWN_POINTER, rrows, PETSC_OWN_POINTER)); 37 /* Collect flags for rows to be zeroed */ 38 PetscCall(PetscSFReduceBegin(sf, MPIU_INT, (PetscInt *)rows, lrows, MPI_LOR)); 39 PetscCall(PetscSFReduceEnd(sf, MPIU_INT, (PetscInt *)rows, lrows, MPI_LOR)); 40 PetscCall(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 PetscFunctionReturn(0); 48 } 49