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 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, (char *)0, 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