1 const char help[] = "Test MatCreateLocalRef()\n\n"; 2 3 #include <petscmat.h> 4 5 static PetscErrorCode GetLocalRef(Mat A, IS isrow, IS iscol, Mat *B) 6 { 7 IS istmp; 8 9 PetscFunctionBegin; 10 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Extracting LocalRef with isrow:\n")); 11 PetscCall(ISOnComm(isrow, PETSC_COMM_WORLD, PETSC_COPY_VALUES, &istmp)); 12 PetscCall(ISView(istmp, PETSC_VIEWER_STDOUT_WORLD)); 13 PetscCall(ISDestroy(&istmp)); 14 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Extracting LocalRef with iscol (only rank=0 shown):\n")); 15 PetscCall(ISOnComm(iscol, PETSC_COMM_WORLD, PETSC_COPY_VALUES, &istmp)); 16 PetscCall(ISView(istmp, PETSC_VIEWER_STDOUT_WORLD)); 17 PetscCall(ISDestroy(&istmp)); 18 19 PetscCall(MatCreateLocalRef(A, isrow, iscol, B)); 20 PetscFunctionReturn(PETSC_SUCCESS); 21 } 22 23 int main(int argc, char *argv[]) 24 { 25 MPI_Comm comm; 26 Mat J, B; 27 PetscInt i, j, k, l, rstart, rend, m, n, top_bs, row_bs, col_bs, nlocblocks, *idx, nrowblocks, ncolblocks, *ridx, *cidx, *icol, *irow; 28 PetscScalar *vals; 29 ISLocalToGlobalMapping brmap; 30 IS is0, is1; 31 PetscBool diag, blocked; 32 33 PetscFunctionBeginUser; 34 PetscCall(PetscInitialize(&argc, &argv, 0, help)); 35 comm = PETSC_COMM_WORLD; 36 37 PetscOptionsBegin(comm, NULL, "LocalRef Test Options", NULL); 38 { 39 top_bs = 2; 40 row_bs = 2; 41 col_bs = 2; 42 diag = PETSC_FALSE; 43 blocked = PETSC_FALSE; 44 PetscCall(PetscOptionsInt("-top_bs", "Block size of top-level matrix", 0, top_bs, &top_bs, NULL)); 45 PetscCall(PetscOptionsInt("-row_bs", "Block size of row map", 0, row_bs, &row_bs, NULL)); 46 PetscCall(PetscOptionsInt("-col_bs", "Block size of col map", 0, col_bs, &col_bs, NULL)); 47 PetscCall(PetscOptionsBool("-diag", "Extract a diagonal black", 0, diag, &diag, NULL)); 48 PetscCall(PetscOptionsBool("-blocked", "Use block insertion", 0, blocked, &blocked, NULL)); 49 } 50 PetscOptionsEnd(); 51 52 PetscCall(MatCreate(comm, &J)); 53 PetscCall(MatSetSizes(J, 6, 6, PETSC_DETERMINE, PETSC_DETERMINE)); 54 PetscCall(MatSetBlockSize(J, top_bs)); 55 PetscCall(MatSetFromOptions(J)); 56 PetscCall(MatSeqBAIJSetPreallocation(J, top_bs, PETSC_DECIDE, 0)); 57 PetscCall(MatMPIBAIJSetPreallocation(J, top_bs, PETSC_DECIDE, 0, PETSC_DECIDE, 0)); 58 PetscCall(MatSetUp(J)); 59 PetscCall(MatGetSize(J, &m, &n)); 60 PetscCall(MatGetOwnershipRange(J, &rstart, &rend)); 61 62 nlocblocks = (rend - rstart) / top_bs + 2; 63 PetscCall(PetscMalloc1(nlocblocks, &idx)); 64 for (i = 0; i < nlocblocks; i++) idx[i] = (rstart / top_bs + i - 1 + m / top_bs) % (m / top_bs); 65 PetscCall(ISLocalToGlobalMappingCreate(comm, top_bs, nlocblocks, idx, PETSC_OWN_POINTER, &brmap)); 66 PetscCall(PetscPrintf(comm, "Block ISLocalToGlobalMapping:\n")); 67 PetscCall(ISLocalToGlobalMappingView(brmap, PETSC_VIEWER_STDOUT_WORLD)); 68 69 PetscCall(MatSetLocalToGlobalMapping(J, brmap, brmap)); 70 PetscCall(ISLocalToGlobalMappingDestroy(&brmap)); 71 72 /* Create index sets for local submatrix */ 73 nrowblocks = (rend - rstart) / row_bs; 74 ncolblocks = (rend - rstart) / col_bs; 75 PetscCall(PetscMalloc2(nrowblocks, &ridx, ncolblocks, &cidx)); 76 for (i = 0; i < nrowblocks; i++) ridx[i] = i + ((i > nrowblocks / 2) ^ !rstart); 77 for (i = 0; i < ncolblocks; i++) cidx[i] = i + ((i < ncolblocks / 2) ^ !!rstart); 78 PetscCall(ISCreateBlock(PETSC_COMM_SELF, row_bs, nrowblocks, ridx, PETSC_COPY_VALUES, &is0)); 79 PetscCall(ISCreateBlock(PETSC_COMM_SELF, col_bs, ncolblocks, cidx, PETSC_COPY_VALUES, &is1)); 80 PetscCall(PetscFree2(ridx, cidx)); 81 82 if (diag) { 83 PetscCall(ISDestroy(&is1)); 84 PetscCall(PetscObjectReference((PetscObject)is0)); 85 is1 = is0; 86 ncolblocks = nrowblocks; 87 } 88 89 PetscCall(GetLocalRef(J, is0, is1, &B)); 90 91 PetscCall(MatZeroEntries(J)); 92 93 PetscCall(PetscMalloc3(row_bs, &irow, col_bs, &icol, row_bs * col_bs, &vals)); 94 for (i = 0; i < nrowblocks; i++) { 95 for (j = 0; j < ncolblocks; j++) { 96 for (k = 0; k < row_bs; k++) { 97 for (l = 0; l < col_bs; l++) vals[k * col_bs + l] = i * 100000 + j * 1000 + k * 10 + l; 98 irow[k] = i * row_bs + k; 99 } 100 for (l = 0; l < col_bs; l++) icol[l] = j * col_bs + l; 101 if (blocked) { 102 PetscCall(MatSetValuesBlockedLocal(B, 1, &i, 1, &j, vals, ADD_VALUES)); 103 } else { 104 PetscCall(MatSetValuesLocal(B, row_bs, irow, col_bs, icol, vals, ADD_VALUES)); 105 } 106 } 107 } 108 PetscCall(PetscFree3(irow, icol, vals)); 109 110 PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY)); 111 PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY)); 112 113 PetscCall(MatView(J, PETSC_VIEWER_STDOUT_WORLD)); 114 115 PetscCall(ISDestroy(&is0)); 116 PetscCall(ISDestroy(&is1)); 117 PetscCall(MatDestroy(&B)); 118 PetscCall(MatDestroy(&J)); 119 PetscCall(PetscFinalize()); 120 return 0; 121 } 122 123 /*TEST 124 125 test: 126 nsize: 2 127 filter: grep -v "type: mpi" 128 args: -blocked {{0 1}} -mat_type {{aij baij}} 129 130 TEST*/ 131