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