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) if (lrows[r] >= 0) lrows[len++] = r;
 43:   }
 44:   if (nr) *nr = len;
 45:   if (olrows) *olrows = lrows;
 46:   return 0;
 47: }