Actual source code: ex123.c

  1: static char help[] = "Test MatSetPreallocationCOO and MatSetValuesCOO\n\n";

  3: #include <petscmat.h>

  5: int main(int argc, char **args)
  6: {
  7:   Mat                    A, At, AAt;
  8:   Vec                    x, y, z;
  9:   ISLocalToGlobalMapping rl2g, cl2g;
 10:   IS                     is;
 11:   PetscLayout            rmap, cmap;
 12:   PetscInt              *it, *jt;
 13:   PetscInt               n1 = 11, n2 = 9;
 14:   PetscInt               i1[] = {7, 6, 2, 0, 4, 1, 1, 0, 2, 2, 1, -1, -1};
 15:   PetscInt               j1[] = {1, 4, 3, 5, 3, 3, 4, 5, 0, 3, 1, -1, -1};
 16:   PetscInt               i2[] = {7, 6, 2, 0, 4, 1, 1, 2, 1, -1, -1};
 17:   PetscInt               j2[] = {1, 4, 3, 5, 3, 3, 4, 0, 1, -1, -1};
 18:   PetscScalar            v1[] = {-1., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., PETSC_MAX_REAL, PETSC_MAX_REAL};
 19:   PetscScalar            v2[] = {1., -1., -2., -3., -4., -5., -6., -7., -8., -9., -10., PETSC_MAX_REAL, PETSC_MAX_REAL};
 20:   PetscInt               N = 6, m = 8, M, rstart, cstart, i;
 21:   PetscMPIInt            size;
 22:   PetscBool              loc      = PETSC_FALSE;
 23:   PetscBool              locdiag  = PETSC_TRUE;
 24:   PetscBool              localapi = PETSC_FALSE;
 25:   PetscBool              neg      = PETSC_FALSE;
 26:   PetscBool              ismatis, ismpiaij;

 28:   PetscFunctionBeginUser;
 29:   PetscCall(PetscInitialize(&argc, &args, (char *)0, help));
 30:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-neg", &neg, NULL));
 31:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-loc", &loc, NULL));
 32:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-locdiag", &locdiag, NULL));
 33:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-localapi", &localapi, NULL));
 34:   PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
 35:   if (loc) {
 36:     if (locdiag) {
 37:       PetscCall(MatSetSizes(A, m, N, PETSC_DECIDE, PETSC_DECIDE));
 38:     } else {
 39:       PetscCall(MatSetSizes(A, m, m + N, PETSC_DECIDE, PETSC_DECIDE));
 40:     }
 41:   } else {
 42:     PetscCall(MatSetSizes(A, m, PETSC_DECIDE, PETSC_DECIDE, N));
 43:   }
 44:   PetscCall(MatSetFromOptions(A));
 45:   PetscCall(MatGetLayouts(A, &rmap, &cmap));
 46:   PetscCall(PetscLayoutSetUp(rmap));
 47:   PetscCall(PetscLayoutSetUp(cmap));
 48:   PetscCall(PetscLayoutGetRange(rmap, &rstart, NULL));
 49:   PetscCall(PetscLayoutGetRange(cmap, &cstart, NULL));
 50:   PetscCall(PetscLayoutGetSize(rmap, &M));
 51:   PetscCall(PetscLayoutGetSize(cmap, &N));

 53:   PetscCall(PetscObjectTypeCompare((PetscObject)A, MATIS, &ismatis));

 55:   /* create fake l2g maps to test the local API */
 56:   PetscCall(ISCreateStride(PETSC_COMM_WORLD, M - rstart, rstart, 1, &is));
 57:   PetscCall(ISLocalToGlobalMappingCreateIS(is, &rl2g));
 58:   PetscCall(ISDestroy(&is));
 59:   PetscCall(ISCreateStride(PETSC_COMM_WORLD, N, 0, 1, &is));
 60:   PetscCall(ISLocalToGlobalMappingCreateIS(is, &cl2g));
 61:   PetscCall(ISDestroy(&is));
 62:   PetscCall(MatSetLocalToGlobalMapping(A, rl2g, cl2g));
 63:   PetscCall(ISLocalToGlobalMappingDestroy(&rl2g));
 64:   PetscCall(ISLocalToGlobalMappingDestroy(&cl2g));

 66:   PetscCall(MatCreateVecs(A, &x, &y));
 67:   PetscCall(MatCreateVecs(A, NULL, &z));
 68:   PetscCall(VecSet(x, 1.));
 69:   PetscCall(VecSet(z, 2.));
 70:   if (!localapi)
 71:     for (i = 0; i < n1; i++) i1[i] += rstart;
 72:   if (!localapi)
 73:     for (i = 0; i < n2; i++) i2[i] += rstart;
 74:   if (loc) {
 75:     if (locdiag) {
 76:       for (i = 0; i < n1; i++) j1[i] += cstart;
 77:       for (i = 0; i < n2; i++) j2[i] += cstart;
 78:     } else {
 79:       for (i = 0; i < n1; i++) j1[i] += cstart + m;
 80:       for (i = 0; i < n2; i++) j2[i] += cstart + m;
 81:     }
 82:   }
 83:   if (neg) {
 84:     n1 += 2;
 85:     n2 += 2;
 86:   }
 87:   /* MatSetPreallocationCOOLocal maps the indices! */
 88:   PetscCall(PetscMalloc2(PetscMax(n1, n2), &it, PetscMax(n1, n2), &jt));
 89:   /* test with repeated entries */
 90:   if (!localapi) {
 91:     PetscCall(MatSetPreallocationCOO(A, n1, i1, j1));
 92:   } else {
 93:     PetscCall(PetscArraycpy(it, i1, n1));
 94:     PetscCall(PetscArraycpy(jt, j1, n1));
 95:     PetscCall(MatSetPreallocationCOOLocal(A, n1, it, jt));
 96:   }
 97:   PetscCall(MatSetValuesCOO(A, v1, ADD_VALUES));
 98:   PetscCall(MatMult(A, x, y));
 99:   PetscCall(MatView(A, NULL));
100:   PetscCall(VecView(y, NULL));
101:   PetscCall(MatSetValuesCOO(A, v2, ADD_VALUES));
102:   PetscCall(MatMultAdd(A, x, y, y));
103:   PetscCall(MatView(A, NULL));
104:   PetscCall(VecView(y, NULL));
105:   PetscCall(MatTranspose(A, MAT_INITIAL_MATRIX, &At));
106:   if (!ismatis) {
107:     PetscCall(MatMatMult(A, At, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &AAt));
108:     PetscCall(MatView(AAt, NULL));
109:     PetscCall(MatDestroy(&AAt));
110:     PetscCall(MatMatMult(At, A, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &AAt));
111:     PetscCall(MatView(AAt, NULL));
112:     PetscCall(MatDestroy(&AAt));
113:   }
114:   PetscCall(MatDestroy(&At));

116:   /* INSERT_VALUES will overwrite matrix entries but
117:      still perform the sum of the repeated entries */
118:   PetscCall(MatSetValuesCOO(A, v2, INSERT_VALUES));
119:   PetscCall(MatView(A, NULL));

121:   /* test with unique entries */
122:   PetscCall(PetscArraycpy(it, i2, n2));
123:   PetscCall(PetscArraycpy(jt, j2, n2));
124:   if (!localapi) {
125:     PetscCall(MatSetPreallocationCOO(A, n2, it, jt));
126:   } else {
127:     PetscCall(MatSetPreallocationCOOLocal(A, n2, it, jt));
128:   }
129:   PetscCall(MatSetValuesCOO(A, v1, ADD_VALUES));
130:   PetscCall(MatMult(A, x, y));
131:   PetscCall(MatView(A, NULL));
132:   PetscCall(VecView(y, NULL));
133:   PetscCall(MatSetValuesCOO(A, v2, ADD_VALUES));
134:   PetscCall(MatMultAdd(A, x, y, z));
135:   PetscCall(MatView(A, NULL));
136:   PetscCall(VecView(z, NULL));
137:   PetscCall(PetscArraycpy(it, i2, n2));
138:   PetscCall(PetscArraycpy(jt, j2, n2));
139:   if (!localapi) {
140:     PetscCall(MatSetPreallocationCOO(A, n2, it, jt));
141:   } else {
142:     PetscCall(MatSetPreallocationCOOLocal(A, n2, it, jt));
143:   }
144:   PetscCall(MatSetValuesCOO(A, v1, INSERT_VALUES));
145:   PetscCall(MatMult(A, x, y));
146:   PetscCall(MatView(A, NULL));
147:   PetscCall(VecView(y, NULL));
148:   PetscCall(MatSetValuesCOO(A, v2, INSERT_VALUES));
149:   PetscCall(MatMultAdd(A, x, y, z));
150:   PetscCall(MatView(A, NULL));
151:   PetscCall(VecView(z, NULL));
152:   PetscCall(MatTranspose(A, MAT_INITIAL_MATRIX, &At));
153:   if (!ismatis) {
154:     PetscCall(MatMatMult(A, At, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &AAt));
155:     PetscCall(MatView(AAt, NULL));
156:     PetscCall(MatDestroy(&AAt));
157:     PetscCall(MatMatMult(At, A, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &AAt));
158:     PetscCall(MatView(AAt, NULL));
159:     PetscCall(MatDestroy(&AAt));
160:   }
161:   PetscCall(MatDestroy(&At));

163:   /* test providing diagonal first, then offdiagonal */
164:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
165:   PetscCall(PetscObjectBaseTypeCompare((PetscObject)A, MATMPIAIJ, &ismpiaij));
166:   if (ismpiaij && size > 1) {
167:     Mat                lA, lB;
168:     const PetscInt    *garray, *iA, *jA, *iB, *jB;
169:     const PetscScalar *vA, *vB;
170:     PetscScalar       *coo_v;
171:     PetscInt          *coo_i, *coo_j;
172:     PetscInt           i, j, nA, nB, nnz;
173:     PetscBool          flg;

175:     PetscCall(MatMPIAIJGetSeqAIJ(A, &lA, &lB, &garray));
176:     PetscCall(MatSeqAIJGetArrayRead(lA, &vA));
177:     PetscCall(MatSeqAIJGetArrayRead(lB, &vB));
178:     PetscCall(MatGetRowIJ(lA, 0, PETSC_FALSE, PETSC_FALSE, &nA, &iA, &jA, &flg));
179:     PetscCall(MatGetRowIJ(lB, 0, PETSC_FALSE, PETSC_FALSE, &nB, &iB, &jB, &flg));
180:     nnz = iA[nA] + iB[nB];
181:     PetscCall(PetscMalloc3(nnz, &coo_i, nnz, &coo_j, nnz, &coo_v));
182:     nnz = 0;
183:     for (i = 0; i < nA; i++) {
184:       for (j = iA[i]; j < iA[i + 1]; j++, nnz++) {
185:         coo_i[nnz] = i + rstart;
186:         coo_j[nnz] = jA[j] + cstart;
187:         coo_v[nnz] = vA[j];
188:       }
189:     }
190:     for (i = 0; i < nB; i++) {
191:       for (j = iB[i]; j < iB[i + 1]; j++, nnz++) {
192:         coo_i[nnz] = i + rstart;
193:         coo_j[nnz] = garray[jB[j]];
194:         coo_v[nnz] = vB[j];
195:       }
196:     }
197:     PetscCall(MatRestoreRowIJ(lA, 0, PETSC_FALSE, PETSC_FALSE, &nA, &iA, &jA, &flg));
198:     PetscCall(MatRestoreRowIJ(lB, 0, PETSC_FALSE, PETSC_FALSE, &nB, &iB, &jB, &flg));
199:     PetscCall(MatSeqAIJRestoreArrayRead(lA, &vA));
200:     PetscCall(MatSeqAIJRestoreArrayRead(lB, &vB));

202:     PetscCall(MatSetPreallocationCOO(A, nnz, coo_i, coo_j));
203:     PetscCall(MatSetValuesCOO(A, coo_v, ADD_VALUES));
204:     PetscCall(MatMult(A, x, y));
205:     PetscCall(MatView(A, NULL));
206:     PetscCall(VecView(y, NULL));
207:     PetscCall(MatSetValuesCOO(A, coo_v, INSERT_VALUES));
208:     PetscCall(MatMult(A, x, y));
209:     PetscCall(MatView(A, NULL));
210:     PetscCall(VecView(y, NULL));
211:     PetscCall(MatTranspose(A, MAT_INITIAL_MATRIX, &At));
212:     PetscCall(MatMatMult(A, At, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &AAt));
213:     PetscCall(MatView(AAt, NULL));
214:     PetscCall(MatDestroy(&AAt));
215:     PetscCall(MatMatMult(At, A, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &AAt));
216:     PetscCall(MatView(AAt, NULL));
217:     PetscCall(MatDestroy(&AAt));
218:     PetscCall(MatDestroy(&At));

220:     PetscCall(PetscFree3(coo_i, coo_j, coo_v));
221:   }
222:   PetscCall(PetscFree2(it, jt));
223:   PetscCall(VecDestroy(&z));
224:   PetscCall(VecDestroy(&x));
225:   PetscCall(VecDestroy(&y));
226:   PetscCall(MatDestroy(&A));
227:   PetscCall(PetscFinalize());
228:   return 0;
229: }

231: /*TEST

233:    test:
234:      suffix: 1
235:      filter: grep -v type
236:      diff_args: -j
237:      args: -mat_type {{seqaij mpiaij}} -localapi {{0 1}} -neg {{0 1}}

239:    test:
240:      requires: cuda
241:      suffix: 1_cuda
242:      filter: grep -v type
243:      diff_args: -j
244:      args: -mat_type {{seqaijcusparse mpiaijcusparse}} -localapi {{0 1}} -neg {{0 1}}
245:      output_file: output/ex123_1.out

247:    test:
248:      requires: kokkos_kernels
249:      suffix: 1_kokkos
250:      filter: grep -v type
251:      diff_args: -j
252:      args: -mat_type {{seqaijkokkos mpiaijkokkos}} -localapi {{0 1}} -neg {{0 1}}
253:      output_file: output/ex123_1.out

255:    test:
256:      suffix: 2
257:      nsize: 7
258:      filter: grep -v type
259:      diff_args: -j
260:      args: -mat_type mpiaij -localapi {{0 1}} -neg {{0 1}}

262:    test:
263:      requires: cuda
264:      suffix: 2_cuda
265:      nsize: 7
266:      filter: grep -v type
267:      diff_args: -j
268:      args: -mat_type mpiaijcusparse -localapi {{0 1}} -neg {{0 1}}
269:      output_file: output/ex123_2.out

271:    test:
272:      requires: kokkos_kernels
273:      suffix: 2_kokkos
274:      nsize: 7
275:      filter: grep -v type
276:      diff_args: -j
277:      args: -mat_type mpiaijkokkos -localapi {{0 1}} -neg {{0 1}}
278:      output_file: output/ex123_2.out

280:    test:
281:      suffix: 3
282:      nsize: 3
283:      filter: grep -v type
284:      diff_args: -j
285:      args: -mat_type mpiaij -loc -localapi {{0 1}} -neg {{0 1}}

287:    test:
288:      requires: cuda
289:      suffix: 3_cuda
290:      nsize: 3
291:      filter: grep -v type
292:      diff_args: -j
293:      args: -mat_type mpiaijcusparse -loc -localapi {{0 1}} -neg {{0 1}}
294:      output_file: output/ex123_3.out

296:    test:
297:      requires: kokkos_kernels
298:      suffix: 3_kokkos
299:      nsize: 3
300:      filter: grep -v type
301:      diff_args: -j
302:      args: -mat_type aijkokkos -loc -localapi {{0 1}} -neg {{0 1}}
303:      output_file: output/ex123_3.out

305:    test:
306:      suffix: 4
307:      nsize: 4
308:      filter: grep -v type
309:      diff_args: -j
310:      args: -mat_type mpiaij -loc -locdiag 0 -localapi {{0 1}} -neg {{0 1}}

312:    test:
313:      requires: cuda
314:      suffix: 4_cuda
315:      nsize: 4
316:      filter: grep -v type
317:      diff_args: -j
318:      args: -mat_type mpiaijcusparse -loc -locdiag 0 -localapi {{0 1}} -neg {{0 1}}
319:      output_file: output/ex123_4.out

321:    test:
322:      requires: kokkos_kernels
323:      suffix: 4_kokkos
324:      nsize: 4
325:      filter: grep -v type
326:      diff_args: -j
327:      args: -mat_type aijkokkos -loc -locdiag 0 -localapi {{0 1}} -neg {{0 1}}
328:      output_file: output/ex123_4.out

330:    test:
331:      suffix: matis
332:      nsize: 3
333:      filter: grep -v type
334:      diff_args: -j
335:      args: -mat_type is -localapi {{0 1}} -neg {{0 1}}

337: TEST*/