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 ISLocalToGlobalMapping rl2g,cl2g; 11 IS is; 12 PetscLayout rmap,cmap; 13 PetscInt *it,*jt; 14 PetscInt n1 = 11, n2 = 9; 15 PetscInt i1[] = { 7, 6, 2, 0, 4, 1, 1, 0, 2, 2, 1 , -1, -1}; 16 PetscInt j1[] = { 1, 4, 3, 5, 3, 3, 4, 5, 0, 3, 1 , -1, -1}; 17 PetscInt i2[] = { 7, 6, 2, 0, 4, 1, 1, 2, 1, -1, -1}; 18 PetscInt j2[] = { 1, 4, 3, 5, 3, 3, 4, 0, 1, -1, -1}; 19 PetscScalar v1[] = { -1., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., PETSC_MAX_REAL, PETSC_MAX_REAL}; 20 PetscScalar v2[] = { 1.,-1.,-2.,-3.,-4.,-5.,-6.,-7.,-8.,-9.,-10., PETSC_MAX_REAL, PETSC_MAX_REAL}; 21 PetscInt N = 6, m = 8, M, rstart, cstart, i; 22 PetscMPIInt size; 23 PetscBool loc = PETSC_FALSE; 24 PetscBool locdiag = PETSC_TRUE; 25 PetscBool localapi = PETSC_FALSE; 26 PetscBool neg = PETSC_FALSE; 27 PetscBool ismatis, ismpiaij; 28 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)); 52 53 PetscCall(PetscObjectTypeCompare((PetscObject)A,MATIS,&ismatis)); 54 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)); 65 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) for (i = 0; i < n1; i++) i1[i] += rstart; 71 if (!localapi) for (i = 0; i < n2; i++) i2[i] += rstart; 72 if (loc) { 73 if (locdiag) { 74 for (i = 0; i < n1; i++) j1[i] += cstart; 75 for (i = 0; i < n2; i++) j2[i] += cstart; 76 } else { 77 for (i = 0; i < n1; i++) j1[i] += cstart + m; 78 for (i = 0; i < n2; i++) j2[i] += cstart + m; 79 } 80 } 81 if (neg) { n1 += 2; n2 += 2; } 82 /* MatSetPreallocationCOOLocal maps the indices! */ 83 PetscCall(PetscMalloc2(PetscMax(n1,n2),&it,PetscMax(n1,n2),&jt)); 84 /* test with repeated entries */ 85 if (!localapi) { 86 PetscCall(MatSetPreallocationCOO(A,n1,i1,j1)); 87 } else { 88 PetscCall(PetscArraycpy(it,i1,n1)); 89 PetscCall(PetscArraycpy(jt,j1,n1)); 90 PetscCall(MatSetPreallocationCOOLocal(A,n1,it,jt)); 91 } 92 PetscCall(MatSetValuesCOO(A,v1,ADD_VALUES)); 93 PetscCall(MatMult(A,x,y)); 94 PetscCall(MyMatView(A,NULL)); 95 PetscCall(MyVecView(y,NULL)); 96 PetscCall(MatSetValuesCOO(A,v2,ADD_VALUES)); 97 PetscCall(MatMultAdd(A,x,y,y)); 98 PetscCall(MyMatView(A,NULL)); 99 PetscCall(MyVecView(y,NULL)); 100 PetscCall(MatTranspose(A,MAT_INITIAL_MATRIX,&At)); 101 if (!ismatis) { 102 PetscCall(MatMatMult(A,At,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt)); 103 PetscCall(MyMatView(AAt,NULL)); 104 PetscCall(MatDestroy(&AAt)); 105 PetscCall(MatMatMult(At,A,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt)); 106 PetscCall(MyMatView(AAt,NULL)); 107 PetscCall(MatDestroy(&AAt)); 108 } 109 PetscCall(MatDestroy(&At)); 110 111 /* INSERT_VALUES will overwrite matrix entries but 112 still perform the sum of the repeated entries */ 113 PetscCall(MatSetValuesCOO(A,v2,INSERT_VALUES)); 114 PetscCall(MyMatView(A,NULL)); 115 116 /* test with unique entries */ 117 PetscCall(PetscArraycpy(it,i2,n2)); 118 PetscCall(PetscArraycpy(jt,j2,n2)); 119 if (!localapi) { 120 PetscCall(MatSetPreallocationCOO(A,n2,it,jt)); 121 } else { 122 PetscCall(MatSetPreallocationCOOLocal(A,n2,it,jt)); 123 } 124 PetscCall(MatSetValuesCOO(A,v1,ADD_VALUES)); 125 PetscCall(MatMult(A,x,y)); 126 PetscCall(MyMatView(A,NULL)); 127 PetscCall(MyVecView(y,NULL)); 128 PetscCall(MatSetValuesCOO(A,v2,ADD_VALUES)); 129 PetscCall(MatMultAdd(A,x,y,z)); 130 PetscCall(MyMatView(A,NULL)); 131 PetscCall(MyVecView(z,NULL)); 132 PetscCall(PetscArraycpy(it,i2,n2)); 133 PetscCall(PetscArraycpy(jt,j2,n2)); 134 if (!localapi) { 135 PetscCall(MatSetPreallocationCOO(A,n2,it,jt)); 136 } else { 137 PetscCall(MatSetPreallocationCOOLocal(A,n2,it,jt)); 138 } 139 PetscCall(MatSetValuesCOO(A,v1,INSERT_VALUES)); 140 PetscCall(MatMult(A,x,y)); 141 PetscCall(MyMatView(A,NULL)); 142 PetscCall(MyVecView(y,NULL)); 143 PetscCall(MatSetValuesCOO(A,v2,INSERT_VALUES)); 144 PetscCall(MatMultAdd(A,x,y,z)); 145 PetscCall(MyMatView(A,NULL)); 146 PetscCall(MyVecView(z,NULL)); 147 PetscCall(MatTranspose(A,MAT_INITIAL_MATRIX,&At)); 148 if (!ismatis) { 149 PetscCall(MatMatMult(A,At,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt)); 150 PetscCall(MyMatView(AAt,NULL)); 151 PetscCall(MatDestroy(&AAt)); 152 PetscCall(MatMatMult(At,A,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt)); 153 PetscCall(MyMatView(AAt,NULL)); 154 PetscCall(MatDestroy(&AAt)); 155 } 156 PetscCall(MatDestroy(&At)); 157 158 /* test providing diagonal first, then offdiagonal */ 159 PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A),&size)); 160 PetscCall(PetscObjectBaseTypeCompare((PetscObject)A,MATMPIAIJ,&ismpiaij)); 161 if (ismpiaij && size > 1) { 162 Mat lA,lB; 163 const PetscInt *garray,*iA,*jA,*iB,*jB; 164 const PetscScalar *vA,*vB; 165 PetscScalar *coo_v; 166 PetscInt *coo_i,*coo_j; 167 PetscInt i,j,nA,nB,nnz; 168 PetscBool flg; 169 170 PetscCall(MatMPIAIJGetSeqAIJ(A,&lA,&lB,&garray)); 171 PetscCall(MatSeqAIJGetArrayRead(lA,&vA)); 172 PetscCall(MatSeqAIJGetArrayRead(lB,&vB)); 173 PetscCall(MatGetRowIJ(lA,0,PETSC_FALSE,PETSC_FALSE,&nA,&iA,&jA,&flg)); 174 PetscCall(MatGetRowIJ(lB,0,PETSC_FALSE,PETSC_FALSE,&nB,&iB,&jB,&flg)); 175 nnz = iA[nA] + iB[nB]; 176 PetscCall(PetscMalloc3(nnz,&coo_i,nnz,&coo_j,nnz,&coo_v)); 177 nnz = 0; 178 for (i=0;i<nA;i++) { 179 for (j=iA[i];j<iA[i+1];j++,nnz++) { 180 coo_i[nnz] = i+rstart; 181 coo_j[nnz] = jA[j]+cstart; 182 coo_v[nnz] = vA[j]; 183 } 184 } 185 for (i=0;i<nB;i++) { 186 for (j=iB[i];j<iB[i+1];j++,nnz++) { 187 coo_i[nnz] = i+rstart; 188 coo_j[nnz] = garray[jB[j]]; 189 coo_v[nnz] = vB[j]; 190 } 191 } 192 PetscCall(MatRestoreRowIJ(lA,0,PETSC_FALSE,PETSC_FALSE,&nA,&iA,&jA,&flg)); 193 PetscCall(MatRestoreRowIJ(lB,0,PETSC_FALSE,PETSC_FALSE,&nB,&iB,&jB,&flg)); 194 PetscCall(MatSeqAIJRestoreArrayRead(lA,&vA)); 195 PetscCall(MatSeqAIJRestoreArrayRead(lB,&vB)); 196 197 PetscCall(MatSetPreallocationCOO(A,nnz,coo_i,coo_j)); 198 PetscCall(MatSetValuesCOO(A,coo_v,ADD_VALUES)); 199 PetscCall(MatMult(A,x,y)); 200 PetscCall(MyMatView(A,NULL)); 201 PetscCall(MyVecView(y,NULL)); 202 PetscCall(MatSetValuesCOO(A,coo_v,INSERT_VALUES)); 203 PetscCall(MatMult(A,x,y)); 204 PetscCall(MyMatView(A,NULL)); 205 PetscCall(MyVecView(y,NULL)); 206 PetscCall(MatTranspose(A,MAT_INITIAL_MATRIX,&At)); 207 PetscCall(MatMatMult(A,At,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt)); 208 PetscCall(MyMatView(AAt,NULL)); 209 PetscCall(MatDestroy(&AAt)); 210 PetscCall(MatMatMult(At,A,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt)); 211 PetscCall(MyMatView(AAt,NULL)); 212 PetscCall(MatDestroy(&AAt)); 213 PetscCall(MatDestroy(&At)); 214 215 PetscCall(PetscFree3(coo_i,coo_j,coo_v)); 216 } 217 PetscCall(PetscFree2(it,jt)); 218 PetscCall(VecDestroy(&z)); 219 PetscCall(VecDestroy(&x)); 220 PetscCall(VecDestroy(&y)); 221 PetscCall(MatDestroy(&A)); 222 PetscCall(PetscFinalize()); 223 return 0; 224 } 225 226 /*TEST 227 228 test: 229 suffix: 1 230 filter: grep -v type 231 diff_args: -j 232 args: -mat_type {{seqaij mpiaij}} -localapi {{0 1}} -neg {{0 1}} 233 234 test: 235 requires: cuda 236 suffix: 1_cuda 237 filter: grep -v type 238 diff_args: -j 239 args: -mat_type {{seqaijcusparse mpiaijcusparse}} -localapi {{0 1}} -neg {{0 1}} 240 output_file: output/ex123_1.out 241 242 test: 243 requires: kokkos_kernels !sycl 244 suffix: 1_kokkos 245 filter: grep -v type 246 diff_args: -j 247 args: -mat_type {{seqaijkokkos mpiaijkokkos}} -localapi {{0 1}} -neg {{0 1}} 248 output_file: output/ex123_1.out 249 250 test: 251 suffix: 2 252 nsize: 7 253 filter: grep -v type 254 diff_args: -j 255 args: -mat_type mpiaij -localapi {{0 1}} -neg {{0 1}} 256 257 test: 258 requires: cuda 259 suffix: 2_cuda 260 nsize: 7 261 filter: grep -v type 262 diff_args: -j 263 args: -mat_type mpiaijcusparse -localapi {{0 1}} -neg {{0 1}} 264 output_file: output/ex123_2.out 265 266 test: 267 requires: kokkos_kernels !sycl 268 suffix: 2_kokkos 269 nsize: 7 270 filter: grep -v type 271 diff_args: -j 272 args: -mat_type mpiaijkokkos -localapi {{0 1}} -neg {{0 1}} 273 output_file: output/ex123_2.out 274 275 test: 276 suffix: 3 277 nsize: 3 278 filter: grep -v type 279 diff_args: -j 280 args: -mat_type mpiaij -loc -localapi {{0 1}} -neg {{0 1}} 281 282 test: 283 requires: cuda 284 suffix: 3_cuda 285 nsize: 3 286 filter: grep -v type 287 diff_args: -j 288 args: -mat_type mpiaijcusparse -loc -localapi {{0 1}} -neg {{0 1}} 289 output_file: output/ex123_3.out 290 291 test: 292 requires: !sycl kokkos_kernels 293 suffix: 3_kokkos 294 nsize: 3 295 filter: grep -v type 296 diff_args: -j 297 args: -mat_type aijkokkos -loc -localapi {{0 1}} -neg {{0 1}} 298 output_file: output/ex123_3.out 299 300 test: 301 suffix: 4 302 nsize: 4 303 filter: grep -v type 304 diff_args: -j 305 args: -mat_type mpiaij -loc -locdiag 0 -localapi {{0 1}} -neg {{0 1}} 306 307 test: 308 requires: cuda 309 suffix: 4_cuda 310 nsize: 4 311 filter: grep -v type 312 diff_args: -j 313 args: -mat_type mpiaijcusparse -loc -locdiag 0 -localapi {{0 1}} -neg {{0 1}} 314 output_file: output/ex123_4.out 315 316 test: 317 requires: !sycl kokkos_kernels 318 suffix: 4_kokkos 319 nsize: 4 320 filter: grep -v type 321 diff_args: -j 322 args: -mat_type aijkokkos -loc -locdiag 0 -localapi {{0 1}} -neg {{0 1}} 323 output_file: output/ex123_4.out 324 325 test: 326 suffix: matis 327 nsize: 3 328 filter: grep -v type 329 diff_args: -j 330 args: -mat_type is -localapi {{0 1}} -neg {{0 1}} 331 332 TEST*/ 333