Lines Matching refs:sf
54 PetscErrorCode PetscSFReset_Basic_NVSHMEM(PetscSF sf) in PetscSFReset_Basic_NVSHMEM() argument
56 PetscSF_Basic *bas = (PetscSF_Basic *)sf->data; in PetscSFReset_Basic_NVSHMEM()
60 PetscCall(PetscSFFree(sf, PETSC_MEMTYPE_CUDA, bas->leafbufdisp_d)); in PetscSFReset_Basic_NVSHMEM()
61 PetscCall(PetscSFFree(sf, PETSC_MEMTYPE_CUDA, bas->leafsigdisp_d)); in PetscSFReset_Basic_NVSHMEM()
62 PetscCall(PetscSFFree(sf, PETSC_MEMTYPE_CUDA, bas->iranks_d)); in PetscSFReset_Basic_NVSHMEM()
63 PetscCall(PetscSFFree(sf, PETSC_MEMTYPE_CUDA, bas->ioffset_d)); in PetscSFReset_Basic_NVSHMEM()
65 PetscCall(PetscFree2(sf->rootsigdisp, sf->rootbufdisp)); in PetscSFReset_Basic_NVSHMEM()
66 PetscCall(PetscSFFree(sf, PETSC_MEMTYPE_CUDA, sf->rootbufdisp_d)); in PetscSFReset_Basic_NVSHMEM()
67 PetscCall(PetscSFFree(sf, PETSC_MEMTYPE_CUDA, sf->rootsigdisp_d)); in PetscSFReset_Basic_NVSHMEM()
68 PetscCall(PetscSFFree(sf, PETSC_MEMTYPE_CUDA, sf->ranks_d)); in PetscSFReset_Basic_NVSHMEM()
69 PetscCall(PetscSFFree(sf, PETSC_MEMTYPE_CUDA, sf->roffset_d)); in PetscSFReset_Basic_NVSHMEM()
74 static PetscErrorCode PetscSFSetUp_Basic_NVSHMEM(PetscSF sf) in PetscSFSetUp_Basic_NVSHMEM() argument
77 PetscSF_Basic *bas = (PetscSF_Basic *)sf->data; in PetscSFSetUp_Basic_NVSHMEM()
85 PetscCall(PetscObjectGetComm((PetscObject)sf, &comm)); in PetscSFSetUp_Basic_NVSHMEM()
86 PetscCall(PetscObjectGetNewTag((PetscObject)sf, &tag)); in PetscSFSetUp_Basic_NVSHMEM()
88 nRemoteRootRanks = sf->nranks - sf->ndranks; in PetscSFSetUp_Basic_NVSHMEM()
90 sf->nRemoteRootRanks = nRemoteRootRanks; in PetscSFSetUp_Basic_NVSHMEM()
96 stmp[1] = sf->leafbuflen[PETSCSF_REMOTE]; in PetscSFSetUp_Basic_NVSHMEM()
102 sf->nRemoteRootRanksMax = rtmp[0]; in PetscSFSetUp_Basic_NVSHMEM()
103 sf->leafbuflen_rmax = rtmp[1]; in PetscSFSetUp_Basic_NVSHMEM()
110 PetscCall(PetscMalloc2(nRemoteRootRanks, &sf->rootsigdisp, nRemoteRootRanks, &sf->rootbufdisp)); in PetscSFSetUp_Basic_NVSHMEM()
111 …i < nRemoteRootRanks; i++) PetscCallMPI(MPIU_Irecv(&sf->rootsigdisp[i], 1, MPIU_INT, sf->ranks[i +… in PetscSFSetUp_Basic_NVSHMEM()
115 …i < nRemoteRootRanks; i++) PetscCallMPI(MPIU_Irecv(&sf->rootbufdisp[i], 1, MPIU_INT, sf->ranks[i +… in PetscSFSetUp_Basic_NVSHMEM()
122 PetscCallCUDA(cudaMalloc((void **)&sf->rootbufdisp_d, nRemoteRootRanks * sizeof(PetscInt))); in PetscSFSetUp_Basic_NVSHMEM()
123 PetscCallCUDA(cudaMalloc((void **)&sf->rootsigdisp_d, nRemoteRootRanks * sizeof(PetscInt))); in PetscSFSetUp_Basic_NVSHMEM()
124 PetscCallCUDA(cudaMalloc((void **)&sf->ranks_d, nRemoteRootRanks * sizeof(PetscMPIInt))); in PetscSFSetUp_Basic_NVSHMEM()
125 PetscCallCUDA(cudaMalloc((void **)&sf->roffset_d, (nRemoteRootRanks + 1) * sizeof(PetscInt))); in PetscSFSetUp_Basic_NVSHMEM()
127 …PetscCallCUDA(cudaMemcpyAsync(sf->rootbufdisp_d, sf->rootbufdisp, nRemoteRootRanks * sizeof(PetscI… in PetscSFSetUp_Basic_NVSHMEM()
128 …PetscCallCUDA(cudaMemcpyAsync(sf->rootsigdisp_d, sf->rootsigdisp, nRemoteRootRanks * sizeof(PetscI… in PetscSFSetUp_Basic_NVSHMEM()
129 …PetscCallCUDA(cudaMemcpyAsync(sf->ranks_d, sf->ranks + sf->ndranks, nRemoteRootRanks * sizeof(Pets… in PetscSFSetUp_Basic_NVSHMEM()
130 …PetscCallCUDA(cudaMemcpyAsync(sf->roffset_d, sf->roffset + sf->ndranks, (nRemoteRootRanks + 1) * s… in PetscSFSetUp_Basic_NVSHMEM()
135 …for (i = 0; i < nRemoteRootRanks; i++) PetscCallMPI(MPI_Send(&i, 1, MPIU_INT, sf->ranks[i + sf->nd… in PetscSFSetUp_Basic_NVSHMEM()
140 tmp = sf->roffset[i + sf->ndranks] - sf->roffset[sf->ndranks]; in PetscSFSetUp_Basic_NVSHMEM()
141 PetscCallMPI(MPI_Send(&tmp, 1, MPIU_INT, sf->ranks[i + sf->ndranks], tag, comm)); in PetscSFSetUp_Basic_NVSHMEM()
159 PetscErrorCode PetscSFLinkNvshmemCheck(PetscSF sf, PetscMemType rootmtype, const void *rootdata, Pe… in PetscSFLinkNvshmemCheck() argument
166 PetscCall(PetscObjectGetComm((PetscObject)sf, &comm)); in PetscSFLinkNvshmemCheck()
170 sf->checked_nvshmem_eligibility = PETSC_TRUE; in PetscSFLinkNvshmemCheck()
171 if (sf->use_nvshmem && !sf->checked_nvshmem_eligibility) { in PetscSFLinkNvshmemCheck()
173 PetscCall(PetscObjectTypeCompare((PetscObject)sf, PETSCSFBASIC, &isBasic)); in PetscSFLinkNvshmemCheck()
175 …if (!isBasic || (result != MPI_IDENT && result != MPI_CONGRUENT)) sf->use_nvshmem = PETSC_FALSE; /… in PetscSFLinkNvshmemCheck()
181 if (sf->use_nvshmem) { in PetscSFLinkNvshmemCheck()
184 if (hasNullRank) sf->use_nvshmem = PETSC_FALSE; in PetscSFLinkNvshmemCheck()
186 sf->checked_nvshmem_eligibility = PETSC_TRUE; /* If eligible, don't do above check again */ in PetscSFLinkNvshmemCheck()
190 if (sf->use_nvshmem) { in PetscSFLinkNvshmemCheck()
199 if (!sf->setup_nvshmem) { /* Set up nvshmem related fields on this SF on-demand */ in PetscSFLinkNvshmemCheck()
200 PetscCall(PetscSFSetUp_Basic_NVSHMEM(sf)); in PetscSFLinkNvshmemCheck()
201 sf->setup_nvshmem = PETSC_TRUE; in PetscSFLinkNvshmemCheck()
214 static PetscErrorCode PetscSFLinkBuildDependenceBegin(PetscSF sf, PetscSFLink link, PetscSFDirectio… in PetscSFLinkBuildDependenceBegin() argument
217 PetscSF_Basic *bas = (PetscSF_Basic *)sf->data; in PetscSFLinkBuildDependenceBegin()
218 …PetscInt buflen = (direction == PETSCSF_ROOT2LEAF) ? bas->rootbuflen[PETSCSF_REMOTE] : sf->l… in PetscSFLinkBuildDependenceBegin()
229 static PetscErrorCode PetscSFLinkBuildDependenceEnd(PetscSF sf, PetscSFLink link, PetscSFDirection … in PetscSFLinkBuildDependenceEnd() argument
232 PetscSF_Basic *bas = (PetscSF_Basic *)sf->data; in PetscSFLinkBuildDependenceEnd()
233 …PetscInt buflen = (direction == PETSCSF_ROOT2LEAF) ? sf->leafbuflen[PETSCSF_REMOTE] : bas->r… in PetscSFLinkBuildDependenceEnd()
304 static PetscErrorCode PetscSFLinkWaitSignalsOfCompletionOfGettingData_NVSHMEM(PetscSF sf, PetscSFLi… in PetscSFLinkWaitSignalsOfCompletionOfGettingData_NVSHMEM() argument
306 PetscSF_Basic *bas = (PetscSF_Basic *)sf->data; in PetscSFLinkWaitSignalsOfCompletionOfGettingData_NVSHMEM()
316 n = sf->nRemoteRootRanks; in PetscSFLinkWaitSignalsOfCompletionOfGettingData_NVSHMEM()
339 static PetscErrorCode PetscSFLinkGetDataBegin_NVSHMEM(PetscSF sf, PetscSFLink link, PetscSFDirectio… in PetscSFLinkGetDataBegin_NVSHMEM() argument
342 PetscSF_Basic *bas = (PetscSF_Basic *)sf->data; in PetscSFLinkGetDataBegin_NVSHMEM()
355 PetscCall(PetscSFLinkBuildDependenceBegin(sf, link, direction)); in PetscSFLinkGetDataBegin_NVSHMEM()
357 nsrcranks = sf->nRemoteRootRanks; in PetscSFLinkGetDataBegin_NVSHMEM()
360 …srcdisp_h = sf->rootbufdisp; /* for my i-th remote root rank, I will access its buf at offset roo… in PetscSFLinkGetDataBegin_NVSHMEM()
361 srcdisp_d = sf->rootbufdisp_d; in PetscSFLinkGetDataBegin_NVSHMEM()
362 srcranks_h = sf->ranks + sf->ndranks; /* my (remote) root ranks */ in PetscSFLinkGetDataBegin_NVSHMEM()
363 srcranks_d = sf->ranks_d; in PetscSFLinkGetDataBegin_NVSHMEM()
368 …dstdisp_h = sf->roffset + sf->ndranks; /* offsets of the local leaf buf. Note dstdisp[0] is not n… in PetscSFLinkGetDataBegin_NVSHMEM()
369 dstdisp_d = sf->roffset_d; in PetscSFLinkGetDataBegin_NVSHMEM()
383 ndstranks = sf->nRemoteRootRanks; in PetscSFLinkGetDataBegin_NVSHMEM()
388 dstranks_d = sf->ranks_d; /* my (remote) root ranks */ in PetscSFLinkGetDataBegin_NVSHMEM()
391 dstsigdisp_d = sf->rootsigdisp_d; in PetscSFLinkGetDataBegin_NVSHMEM()
435 static PetscErrorCode PetscSFLinkGetDataEnd_NVSHMEM(PetscSF sf, PetscSFLink link, PetscSFDirection … in PetscSFLinkGetDataEnd_NVSHMEM() argument
438 PetscSF_Basic *bas = (PetscSF_Basic *)sf->data; in PetscSFLinkGetDataEnd_NVSHMEM()
445 nsrcranks = sf->nRemoteRootRanks; in PetscSFLinkGetDataEnd_NVSHMEM()
447 srcsigdisp = sf->rootsigdisp_d; /* offset of each root signal */ in PetscSFLinkGetDataEnd_NVSHMEM()
448 srcranks = sf->ranks_d; /* ranks of the n root ranks */ in PetscSFLinkGetDataEnd_NVSHMEM()
462 PetscCall(PetscSFLinkBuildDependenceEnd(sf, link, direction)); in PetscSFLinkGetDataEnd_NVSHMEM()
514 static PetscErrorCode PetscSFLinkPutDataBegin_NVSHMEM(PetscSF sf, PetscSFLink link, PetscSFDirectio… in PetscSFLinkPutDataBegin_NVSHMEM() argument
517 PetscSF_Basic *bas = (PetscSF_Basic *)sf->data; in PetscSFLinkPutDataBegin_NVSHMEM()
527 PetscCall(PetscSFLinkBuildDependenceBegin(sf, link, direction)); in PetscSFLinkPutDataBegin_NVSHMEM()
542 ndstranks = sf->nRemoteRootRanks; in PetscSFLinkPutDataBegin_NVSHMEM()
546 srcdisp_h = sf->roffset + sf->ndranks; /* offsets of leafbuf */ in PetscSFLinkPutDataBegin_NVSHMEM()
547 srcdisp_d = sf->roffset_d; in PetscSFLinkPutDataBegin_NVSHMEM()
550 …dstdisp_h = sf->rootbufdisp; /* for my i-th remote root rank, I will access its root buf at offse… in PetscSFLinkPutDataBegin_NVSHMEM()
551 dstdisp_d = sf->rootbufdisp_d; in PetscSFLinkPutDataBegin_NVSHMEM()
552 dstranks_h = sf->ranks + sf->ndranks; /* remote root ranks */ in PetscSFLinkPutDataBegin_NVSHMEM()
553 dstranks_d = sf->ranks_d; in PetscSFLinkPutDataBegin_NVSHMEM()
606 static PetscErrorCode PetscSFLinkPutDataEnd_NVSHMEM(PetscSF sf, PetscSFLink link, PetscSFDirection … in PetscSFLinkPutDataEnd_NVSHMEM() argument
609 PetscSF_Basic *bas = (PetscSF_Basic *)sf->data; in PetscSFLinkPutDataEnd_NVSHMEM()
616 nsrcranks = sf->nRemoteRootRanks; in PetscSFLinkPutDataEnd_NVSHMEM()
625 ndstranks = sf->nRemoteRootRanks; in PetscSFLinkPutDataEnd_NVSHMEM()
626 dstranks = sf->ranks_d; in PetscSFLinkPutDataEnd_NVSHMEM()
628 dstsigdisp = sf->rootsigdisp_d; in PetscSFLinkPutDataEnd_NVSHMEM()
635 PetscCall(PetscSFLinkBuildDependenceEnd(sf, link, direction)); in PetscSFLinkPutDataEnd_NVSHMEM()
640 static PetscErrorCode PetscSFLinkSendSignalsToAllowPuttingData_NVSHMEM(PetscSF sf, PetscSFLink link… in PetscSFLinkSendSignalsToAllowPuttingData_NVSHMEM() argument
642 PetscSF_Basic *bas = (PetscSF_Basic *)sf->data; in PetscSFLinkSendSignalsToAllowPuttingData_NVSHMEM()
649 nsrcranks = sf->nRemoteRootRanks; in PetscSFLinkSendSignalsToAllowPuttingData_NVSHMEM()
651 srcsigdisp_d = sf->rootsigdisp_d; /* offset of each root signal */ in PetscSFLinkSendSignalsToAllowPuttingData_NVSHMEM()
652 srcranks_d = sf->ranks_d; /* ranks of the n root ranks */ in PetscSFLinkSendSignalsToAllowPuttingData_NVSHMEM()
668 static PetscErrorCode PetscSFLinkDestroy_NVSHMEM(PetscSF sf, PetscSFLink link) in PetscSFLinkDestroy_NVSHMEM() argument
687 PetscErrorCode PetscSFLinkCreate_NVSHMEM(PetscSF sf, MPI_Datatype unit, PetscMemType rootmtype, con… in PetscSFLinkCreate_NVSHMEM() argument
690 PetscSF_Basic *bas = (PetscSF_Basic *)sf->data; in PetscSFLinkCreate_NVSHMEM()
700 if (sf->use_nvshmem_get) { in PetscSFLinkCreate_NVSHMEM()
702 …leafdirect[PETSCSF_REMOTE] = (PetscMemTypeNVSHMEM(leafmtype) && sf->leafcontig[PETSCSF_REMOTE] && … in PetscSFLinkCreate_NVSHMEM()
708 if (sf->use_nvshmem_get) { in PetscSFLinkCreate_NVSHMEM()
713 …leafdirect[PETSCSF_REMOTE] = (PetscMemTypeNVSHMEM(leafmtype) && sf->leafcontig[PETSCSF_REMOTE]) ? … in PetscSFLinkCreate_NVSHMEM()
731 …PetscCall(PetscSFLinkSetUp_Host(sf, link, unit)); /* Comp… in PetscSFLinkCreate_NVSHMEM()
732 …if (sf->backend == PETSCSF_BACKEND_CUDA) PetscCall(PetscSFLinkSetUp_CUDA(sf, link, unit)); /* Setu… in PetscSFLinkCreate_NVSHMEM()
734 …else if (sf->backend == PETSCSF_BACKEND_KOKKOS) PetscCall(PetscSFLinkSetUp_Kokkos(sf, link, unit)); in PetscSFLinkCreate_NVSHMEM()
743 …if (!link->leafSendSig) PetscCall(PetscNvshmemCalloc(sf->nRemoteRootRanksMax * sizeof(uint64_t), (… in PetscSFLinkCreate_NVSHMEM()
744 …if (!link->leafRecvSig) PetscCall(PetscNvshmemCalloc(sf->nRemoteRootRanksMax * sizeof(uint64_t), (… in PetscSFLinkCreate_NVSHMEM()
751 if (sf->use_nvshmem_get) { /* get-based protocol */ in PetscSFLinkCreate_NVSHMEM()
776 …link->leafbuf[PETSCSF_REMOTE][PETSC_MEMTYPE_DEVICE] = (char *)leafdata + sf->leafstart[PETSCSF_REM… in PetscSFLinkCreate_NVSHMEM()
778 …if (!link->leafbuf_alloc[PETSCSF_REMOTE][PETSC_MEMTYPE_DEVICE]) PetscCall(PetscNvshmemMalloc(sf->l… in PetscSFLinkCreate_NVSHMEM()