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