1 static char help[] = "Test MatSetPreallocationCOO and MatSetValuesCOO\n\n"; 2 3 #include <petscmat.h> 4 #define MyMatView(a,b) PetscPrintf(PetscObjectComm((PetscObject)(a)),"LINE %d\n",__LINE__),MatView(a,b); 5 #define MyVecView(a,b) PetscPrintf(PetscObjectComm((PetscObject)(a)),"LINE %d\n",__LINE__),VecView(a,b); 6 int main(int argc,char **args) 7 { 8 Mat A,At,AAt; 9 Vec x,y,z; 10 PetscLayout rmap,cmap; 11 PetscInt n1 = 11, n2 = 9; 12 PetscInt i1[] = { 7, 6, 2, 0, 4, 1, 1, 0, 2, 2, 1 }; 13 PetscInt j1[] = { 1, 4, 3, 5, 3, 3, 4, 5, 0, 3, 1 }; 14 PetscInt i2[] = { 7, 6, 2, 0, 4, 1, 1, 2, 1 }; 15 PetscInt j2[] = { 1, 4, 3, 5, 3, 3, 4, 0, 1 }; 16 PetscScalar v1[] = { -1., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10.}; 17 PetscScalar v2[] = { 1.,-1.,-2.,-3.,-4.,-5.,-6.,-7.,-8.,-9.,-10.}; 18 PetscInt N = 6, m = 8, rstart, cstart, i; 19 PetscMPIInt size; 20 PetscBool loc = PETSC_FALSE; 21 PetscBool locdiag = PETSC_TRUE, ismpiaij; 22 PetscErrorCode ierr; 23 24 ierr = PetscInitialize(&argc,&args,(char*)0,help);if (ierr) return ierr; 25 ierr = PetscOptionsGetBool(NULL,NULL,"-loc",&loc,NULL);CHKERRQ(ierr); 26 ierr = PetscOptionsGetBool(NULL,NULL,"-locdiag",&locdiag,NULL);CHKERRQ(ierr); 27 28 ierr = MatCreate(PETSC_COMM_WORLD,&A);CHKERRQ(ierr); 29 if (loc) { 30 if (locdiag) { 31 ierr = MatSetSizes(A,m,N,PETSC_DECIDE,PETSC_DECIDE);CHKERRQ(ierr); 32 } else { 33 ierr = MatSetSizes(A,m,m+N,PETSC_DECIDE,PETSC_DECIDE);CHKERRQ(ierr); 34 } 35 } else { 36 ierr = MatSetSizes(A,m,PETSC_DECIDE,PETSC_DECIDE,N);CHKERRQ(ierr); 37 } 38 ierr = MatSetFromOptions(A);CHKERRQ(ierr); 39 ierr = MatGetLayouts(A,&rmap,&cmap);CHKERRQ(ierr); 40 ierr = PetscLayoutSetUp(rmap);CHKERRQ(ierr); 41 ierr = PetscLayoutSetUp(cmap);CHKERRQ(ierr); 42 ierr = MatCreateVecs(A,&x,&y);CHKERRQ(ierr); 43 ierr = MatCreateVecs(A,NULL,&z);CHKERRQ(ierr); 44 ierr = VecSet(x,1.);CHKERRQ(ierr); 45 ierr = VecSet(z,2.);CHKERRQ(ierr); 46 ierr = PetscLayoutGetRange(rmap,&rstart,NULL);CHKERRQ(ierr); 47 ierr = PetscLayoutGetRange(cmap,&cstart,NULL);CHKERRQ(ierr); 48 for (i = 0; i < n1; i++) i1[i] += rstart; 49 for (i = 0; i < n2; i++) i2[i] += rstart; 50 if (loc) { 51 if (locdiag) { 52 for (i = 0; i < n1; i++) j1[i] += cstart; 53 for (i = 0; i < n2; i++) j2[i] += cstart; 54 } else { 55 for (i = 0; i < n1; i++) j1[i] += cstart + m; 56 for (i = 0; i < n2; i++) j2[i] += cstart + m; 57 } 58 } 59 60 /* test with repeated entries */ 61 ierr = MatSetPreallocationCOO(A,n1,i1,j1);CHKERRQ(ierr); 62 ierr = MatSetValuesCOO(A,v1,ADD_VALUES);CHKERRQ(ierr); 63 ierr = MyMatView(A,NULL);CHKERRQ(ierr); 64 ierr = MatMult(A,x,y);CHKERRQ(ierr); 65 ierr = MyVecView(y,NULL);CHKERRQ(ierr); 66 ierr = MatSetValuesCOO(A,v2,ADD_VALUES);CHKERRQ(ierr); 67 ierr = MyMatView(A,NULL);CHKERRQ(ierr); 68 ierr = MatMultAdd(A,x,y,y);CHKERRQ(ierr); 69 ierr = MyVecView(y,NULL);CHKERRQ(ierr); 70 ierr = MatTranspose(A,MAT_INITIAL_MATRIX,&At);CHKERRQ(ierr); 71 ierr = MatMatMult(A,At,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr); 72 ierr = MyMatView(AAt,NULL);CHKERRQ(ierr); 73 ierr = MatDestroy(&AAt);CHKERRQ(ierr); 74 ierr = MatMatMult(At,A,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr); 75 ierr = MyMatView(AAt,NULL);CHKERRQ(ierr); 76 ierr = MatDestroy(&AAt);CHKERRQ(ierr); 77 ierr = MatDestroy(&At);CHKERRQ(ierr); 78 79 /* test with unique entries */ 80 ierr = MatSetPreallocationCOO(A,n2,i2,j2);CHKERRQ(ierr); 81 ierr = MatSetValuesCOO(A,v1,ADD_VALUES);CHKERRQ(ierr); 82 ierr = MyMatView(A,NULL);CHKERRQ(ierr); 83 ierr = MatMult(A,x,y);CHKERRQ(ierr); 84 ierr = MyVecView(y,NULL);CHKERRQ(ierr); 85 ierr = MatSetValuesCOO(A,v2,ADD_VALUES);CHKERRQ(ierr); 86 ierr = MyMatView(A,NULL);CHKERRQ(ierr); 87 ierr = MatMultAdd(A,x,y,z);CHKERRQ(ierr); 88 ierr = MyVecView(z,NULL);CHKERRQ(ierr); 89 ierr = MatSetPreallocationCOO(A,n2,i2,j2);CHKERRQ(ierr); 90 ierr = MatSetValuesCOO(A,v1,INSERT_VALUES);CHKERRQ(ierr); 91 ierr = MyMatView(A,NULL);CHKERRQ(ierr); 92 ierr = MatMult(A,x,y);CHKERRQ(ierr); 93 ierr = MyVecView(y,NULL);CHKERRQ(ierr); 94 ierr = MatSetValuesCOO(A,v2,INSERT_VALUES);CHKERRQ(ierr); 95 ierr = MyMatView(A,NULL);CHKERRQ(ierr); 96 ierr = MatMultAdd(A,x,y,z);CHKERRQ(ierr); 97 ierr = MyVecView(z,NULL);CHKERRQ(ierr); 98 ierr = MatTranspose(A,MAT_INITIAL_MATRIX,&At);CHKERRQ(ierr); 99 ierr = MatMatMult(A,At,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr); 100 ierr = MyMatView(AAt,NULL);CHKERRQ(ierr); 101 ierr = MatDestroy(&AAt);CHKERRQ(ierr); 102 ierr = MatMatMult(At,A,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr); 103 ierr = MyMatView(AAt,NULL);CHKERRQ(ierr); 104 ierr = MatDestroy(&AAt);CHKERRQ(ierr); 105 ierr = MatDestroy(&At);CHKERRQ(ierr); 106 107 /* test providing diagonal first, the offdiagonal */ 108 ierr = MPI_Comm_size(PetscObjectComm((PetscObject)A),&size);CHKERRQ(ierr); 109 ierr = PetscObjectBaseTypeCompare((PetscObject)A,MATMPIAIJ,&ismpiaij);CHKERRQ(ierr); 110 if (ismpiaij && size > 1) { 111 Mat lA,lB; 112 const PetscInt *garray,*iA,*jA,*iB,*jB; 113 const PetscScalar *vA,*vB; 114 PetscScalar *coo_v; 115 PetscInt *coo_i,*coo_j; 116 PetscInt i,j,nA,nB,nnz; 117 PetscBool flg; 118 119 ierr = MatMPIAIJGetSeqAIJ(A,&lA,&lB,&garray);CHKERRQ(ierr); 120 ierr = MatSeqAIJGetArrayRead(lA,&vA);CHKERRQ(ierr); 121 ierr = MatSeqAIJGetArrayRead(lB,&vB);CHKERRQ(ierr); 122 ierr = MatGetRowIJ(lA,0,PETSC_FALSE,PETSC_FALSE,&nA,&iA,&jA,&flg);CHKERRQ(ierr); 123 ierr = MatGetRowIJ(lB,0,PETSC_FALSE,PETSC_FALSE,&nB,&iB,&jB,&flg);CHKERRQ(ierr); 124 nnz = iA[nA] + iB[nB]; 125 ierr = PetscMalloc3(nnz,&coo_i,nnz,&coo_j,nnz,&coo_v);CHKERRQ(ierr); 126 nnz = 0; 127 for (i=0;i<nA;i++) { 128 for (j=iA[i];j<iA[i+1];j++,nnz++) { 129 coo_i[nnz] = i+rstart; 130 coo_j[nnz] = jA[j]+cstart; 131 coo_v[nnz] = vA[j]; 132 } 133 } 134 for (i=0;i<nB;i++) { 135 for (j=iB[i];j<iB[i+1];j++,nnz++) { 136 coo_i[nnz] = i+rstart; 137 coo_j[nnz] = garray[jB[j]]; 138 coo_v[nnz] = vB[j]; 139 } 140 } 141 ierr = MatRestoreRowIJ(lA,0,PETSC_FALSE,PETSC_FALSE,&nA,&iA,&jA,&flg);CHKERRQ(ierr); 142 ierr = MatRestoreRowIJ(lB,0,PETSC_FALSE,PETSC_FALSE,&nB,&iB,&jB,&flg);CHKERRQ(ierr); 143 ierr = MatSeqAIJRestoreArrayRead(lA,&vA);CHKERRQ(ierr); 144 ierr = MatSeqAIJRestoreArrayRead(lB,&vB);CHKERRQ(ierr); 145 146 ierr = MatSetPreallocationCOO(A,nnz,coo_i,coo_j);CHKERRQ(ierr); 147 ierr = MatSetValuesCOO(A,coo_v,ADD_VALUES);CHKERRQ(ierr); 148 ierr = MyMatView(A,NULL);CHKERRQ(ierr); 149 ierr = MatMult(A,x,y);CHKERRQ(ierr); 150 ierr = MyVecView(y,NULL);CHKERRQ(ierr); 151 ierr = MatSetValuesCOO(A,coo_v,INSERT_VALUES);CHKERRQ(ierr); 152 ierr = MyMatView(A,NULL);CHKERRQ(ierr); 153 ierr = MatMult(A,x,y);CHKERRQ(ierr); 154 ierr = MyVecView(y,NULL);CHKERRQ(ierr); 155 ierr = MatTranspose(A,MAT_INITIAL_MATRIX,&At);CHKERRQ(ierr); 156 ierr = MatMatMult(A,At,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr); 157 ierr = MyMatView(AAt,NULL);CHKERRQ(ierr); 158 ierr = MatDestroy(&AAt);CHKERRQ(ierr); 159 ierr = MatMatMult(At,A,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr); 160 ierr = MyMatView(AAt,NULL);CHKERRQ(ierr); 161 ierr = MatDestroy(&AAt);CHKERRQ(ierr); 162 ierr = MatDestroy(&At);CHKERRQ(ierr); 163 164 ierr = PetscFree3(coo_i,coo_j,coo_v);CHKERRQ(ierr); 165 } 166 ierr = VecDestroy(&z);CHKERRQ(ierr); 167 ierr = VecDestroy(&x);CHKERRQ(ierr); 168 ierr = VecDestroy(&y);CHKERRQ(ierr); 169 ierr = MatDestroy(&A);CHKERRQ(ierr); 170 ierr = PetscFinalize(); 171 return ierr; 172 } 173 174 175 /*TEST 176 177 test: 178 suffix: 1 179 filter: grep -v type 180 diff_args: -j 181 args: -mat_type {{seqaij mpiaij}} 182 183 test: 184 requires: cuda 185 suffix: 1_cuda 186 filter: grep -v type 187 diff_args: -j 188 args: -mat_type {{seqaijcusparse mpiaijcusparse}} 189 output_file: output/ex123_1.out 190 191 test: 192 suffix: 2 193 nsize: 7 194 filter: grep -v type 195 diff_args: -j 196 args: -mat_type mpiaij 197 198 test: 199 requires: cuda 200 suffix: 2_cuda 201 nsize: 7 202 filter: grep -v type 203 diff_args: -j 204 args: -mat_type mpiaijcusparse 205 output_file: output/ex123_2.out 206 207 test: 208 suffix: 3 209 nsize: 3 210 filter: grep -v type 211 diff_args: -j 212 args: -mat_type mpiaij -loc 213 214 test: 215 requires: cuda 216 suffix: 3_cuda 217 nsize: 3 218 filter: grep -v type 219 diff_args: -j 220 args: -mat_type mpiaijcusparse -loc 221 output_file: output/ex123_3.out 222 223 test: 224 suffix: 4 225 nsize: 4 226 filter: grep -v type 227 diff_args: -j 228 args: -mat_type mpiaij -loc -locdiag 0 229 230 test: 231 requires: cuda 232 suffix: 4_cuda 233 nsize: 4 234 filter: grep -v type 235 diff_args: -j 236 args: -mat_type mpiaijcusparse -loc -locdiag 0 237 output_file: output/ex123_4.out 238 239 TEST*/ 240