1 #include <../src/mat/impls/aij/mpi/mpiaij.h> 2 3 PetscErrorCode MatGetMultiProcBlock_MPIAIJ(Mat mat, MPI_Comm subComm, MatReuse scall, Mat *subMat) 4 { 5 Mat_MPIAIJ *aij = (Mat_MPIAIJ *)mat->data; 6 Mat_SeqAIJ *aijB = (Mat_SeqAIJ *)aij->B->data; 7 PetscMPIInt subCommSize, subCommRank; 8 PetscMPIInt *commRankMap, subRank, rank, commRank; 9 PetscInt *garrayCMap, col, i, j, *nnz, newRow, newCol; 10 11 PetscFunctionBegin; 12 PetscCallMPI(MPI_Comm_size(subComm, &subCommSize)); 13 PetscCallMPI(MPI_Comm_rank(subComm, &subCommRank)); 14 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)mat), &commRank)); 15 16 /* create subMat object with the relevant layout */ 17 if (scall == MAT_INITIAL_MATRIX) { 18 PetscCall(MatCreate(subComm, subMat)); 19 PetscCall(MatSetType(*subMat, MATMPIAIJ)); 20 PetscCall(MatSetSizes(*subMat, mat->rmap->n, mat->cmap->n, PETSC_DECIDE, PETSC_DECIDE)); 21 PetscCall(MatSetBlockSizesFromMats(*subMat, mat, mat)); 22 23 /* need to setup rmap and cmap before Preallocation */ 24 PetscCall(PetscLayoutSetUp((*subMat)->rmap)); 25 PetscCall(PetscLayoutSetUp((*subMat)->cmap)); 26 } 27 28 /* create a map of comm_rank from subComm to comm - should commRankMap and garrayCMap be kept for reused? */ 29 PetscCall(PetscMalloc1(subCommSize, &commRankMap)); 30 PetscCallMPI(MPI_Allgather(&commRank, 1, MPI_INT, commRankMap, 1, MPI_INT, subComm)); 31 32 /* Traverse garray and identify column indices [of offdiag mat] that 33 should be discarded. For the ones not discarded, store the newCol+1 34 value in garrayCMap */ 35 PetscCall(PetscCalloc1(aij->B->cmap->n, &garrayCMap)); 36 for (i = 0; i < aij->B->cmap->n; i++) { 37 col = aij->garray[i]; 38 for (subRank = 0; subRank < subCommSize; subRank++) { 39 rank = commRankMap[subRank]; 40 if ((col >= mat->cmap->range[rank]) && (col < mat->cmap->range[rank + 1])) { 41 garrayCMap[i] = (*subMat)->cmap->range[subRank] + col - mat->cmap->range[rank] + 1; 42 break; 43 } 44 } 45 } 46 47 if (scall == MAT_INITIAL_MATRIX) { 48 /* Compute preallocation for the offdiag mat */ 49 PetscCall(PetscCalloc1(aij->B->rmap->n, &nnz)); 50 for (i = 0; i < aij->B->rmap->n; i++) { 51 for (j = aijB->i[i]; j < aijB->i[i + 1]; j++) { 52 if (garrayCMap[aijB->j[j]]) nnz[i]++; 53 } 54 } 55 PetscCall(MatMPIAIJSetPreallocation(*(subMat), 0, NULL, 0, nnz)); 56 57 /* reuse diag block with the new submat */ 58 PetscCall(MatDestroy(&((Mat_MPIAIJ *)((*subMat)->data))->A)); 59 ((Mat_MPIAIJ *)((*subMat)->data))->A = aij->A; 60 PetscCall(PetscObjectReference((PetscObject)aij->A)); 61 } else if (((Mat_MPIAIJ *)(*subMat)->data)->A != aij->A) { 62 PetscObject obj = (PetscObject)((Mat_MPIAIJ *)((*subMat)->data))->A; 63 PetscCall(PetscObjectReference((PetscObject)obj)); 64 ((Mat_MPIAIJ *)((*subMat)->data))->A = aij->A; 65 PetscCall(PetscObjectReference((PetscObject)aij->A)); 66 } 67 68 /* Traverse aij->B and insert values into subMat */ 69 if ((*subMat)->assembled) { 70 (*subMat)->was_assembled = PETSC_TRUE; 71 (*subMat)->assembled = PETSC_FALSE; 72 } 73 for (i = 0; i < aij->B->rmap->n; i++) { 74 newRow = (*subMat)->rmap->range[subCommRank] + i; 75 for (j = aijB->i[i]; j < aijB->i[i + 1]; j++) { 76 newCol = garrayCMap[aijB->j[j]]; 77 if (newCol) { 78 newCol--; /* remove the increment */ 79 PetscCall(MatSetValues_MPIAIJ(*subMat, 1, &newRow, 1, &newCol, (aijB->a + j), INSERT_VALUES)); 80 } 81 } 82 } 83 PetscCall(MatAssemblyBegin(*subMat, MAT_FINAL_ASSEMBLY)); 84 PetscCall(MatAssemblyEnd(*subMat, MAT_FINAL_ASSEMBLY)); 85 86 /* deallocate temporary data */ 87 PetscCall(PetscFree(commRankMap)); 88 PetscCall(PetscFree(garrayCMap)); 89 if (scall == MAT_INITIAL_MATRIX) PetscCall(PetscFree(nnz)); 90 PetscFunctionReturn(PETSC_SUCCESS); 91 } 92