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 /* INSERT_VALUES will overwrite matrix entries but 79 still perform the sum of the repeated entries */ 80 ierr = MatSetValuesCOO(A,v2,INSERT_VALUES);CHKERRQ(ierr); 81 ierr = MyMatView(A,NULL);CHKERRQ(ierr); 82 83 /* test with unique entries */ 84 ierr = MatSetPreallocationCOO(A,n2,i2,j2);CHKERRQ(ierr); 85 ierr = MatSetValuesCOO(A,v1,ADD_VALUES);CHKERRQ(ierr); 86 ierr = MyMatView(A,NULL);CHKERRQ(ierr); 87 ierr = MatMult(A,x,y);CHKERRQ(ierr); 88 ierr = MyVecView(y,NULL);CHKERRQ(ierr); 89 ierr = MatSetValuesCOO(A,v2,ADD_VALUES);CHKERRQ(ierr); 90 ierr = MyMatView(A,NULL);CHKERRQ(ierr); 91 ierr = MatMultAdd(A,x,y,z);CHKERRQ(ierr); 92 ierr = MyVecView(z,NULL);CHKERRQ(ierr); 93 ierr = MatSetPreallocationCOO(A,n2,i2,j2);CHKERRQ(ierr); 94 ierr = MatSetValuesCOO(A,v1,INSERT_VALUES);CHKERRQ(ierr); 95 ierr = MyMatView(A,NULL);CHKERRQ(ierr); 96 ierr = MatMult(A,x,y);CHKERRQ(ierr); 97 ierr = MyVecView(y,NULL);CHKERRQ(ierr); 98 ierr = MatSetValuesCOO(A,v2,INSERT_VALUES);CHKERRQ(ierr); 99 ierr = MyMatView(A,NULL);CHKERRQ(ierr); 100 ierr = MatMultAdd(A,x,y,z);CHKERRQ(ierr); 101 ierr = MyVecView(z,NULL);CHKERRQ(ierr); 102 ierr = MatTranspose(A,MAT_INITIAL_MATRIX,&At);CHKERRQ(ierr); 103 ierr = MatMatMult(A,At,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr); 104 ierr = MyMatView(AAt,NULL);CHKERRQ(ierr); 105 ierr = MatDestroy(&AAt);CHKERRQ(ierr); 106 ierr = MatMatMult(At,A,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr); 107 ierr = MyMatView(AAt,NULL);CHKERRQ(ierr); 108 ierr = MatDestroy(&AAt);CHKERRQ(ierr); 109 ierr = MatDestroy(&At);CHKERRQ(ierr); 110 111 /* test providing diagonal first, the offdiagonal */ 112 ierr = MPI_Comm_size(PetscObjectComm((PetscObject)A),&size);CHKERRQ(ierr); 113 ierr = PetscObjectBaseTypeCompare((PetscObject)A,MATMPIAIJ,&ismpiaij);CHKERRQ(ierr); 114 if (ismpiaij && size > 1) { 115 Mat lA,lB; 116 const PetscInt *garray,*iA,*jA,*iB,*jB; 117 const PetscScalar *vA,*vB; 118 PetscScalar *coo_v; 119 PetscInt *coo_i,*coo_j; 120 PetscInt i,j,nA,nB,nnz; 121 PetscBool flg; 122 123 ierr = MatMPIAIJGetSeqAIJ(A,&lA,&lB,&garray);CHKERRQ(ierr); 124 ierr = MatSeqAIJGetArrayRead(lA,&vA);CHKERRQ(ierr); 125 ierr = MatSeqAIJGetArrayRead(lB,&vB);CHKERRQ(ierr); 126 ierr = MatGetRowIJ(lA,0,PETSC_FALSE,PETSC_FALSE,&nA,&iA,&jA,&flg);CHKERRQ(ierr); 127 ierr = MatGetRowIJ(lB,0,PETSC_FALSE,PETSC_FALSE,&nB,&iB,&jB,&flg);CHKERRQ(ierr); 128 nnz = iA[nA] + iB[nB]; 129 ierr = PetscMalloc3(nnz,&coo_i,nnz,&coo_j,nnz,&coo_v);CHKERRQ(ierr); 130 nnz = 0; 131 for (i=0;i<nA;i++) { 132 for (j=iA[i];j<iA[i+1];j++,nnz++) { 133 coo_i[nnz] = i+rstart; 134 coo_j[nnz] = jA[j]+cstart; 135 coo_v[nnz] = vA[j]; 136 } 137 } 138 for (i=0;i<nB;i++) { 139 for (j=iB[i];j<iB[i+1];j++,nnz++) { 140 coo_i[nnz] = i+rstart; 141 coo_j[nnz] = garray[jB[j]]; 142 coo_v[nnz] = vB[j]; 143 } 144 } 145 ierr = MatRestoreRowIJ(lA,0,PETSC_FALSE,PETSC_FALSE,&nA,&iA,&jA,&flg);CHKERRQ(ierr); 146 ierr = MatRestoreRowIJ(lB,0,PETSC_FALSE,PETSC_FALSE,&nB,&iB,&jB,&flg);CHKERRQ(ierr); 147 ierr = MatSeqAIJRestoreArrayRead(lA,&vA);CHKERRQ(ierr); 148 ierr = MatSeqAIJRestoreArrayRead(lB,&vB);CHKERRQ(ierr); 149 150 ierr = MatSetPreallocationCOO(A,nnz,coo_i,coo_j);CHKERRQ(ierr); 151 ierr = MatSetValuesCOO(A,coo_v,ADD_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 = MatSetValuesCOO(A,coo_v,INSERT_VALUES);CHKERRQ(ierr); 156 ierr = MyMatView(A,NULL);CHKERRQ(ierr); 157 ierr = MatMult(A,x,y);CHKERRQ(ierr); 158 ierr = MyVecView(y,NULL);CHKERRQ(ierr); 159 ierr = MatTranspose(A,MAT_INITIAL_MATRIX,&At);CHKERRQ(ierr); 160 ierr = MatMatMult(A,At,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr); 161 ierr = MyMatView(AAt,NULL);CHKERRQ(ierr); 162 ierr = MatDestroy(&AAt);CHKERRQ(ierr); 163 ierr = MatMatMult(At,A,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr); 164 ierr = MyMatView(AAt,NULL);CHKERRQ(ierr); 165 ierr = MatDestroy(&AAt);CHKERRQ(ierr); 166 ierr = MatDestroy(&At);CHKERRQ(ierr); 167 168 ierr = PetscFree3(coo_i,coo_j,coo_v);CHKERRQ(ierr); 169 } 170 ierr = VecDestroy(&z);CHKERRQ(ierr); 171 ierr = VecDestroy(&x);CHKERRQ(ierr); 172 ierr = VecDestroy(&y);CHKERRQ(ierr); 173 ierr = MatDestroy(&A);CHKERRQ(ierr); 174 ierr = PetscFinalize(); 175 return ierr; 176 } 177 178 179 /*TEST 180 181 test: 182 suffix: 1 183 filter: grep -v type 184 diff_args: -j 185 args: -mat_type {{seqaij mpiaij}} 186 187 test: 188 requires: cuda 189 suffix: 1_cuda 190 filter: grep -v type 191 diff_args: -j 192 args: -mat_type {{seqaijcusparse mpiaijcusparse}} 193 output_file: output/ex123_1.out 194 195 test: 196 suffix: 2 197 nsize: 7 198 filter: grep -v type 199 diff_args: -j 200 args: -mat_type mpiaij 201 202 test: 203 requires: cuda 204 suffix: 2_cuda 205 nsize: 7 206 filter: grep -v type 207 diff_args: -j 208 args: -mat_type mpiaijcusparse 209 output_file: output/ex123_2.out 210 211 test: 212 suffix: 3 213 nsize: 3 214 filter: grep -v type 215 diff_args: -j 216 args: -mat_type mpiaij -loc 217 218 test: 219 requires: cuda 220 suffix: 3_cuda 221 nsize: 3 222 filter: grep -v type 223 diff_args: -j 224 args: -mat_type mpiaijcusparse -loc 225 output_file: output/ex123_3.out 226 227 test: 228 suffix: 4 229 nsize: 4 230 filter: grep -v type 231 diff_args: -j 232 args: -mat_type mpiaij -loc -locdiag 0 233 234 test: 235 requires: cuda 236 suffix: 4_cuda 237 nsize: 4 238 filter: grep -v type 239 diff_args: -j 240 args: -mat_type mpiaijcusparse -loc -locdiag 0 241 output_file: output/ex123_4.out 242 243 TEST*/ 244