1 2 #include <../src/vec/is/sf/impls/basic/gatherv/sfgatherv.h> 3 4 /* Reuse the type. The difference is some fields (displs, recvcounts) are only significant 5 on rank 0 in Gatherv. On other ranks they are harmless NULL. 6 */ 7 typedef PetscSF_Allgatherv PetscSF_Gatherv; 8 9 PETSC_INTERN PetscErrorCode PetscSFBcastBegin_Gatherv(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata,MPI_Op op) 10 { 11 PetscErrorCode ierr; 12 PetscSFLink link; 13 PetscMPIInt sendcount; 14 MPI_Comm comm; 15 PetscSF_Gatherv *dat = (PetscSF_Gatherv*)sf->data; 16 void *rootbuf = NULL,*leafbuf = NULL; /* buffer seen by MPI */ 17 MPI_Request *req; 18 19 PetscFunctionBegin; 20 ierr = PetscSFLinkCreate(sf,unit,rootmtype,rootdata,leafmtype,leafdata,op,PETSCSF_BCAST,&link);CHKERRQ(ierr); 21 ierr = PetscSFLinkPackRootData(sf,link,PETSCSF_REMOTE,rootdata);CHKERRQ(ierr); 22 ierr = PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf,link,PETSC_TRUE/* device2host before sending */);CHKERRQ(ierr); 23 ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr); 24 ierr = PetscMPIIntCast(sf->nroots,&sendcount);CHKERRQ(ierr); 25 ierr = PetscSFLinkGetMPIBuffersAndRequests(sf,link,PETSCSF_ROOT2LEAF,&rootbuf,&leafbuf,&req,NULL);CHKERRQ(ierr); 26 ierr = PetscSFLinkSyncStreamBeforeCallMPI(sf,link,PETSCSF_ROOT2LEAF);CHKERRQ(ierr); 27 ierr = MPIU_Igatherv(rootbuf,sendcount,unit,leafbuf,dat->recvcounts,dat->displs,unit,0/*rank 0*/,comm,req);CHKERRMPI(ierr); 28 PetscFunctionReturn(0); 29 } 30 31 static PetscErrorCode PetscSFReduceBegin_Gatherv(PetscSF sf,MPI_Datatype unit,PetscMemType leafmtype,const void *leafdata,PetscMemType rootmtype,void *rootdata,MPI_Op op) 32 { 33 PetscErrorCode ierr; 34 PetscSFLink link; 35 PetscMPIInt recvcount; 36 MPI_Comm comm; 37 PetscSF_Gatherv *dat = (PetscSF_Gatherv*)sf->data; 38 void *rootbuf = NULL,*leafbuf = NULL; /* buffer seen by MPI */ 39 MPI_Request *req; 40 41 PetscFunctionBegin; 42 ierr = PetscSFLinkCreate(sf,unit,rootmtype,rootdata,leafmtype,leafdata,op,PETSCSF_REDUCE,&link);CHKERRQ(ierr); 43 ierr = PetscSFLinkPackLeafData(sf,link,PETSCSF_REMOTE,leafdata);CHKERRQ(ierr); 44 ierr = PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf,link,PETSC_TRUE/* device2host before sending */);CHKERRQ(ierr); 45 ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr); 46 ierr = PetscMPIIntCast(sf->nroots,&recvcount);CHKERRQ(ierr); 47 ierr = PetscSFLinkGetMPIBuffersAndRequests(sf,link,PETSCSF_LEAF2ROOT,&rootbuf,&leafbuf,&req,NULL);CHKERRQ(ierr); 48 ierr = PetscSFLinkSyncStreamBeforeCallMPI(sf,link,PETSCSF_LEAF2ROOT);CHKERRQ(ierr); 49 ierr = MPIU_Iscatterv(leafbuf,dat->recvcounts,dat->displs,unit,rootbuf,recvcount,unit,0,comm,req);CHKERRMPI(ierr); 50 PetscFunctionReturn(0); 51 } 52 53 PETSC_INTERN PetscErrorCode PetscSFFetchAndOpBegin_Gatherv(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,void *rootdata,PetscMemType leafmtype,const void *leafdata,void *leafupdate,MPI_Op op) 54 { 55 PetscErrorCode ierr; 56 57 PetscFunctionBegin; 58 /* In Gatherv, each root only has one leaf. So we just need to bcast rootdata to leafupdate and then reduce leafdata to rootdata */ 59 ierr = PetscSFBcastBegin(sf,unit,rootdata,leafupdate,MPI_REPLACE);CHKERRQ(ierr); 60 ierr = PetscSFBcastEnd(sf,unit,rootdata,leafupdate,MPI_REPLACE);CHKERRQ(ierr); 61 ierr = PetscSFReduceBegin(sf,unit,leafdata,rootdata,op);CHKERRQ(ierr); 62 PetscFunctionReturn(0); 63 } 64 65 PETSC_INTERN PetscErrorCode PetscSFCreate_Gatherv(PetscSF sf) 66 { 67 PetscErrorCode ierr; 68 PetscSF_Gatherv *dat = (PetscSF_Gatherv*)sf->data; 69 70 PetscFunctionBegin; 71 sf->ops->BcastEnd = PetscSFBcastEnd_Basic; 72 sf->ops->ReduceEnd = PetscSFReduceEnd_Basic; 73 74 /* Inherit from Allgatherv */ 75 sf->ops->SetUp = PetscSFSetUp_Allgatherv; 76 sf->ops->Reset = PetscSFReset_Allgatherv; 77 sf->ops->Destroy = PetscSFDestroy_Allgatherv; 78 sf->ops->GetGraph = PetscSFGetGraph_Allgatherv; 79 sf->ops->GetLeafRanks = PetscSFGetLeafRanks_Allgatherv; 80 sf->ops->GetRootRanks = PetscSFGetRootRanks_Allgatherv; 81 sf->ops->FetchAndOpEnd = PetscSFFetchAndOpEnd_Allgatherv; 82 sf->ops->CreateLocalSF = PetscSFCreateLocalSF_Allgatherv; 83 84 /* Gatherv stuff */ 85 sf->ops->BcastBegin = PetscSFBcastBegin_Gatherv; 86 sf->ops->ReduceBegin = PetscSFReduceBegin_Gatherv; 87 sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Gatherv; 88 89 ierr = PetscNewLog(sf,&dat);CHKERRQ(ierr); 90 sf->data = (void*)dat; 91 PetscFunctionReturn(0); 92 } 93