16e520ac8SStefano Zampini #include <petsc/private/matimpl.h> 26e520ac8SStefano Zampini #include <petscsf.h> 36e520ac8SStefano Zampini 46e520ac8SStefano Zampini /* this function maps rows to locally owned rows */ 5d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode MatZeroRowsMapLocal_Private(Mat A, PetscInt N, const PetscInt *rows, PetscInt *nr, PetscInt **olrows) 6d71ae5a4SJacob Faibussowitsch { 76e520ac8SStefano Zampini PetscInt *owners = A->rmap->range; 86e520ac8SStefano Zampini PetscInt n = A->rmap->n; 96e520ac8SStefano Zampini PetscSF sf; 106e520ac8SStefano Zampini PetscInt *lrows; 116e520ac8SStefano Zampini PetscSFNode *rrows; 12131c27b5Sprj- PetscMPIInt rank, p = 0; 13131c27b5Sprj- PetscInt r, len = 0; 146e520ac8SStefano Zampini 156e520ac8SStefano Zampini PetscFunctionBegin; 166e520ac8SStefano Zampini /* Create SF where leaves are input rows and roots are owned rows */ 179566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)A), &rank)); 189566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(n, &lrows)); 196e520ac8SStefano Zampini for (r = 0; r < n; ++r) lrows[r] = -1; 209566063dSJacob Faibussowitsch if (!A->nooffproczerorows) PetscCall(PetscMalloc1(N, &rrows)); 216e520ac8SStefano Zampini for (r = 0; r < N; ++r) { 226e520ac8SStefano Zampini const PetscInt idx = rows[r]; 23aed4548fSBarry Smith 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); 246e520ac8SStefano Zampini if (idx < owners[p] || owners[p + 1] <= idx) { /* short-circuit the search if the last p owns this row too */ 259566063dSJacob Faibussowitsch PetscCall(PetscLayoutFindOwner(A->rmap, idx, &p)); 266e520ac8SStefano Zampini } 276e520ac8SStefano Zampini if (A->nooffproczerorows) { 2808401ef6SPierre Jolivet 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); 296e520ac8SStefano Zampini lrows[len++] = idx - owners[p]; 306e520ac8SStefano Zampini } else { 316e520ac8SStefano Zampini rrows[r].rank = p; 326e520ac8SStefano Zampini rrows[r].index = rows[r] - owners[p]; 336e520ac8SStefano Zampini } 346e520ac8SStefano Zampini } 356e520ac8SStefano Zampini if (!A->nooffproczerorows) { 369566063dSJacob Faibussowitsch PetscCall(PetscSFCreate(PetscObjectComm((PetscObject)A), &sf)); 379566063dSJacob Faibussowitsch PetscCall(PetscSFSetGraph(sf, n, N, NULL, PETSC_OWN_POINTER, rrows, PETSC_OWN_POINTER)); 386e520ac8SStefano Zampini /* Collect flags for rows to be zeroed */ 399566063dSJacob Faibussowitsch PetscCall(PetscSFReduceBegin(sf, MPIU_INT, (PetscInt *)rows, lrows, MPI_LOR)); 409566063dSJacob Faibussowitsch PetscCall(PetscSFReduceEnd(sf, MPIU_INT, (PetscInt *)rows, lrows, MPI_LOR)); 419566063dSJacob Faibussowitsch PetscCall(PetscSFDestroy(&sf)); 426e520ac8SStefano Zampini /* Compress and put in row numbers */ 439371c9d4SSatish Balay for (r = 0; r < n; ++r) 449371c9d4SSatish Balay if (lrows[r] >= 0) lrows[len++] = r; 456e520ac8SStefano Zampini } 466e520ac8SStefano Zampini if (nr) *nr = len; 476e520ac8SStefano Zampini if (olrows) *olrows = lrows; 48*3ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 496e520ac8SStefano Zampini } 50