1 2 static char help[] = "Test PetscSFConcatenate()\n\n"; 3 4 #include <petscsf.h> 5 6 typedef struct { 7 MPI_Comm comm; 8 PetscMPIInt rank, size; 9 PetscInt leaveStep, nsfs, nLeavesPerRank; 10 PetscBool shareRoots, sparseLeaves; 11 } AppCtx; 12 13 static PetscErrorCode GetOptions(MPI_Comm comm, AppCtx *ctx) { 14 PetscFunctionBegin; 15 ctx->comm = comm; 16 ctx->nsfs = 3; 17 ctx->nLeavesPerRank = 4; 18 ctx->leaveStep = 1; 19 ctx->shareRoots = PETSC_FALSE; 20 ctx->sparseLeaves = PETSC_FALSE; 21 PetscCall(PetscOptionsGetInt(NULL, NULL, "-nsfs", &ctx->nsfs, NULL)); 22 PetscCall(PetscOptionsGetInt(NULL, NULL, "-n_leaves_per_rank", &ctx->nLeavesPerRank, NULL)); 23 PetscCall(PetscOptionsGetInt(NULL, NULL, "-leave_step", &ctx->leaveStep, NULL)); 24 PetscCall(PetscOptionsGetBool(NULL, NULL, "-share_roots", &ctx->shareRoots, NULL)); 25 ctx->sparseLeaves = (PetscBool)(ctx->leaveStep != 1); 26 PetscCallMPI(MPI_Comm_size(comm, &ctx->size)); 27 PetscCallMPI(MPI_Comm_rank(comm, &ctx->rank)); 28 PetscFunctionReturn(0); 29 } 30 31 static PetscErrorCode PetscSFCheckEqual_Private(PetscSF sf0, PetscSF sf1) { 32 PetscInt nRoot, nLeave; 33 Vec vecRoot0, vecLeave0, vecRoot1, vecLeave1; 34 MPI_Comm comm; 35 PetscBool flg; 36 37 PetscFunctionBegin; 38 PetscCall(PetscObjectGetComm((PetscObject)sf0, &comm)); 39 PetscCall(PetscSFGetGraph(sf0, &nRoot, NULL, NULL, NULL)); 40 PetscCall(PetscSFGetLeafRange(sf0, NULL, &nLeave)); 41 nLeave++; 42 PetscCall(VecCreateMPI(comm, nRoot, PETSC_DECIDE, &vecRoot0)); 43 PetscCall(VecCreateMPI(comm, nLeave, PETSC_DECIDE, &vecLeave0)); 44 PetscCall(VecDuplicate(vecRoot0, &vecRoot1)); 45 PetscCall(VecDuplicate(vecLeave0, &vecLeave1)); 46 { 47 PetscRandom rand; 48 49 PetscCall(PetscRandomCreate(comm, &rand)); 50 PetscCall(PetscRandomSetFromOptions(rand)); 51 PetscCall(VecSetRandom(vecRoot0, rand)); 52 PetscCall(VecSetRandom(vecLeave0, rand)); 53 PetscCall(VecCopy(vecRoot0, vecRoot1)); 54 PetscCall(VecCopy(vecLeave0, vecLeave1)); 55 PetscCall(PetscRandomDestroy(&rand)); 56 } 57 58 PetscCall(VecScatterBegin(sf0, vecRoot0, vecLeave0, ADD_VALUES, SCATTER_FORWARD)); 59 PetscCall(VecScatterEnd(sf0, vecRoot0, vecLeave0, ADD_VALUES, SCATTER_FORWARD)); 60 PetscCall(VecScatterBegin(sf1, vecRoot1, vecLeave1, ADD_VALUES, SCATTER_FORWARD)); 61 PetscCall(VecScatterEnd(sf1, vecRoot1, vecLeave1, ADD_VALUES, SCATTER_FORWARD)); 62 PetscCall(VecEqual(vecLeave0, vecLeave1, &flg)); 63 PetscCheck(flg, comm, PETSC_ERR_PLIB, "leave vectors differ"); 64 65 PetscCall(VecScatterBegin(sf0, vecLeave0, vecRoot0, ADD_VALUES, SCATTER_REVERSE)); 66 PetscCall(VecScatterEnd(sf0, vecLeave0, vecRoot0, ADD_VALUES, SCATTER_REVERSE)); 67 PetscCall(VecScatterBegin(sf1, vecLeave1, vecRoot1, ADD_VALUES, SCATTER_REVERSE)); 68 PetscCall(VecScatterEnd(sf1, vecLeave1, vecRoot1, ADD_VALUES, SCATTER_REVERSE)); 69 PetscCall(VecEqual(vecRoot0, vecRoot1, &flg)); 70 PetscCheck(flg, comm, PETSC_ERR_PLIB, "root vectors differ"); 71 72 PetscCall(VecDestroy(&vecRoot0)); 73 PetscCall(VecDestroy(&vecRoot1)); 74 PetscCall(VecDestroy(&vecLeave0)); 75 PetscCall(VecDestroy(&vecLeave1)); 76 PetscFunctionReturn(0); 77 } 78 79 PetscErrorCode CreateReferenceSF(AppCtx *ctx, PetscSF *refSF) { 80 PetscInt i, j, k, r; 81 PetscInt *ilocal = NULL; 82 PetscSFNode *iremote; 83 PetscInt nLeaves = ctx->nsfs * ctx->nLeavesPerRank * ctx->size; 84 PetscInt nroots = ctx->nLeavesPerRank * ctx->nsfs; 85 PetscSF sf; 86 87 PetscFunctionBegin; 88 ilocal = NULL; 89 if (ctx->sparseLeaves) { PetscCall(PetscCalloc1(nLeaves + 1, &ilocal)); } 90 PetscCall(PetscMalloc1(nLeaves, &iremote)); 91 PetscCall(PetscSFCreate(ctx->comm, &sf)); 92 for (i = 0, j = 0; i < ctx->nsfs; i++) { 93 for (r = 0; r < ctx->size; r++) { 94 for (k = 0; k < ctx->nLeavesPerRank; k++, j++) { 95 if (ctx->sparseLeaves) { ilocal[j + 1] = ilocal[j] + ctx->leaveStep; } 96 iremote[j].rank = r; 97 iremote[j].index = k + i * ctx->nLeavesPerRank; 98 } 99 } 100 } 101 PetscCall(PetscSFSetGraph(sf, nroots, nLeaves, ilocal, PETSC_OWN_POINTER, iremote, PETSC_OWN_POINTER)); 102 *refSF = sf; 103 PetscFunctionReturn(0); 104 } 105 106 PetscErrorCode CreateSFs(AppCtx *ctx, PetscSF *newSFs[], PetscInt *leafOffsets[]) { 107 PetscInt i; 108 PetscInt *lOffsets = NULL; 109 PetscSF *sfs; 110 PetscInt nLeaves = ctx->nLeavesPerRank * ctx->size; 111 PetscInt nroots = ctx->shareRoots ? ctx->nLeavesPerRank * ctx->nsfs : ctx->nLeavesPerRank; 112 113 PetscFunctionBegin; 114 if (ctx->sparseLeaves) { PetscCall(PetscCalloc1(ctx->nsfs + 1, &lOffsets)); } 115 PetscCall(PetscMalloc1(ctx->nsfs, &sfs)); 116 for (i = 0; i < ctx->nsfs; i++) { 117 PetscInt j, k; 118 PetscMPIInt r; 119 PetscInt *ilocal = NULL; 120 PetscSFNode *iremote; 121 122 if (ctx->sparseLeaves) { PetscCall(PetscCalloc1(nLeaves + 1, &ilocal)); } 123 PetscCall(PetscMalloc1(nLeaves, &iremote)); 124 for (r = 0, j = 0; r < ctx->size; r++) { 125 for (k = 0; k < ctx->nLeavesPerRank; k++, j++) { 126 if (ctx->sparseLeaves) { ilocal[j + 1] = ilocal[j] + ctx->leaveStep; } 127 iremote[j].rank = r; 128 iremote[j].index = ctx->shareRoots ? k + i * ctx->nLeavesPerRank : k; 129 } 130 } 131 if (ctx->sparseLeaves) lOffsets[i + 1] = lOffsets[i] + ilocal[j]; 132 133 PetscCall(PetscSFCreate(ctx->comm, &sfs[i])); 134 PetscCall(PetscSFSetGraph(sfs[i], nroots, nLeaves, ilocal, PETSC_OWN_POINTER, iremote, PETSC_OWN_POINTER)); 135 } 136 *newSFs = sfs; 137 *leafOffsets = lOffsets; 138 PetscFunctionReturn(0); 139 } 140 141 PetscErrorCode DestroySFs(AppCtx *ctx, PetscSF *sfs[]) { 142 PetscInt i; 143 144 PetscFunctionBegin; 145 for (i = 0; i < ctx->nsfs; i++) { PetscCall(PetscSFDestroy(&(*sfs)[i])); } 146 PetscCall(PetscFree(*sfs)); 147 PetscFunctionReturn(0); 148 } 149 150 int main(int argc, char **argv) { 151 AppCtx ctx; 152 PetscSF sf, sfRef; 153 PetscSF *sfs; 154 PetscInt *leafOffsets; 155 MPI_Comm comm; 156 157 PetscFunctionBeginUser; 158 PetscCall(PetscInitialize(&argc, &argv, NULL, help)); 159 comm = PETSC_COMM_WORLD; 160 PetscCall(GetOptions(comm, &ctx)); 161 162 PetscCall(CreateSFs(&ctx, &sfs, &leafOffsets)); 163 PetscCall(PetscSFConcatenate(comm, ctx.nsfs, sfs, ctx.shareRoots, leafOffsets, &sf)); 164 PetscCall(CreateReferenceSF(&ctx, &sfRef)); 165 PetscCall(PetscSFCheckEqual_Private(sf, sfRef)); 166 167 PetscCall(DestroySFs(&ctx, &sfs)); 168 PetscCall(PetscFree(leafOffsets)); 169 PetscCall(PetscSFDestroy(&sf)); 170 PetscCall(PetscSFDestroy(&sfRef)); 171 PetscCall(PetscFinalize()); 172 return 0; 173 } 174 175 /*TEST 176 test: 177 nsize: {{1 3}} 178 args: -nsfs {{1 3}} -n_leaves_per_rank {{0 1 5}} -leave_step {{1 3}} -share_roots {{true false}} 179 TEST*/ 180