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