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:   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: }