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;

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