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