1 #include <petsc/private/matimpl.h> 2 #include <petscsf.h> 3 4 /* this function maps rows to locally owned rows */ 5 #undef __FUNCT__ 6 #define __FUNCT__ "MatZeroRowsMapLocal_Private" 7 PETSC_INTERN PetscErrorCode MatZeroRowsMapLocal_Private(Mat A,PetscInt N,const PetscInt *rows,PetscInt *nr,PetscInt **olrows) 8 { 9 PetscInt *owners = A->rmap->range; 10 PetscInt n = A->rmap->n; 11 PetscSF sf; 12 PetscInt *lrows; 13 PetscSFNode *rrows; 14 PetscMPIInt rank; 15 PetscInt r, p = 0, len = 0; 16 PetscErrorCode ierr; 17 18 PetscFunctionBegin; 19 /* Create SF where leaves are input rows and roots are owned rows */ 20 ierr = MPI_Comm_rank(PetscObjectComm((PetscObject)A),&rank);CHKERRQ(ierr); 21 ierr = PetscMalloc1(n, &lrows);CHKERRQ(ierr); 22 for (r = 0; r < n; ++r) lrows[r] = -1; 23 if (!A->nooffproczerorows) {ierr = PetscMalloc1(N, &rrows);CHKERRQ(ierr);} 24 for (r = 0; r < N; ++r) { 25 const PetscInt idx = rows[r]; 26 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); 27 if (idx < owners[p] || owners[p+1] <= idx) { /* short-circuit the search if the last p owns this row too */ 28 ierr = PetscLayoutFindOwner(A->rmap,idx,&p);CHKERRQ(ierr); 29 } 30 if (A->nooffproczerorows) { 31 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); 32 lrows[len++] = idx - owners[p]; 33 } else { 34 rrows[r].rank = p; 35 rrows[r].index = rows[r] - owners[p]; 36 } 37 } 38 if (!A->nooffproczerorows) { 39 ierr = PetscSFCreate(PetscObjectComm((PetscObject) A), &sf);CHKERRQ(ierr); 40 ierr = PetscSFSetGraph(sf, n, N, NULL, PETSC_OWN_POINTER, rrows, PETSC_OWN_POINTER);CHKERRQ(ierr); 41 /* Collect flags for rows to be zeroed */ 42 ierr = PetscSFReduceBegin(sf, MPIU_INT, (PetscInt*)rows, lrows, MPI_LOR);CHKERRQ(ierr); 43 ierr = PetscSFReduceEnd(sf, MPIU_INT, (PetscInt*)rows, lrows, MPI_LOR);CHKERRQ(ierr); 44 ierr = PetscSFDestroy(&sf);CHKERRQ(ierr); 45 /* Compress and put in row numbers */ 46 for (r = 0; r < n; ++r) if (lrows[r] >= 0) lrows[len++] = r; 47 } 48 if (nr) *nr = len; 49 if (olrows) *olrows = lrows; 50 PetscFunctionReturn(0); 51 } 52 53