1 const char help[] = "Test MatCreateLocalRef()\n\n";
2
3 #include <petscmat.h>
4
GetLocalRef(Mat A,IS isrow,IS iscol,Mat * B)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
main(int argc,char * argv[])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