1 static char help[] = "Tests MatMPIBAIJSetPreallocationCSR()\n\n";
2
3 /*
4 Include "petscmat.h" so that we can use matrices. Note that this file
5 automatically includes:
6 petscsys.h - base PETSc routines petscvec.h - vectors
7 petscmat.h - matrices
8 petscis.h - index sets
9 petscviewer.h - viewers
10 */
11 #include <petscmat.h>
12
main(int argc,char ** args)13 int main(int argc, char **args)
14 {
15 Mat A;
16 PetscInt *ia, *ja, bs = 2;
17 PetscInt N = 9, n;
18 PetscInt rstart, rend, row, col;
19 PetscInt i;
20 PetscMPIInt rank, size;
21 Vec v;
22
23 PetscFunctionBeginUser;
24 PetscCall(PetscInitialize(&argc, &args, NULL, help));
25 PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
26 PetscCheck(size < 5, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "Can only use at most 4 processors.");
27 PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
28
29 /* Get a partition range based on the vector size */
30 PetscCall(VecCreateFromOptions(PETSC_COMM_WORLD, NULL, 1, PETSC_DECIDE, N, &v));
31 PetscCall(VecGetLocalSize(v, &n));
32 PetscCall(VecGetOwnershipRange(v, &rstart, &rend));
33 PetscCall(VecDestroy(&v));
34
35 PetscCall(PetscMalloc1(n + 1, &ia));
36 PetscCall(PetscMalloc1(3 * n, &ja));
37
38 /* Construct a tri-diagonal CSR indexing */
39 i = 1;
40 ia[0] = 0;
41 for (row = rstart; row < rend; row++) {
42 ia[i] = ia[i - 1];
43
44 /* diagonal */
45 col = row;
46 {
47 ja[ia[i]] = col;
48 ia[i]++;
49 }
50
51 /* lower diagonal */
52 col = row - 1;
53 if (col >= 0) {
54 ja[ia[i]] = col;
55 ia[i]++;
56 }
57
58 /* upper diagonal */
59 col = row + 1;
60 if (col < N) {
61 ja[ia[i]] = col;
62 ia[i]++;
63 }
64 i++;
65 }
66
67 PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
68 PetscCall(MatSetSizes(A, n, n, PETSC_DETERMINE, PETSC_DETERMINE));
69 PetscCall(MatSetType(A, MATMPIAIJ));
70 PetscCall(MatMPIAIJSetPreallocationCSR(A, ia, ja, NULL));
71 PetscCall(MatView(A, PETSC_VIEWER_STDOUT_WORLD));
72 PetscCall(MatDestroy(&A));
73
74 PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
75 PetscCall(MatSetSizes(A, bs * n, bs * n, PETSC_DETERMINE, PETSC_DETERMINE));
76 PetscCall(MatSetType(A, MATMPIBAIJ));
77 PetscCall(MatMPIBAIJSetPreallocationCSR(A, bs, ia, ja, NULL));
78 PetscCall(MatView(A, PETSC_VIEWER_STDOUT_WORLD));
79 PetscCall(MatDestroy(&A));
80
81 PetscCall(PetscFree(ia));
82 PetscCall(PetscFree(ja));
83 PetscCall(PetscFinalize());
84 return 0;
85 }
86
87 /*TEST
88
89 test:
90 nsize: 4
91
92 TEST*/
93