140e23c03SJunchao Zhang #include <../src/vec/is/sf/impls/basic/sfpack.h> 240e23c03SJunchao Zhang #include <../src/vec/is/sf/impls/basic/sfbasic.h> 340e23c03SJunchao Zhang 4eb02082bSJunchao Zhang #if defined(PETSC_HAVE_CUDA) 5eb02082bSJunchao Zhang #include <cuda_runtime.h> 6eb02082bSJunchao Zhang #endif 740e23c03SJunchao Zhang /* 840e23c03SJunchao Zhang * MPI_Reduce_local is not really useful because it can't handle sparse data and it vectorizes "in the wrong direction", 940e23c03SJunchao Zhang * therefore we pack data types manually. This file defines packing routines for the standard data types. 1040e23c03SJunchao Zhang */ 1140e23c03SJunchao Zhang 12cd620004SJunchao Zhang #define CPPJoin4(a,b,c,d) a##_##b##_##c##_##d 1340e23c03SJunchao Zhang 14cd620004SJunchao Zhang /* Operations working like s += t */ 15cd620004SJunchao Zhang #define OP_BINARY(op,s,t) do {(s) = (s) op (t); } while(0) /* binary ops in the middle such as +, *, && etc. */ 16cd620004SJunchao Zhang #define OP_FUNCTION(op,s,t) do {(s) = op((s),(t)); } while(0) /* ops like a function, such as PetscMax, PetscMin */ 17cd620004SJunchao Zhang #define OP_LXOR(op,s,t) do {(s) = (!(s)) != (!(t));} while(0) /* logical exclusive OR */ 18cd620004SJunchao Zhang #define OP_ASSIGN(op,s,t) do {(s) = (t);} while(0) 19cd620004SJunchao Zhang /* Ref MPI MAXLOC */ 20cd620004SJunchao Zhang #define OP_XLOC(op,s,t) \ 21cd620004SJunchao Zhang do { \ 22cd620004SJunchao Zhang if ((s).u == (t).u) (s).i = PetscMin((s).i,(t).i); \ 23cd620004SJunchao Zhang else if (!((s).u op (t).u)) s = t; \ 24cd620004SJunchao Zhang } while(0) 2540e23c03SJunchao Zhang 2640e23c03SJunchao Zhang /* DEF_PackFunc - macro defining a Pack routine 2740e23c03SJunchao Zhang 2840e23c03SJunchao Zhang Arguments of the macro: 29b23bfdefSJunchao Zhang +Type Type of the basic data in an entry, i.e., int, PetscInt, PetscReal etc. It is not the type of an entry. 30fcc7397dSJunchao Zhang .BS Block size for vectorization. It is a factor of bsz. 31b23bfdefSJunchao Zhang -EQ (bs == BS) ? 1 : 0. EQ is a compile-time const to help compiler optimizations. See below. 3240e23c03SJunchao Zhang 3340e23c03SJunchao Zhang Arguments of the Pack routine: 34cd620004SJunchao Zhang +count Number of indices in idx[]. 35fcc7397dSJunchao Zhang .start When opt and idx are NULL, it means indices are contiguous & start is the first index; otherwise, not used. 36fcc7397dSJunchao Zhang .opt Per-pack optimization plan. NULL means no such plan. 37fcc7397dSJunchao Zhang .idx Indices of entries to packed. 38eb02082bSJunchao Zhang .link Provide a context for the current call, such as link->bs, number of basic types in an entry. Ex. if unit is MPI_2INT, then bs=2 and the basic type is int. 39cd620004SJunchao Zhang .unpacked Address of the unpacked data. The entries will be packed are unpacked[idx[i]],for i in [0,count). 40cd620004SJunchao Zhang -packed Address of the packed data. 4140e23c03SJunchao Zhang */ 42b23bfdefSJunchao Zhang #define DEF_PackFunc(Type,BS,EQ) \ 43fcc7397dSJunchao Zhang static PetscErrorCode CPPJoin4(Pack,Type,BS,EQ)(PetscSFLink link,PetscInt count,PetscInt start,PetscSFPackOpt opt,const PetscInt *idx,const void *unpacked,void *packed) \ 44b23bfdefSJunchao Zhang { \ 4540e23c03SJunchao Zhang PetscErrorCode ierr; \ 46b23bfdefSJunchao Zhang const Type *u = (const Type*)unpacked,*u2; \ 47b23bfdefSJunchao Zhang Type *p = (Type*)packed,*p2; \ 48fcc7397dSJunchao Zhang PetscInt i,j,k,X,Y,r,bs=link->bs; \ 49fcc7397dSJunchao Zhang const PetscInt M = (EQ) ? 1 : bs/BS; /* If EQ, then M=1 enables compiler's const-propagation */ \ 50b23bfdefSJunchao Zhang const PetscInt MBS = M*BS; /* MBS=bs. We turn MBS into a compile time const when EQ=1. */ \ 5140e23c03SJunchao Zhang PetscFunctionBegin; \ 52fcc7397dSJunchao Zhang if (!idx) {ierr = PetscArraycpy(p,u+start*MBS,MBS*count);CHKERRQ(ierr);}/* idx[] are contiguous */ \ 53fcc7397dSJunchao Zhang else if (opt) { /* has optimizations available */ \ 54fcc7397dSJunchao Zhang p2 = p; \ 55fcc7397dSJunchao Zhang for (r=0; r<opt->n; r++) { \ 56fcc7397dSJunchao Zhang u2 = u + opt->start[r]*MBS; \ 57fcc7397dSJunchao Zhang X = opt->X[r]; \ 58fcc7397dSJunchao Zhang Y = opt->Y[r]; \ 59fcc7397dSJunchao Zhang for (k=0; k<opt->dz[r]; k++) \ 60fcc7397dSJunchao Zhang for (j=0; j<opt->dy[r]; j++) { \ 61fcc7397dSJunchao Zhang ierr = PetscArraycpy(p2,u2+(X*Y*k+X*j)*MBS,opt->dx[r]*MBS);CHKERRQ(ierr); \ 62fcc7397dSJunchao Zhang p2 += opt->dx[r]*MBS; \ 63fcc7397dSJunchao Zhang } \ 64fcc7397dSJunchao Zhang } \ 65fcc7397dSJunchao Zhang } else { \ 66b23bfdefSJunchao Zhang for (i=0; i<count; i++) \ 67eb02082bSJunchao Zhang for (j=0; j<M; j++) /* Decent compilers should eliminate this loop when M = const 1 */ \ 68eb02082bSJunchao Zhang for (k=0; k<BS; k++) /* Compiler either unrolls (BS=1) or vectorizes (BS=2,4,8,etc) this loop */ \ 69b23bfdefSJunchao Zhang p[i*MBS+j*BS+k] = u[idx[i]*MBS+j*BS+k]; \ 7040e23c03SJunchao Zhang } \ 7140e23c03SJunchao Zhang PetscFunctionReturn(0); \ 7240e23c03SJunchao Zhang } 7340e23c03SJunchao Zhang 74cd620004SJunchao Zhang /* DEF_Action - macro defining a UnpackAndInsert routine that unpacks data from a contiguous buffer 75cd620004SJunchao Zhang and inserts into a sparse array. 7640e23c03SJunchao Zhang 7740e23c03SJunchao Zhang Arguments: 78b23bfdefSJunchao Zhang .Type Type of the data 7940e23c03SJunchao Zhang .BS Block size for vectorization 80b23bfdefSJunchao Zhang .EQ (bs == BS) ? 1 : 0. EQ is a compile-time const. 8140e23c03SJunchao Zhang 8240e23c03SJunchao Zhang Notes: 8340e23c03SJunchao Zhang This macro is not combined with DEF_ActionAndOp because we want to use memcpy in this macro. 8440e23c03SJunchao Zhang */ 85cd620004SJunchao Zhang #define DEF_UnpackFunc(Type,BS,EQ) \ 86fcc7397dSJunchao Zhang static PetscErrorCode CPPJoin4(UnpackAndInsert,Type,BS,EQ)(PetscSFLink link,PetscInt count,PetscInt start,PetscSFPackOpt opt,const PetscInt *idx,void *unpacked,const void *packed) \ 87b23bfdefSJunchao Zhang { \ 8840e23c03SJunchao Zhang PetscErrorCode ierr; \ 89b23bfdefSJunchao Zhang Type *u = (Type*)unpacked,*u2; \ 90fcc7397dSJunchao Zhang const Type *p = (const Type*)packed; \ 91fcc7397dSJunchao Zhang PetscInt i,j,k,X,Y,r,bs=link->bs; \ 92fcc7397dSJunchao Zhang const PetscInt M = (EQ) ? 1 : bs/BS; /* If EQ, then M=1 enables compiler's const-propagation */ \ 93b23bfdefSJunchao Zhang const PetscInt MBS = M*BS; /* MBS=bs. We turn MBS into a compile time const when EQ=1. */ \ 9440e23c03SJunchao Zhang PetscFunctionBegin; \ 95b23bfdefSJunchao Zhang if (!idx) { \ 96fcc7397dSJunchao Zhang u += start*MBS; \ 97fcc7397dSJunchao Zhang if (u != p) {ierr = PetscArraycpy(u,p,count*MBS);CHKERRQ(ierr);} \ 98fcc7397dSJunchao Zhang } else if (opt) { /* has optimizations available */ \ 99fcc7397dSJunchao Zhang for (r=0; r<opt->n; r++) { \ 100fcc7397dSJunchao Zhang u2 = u + opt->start[r]*MBS; \ 101fcc7397dSJunchao Zhang X = opt->X[r]; \ 102fcc7397dSJunchao Zhang Y = opt->Y[r]; \ 103fcc7397dSJunchao Zhang for (k=0; k<opt->dz[r]; k++) \ 104fcc7397dSJunchao Zhang for (j=0; j<opt->dy[r]; j++) { \ 105fcc7397dSJunchao Zhang ierr = PetscArraycpy(u2+(X*Y*k+X*j)*MBS,p,opt->dx[r]*MBS);CHKERRQ(ierr); \ 106fcc7397dSJunchao Zhang p += opt->dx[r]*MBS; \ 107fcc7397dSJunchao Zhang } \ 108fcc7397dSJunchao Zhang } \ 109fcc7397dSJunchao Zhang } else { \ 110b23bfdefSJunchao Zhang for (i=0; i<count; i++) \ 111b23bfdefSJunchao Zhang for (j=0; j<M; j++) \ 112cd620004SJunchao Zhang for (k=0; k<BS; k++) u[idx[i]*MBS+j*BS+k] = p[i*MBS+j*BS+k]; \ 11340e23c03SJunchao Zhang } \ 11440e23c03SJunchao Zhang PetscFunctionReturn(0); \ 11540e23c03SJunchao Zhang } 11640e23c03SJunchao Zhang 117cd620004SJunchao Zhang /* DEF_UnpackAndOp - macro defining a UnpackAndOp routine where Op should not be Insert 11840e23c03SJunchao Zhang 11940e23c03SJunchao Zhang Arguments: 120cd620004SJunchao Zhang +Opname Name of the Op, such as Add, Mult, LAND, etc. 121b23bfdefSJunchao Zhang .Type Type of the data 12240e23c03SJunchao Zhang .BS Block size for vectorization 123b23bfdefSJunchao Zhang .EQ (bs == BS) ? 1 : 0. EQ is a compile-time const. 124cd620004SJunchao Zhang .Op Operator for the op, such as +, *, &&, ||, PetscMax, PetscMin, etc. 125cd620004SJunchao Zhang .OpApply Macro defining application of the op. Could be OP_BINARY, OP_FUNCTION, OP_LXOR 12640e23c03SJunchao Zhang */ 127cd620004SJunchao Zhang #define DEF_UnpackAndOp(Type,BS,EQ,Opname,Op,OpApply) \ 128fcc7397dSJunchao Zhang static PetscErrorCode CPPJoin4(UnpackAnd##Opname,Type,BS,EQ)(PetscSFLink link,PetscInt count,PetscInt start,PetscSFPackOpt opt,const PetscInt *idx,void *unpacked,const void *packed) \ 129b23bfdefSJunchao Zhang { \ 130cd620004SJunchao Zhang Type *u = (Type*)unpacked,*u2; \ 131fcc7397dSJunchao Zhang const Type *p = (const Type*)packed; \ 132fcc7397dSJunchao Zhang PetscInt i,j,k,X,Y,r,bs=link->bs; \ 133fcc7397dSJunchao Zhang const PetscInt M = (EQ) ? 1 : bs/BS; /* If EQ, then M=1 enables compiler's const-propagation */ \ 134b23bfdefSJunchao Zhang const PetscInt MBS = M*BS; /* MBS=bs. We turn MBS into a compile time const when EQ=1. */ \ 13540e23c03SJunchao Zhang PetscFunctionBegin; \ 136b23bfdefSJunchao Zhang if (!idx) { \ 137fcc7397dSJunchao Zhang u += start*MBS; \ 138cd620004SJunchao Zhang for (i=0; i<count; i++) \ 139cd620004SJunchao Zhang for (j=0; j<M; j++) \ 140cd620004SJunchao Zhang for (k=0; k<BS; k++) \ 141cd620004SJunchao Zhang OpApply(Op,u[i*MBS+j*BS+k],p[i*MBS+j*BS+k]); \ 142fcc7397dSJunchao Zhang } else if (opt) { /* idx[] has patterns */ \ 143fcc7397dSJunchao Zhang for (r=0; r<opt->n; r++) { \ 144fcc7397dSJunchao Zhang u2 = u + opt->start[r]*MBS; \ 145fcc7397dSJunchao Zhang X = opt->X[r]; \ 146fcc7397dSJunchao Zhang Y = opt->Y[r]; \ 147fcc7397dSJunchao Zhang for (k=0; k<opt->dz[r]; k++) \ 148fcc7397dSJunchao Zhang for (j=0; j<opt->dy[r]; j++) { \ 149fcc7397dSJunchao Zhang for (i=0; i<opt->dx[r]*MBS; i++) OpApply(Op,u2[(X*Y*k+X*j)*MBS+i],p[i]); \ 150fcc7397dSJunchao Zhang p += opt->dx[r]*MBS; \ 151fcc7397dSJunchao Zhang } \ 152fcc7397dSJunchao Zhang } \ 153fcc7397dSJunchao Zhang } else { \ 154cd620004SJunchao Zhang for (i=0; i<count; i++) \ 155cd620004SJunchao Zhang for (j=0; j<M; j++) \ 156cd620004SJunchao Zhang for (k=0; k<BS; k++) \ 157cd620004SJunchao Zhang OpApply(Op,u[idx[i]*MBS+j*BS+k],p[i*MBS+j*BS+k]); \ 158cd620004SJunchao Zhang } \ 159cd620004SJunchao Zhang PetscFunctionReturn(0); \ 160cd620004SJunchao Zhang } 161cd620004SJunchao Zhang 162cd620004SJunchao Zhang #define DEF_FetchAndOp(Type,BS,EQ,Opname,Op,OpApply) \ 163fcc7397dSJunchao Zhang static PetscErrorCode CPPJoin4(FetchAnd##Opname,Type,BS,EQ)(PetscSFLink link,PetscInt count,PetscInt start,PetscSFPackOpt opt,const PetscInt *idx,void *unpacked,void *packed) \ 164cd620004SJunchao Zhang { \ 165fcc7397dSJunchao Zhang Type *u = (Type*)unpacked,*p = (Type*)packed,tmp; \ 166fcc7397dSJunchao Zhang PetscInt i,j,k,r,l,bs=link->bs; \ 167fcc7397dSJunchao Zhang const PetscInt M = (EQ) ? 1 : bs/BS; \ 168fcc7397dSJunchao Zhang const PetscInt MBS = M*BS; \ 169cd620004SJunchao Zhang PetscFunctionBegin; \ 170fcc7397dSJunchao Zhang for (i=0; i<count; i++) { \ 171fcc7397dSJunchao Zhang r = (!idx ? start+i : idx[i])*MBS; \ 172fcc7397dSJunchao Zhang l = i*MBS; \ 173b23bfdefSJunchao Zhang for (j=0; j<M; j++) \ 174b23bfdefSJunchao Zhang for (k=0; k<BS; k++) { \ 175fcc7397dSJunchao Zhang tmp = u[r+j*BS+k]; \ 176fcc7397dSJunchao Zhang OpApply(Op,u[r+j*BS+k],p[l+j*BS+k]); \ 177fcc7397dSJunchao Zhang p[l+j*BS+k] = tmp; \ 178cd620004SJunchao Zhang } \ 179cd620004SJunchao Zhang } \ 180cd620004SJunchao Zhang PetscFunctionReturn(0); \ 181cd620004SJunchao Zhang } 182cd620004SJunchao Zhang 183cd620004SJunchao Zhang #define DEF_ScatterAndOp(Type,BS,EQ,Opname,Op,OpApply) \ 184fcc7397dSJunchao Zhang static PetscErrorCode CPPJoin4(ScatterAnd##Opname,Type,BS,EQ)(PetscSFLink link,PetscInt count,PetscInt srcStart,PetscSFPackOpt srcOpt,const PetscInt *srcIdx,const void *src,PetscInt dstStart,PetscSFPackOpt dstOpt,const PetscInt *dstIdx,void *dst) \ 185cd620004SJunchao Zhang { \ 186fcc7397dSJunchao Zhang PetscErrorCode ierr; \ 187fcc7397dSJunchao Zhang const Type *u = (const Type*)src; \ 188fcc7397dSJunchao Zhang Type *v = (Type*)dst; \ 189fcc7397dSJunchao Zhang PetscInt i,j,k,s,t,X,Y,bs = link->bs; \ 190cd620004SJunchao Zhang const PetscInt M = (EQ) ? 1 : bs/BS; \ 191cd620004SJunchao Zhang const PetscInt MBS = M*BS; \ 192cd620004SJunchao Zhang PetscFunctionBegin; \ 193fcc7397dSJunchao Zhang if (!srcIdx) { /* src is contiguous */ \ 194fcc7397dSJunchao Zhang u += srcStart*MBS; \ 195fcc7397dSJunchao Zhang ierr = CPPJoin4(UnpackAnd##Opname,Type,BS,EQ)(link,count,dstStart,dstOpt,dstIdx,dst,u);CHKERRQ(ierr); \ 196fcc7397dSJunchao Zhang } else if (srcOpt && !dstIdx) { /* src is 3D, dst is contiguous */ \ 197fcc7397dSJunchao Zhang u += srcOpt->start[0]*MBS; \ 198fcc7397dSJunchao Zhang v += dstStart*MBS; \ 199fcc7397dSJunchao Zhang X = srcOpt->X[0]; Y = srcOpt->Y[0]; \ 200fcc7397dSJunchao Zhang for (k=0; k<srcOpt->dz[0]; k++) \ 201fcc7397dSJunchao Zhang for (j=0; j<srcOpt->dy[0]; j++) { \ 202fcc7397dSJunchao Zhang for (i=0; i<srcOpt->dx[0]*MBS; i++) OpApply(Op,v[i],u[(X*Y*k+X*j)*MBS+i]); \ 203fcc7397dSJunchao Zhang v += srcOpt->dx[0]*MBS; \ 204fcc7397dSJunchao Zhang } \ 205fcc7397dSJunchao Zhang } else { /* all other cases */ \ 206fcc7397dSJunchao Zhang for (i=0; i<count; i++) { \ 207fcc7397dSJunchao Zhang s = (!srcIdx ? srcStart+i : srcIdx[i])*MBS; \ 208fcc7397dSJunchao Zhang t = (!dstIdx ? dstStart+i : dstIdx[i])*MBS; \ 209cd620004SJunchao Zhang for (j=0; j<M; j++) \ 210fcc7397dSJunchao Zhang for (k=0; k<BS; k++) OpApply(Op,v[t+j*BS+k],u[s+j*BS+k]); \ 211fcc7397dSJunchao Zhang } \ 212cd620004SJunchao Zhang } \ 213cd620004SJunchao Zhang PetscFunctionReturn(0); \ 214cd620004SJunchao Zhang } 215cd620004SJunchao Zhang 216cd620004SJunchao Zhang #define DEF_FetchAndOpLocal(Type,BS,EQ,Opname,Op,OpApply) \ 217fcc7397dSJunchao Zhang static PetscErrorCode CPPJoin4(FetchAnd##Opname##Local,Type,BS,EQ)(PetscSFLink link,PetscInt count,PetscInt rootstart,PetscSFPackOpt rootopt,const PetscInt *rootidx,void *rootdata,PetscInt leafstart,PetscSFPackOpt leafopt,const PetscInt *leafidx,const void *leafdata,void *leafupdate) \ 218cd620004SJunchao Zhang { \ 219fcc7397dSJunchao Zhang Type *rdata = (Type*)rootdata,*lupdate = (Type*)leafupdate; \ 220fcc7397dSJunchao Zhang const Type *ldata = (const Type*)leafdata; \ 221fcc7397dSJunchao Zhang PetscInt i,j,k,r,l,bs = link->bs; \ 222cd620004SJunchao Zhang const PetscInt M = (EQ) ? 1 : bs/BS; \ 223cd620004SJunchao Zhang const PetscInt MBS = M*BS; \ 224cd620004SJunchao Zhang PetscFunctionBegin; \ 225fcc7397dSJunchao Zhang for (i=0; i<count; i++) { \ 226fcc7397dSJunchao Zhang r = (rootidx ? rootidx[i] : rootstart+i)*MBS; \ 227fcc7397dSJunchao Zhang l = (leafidx ? leafidx[i] : leafstart+i)*MBS; \ 228cd620004SJunchao Zhang for (j=0; j<M; j++) \ 229cd620004SJunchao Zhang for (k=0; k<BS; k++) { \ 230fcc7397dSJunchao Zhang lupdate[l+j*BS+k] = rdata[r+j*BS+k]; \ 231fcc7397dSJunchao Zhang OpApply(Op,rdata[r+j*BS+k],ldata[l+j*BS+k]); \ 23240e23c03SJunchao Zhang } \ 23340e23c03SJunchao Zhang } \ 23440e23c03SJunchao Zhang PetscFunctionReturn(0); \ 23540e23c03SJunchao Zhang } 23640e23c03SJunchao Zhang 237b23bfdefSJunchao Zhang /* Pack, Unpack/Fetch ops */ 238b23bfdefSJunchao Zhang #define DEF_Pack(Type,BS,EQ) \ 239b23bfdefSJunchao Zhang DEF_PackFunc(Type,BS,EQ) \ 240cd620004SJunchao Zhang DEF_UnpackFunc(Type,BS,EQ) \ 241cd620004SJunchao Zhang DEF_ScatterAndOp(Type,BS,EQ,Insert,=,OP_ASSIGN) \ 242cd620004SJunchao Zhang static void CPPJoin4(PackInit_Pack,Type,BS,EQ)(PetscSFLink link) { \ 243eb02082bSJunchao Zhang link->h_Pack = CPPJoin4(Pack, Type,BS,EQ); \ 244eb02082bSJunchao Zhang link->h_UnpackAndInsert = CPPJoin4(UnpackAndInsert,Type,BS,EQ); \ 245cd620004SJunchao Zhang link->h_ScatterAndInsert= CPPJoin4(ScatterAndInsert,Type,BS,EQ); \ 24640e23c03SJunchao Zhang } 24740e23c03SJunchao Zhang 248b23bfdefSJunchao Zhang /* Add, Mult ops */ 249b23bfdefSJunchao Zhang #define DEF_Add(Type,BS,EQ) \ 250cd620004SJunchao Zhang DEF_UnpackAndOp (Type,BS,EQ,Add, +,OP_BINARY) \ 251cd620004SJunchao Zhang DEF_UnpackAndOp (Type,BS,EQ,Mult,*,OP_BINARY) \ 252cd620004SJunchao Zhang DEF_FetchAndOp (Type,BS,EQ,Add, +,OP_BINARY) \ 253cd620004SJunchao Zhang DEF_ScatterAndOp (Type,BS,EQ,Add, +,OP_BINARY) \ 254cd620004SJunchao Zhang DEF_ScatterAndOp (Type,BS,EQ,Mult,*,OP_BINARY) \ 255cd620004SJunchao Zhang DEF_FetchAndOpLocal(Type,BS,EQ,Add, +,OP_BINARY) \ 256cd620004SJunchao Zhang static void CPPJoin4(PackInit_Add,Type,BS,EQ)(PetscSFLink link) { \ 257eb02082bSJunchao Zhang link->h_UnpackAndAdd = CPPJoin4(UnpackAndAdd, Type,BS,EQ); \ 258eb02082bSJunchao Zhang link->h_UnpackAndMult = CPPJoin4(UnpackAndMult, Type,BS,EQ); \ 259eb02082bSJunchao Zhang link->h_FetchAndAdd = CPPJoin4(FetchAndAdd, Type,BS,EQ); \ 260cd620004SJunchao Zhang link->h_ScatterAndAdd = CPPJoin4(ScatterAndAdd, Type,BS,EQ); \ 261cd620004SJunchao Zhang link->h_ScatterAndMult = CPPJoin4(ScatterAndMult, Type,BS,EQ); \ 262cd620004SJunchao Zhang link->h_FetchAndAddLocal = CPPJoin4(FetchAndAddLocal,Type,BS,EQ); \ 26340e23c03SJunchao Zhang } 26440e23c03SJunchao Zhang 265b23bfdefSJunchao Zhang /* Max, Min ops */ 266b23bfdefSJunchao Zhang #define DEF_Cmp(Type,BS,EQ) \ 267cd620004SJunchao Zhang DEF_UnpackAndOp (Type,BS,EQ,Max,PetscMax,OP_FUNCTION) \ 268cd620004SJunchao Zhang DEF_UnpackAndOp (Type,BS,EQ,Min,PetscMin,OP_FUNCTION) \ 269cd620004SJunchao Zhang DEF_ScatterAndOp(Type,BS,EQ,Max,PetscMax,OP_FUNCTION) \ 270cd620004SJunchao Zhang DEF_ScatterAndOp(Type,BS,EQ,Min,PetscMin,OP_FUNCTION) \ 271cd620004SJunchao Zhang static void CPPJoin4(PackInit_Compare,Type,BS,EQ)(PetscSFLink link) { \ 272eb02082bSJunchao Zhang link->h_UnpackAndMax = CPPJoin4(UnpackAndMax, Type,BS,EQ); \ 273eb02082bSJunchao Zhang link->h_UnpackAndMin = CPPJoin4(UnpackAndMin, Type,BS,EQ); \ 274cd620004SJunchao Zhang link->h_ScatterAndMax = CPPJoin4(ScatterAndMax, Type,BS,EQ); \ 275cd620004SJunchao Zhang link->h_ScatterAndMin = CPPJoin4(ScatterAndMin, Type,BS,EQ); \ 276b23bfdefSJunchao Zhang } 277b23bfdefSJunchao Zhang 278b23bfdefSJunchao Zhang /* Logical ops. 279cd620004SJunchao Zhang The operator in OP_LXOR should be empty but is ||. It is not used. Put here to avoid 28040e23c03SJunchao Zhang the compilation warning "empty macro arguments are undefined in ISO C90" 28140e23c03SJunchao Zhang */ 282b23bfdefSJunchao Zhang #define DEF_Log(Type,BS,EQ) \ 283cd620004SJunchao Zhang DEF_UnpackAndOp (Type,BS,EQ,LAND,&&,OP_BINARY) \ 284cd620004SJunchao Zhang DEF_UnpackAndOp (Type,BS,EQ,LOR, ||,OP_BINARY) \ 285cd620004SJunchao Zhang DEF_UnpackAndOp (Type,BS,EQ,LXOR,||, OP_LXOR) \ 286cd620004SJunchao Zhang DEF_ScatterAndOp(Type,BS,EQ,LAND,&&,OP_BINARY) \ 287cd620004SJunchao Zhang DEF_ScatterAndOp(Type,BS,EQ,LOR, ||,OP_BINARY) \ 288cd620004SJunchao Zhang DEF_ScatterAndOp(Type,BS,EQ,LXOR,||, OP_LXOR) \ 289cd620004SJunchao Zhang static void CPPJoin4(PackInit_Logical,Type,BS,EQ)(PetscSFLink link) { \ 290eb02082bSJunchao Zhang link->h_UnpackAndLAND = CPPJoin4(UnpackAndLAND, Type,BS,EQ); \ 291eb02082bSJunchao Zhang link->h_UnpackAndLOR = CPPJoin4(UnpackAndLOR, Type,BS,EQ); \ 292eb02082bSJunchao Zhang link->h_UnpackAndLXOR = CPPJoin4(UnpackAndLXOR, Type,BS,EQ); \ 293cd620004SJunchao Zhang link->h_ScatterAndLAND = CPPJoin4(ScatterAndLAND,Type,BS,EQ); \ 294cd620004SJunchao Zhang link->h_ScatterAndLOR = CPPJoin4(ScatterAndLOR, Type,BS,EQ); \ 295cd620004SJunchao Zhang link->h_ScatterAndLXOR = CPPJoin4(ScatterAndLXOR,Type,BS,EQ); \ 29640e23c03SJunchao Zhang } 29740e23c03SJunchao Zhang 298b23bfdefSJunchao Zhang /* Bitwise ops */ 299b23bfdefSJunchao Zhang #define DEF_Bit(Type,BS,EQ) \ 300cd620004SJunchao Zhang DEF_UnpackAndOp (Type,BS,EQ,BAND,&,OP_BINARY) \ 301cd620004SJunchao Zhang DEF_UnpackAndOp (Type,BS,EQ,BOR, |,OP_BINARY) \ 302cd620004SJunchao Zhang DEF_UnpackAndOp (Type,BS,EQ,BXOR,^,OP_BINARY) \ 303cd620004SJunchao Zhang DEF_ScatterAndOp(Type,BS,EQ,BAND,&,OP_BINARY) \ 304cd620004SJunchao Zhang DEF_ScatterAndOp(Type,BS,EQ,BOR, |,OP_BINARY) \ 305cd620004SJunchao Zhang DEF_ScatterAndOp(Type,BS,EQ,BXOR,^,OP_BINARY) \ 306cd620004SJunchao Zhang static void CPPJoin4(PackInit_Bitwise,Type,BS,EQ)(PetscSFLink link) { \ 307eb02082bSJunchao Zhang link->h_UnpackAndBAND = CPPJoin4(UnpackAndBAND, Type,BS,EQ); \ 308eb02082bSJunchao Zhang link->h_UnpackAndBOR = CPPJoin4(UnpackAndBOR, Type,BS,EQ); \ 309eb02082bSJunchao Zhang link->h_UnpackAndBXOR = CPPJoin4(UnpackAndBXOR, Type,BS,EQ); \ 310cd620004SJunchao Zhang link->h_ScatterAndBAND = CPPJoin4(ScatterAndBAND,Type,BS,EQ); \ 311cd620004SJunchao Zhang link->h_ScatterAndBOR = CPPJoin4(ScatterAndBOR, Type,BS,EQ); \ 312cd620004SJunchao Zhang link->h_ScatterAndBXOR = CPPJoin4(ScatterAndBXOR,Type,BS,EQ); \ 31340e23c03SJunchao Zhang } 31440e23c03SJunchao Zhang 315cd620004SJunchao Zhang /* Maxloc, Minloc ops */ 316cd620004SJunchao Zhang #define DEF_Xloc(Type,BS,EQ) \ 317cd620004SJunchao Zhang DEF_UnpackAndOp (Type,BS,EQ,Max,>,OP_XLOC) \ 318cd620004SJunchao Zhang DEF_UnpackAndOp (Type,BS,EQ,Min,<,OP_XLOC) \ 319cd620004SJunchao Zhang DEF_ScatterAndOp(Type,BS,EQ,Max,>,OP_XLOC) \ 320cd620004SJunchao Zhang DEF_ScatterAndOp(Type,BS,EQ,Min,<,OP_XLOC) \ 321cd620004SJunchao Zhang static void CPPJoin4(PackInit_Xloc,Type,BS,EQ)(PetscSFLink link) { \ 322cd620004SJunchao Zhang link->h_UnpackAndMaxloc = CPPJoin4(UnpackAndMax, Type,BS,EQ); \ 323cd620004SJunchao Zhang link->h_UnpackAndMinloc = CPPJoin4(UnpackAndMin, Type,BS,EQ); \ 324cd620004SJunchao Zhang link->h_ScatterAndMaxloc = CPPJoin4(ScatterAndMax,Type,BS,EQ); \ 325cd620004SJunchao Zhang link->h_ScatterAndMinloc = CPPJoin4(ScatterAndMin,Type,BS,EQ); \ 32640e23c03SJunchao Zhang } 32740e23c03SJunchao Zhang 328b23bfdefSJunchao Zhang #define DEF_IntegerType(Type,BS,EQ) \ 329b23bfdefSJunchao Zhang DEF_Pack(Type,BS,EQ) \ 330b23bfdefSJunchao Zhang DEF_Add(Type,BS,EQ) \ 331b23bfdefSJunchao Zhang DEF_Cmp(Type,BS,EQ) \ 332b23bfdefSJunchao Zhang DEF_Log(Type,BS,EQ) \ 333b23bfdefSJunchao Zhang DEF_Bit(Type,BS,EQ) \ 334cd620004SJunchao Zhang static void CPPJoin4(PackInit_IntegerType,Type,BS,EQ)(PetscSFLink link) { \ 335b23bfdefSJunchao Zhang CPPJoin4(PackInit_Pack,Type,BS,EQ)(link); \ 336b23bfdefSJunchao Zhang CPPJoin4(PackInit_Add,Type,BS,EQ)(link); \ 337b23bfdefSJunchao Zhang CPPJoin4(PackInit_Compare,Type,BS,EQ)(link); \ 338b23bfdefSJunchao Zhang CPPJoin4(PackInit_Logical,Type,BS,EQ)(link); \ 339b23bfdefSJunchao Zhang CPPJoin4(PackInit_Bitwise,Type,BS,EQ)(link); \ 34040e23c03SJunchao Zhang } 34140e23c03SJunchao Zhang 342b23bfdefSJunchao Zhang #define DEF_RealType(Type,BS,EQ) \ 343b23bfdefSJunchao Zhang DEF_Pack(Type,BS,EQ) \ 344b23bfdefSJunchao Zhang DEF_Add(Type,BS,EQ) \ 345b23bfdefSJunchao Zhang DEF_Cmp(Type,BS,EQ) \ 346cd620004SJunchao Zhang static void CPPJoin4(PackInit_RealType,Type,BS,EQ)(PetscSFLink link) { \ 347b23bfdefSJunchao Zhang CPPJoin4(PackInit_Pack,Type,BS,EQ)(link); \ 348b23bfdefSJunchao Zhang CPPJoin4(PackInit_Add,Type,BS,EQ)(link); \ 349b23bfdefSJunchao Zhang CPPJoin4(PackInit_Compare,Type,BS,EQ)(link); \ 350b23bfdefSJunchao Zhang } 35140e23c03SJunchao Zhang 35240e23c03SJunchao Zhang #if defined(PETSC_HAVE_COMPLEX) 353b23bfdefSJunchao Zhang #define DEF_ComplexType(Type,BS,EQ) \ 354b23bfdefSJunchao Zhang DEF_Pack(Type,BS,EQ) \ 355b23bfdefSJunchao Zhang DEF_Add(Type,BS,EQ) \ 356cd620004SJunchao Zhang static void CPPJoin4(PackInit_ComplexType,Type,BS,EQ)(PetscSFLink link) { \ 357b23bfdefSJunchao Zhang CPPJoin4(PackInit_Pack,Type,BS,EQ)(link); \ 358b23bfdefSJunchao Zhang CPPJoin4(PackInit_Add,Type,BS,EQ)(link); \ 359b23bfdefSJunchao Zhang } 36040e23c03SJunchao Zhang #endif 361b23bfdefSJunchao Zhang 362b23bfdefSJunchao Zhang #define DEF_DumbType(Type,BS,EQ) \ 363b23bfdefSJunchao Zhang DEF_Pack(Type,BS,EQ) \ 364cd620004SJunchao Zhang static void CPPJoin4(PackInit_DumbType,Type,BS,EQ)(PetscSFLink link) { \ 365b23bfdefSJunchao Zhang CPPJoin4(PackInit_Pack,Type,BS,EQ)(link); \ 366b23bfdefSJunchao Zhang } 367b23bfdefSJunchao Zhang 368b23bfdefSJunchao Zhang /* Maxloc, Minloc */ 369cd620004SJunchao Zhang #define DEF_PairType(Type,BS,EQ) \ 370cd620004SJunchao Zhang DEF_Pack(Type,BS,EQ) \ 371cd620004SJunchao Zhang DEF_Xloc(Type,BS,EQ) \ 372cd620004SJunchao Zhang static void CPPJoin4(PackInit_PairType,Type,BS,EQ)(PetscSFLink link) { \ 373cd620004SJunchao Zhang CPPJoin4(PackInit_Pack,Type,BS,EQ)(link); \ 374cd620004SJunchao Zhang CPPJoin4(PackInit_Xloc,Type,BS,EQ)(link); \ 375b23bfdefSJunchao Zhang } 376b23bfdefSJunchao Zhang 377b23bfdefSJunchao Zhang DEF_IntegerType(PetscInt,1,1) /* unit = 1 MPIU_INT */ 378b23bfdefSJunchao Zhang DEF_IntegerType(PetscInt,2,1) /* unit = 2 MPIU_INTs */ 379b23bfdefSJunchao Zhang DEF_IntegerType(PetscInt,4,1) /* unit = 4 MPIU_INTs */ 380b23bfdefSJunchao Zhang DEF_IntegerType(PetscInt,8,1) /* unit = 8 MPIU_INTs */ 381b23bfdefSJunchao Zhang DEF_IntegerType(PetscInt,1,0) /* unit = 1*n MPIU_INTs, n>1 */ 382b23bfdefSJunchao Zhang DEF_IntegerType(PetscInt,2,0) /* unit = 2*n MPIU_INTs, n>1 */ 383b23bfdefSJunchao Zhang DEF_IntegerType(PetscInt,4,0) /* unit = 4*n MPIU_INTs, n>1 */ 384b23bfdefSJunchao Zhang DEF_IntegerType(PetscInt,8,0) /* unit = 8*n MPIU_INTs, n>1. Routines with bigger BS are tried first. */ 385b23bfdefSJunchao Zhang 386b23bfdefSJunchao Zhang #if defined(PETSC_USE_64BIT_INDICES) /* Do not need (though it is OK) to generate redundant functions if PetscInt is int */ 387b23bfdefSJunchao Zhang DEF_IntegerType(int,1,1) 388b23bfdefSJunchao Zhang DEF_IntegerType(int,2,1) 389b23bfdefSJunchao Zhang DEF_IntegerType(int,4,1) 390b23bfdefSJunchao Zhang DEF_IntegerType(int,8,1) 391b23bfdefSJunchao Zhang DEF_IntegerType(int,1,0) 392b23bfdefSJunchao Zhang DEF_IntegerType(int,2,0) 393b23bfdefSJunchao Zhang DEF_IntegerType(int,4,0) 394b23bfdefSJunchao Zhang DEF_IntegerType(int,8,0) 395b23bfdefSJunchao Zhang #endif 396b23bfdefSJunchao Zhang 397b23bfdefSJunchao Zhang /* The typedefs are used to get a typename without space that CPPJoin can handle */ 398b23bfdefSJunchao Zhang typedef signed char SignedChar; 399b23bfdefSJunchao Zhang DEF_IntegerType(SignedChar,1,1) 400b23bfdefSJunchao Zhang DEF_IntegerType(SignedChar,2,1) 401b23bfdefSJunchao Zhang DEF_IntegerType(SignedChar,4,1) 402b23bfdefSJunchao Zhang DEF_IntegerType(SignedChar,8,1) 403b23bfdefSJunchao Zhang DEF_IntegerType(SignedChar,1,0) 404b23bfdefSJunchao Zhang DEF_IntegerType(SignedChar,2,0) 405b23bfdefSJunchao Zhang DEF_IntegerType(SignedChar,4,0) 406b23bfdefSJunchao Zhang DEF_IntegerType(SignedChar,8,0) 407b23bfdefSJunchao Zhang 408b23bfdefSJunchao Zhang typedef unsigned char UnsignedChar; 409b23bfdefSJunchao Zhang DEF_IntegerType(UnsignedChar,1,1) 410b23bfdefSJunchao Zhang DEF_IntegerType(UnsignedChar,2,1) 411b23bfdefSJunchao Zhang DEF_IntegerType(UnsignedChar,4,1) 412b23bfdefSJunchao Zhang DEF_IntegerType(UnsignedChar,8,1) 413b23bfdefSJunchao Zhang DEF_IntegerType(UnsignedChar,1,0) 414b23bfdefSJunchao Zhang DEF_IntegerType(UnsignedChar,2,0) 415b23bfdefSJunchao Zhang DEF_IntegerType(UnsignedChar,4,0) 416b23bfdefSJunchao Zhang DEF_IntegerType(UnsignedChar,8,0) 417b23bfdefSJunchao Zhang 418b23bfdefSJunchao Zhang DEF_RealType(PetscReal,1,1) 419b23bfdefSJunchao Zhang DEF_RealType(PetscReal,2,1) 420b23bfdefSJunchao Zhang DEF_RealType(PetscReal,4,1) 421b23bfdefSJunchao Zhang DEF_RealType(PetscReal,8,1) 422b23bfdefSJunchao Zhang DEF_RealType(PetscReal,1,0) 423b23bfdefSJunchao Zhang DEF_RealType(PetscReal,2,0) 424b23bfdefSJunchao Zhang DEF_RealType(PetscReal,4,0) 425b23bfdefSJunchao Zhang DEF_RealType(PetscReal,8,0) 426b23bfdefSJunchao Zhang 427b23bfdefSJunchao Zhang #if defined(PETSC_HAVE_COMPLEX) 428b23bfdefSJunchao Zhang DEF_ComplexType(PetscComplex,1,1) 429b23bfdefSJunchao Zhang DEF_ComplexType(PetscComplex,2,1) 430b23bfdefSJunchao Zhang DEF_ComplexType(PetscComplex,4,1) 431b23bfdefSJunchao Zhang DEF_ComplexType(PetscComplex,8,1) 432b23bfdefSJunchao Zhang DEF_ComplexType(PetscComplex,1,0) 433b23bfdefSJunchao Zhang DEF_ComplexType(PetscComplex,2,0) 434b23bfdefSJunchao Zhang DEF_ComplexType(PetscComplex,4,0) 435b23bfdefSJunchao Zhang DEF_ComplexType(PetscComplex,8,0) 436b23bfdefSJunchao Zhang #endif 437b23bfdefSJunchao Zhang 438cd620004SJunchao Zhang #define PairType(Type1,Type2) Type1##_##Type2 439cd620004SJunchao Zhang typedef struct {int u; int i;} PairType(int,int); 440cd620004SJunchao Zhang typedef struct {PetscInt u; PetscInt i;} PairType(PetscInt,PetscInt); 441cd620004SJunchao Zhang DEF_PairType(PairType(int,int),1,1) 442cd620004SJunchao Zhang DEF_PairType(PairType(PetscInt,PetscInt),1,1) 443b23bfdefSJunchao Zhang 444b23bfdefSJunchao Zhang /* If we don't know the basic type, we treat it as a stream of chars or ints */ 445b23bfdefSJunchao Zhang DEF_DumbType(char,1,1) 446b23bfdefSJunchao Zhang DEF_DumbType(char,2,1) 447b23bfdefSJunchao Zhang DEF_DumbType(char,4,1) 448b23bfdefSJunchao Zhang DEF_DumbType(char,1,0) 449b23bfdefSJunchao Zhang DEF_DumbType(char,2,0) 450b23bfdefSJunchao Zhang DEF_DumbType(char,4,0) 451b23bfdefSJunchao Zhang 452eb02082bSJunchao Zhang typedef int DumbInt; /* To have a different name than 'int' used above. The name is used to make routine names. */ 453b23bfdefSJunchao Zhang DEF_DumbType(DumbInt,1,1) 454b23bfdefSJunchao Zhang DEF_DumbType(DumbInt,2,1) 455b23bfdefSJunchao Zhang DEF_DumbType(DumbInt,4,1) 456b23bfdefSJunchao Zhang DEF_DumbType(DumbInt,8,1) 457b23bfdefSJunchao Zhang DEF_DumbType(DumbInt,1,0) 458b23bfdefSJunchao Zhang DEF_DumbType(DumbInt,2,0) 459b23bfdefSJunchao Zhang DEF_DumbType(DumbInt,4,0) 460b23bfdefSJunchao Zhang DEF_DumbType(DumbInt,8,0) 46140e23c03SJunchao Zhang 46240e23c03SJunchao Zhang #if !defined(PETSC_HAVE_MPI_TYPE_DUP) 46340e23c03SJunchao Zhang PETSC_STATIC_INLINE int MPI_Type_dup(MPI_Datatype datatype,MPI_Datatype *newtype) 46440e23c03SJunchao Zhang { 46540e23c03SJunchao Zhang int ierr; 46640e23c03SJunchao Zhang ierr = MPI_Type_contiguous(1,datatype,newtype); if (ierr) return ierr; 46740e23c03SJunchao Zhang ierr = MPI_Type_commit(newtype); if (ierr) return ierr; 46840e23c03SJunchao Zhang return MPI_SUCCESS; 46940e23c03SJunchao Zhang } 47040e23c03SJunchao Zhang #endif 47140e23c03SJunchao Zhang 472cd620004SJunchao Zhang /* 473cd620004SJunchao Zhang The routine Creates a communication link for the given operation. It first looks up its link cache. If 474cd620004SJunchao Zhang there is a free & suitable one, it uses it. Otherwise it creates a new one. 475cd620004SJunchao Zhang 476cd620004SJunchao Zhang A link contains buffers and MPI requests for send/recv. It also contains pack/unpack routines to pack/unpack 477cd620004SJunchao Zhang root/leafdata to/from these buffers. Buffers are allocated at our discretion. When we find root/leafata 478cd620004SJunchao Zhang can be directly passed to MPI, we won't allocate them. Even we allocate buffers, we only allocate 479cd620004SJunchao Zhang those that are needed by the given `sfop` and `op`, in other words, we do lazy memory-allocation. 480cd620004SJunchao Zhang 481cd620004SJunchao Zhang The routine also allocates buffers on CPU when one does not use gpu-aware MPI but data is on GPU. 482cd620004SJunchao Zhang 483cd620004SJunchao Zhang In SFBasic, MPI requests are persistent. They are init'ed until we try to get requests from a link. 484cd620004SJunchao Zhang 485cd620004SJunchao Zhang The routine is shared by SFBasic and SFNeighbor based on the fact they all deal with sparse graphs and 486cd620004SJunchao Zhang need pack/unpack data. 487cd620004SJunchao Zhang */ 488cd620004SJunchao Zhang PetscErrorCode PetscSFLinkCreate(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,const void *leafdata,MPI_Op op,PetscSFOperation sfop,PetscSFLink *mylink) 48940e23c03SJunchao Zhang { 49040e23c03SJunchao Zhang PetscErrorCode ierr; 491cd620004SJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 492cd620004SJunchao Zhang PetscInt i,j,k,nrootreqs,nleafreqs,nreqs; 493cd620004SJunchao Zhang PetscSFLink *p,link; 494cd620004SJunchao Zhang PetscSFDirection direction; 495cd620004SJunchao Zhang MPI_Request *reqs = NULL; 496cd620004SJunchao Zhang PetscBool match,rootdirect[2],leafdirect[2]; 497cd620004SJunchao Zhang PetscMemType rootmtype_mpi,leafmtype_mpi; /* mtypes seen by MPI */ 498cd620004SJunchao Zhang PetscInt rootdirect_mpi,leafdirect_mpi; /* root/leafdirect seen by MPI*/ 499cd620004SJunchao Zhang 500cd620004SJunchao Zhang PetscFunctionBegin; 501cd620004SJunchao Zhang ierr = PetscSFSetErrorOnUnsupportedOverlap(sf,unit,rootdata,leafdata);CHKERRQ(ierr); 502cd620004SJunchao Zhang 503cd620004SJunchao Zhang /* Can we directly use root/leafdirect with the given sf, sfop and op? */ 504cd620004SJunchao Zhang for (i=PETSCSF_LOCAL; i<=PETSCSF_REMOTE; i++) { 505cd620004SJunchao Zhang if (sfop == PETSCSF_BCAST) { 506cd620004SJunchao Zhang rootdirect[i] = bas->rootcontig[i]; /* Pack roots */ 507cd620004SJunchao Zhang leafdirect[i] = (sf->leafcontig[i] && op == MPIU_REPLACE) ? PETSC_TRUE : PETSC_FALSE; /* Unpack leaves */ 508cd620004SJunchao Zhang } else if (sfop == PETSCSF_REDUCE) { 509cd620004SJunchao Zhang leafdirect[i] = sf->leafcontig[i]; /* Pack leaves */ 510cd620004SJunchao Zhang rootdirect[i] = (bas->rootcontig[i] && op == MPIU_REPLACE) ? PETSC_TRUE : PETSC_FALSE; /* Unpack roots */ 511cd620004SJunchao Zhang } else { /* PETSCSF_FETCH */ 512cd620004SJunchao Zhang rootdirect[i] = PETSC_FALSE; /* FETCH always need a separate rootbuf */ 513cd620004SJunchao Zhang leafdirect[i] = PETSC_FALSE; /* We also force allocating a separate leafbuf so that leafdata and leafupdate can share mpi requests */ 514cd620004SJunchao Zhang } 515cd620004SJunchao Zhang } 516cd620004SJunchao Zhang 517c2a741eeSJunchao Zhang if (sf->use_gpu_aware_mpi) { 518cd620004SJunchao Zhang rootmtype_mpi = rootmtype; 519cd620004SJunchao Zhang leafmtype_mpi = leafmtype; 520cd620004SJunchao Zhang } else { 521cd620004SJunchao Zhang rootmtype_mpi = leafmtype_mpi = PETSC_MEMTYPE_HOST; 522cd620004SJunchao Zhang } 523cd620004SJunchao Zhang /* Will root/leafdata be directly accessed by MPI? Without use_gpu_aware_mpi, device data is bufferred on host and then passed to MPI */ 524cd620004SJunchao Zhang rootdirect_mpi = rootdirect[PETSCSF_REMOTE] && (rootmtype_mpi == rootmtype)? 1 : 0; 525cd620004SJunchao Zhang leafdirect_mpi = leafdirect[PETSCSF_REMOTE] && (leafmtype_mpi == leafmtype)? 1 : 0; 526cd620004SJunchao Zhang 527cd620004SJunchao Zhang direction = (sfop == PETSCSF_BCAST)? PETSCSF_ROOT2LEAF : PETSCSF_LEAF2ROOT; 528cd620004SJunchao Zhang nrootreqs = bas->nrootreqs; 529cd620004SJunchao Zhang nleafreqs = sf->nleafreqs; 530cd620004SJunchao Zhang 531cd620004SJunchao Zhang /* Look for free links in cache */ 532cd620004SJunchao Zhang for (p=&bas->avail; (link=*p); p=&link->next) { 533cd620004SJunchao Zhang ierr = MPIPetsc_Type_compare(unit,link->unit,&match);CHKERRQ(ierr); 534cd620004SJunchao Zhang if (match) { 535cd620004SJunchao Zhang /* If root/leafdata will be directly passed to MPI, test if the data used to initialized the MPI requests matches with current. 536cd620004SJunchao Zhang If not, free old requests. New requests will be lazily init'ed until one calls PetscSFLinkGetMPIBuffersAndRequests(). 537cd620004SJunchao Zhang */ 538cd620004SJunchao Zhang if (rootdirect_mpi && sf->persistent && link->rootreqsinited[direction][rootmtype][1] && link->rootdatadirect[direction][rootmtype] != rootdata) { 539cd620004SJunchao Zhang reqs = link->rootreqs[direction][rootmtype][1]; /* Here, rootmtype = rootmtype_mpi */ 540cd620004SJunchao Zhang for (i=0; i<nrootreqs; i++) {if (reqs[i] != MPI_REQUEST_NULL) {ierr = MPI_Request_free(&reqs[i]);CHKERRQ(ierr);}} 541cd620004SJunchao Zhang link->rootreqsinited[direction][rootmtype][1] = PETSC_FALSE; 542cd620004SJunchao Zhang } 543cd620004SJunchao Zhang if (leafdirect_mpi && sf->persistent && link->leafreqsinited[direction][leafmtype][1] && link->leafdatadirect[direction][leafmtype] != leafdata) { 544cd620004SJunchao Zhang reqs = link->leafreqs[direction][leafmtype][1]; 545cd620004SJunchao Zhang for (i=0; i<nleafreqs; i++) {if (reqs[i] != MPI_REQUEST_NULL) {ierr = MPI_Request_free(&reqs[i]);CHKERRQ(ierr);}} 546cd620004SJunchao Zhang link->leafreqsinited[direction][leafmtype][1] = PETSC_FALSE; 547cd620004SJunchao Zhang } 548cd620004SJunchao Zhang *p = link->next; /* Remove from available list */ 549cd620004SJunchao Zhang goto found; 550cd620004SJunchao Zhang } 551cd620004SJunchao Zhang } 552cd620004SJunchao Zhang 553cd620004SJunchao Zhang ierr = PetscNew(&link);CHKERRQ(ierr); 554cd620004SJunchao Zhang ierr = PetscSFLinkSetUp_Host(sf,link,unit);CHKERRQ(ierr); 555cd620004SJunchao Zhang ierr = PetscCommGetNewTag(PetscObjectComm((PetscObject)sf),&link->tag);CHKERRQ(ierr); /* One tag per link */ 556cd620004SJunchao Zhang 557cd620004SJunchao Zhang nreqs = (nrootreqs+nleafreqs)*8; 558cd620004SJunchao Zhang ierr = PetscMalloc1(nreqs,&link->reqs);CHKERRQ(ierr); 559cd620004SJunchao Zhang for (i=0; i<nreqs; i++) link->reqs[i] = MPI_REQUEST_NULL; /* Initialized to NULL so that we know which need to be freed in Destroy */ 560cd620004SJunchao Zhang 561cd620004SJunchao Zhang for (i=0; i<2; i++) { /* Two communication directions */ 562cd620004SJunchao Zhang for (j=0; j<2; j++) { /* Two memory types */ 563cd620004SJunchao Zhang for (k=0; k<2; k++) { /* root/leafdirect 0 or 1 */ 564cd620004SJunchao Zhang link->rootreqs[i][j][k] = link->reqs + nrootreqs*(4*i+2*j+k); 565cd620004SJunchao Zhang link->leafreqs[i][j][k] = link->reqs + nrootreqs*8 + nleafreqs*(4*i+2*j+k); 566cd620004SJunchao Zhang } 567cd620004SJunchao Zhang } 568cd620004SJunchao Zhang } 569cd620004SJunchao Zhang 570cd620004SJunchao Zhang found: 571cd620004SJunchao Zhang if ((rootmtype == PETSC_MEMTYPE_DEVICE || leafmtype == PETSC_MEMTYPE_DEVICE) && !link->deviceinited) {ierr = PetscSFLinkSetUp_Device(sf,link,unit);CHKERRQ(ierr);} 572cd620004SJunchao Zhang 573cd620004SJunchao Zhang /* Allocate buffers along root/leafdata */ 574cd620004SJunchao Zhang for (i=PETSCSF_LOCAL; i<=PETSCSF_REMOTE; i++) { 575cd620004SJunchao Zhang /* For local communication, buffers are only needed when roots and leaves have different mtypes */ 576cd620004SJunchao Zhang if (i == PETSCSF_LOCAL && rootmtype == leafmtype) continue; 577cd620004SJunchao Zhang if (bas->rootbuflen[i]) { 578cd620004SJunchao Zhang if (rootdirect[i]) { /* Aha, we disguise rootdata as rootbuf */ 579cd620004SJunchao Zhang link->rootbuf[i][rootmtype] = (char*)rootdata + bas->rootstart[i]*link->unitbytes; 580cd620004SJunchao Zhang } else { /* Have to have a separate rootbuf */ 581cd620004SJunchao Zhang if (!link->rootbuf_alloc[i][rootmtype]) { 582cd620004SJunchao Zhang ierr = PetscMallocWithMemType(rootmtype,bas->rootbuflen[i]*link->unitbytes,(void**)&link->rootbuf_alloc[i][rootmtype]);CHKERRQ(ierr); 583cd620004SJunchao Zhang } 584cd620004SJunchao Zhang link->rootbuf[i][rootmtype] = link->rootbuf_alloc[i][rootmtype]; 585cd620004SJunchao Zhang } 586cd620004SJunchao Zhang } 587cd620004SJunchao Zhang 588cd620004SJunchao Zhang if (sf->leafbuflen[i]) { 589cd620004SJunchao Zhang if (leafdirect[i]) { 590cd620004SJunchao Zhang link->leafbuf[i][leafmtype] = (char*)leafdata + sf->leafstart[i]*link->unitbytes; 591cd620004SJunchao Zhang } else { 592cd620004SJunchao Zhang if (!link->leafbuf_alloc[i][leafmtype]) { 593cd620004SJunchao Zhang ierr = PetscMallocWithMemType(leafmtype,sf->leafbuflen[i]*link->unitbytes,(void**)&link->leafbuf_alloc[i][leafmtype]);CHKERRQ(ierr); 594cd620004SJunchao Zhang } 595cd620004SJunchao Zhang link->leafbuf[i][leafmtype] = link->leafbuf_alloc[i][leafmtype]; 596cd620004SJunchao Zhang } 597cd620004SJunchao Zhang } 598cd620004SJunchao Zhang } 599cd620004SJunchao Zhang 600cd620004SJunchao Zhang /* Allocate buffers on host for buffering data on device in cast not use_gpu_aware_mpi */ 601cd620004SJunchao Zhang if (rootmtype == PETSC_MEMTYPE_DEVICE && rootmtype_mpi == PETSC_MEMTYPE_HOST) { 602cd620004SJunchao Zhang if(!link->rootbuf_alloc[PETSCSF_REMOTE][PETSC_MEMTYPE_HOST]) { 603cd620004SJunchao Zhang ierr = PetscMalloc(bas->rootbuflen[PETSCSF_REMOTE]*link->unitbytes,&link->rootbuf_alloc[PETSCSF_REMOTE][PETSC_MEMTYPE_HOST]);CHKERRQ(ierr); 604cd620004SJunchao Zhang } 605cd620004SJunchao Zhang link->rootbuf[PETSCSF_REMOTE][PETSC_MEMTYPE_HOST] = link->rootbuf_alloc[PETSCSF_REMOTE][PETSC_MEMTYPE_HOST]; 606cd620004SJunchao Zhang } 607cd620004SJunchao Zhang if (leafmtype == PETSC_MEMTYPE_DEVICE && leafmtype_mpi == PETSC_MEMTYPE_HOST) { 608cd620004SJunchao Zhang if (!link->leafbuf_alloc[PETSCSF_REMOTE][PETSC_MEMTYPE_HOST]) { 609cd620004SJunchao Zhang ierr = PetscMalloc(sf->leafbuflen[PETSCSF_REMOTE]*link->unitbytes,&link->leafbuf_alloc[PETSCSF_REMOTE][PETSC_MEMTYPE_HOST]);CHKERRQ(ierr); 610cd620004SJunchao Zhang } 611cd620004SJunchao Zhang link->leafbuf[PETSCSF_REMOTE][PETSC_MEMTYPE_HOST] = link->leafbuf_alloc[PETSCSF_REMOTE][PETSC_MEMTYPE_HOST]; 612cd620004SJunchao Zhang } 613cd620004SJunchao Zhang 614cd620004SJunchao Zhang /* Set `current` state of the link. They may change between different SF invocations with the same link */ 615cd620004SJunchao Zhang if (sf->persistent) { /* If data is directly passed to MPI and inits MPI requests, record the data for comparison on future invocations */ 616cd620004SJunchao Zhang if (rootdirect_mpi) link->rootdatadirect[direction][rootmtype] = rootdata; 617cd620004SJunchao Zhang if (leafdirect_mpi) link->leafdatadirect[direction][leafmtype] = leafdata; 618cd620004SJunchao Zhang } 619cd620004SJunchao Zhang 620cd620004SJunchao Zhang link->rootdata = rootdata; /* root/leafdata are keys to look up links in PetscSFXxxEnd */ 621cd620004SJunchao Zhang link->leafdata = leafdata; 622cd620004SJunchao Zhang for (i=PETSCSF_LOCAL; i<=PETSCSF_REMOTE; i++) { 623cd620004SJunchao Zhang link->rootdirect[i] = rootdirect[i]; 624cd620004SJunchao Zhang link->leafdirect[i] = leafdirect[i]; 625cd620004SJunchao Zhang } 626cd620004SJunchao Zhang link->rootdirect_mpi = rootdirect_mpi; 627cd620004SJunchao Zhang link->leafdirect_mpi = leafdirect_mpi; 628cd620004SJunchao Zhang link->rootmtype = rootmtype; 629cd620004SJunchao Zhang link->leafmtype = leafmtype; 630cd620004SJunchao Zhang link->rootmtype_mpi = rootmtype_mpi; 631cd620004SJunchao Zhang link->leafmtype_mpi = leafmtype_mpi; 632cd620004SJunchao Zhang 633cd620004SJunchao Zhang link->next = bas->inuse; 634cd620004SJunchao Zhang bas->inuse = link; 635cd620004SJunchao Zhang *mylink = link; 636cd620004SJunchao Zhang PetscFunctionReturn(0); 637cd620004SJunchao Zhang } 638cd620004SJunchao Zhang 639cd620004SJunchao Zhang /* Return root/leaf buffers and MPI requests attached to the link for MPI communication in the given direction. 640cd620004SJunchao Zhang If the sf uses persistent requests and the requests have not been initialized, then initialize them. 641cd620004SJunchao Zhang */ 642cd620004SJunchao Zhang PetscErrorCode PetscSFLinkGetMPIBuffersAndRequests(PetscSF sf,PetscSFLink link,PetscSFDirection direction,void **rootbuf, void **leafbuf,MPI_Request **rootreqs,MPI_Request **leafreqs) 643cd620004SJunchao Zhang { 644cd620004SJunchao Zhang PetscErrorCode ierr; 645cd620004SJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 646cd620004SJunchao Zhang PetscInt i,j,nrootranks,ndrootranks,nleafranks,ndleafranks; 647cd620004SJunchao Zhang const PetscInt *rootoffset,*leafoffset; 648cd620004SJunchao Zhang PetscMPIInt n; 649cd620004SJunchao Zhang MPI_Aint disp; 650cd620004SJunchao Zhang MPI_Comm comm = PetscObjectComm((PetscObject)sf); 651cd620004SJunchao Zhang MPI_Datatype unit = link->unit; 652cd620004SJunchao Zhang const PetscMemType rootmtype_mpi = link->rootmtype_mpi,leafmtype_mpi = link->leafmtype_mpi; /* Used to select buffers passed to MPI */ 653cd620004SJunchao Zhang const PetscInt rootdirect_mpi = link->rootdirect_mpi,leafdirect_mpi = link->leafdirect_mpi; 654cd620004SJunchao Zhang 655cd620004SJunchao Zhang PetscFunctionBegin; 656cd620004SJunchao Zhang /* Init persistent MPI requests if not yet. Currently only SFBasic uses persistent MPI */ 657cd620004SJunchao Zhang if (sf->persistent) { 658cd620004SJunchao Zhang if (rootreqs && bas->rootbuflen[PETSCSF_REMOTE] && !link->rootreqsinited[direction][rootmtype_mpi][rootdirect_mpi]) { 659cd620004SJunchao Zhang ierr = PetscSFGetRootInfo_Basic(sf,&nrootranks,&ndrootranks,NULL,&rootoffset,NULL);CHKERRQ(ierr); 660cd620004SJunchao Zhang if (direction == PETSCSF_LEAF2ROOT) { 661cd620004SJunchao Zhang for (i=ndrootranks,j=0; i<nrootranks; i++,j++) { 662cd620004SJunchao Zhang disp = (rootoffset[i] - rootoffset[ndrootranks])*link->unitbytes; 663cd620004SJunchao Zhang ierr = PetscMPIIntCast(rootoffset[i+1]-rootoffset[i],&n);CHKERRQ(ierr); 664cd620004SJunchao Zhang ierr = MPI_Recv_init(link->rootbuf[PETSCSF_REMOTE][rootmtype_mpi]+disp,n,unit,bas->iranks[i],link->tag,comm,link->rootreqs[direction][rootmtype_mpi][rootdirect_mpi]+j);CHKERRQ(ierr); 665cd620004SJunchao Zhang } 666cd620004SJunchao Zhang } else { /* PETSCSF_ROOT2LEAF */ 667cd620004SJunchao Zhang for (i=ndrootranks,j=0; i<nrootranks; i++,j++) { 668cd620004SJunchao Zhang disp = (rootoffset[i] - rootoffset[ndrootranks])*link->unitbytes; 669cd620004SJunchao Zhang ierr = PetscMPIIntCast(rootoffset[i+1]-rootoffset[i],&n);CHKERRQ(ierr); 670cd620004SJunchao Zhang ierr = MPI_Send_init(link->rootbuf[PETSCSF_REMOTE][rootmtype_mpi]+disp,n,unit,bas->iranks[i],link->tag,comm,link->rootreqs[direction][rootmtype_mpi][rootdirect_mpi]+j);CHKERRQ(ierr); 671cd620004SJunchao Zhang } 672cd620004SJunchao Zhang } 673cd620004SJunchao Zhang link->rootreqsinited[direction][rootmtype_mpi][rootdirect_mpi] = PETSC_TRUE; 674cd620004SJunchao Zhang } 675cd620004SJunchao Zhang 676cd620004SJunchao Zhang if (leafreqs && sf->leafbuflen[PETSCSF_REMOTE] && !link->leafreqsinited[direction][leafmtype_mpi][leafdirect_mpi]) { 677cd620004SJunchao Zhang ierr = PetscSFGetLeafInfo_Basic(sf,&nleafranks,&ndleafranks,NULL,&leafoffset,NULL,NULL);CHKERRQ(ierr); 678cd620004SJunchao Zhang if (direction == PETSCSF_LEAF2ROOT) { 679cd620004SJunchao Zhang for (i=ndleafranks,j=0; i<nleafranks; i++,j++) { 680cd620004SJunchao Zhang disp = (leafoffset[i] - leafoffset[ndleafranks])*link->unitbytes; 681cd620004SJunchao Zhang ierr = PetscMPIIntCast(leafoffset[i+1]-leafoffset[i],&n);CHKERRQ(ierr); 682cd620004SJunchao Zhang ierr = MPI_Send_init(link->leafbuf[PETSCSF_REMOTE][leafmtype_mpi]+disp,n,unit,sf->ranks[i],link->tag,comm,link->leafreqs[direction][leafmtype_mpi][leafdirect_mpi]+j);CHKERRQ(ierr); 683cd620004SJunchao Zhang } 684cd620004SJunchao Zhang } else { /* PETSCSF_ROOT2LEAF */ 685cd620004SJunchao Zhang for (i=ndleafranks,j=0; i<nleafranks; i++,j++) { 686cd620004SJunchao Zhang disp = (leafoffset[i] - leafoffset[ndleafranks])*link->unitbytes; 687cd620004SJunchao Zhang ierr = PetscMPIIntCast(leafoffset[i+1]-leafoffset[i],&n);CHKERRQ(ierr); 688cd620004SJunchao Zhang ierr = MPI_Recv_init(link->leafbuf[PETSCSF_REMOTE][leafmtype_mpi]+disp,n,unit,sf->ranks[i],link->tag,comm,link->leafreqs[direction][leafmtype_mpi][leafdirect_mpi]+j);CHKERRQ(ierr); 689cd620004SJunchao Zhang } 690cd620004SJunchao Zhang } 691cd620004SJunchao Zhang link->leafreqsinited[direction][leafmtype_mpi][leafdirect_mpi] = PETSC_TRUE; 692cd620004SJunchao Zhang } 693cd620004SJunchao Zhang } 694cd620004SJunchao Zhang if (rootbuf) *rootbuf = link->rootbuf[PETSCSF_REMOTE][rootmtype_mpi]; 695cd620004SJunchao Zhang if (leafbuf) *leafbuf = link->leafbuf[PETSCSF_REMOTE][leafmtype_mpi]; 696cd620004SJunchao Zhang if (rootreqs) *rootreqs = link->rootreqs[direction][rootmtype_mpi][rootdirect_mpi]; 697cd620004SJunchao Zhang if (leafreqs) *leafreqs = link->leafreqs[direction][leafmtype_mpi][leafdirect_mpi]; 698cd620004SJunchao Zhang PetscFunctionReturn(0); 699cd620004SJunchao Zhang } 700cd620004SJunchao Zhang 701cd620004SJunchao Zhang 702cd620004SJunchao Zhang PetscErrorCode PetscSFLinkGetInUse(PetscSF sf,MPI_Datatype unit,const void *rootdata,const void *leafdata,PetscCopyMode cmode,PetscSFLink *mylink) 703cd620004SJunchao Zhang { 704cd620004SJunchao Zhang PetscErrorCode ierr; 705cd620004SJunchao Zhang PetscSFLink link,*p; 70640e23c03SJunchao Zhang PetscSF_Basic *bas=(PetscSF_Basic*)sf->data; 70740e23c03SJunchao Zhang 70840e23c03SJunchao Zhang PetscFunctionBegin; 70940e23c03SJunchao Zhang /* Look for types in cache */ 71040e23c03SJunchao Zhang for (p=&bas->inuse; (link=*p); p=&link->next) { 71140e23c03SJunchao Zhang PetscBool match; 71240e23c03SJunchao Zhang ierr = MPIPetsc_Type_compare(unit,link->unit,&match);CHKERRQ(ierr); 713637e6665SJunchao Zhang if (match && (rootdata == link->rootdata) && (leafdata == link->leafdata)) { 71440e23c03SJunchao Zhang switch (cmode) { 71540e23c03SJunchao Zhang case PETSC_OWN_POINTER: *p = link->next; break; /* Remove from inuse list */ 71640e23c03SJunchao Zhang case PETSC_USE_POINTER: break; 71740e23c03SJunchao Zhang default: SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_INCOMP,"invalid cmode"); 71840e23c03SJunchao Zhang } 71940e23c03SJunchao Zhang *mylink = link; 72040e23c03SJunchao Zhang PetscFunctionReturn(0); 72140e23c03SJunchao Zhang } 72240e23c03SJunchao Zhang } 72340e23c03SJunchao Zhang SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Could not find pack"); 72440e23c03SJunchao Zhang PetscFunctionReturn(0); 72540e23c03SJunchao Zhang } 72640e23c03SJunchao Zhang 727cd620004SJunchao Zhang PetscErrorCode PetscSFLinkReclaim(PetscSF sf,PetscSFLink *link) 72840e23c03SJunchao Zhang { 72940e23c03SJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 73040e23c03SJunchao Zhang 73140e23c03SJunchao Zhang PetscFunctionBegin; 732637e6665SJunchao Zhang (*link)->rootdata = NULL; 733637e6665SJunchao Zhang (*link)->leafdata = NULL; 73440e23c03SJunchao Zhang (*link)->next = bas->avail; 73540e23c03SJunchao Zhang bas->avail = *link; 73640e23c03SJunchao Zhang *link = NULL; 73740e23c03SJunchao Zhang PetscFunctionReturn(0); 73840e23c03SJunchao Zhang } 73940e23c03SJunchao Zhang 740cd620004SJunchao Zhang /* Destroy all links chained in 'avail' */ 741cd620004SJunchao Zhang PetscErrorCode PetscSFLinkDestroy(PetscSF sf,PetscSFLink *avail) 742eb02082bSJunchao Zhang { 743eb02082bSJunchao Zhang PetscErrorCode ierr; 744cd620004SJunchao Zhang PetscSFLink link = *avail,next; 745cd620004SJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 746cd620004SJunchao Zhang PetscInt i,nreqs = (bas->nrootreqs+sf->nleafreqs)*8; 747eb02082bSJunchao Zhang 748eb02082bSJunchao Zhang PetscFunctionBegin; 749eb02082bSJunchao Zhang for (; link; link=next) { 750eb02082bSJunchao Zhang next = link->next; 751eb02082bSJunchao Zhang if (!link->isbuiltin) {ierr = MPI_Type_free(&link->unit);CHKERRQ(ierr);} 752cd620004SJunchao Zhang for (i=0; i<nreqs; i++) { /* Persistent reqs must be freed. */ 753eb02082bSJunchao Zhang if (link->reqs[i] != MPI_REQUEST_NULL) {ierr = MPI_Request_free(&link->reqs[i]);CHKERRQ(ierr);} 754eb02082bSJunchao Zhang } 755eb02082bSJunchao Zhang ierr = PetscFree(link->reqs);CHKERRQ(ierr); 756cd620004SJunchao Zhang for (i=PETSCSF_LOCAL; i<=PETSCSF_REMOTE; i++) { 75751ccb202SJunchao Zhang #if defined(PETSC_HAVE_CUDA) 758cd620004SJunchao Zhang ierr = PetscFreeWithMemType(PETSC_MEMTYPE_DEVICE,link->rootbuf_alloc[i][PETSC_MEMTYPE_DEVICE]);CHKERRQ(ierr); 759cd620004SJunchao Zhang ierr = PetscFreeWithMemType(PETSC_MEMTYPE_DEVICE,link->leafbuf_alloc[i][PETSC_MEMTYPE_DEVICE]);CHKERRQ(ierr); 760eb02082bSJunchao Zhang if (link->stream) {cudaError_t err = cudaStreamDestroy(link->stream);CHKERRCUDA(err); link->stream = NULL;} 761eb02082bSJunchao Zhang #endif 762cd620004SJunchao Zhang ierr = PetscFree(link->rootbuf_alloc[i][PETSC_MEMTYPE_HOST]);CHKERRQ(ierr); 763cd620004SJunchao Zhang ierr = PetscFree(link->leafbuf_alloc[i][PETSC_MEMTYPE_HOST]);CHKERRQ(ierr); 764cd620004SJunchao Zhang } 765eb02082bSJunchao Zhang ierr = PetscFree(link);CHKERRQ(ierr); 766eb02082bSJunchao Zhang } 767eb02082bSJunchao Zhang *avail = NULL; 768eb02082bSJunchao Zhang PetscFunctionReturn(0); 769eb02082bSJunchao Zhang } 770eb02082bSJunchao Zhang 7719d1c8addSJunchao Zhang /* Error out on unsupported overlapped communications */ 772cd620004SJunchao Zhang PetscErrorCode PetscSFSetErrorOnUnsupportedOverlap(PetscSF sf,MPI_Datatype unit,const void *rootdata,const void *leafdata) 7739d1c8addSJunchao Zhang { 7749d1c8addSJunchao Zhang PetscErrorCode ierr; 775cd620004SJunchao Zhang PetscSFLink link,*p; 7769d1c8addSJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 7779d1c8addSJunchao Zhang PetscBool match; 7789d1c8addSJunchao Zhang 7799d1c8addSJunchao Zhang PetscFunctionBegin; 780*4e1ad211SJed Brown if (!PetscDefined(USE_DEBUG)) PetscFunctionReturn(0); 78118fb5014SJunchao Zhang /* Look up links in use and error out if there is a match. When both rootdata and leafdata are NULL, ignore 78218fb5014SJunchao Zhang the potential overlapping since this process does not participate in communication. Overlapping is harmless. 78318fb5014SJunchao Zhang */ 784637e6665SJunchao Zhang if (rootdata || leafdata) { 7859d1c8addSJunchao Zhang for (p=&bas->inuse; (link=*p); p=&link->next) { 7869d1c8addSJunchao Zhang ierr = MPIPetsc_Type_compare(unit,link->unit,&match);CHKERRQ(ierr); 787cd620004SJunchao Zhang if (match && (rootdata == link->rootdata) && (leafdata == link->leafdata)) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_SUP,"Overlapped PetscSF with the same rootdata(%p), leafdata(%p) and data type. Undo the overlapping to avoid the error.",rootdata,leafdata); 7889d1c8addSJunchao Zhang } 78918fb5014SJunchao Zhang } 7909d1c8addSJunchao Zhang PetscFunctionReturn(0); 7919d1c8addSJunchao Zhang } 7929d1c8addSJunchao Zhang 793cd620004SJunchao Zhang PetscErrorCode PetscSFLinkSetUp_Host(PetscSF sf,PetscSFLink link,MPI_Datatype unit) 79440e23c03SJunchao Zhang { 79540e23c03SJunchao Zhang PetscErrorCode ierr; 796b23bfdefSJunchao Zhang PetscInt nSignedChar=0,nUnsignedChar=0,nInt=0,nPetscInt=0,nPetscReal=0; 797b23bfdefSJunchao Zhang PetscBool is2Int,is2PetscInt; 79840e23c03SJunchao Zhang PetscMPIInt ni,na,nd,combiner; 79940e23c03SJunchao Zhang #if defined(PETSC_HAVE_COMPLEX) 800b23bfdefSJunchao Zhang PetscInt nPetscComplex=0; 80140e23c03SJunchao Zhang #endif 80240e23c03SJunchao Zhang 80340e23c03SJunchao Zhang PetscFunctionBegin; 804b23bfdefSJunchao Zhang ierr = MPIPetsc_Type_compare_contig(unit,MPI_SIGNED_CHAR, &nSignedChar);CHKERRQ(ierr); 805b23bfdefSJunchao Zhang ierr = MPIPetsc_Type_compare_contig(unit,MPI_UNSIGNED_CHAR,&nUnsignedChar);CHKERRQ(ierr); 806b23bfdefSJunchao Zhang /* MPI_CHAR is treated below as a dumb type that does not support reduction according to MPI standard */ 807b23bfdefSJunchao Zhang ierr = MPIPetsc_Type_compare_contig(unit,MPI_INT, &nInt);CHKERRQ(ierr); 808b23bfdefSJunchao Zhang ierr = MPIPetsc_Type_compare_contig(unit,MPIU_INT, &nPetscInt);CHKERRQ(ierr); 809b23bfdefSJunchao Zhang ierr = MPIPetsc_Type_compare_contig(unit,MPIU_REAL,&nPetscReal);CHKERRQ(ierr); 81040e23c03SJunchao Zhang #if defined(PETSC_HAVE_COMPLEX) 811b23bfdefSJunchao Zhang ierr = MPIPetsc_Type_compare_contig(unit,MPIU_COMPLEX,&nPetscComplex);CHKERRQ(ierr); 81240e23c03SJunchao Zhang #endif 81340e23c03SJunchao Zhang ierr = MPIPetsc_Type_compare(unit,MPI_2INT,&is2Int);CHKERRQ(ierr); 81440e23c03SJunchao Zhang ierr = MPIPetsc_Type_compare(unit,MPIU_2INT,&is2PetscInt);CHKERRQ(ierr); 815b23bfdefSJunchao Zhang /* TODO: shaell we also handle Fortran MPI_2REAL? */ 81640e23c03SJunchao Zhang ierr = MPI_Type_get_envelope(unit,&ni,&na,&nd,&combiner);CHKERRQ(ierr); 8175ad15460SJunchao Zhang link->isbuiltin = (combiner == MPI_COMBINER_NAMED) ? PETSC_TRUE : PETSC_FALSE; /* unit is MPI builtin */ 818b23bfdefSJunchao Zhang link->bs = 1; /* default */ 81940e23c03SJunchao Zhang 820eb02082bSJunchao Zhang if (is2Int) { 821cd620004SJunchao Zhang PackInit_PairType_int_int_1_1(link); 822eb02082bSJunchao Zhang link->bs = 1; 823eb02082bSJunchao Zhang link->unitbytes = 2*sizeof(int); 8245ad15460SJunchao Zhang link->isbuiltin = PETSC_TRUE; /* unit is PETSc builtin */ 825eb02082bSJunchao Zhang link->basicunit = MPI_2INT; 8265ad15460SJunchao Zhang link->unit = MPI_2INT; 827eb02082bSJunchao Zhang } else if (is2PetscInt) { /* TODO: when is2PetscInt and nPetscInt=2, we don't know which path to take. The two paths support different ops. */ 828cd620004SJunchao Zhang PackInit_PairType_PetscInt_PetscInt_1_1(link); 829eb02082bSJunchao Zhang link->bs = 1; 830eb02082bSJunchao Zhang link->unitbytes = 2*sizeof(PetscInt); 831eb02082bSJunchao Zhang link->basicunit = MPIU_2INT; 8325ad15460SJunchao Zhang link->isbuiltin = PETSC_TRUE; /* unit is PETSc builtin */ 8335ad15460SJunchao Zhang link->unit = MPIU_2INT; 834eb02082bSJunchao Zhang } else if (nPetscReal) { 835b23bfdefSJunchao Zhang if (nPetscReal == 8) PackInit_RealType_PetscReal_8_1(link); else if (nPetscReal%8 == 0) PackInit_RealType_PetscReal_8_0(link); 836b23bfdefSJunchao Zhang else if (nPetscReal == 4) PackInit_RealType_PetscReal_4_1(link); else if (nPetscReal%4 == 0) PackInit_RealType_PetscReal_4_0(link); 837b23bfdefSJunchao Zhang else if (nPetscReal == 2) PackInit_RealType_PetscReal_2_1(link); else if (nPetscReal%2 == 0) PackInit_RealType_PetscReal_2_0(link); 838b23bfdefSJunchao Zhang else if (nPetscReal == 1) PackInit_RealType_PetscReal_1_1(link); else if (nPetscReal%1 == 0) PackInit_RealType_PetscReal_1_0(link); 839b23bfdefSJunchao Zhang link->bs = nPetscReal; 840eb02082bSJunchao Zhang link->unitbytes = nPetscReal*sizeof(PetscReal); 84140e23c03SJunchao Zhang link->basicunit = MPIU_REAL; 8425ad15460SJunchao Zhang if (link->bs == 1) {link->isbuiltin = PETSC_TRUE; link->unit = MPIU_REAL;} 843b23bfdefSJunchao Zhang } else if (nPetscInt) { 844b23bfdefSJunchao Zhang if (nPetscInt == 8) PackInit_IntegerType_PetscInt_8_1(link); else if (nPetscInt%8 == 0) PackInit_IntegerType_PetscInt_8_0(link); 845b23bfdefSJunchao Zhang else if (nPetscInt == 4) PackInit_IntegerType_PetscInt_4_1(link); else if (nPetscInt%4 == 0) PackInit_IntegerType_PetscInt_4_0(link); 846b23bfdefSJunchao Zhang else if (nPetscInt == 2) PackInit_IntegerType_PetscInt_2_1(link); else if (nPetscInt%2 == 0) PackInit_IntegerType_PetscInt_2_0(link); 847b23bfdefSJunchao Zhang else if (nPetscInt == 1) PackInit_IntegerType_PetscInt_1_1(link); else if (nPetscInt%1 == 0) PackInit_IntegerType_PetscInt_1_0(link); 848b23bfdefSJunchao Zhang link->bs = nPetscInt; 849eb02082bSJunchao Zhang link->unitbytes = nPetscInt*sizeof(PetscInt); 850b23bfdefSJunchao Zhang link->basicunit = MPIU_INT; 8515ad15460SJunchao Zhang if (link->bs == 1) {link->isbuiltin = PETSC_TRUE; link->unit = MPIU_INT;} 852b23bfdefSJunchao Zhang #if defined(PETSC_USE_64BIT_INDICES) 853b23bfdefSJunchao Zhang } else if (nInt) { 854b23bfdefSJunchao Zhang if (nInt == 8) PackInit_IntegerType_int_8_1(link); else if (nInt%8 == 0) PackInit_IntegerType_int_8_0(link); 855b23bfdefSJunchao Zhang else if (nInt == 4) PackInit_IntegerType_int_4_1(link); else if (nInt%4 == 0) PackInit_IntegerType_int_4_0(link); 856b23bfdefSJunchao Zhang else if (nInt == 2) PackInit_IntegerType_int_2_1(link); else if (nInt%2 == 0) PackInit_IntegerType_int_2_0(link); 857b23bfdefSJunchao Zhang else if (nInt == 1) PackInit_IntegerType_int_1_1(link); else if (nInt%1 == 0) PackInit_IntegerType_int_1_0(link); 858b23bfdefSJunchao Zhang link->bs = nInt; 859eb02082bSJunchao Zhang link->unitbytes = nInt*sizeof(int); 860b23bfdefSJunchao Zhang link->basicunit = MPI_INT; 8615ad15460SJunchao Zhang if (link->bs == 1) {link->isbuiltin = PETSC_TRUE; link->unit = MPI_INT;} 862b23bfdefSJunchao Zhang #endif 863b23bfdefSJunchao Zhang } else if (nSignedChar) { 864b23bfdefSJunchao Zhang if (nSignedChar == 8) PackInit_IntegerType_SignedChar_8_1(link); else if (nSignedChar%8 == 0) PackInit_IntegerType_SignedChar_8_0(link); 865b23bfdefSJunchao Zhang else if (nSignedChar == 4) PackInit_IntegerType_SignedChar_4_1(link); else if (nSignedChar%4 == 0) PackInit_IntegerType_SignedChar_4_0(link); 866b23bfdefSJunchao Zhang else if (nSignedChar == 2) PackInit_IntegerType_SignedChar_2_1(link); else if (nSignedChar%2 == 0) PackInit_IntegerType_SignedChar_2_0(link); 867b23bfdefSJunchao Zhang else if (nSignedChar == 1) PackInit_IntegerType_SignedChar_1_1(link); else if (nSignedChar%1 == 0) PackInit_IntegerType_SignedChar_1_0(link); 868b23bfdefSJunchao Zhang link->bs = nSignedChar; 869eb02082bSJunchao Zhang link->unitbytes = nSignedChar*sizeof(SignedChar); 870b23bfdefSJunchao Zhang link->basicunit = MPI_SIGNED_CHAR; 8715ad15460SJunchao Zhang if (link->bs == 1) {link->isbuiltin = PETSC_TRUE; link->unit = MPI_SIGNED_CHAR;} 872b23bfdefSJunchao Zhang } else if (nUnsignedChar) { 873b23bfdefSJunchao Zhang if (nUnsignedChar == 8) PackInit_IntegerType_UnsignedChar_8_1(link); else if (nUnsignedChar%8 == 0) PackInit_IntegerType_UnsignedChar_8_0(link); 874b23bfdefSJunchao Zhang else if (nUnsignedChar == 4) PackInit_IntegerType_UnsignedChar_4_1(link); else if (nUnsignedChar%4 == 0) PackInit_IntegerType_UnsignedChar_4_0(link); 875b23bfdefSJunchao Zhang else if (nUnsignedChar == 2) PackInit_IntegerType_UnsignedChar_2_1(link); else if (nUnsignedChar%2 == 0) PackInit_IntegerType_UnsignedChar_2_0(link); 876b23bfdefSJunchao Zhang else if (nUnsignedChar == 1) PackInit_IntegerType_UnsignedChar_1_1(link); else if (nUnsignedChar%1 == 0) PackInit_IntegerType_UnsignedChar_1_0(link); 877b23bfdefSJunchao Zhang link->bs = nUnsignedChar; 878eb02082bSJunchao Zhang link->unitbytes = nUnsignedChar*sizeof(UnsignedChar); 879b23bfdefSJunchao Zhang link->basicunit = MPI_UNSIGNED_CHAR; 8805ad15460SJunchao Zhang if (link->bs == 1) {link->isbuiltin = PETSC_TRUE; link->unit = MPI_UNSIGNED_CHAR;} 88140e23c03SJunchao Zhang #if defined(PETSC_HAVE_COMPLEX) 882b23bfdefSJunchao Zhang } else if (nPetscComplex) { 883b23bfdefSJunchao Zhang if (nPetscComplex == 8) PackInit_ComplexType_PetscComplex_8_1(link); else if (nPetscComplex%8 == 0) PackInit_ComplexType_PetscComplex_8_0(link); 884b23bfdefSJunchao Zhang else if (nPetscComplex == 4) PackInit_ComplexType_PetscComplex_4_1(link); else if (nPetscComplex%4 == 0) PackInit_ComplexType_PetscComplex_4_0(link); 885b23bfdefSJunchao Zhang else if (nPetscComplex == 2) PackInit_ComplexType_PetscComplex_2_1(link); else if (nPetscComplex%2 == 0) PackInit_ComplexType_PetscComplex_2_0(link); 886b23bfdefSJunchao Zhang else if (nPetscComplex == 1) PackInit_ComplexType_PetscComplex_1_1(link); else if (nPetscComplex%1 == 0) PackInit_ComplexType_PetscComplex_1_0(link); 887b23bfdefSJunchao Zhang link->bs = nPetscComplex; 888eb02082bSJunchao Zhang link->unitbytes = nPetscComplex*sizeof(PetscComplex); 88940e23c03SJunchao Zhang link->basicunit = MPIU_COMPLEX; 8905ad15460SJunchao Zhang if (link->bs == 1) {link->isbuiltin = PETSC_TRUE; link->unit = MPIU_COMPLEX;} 89140e23c03SJunchao Zhang #endif 89240e23c03SJunchao Zhang } else { 893b23bfdefSJunchao Zhang MPI_Aint lb,nbyte; 894b23bfdefSJunchao Zhang ierr = MPI_Type_get_extent(unit,&lb,&nbyte);CHKERRQ(ierr); 89540e23c03SJunchao Zhang if (lb != 0) SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_SUP,"Datatype with nonzero lower bound %ld\n",(long)lb); 896eb02082bSJunchao Zhang if (nbyte % sizeof(int)) { /* If the type size is not multiple of int */ 897eb02082bSJunchao Zhang if (nbyte == 4) PackInit_DumbType_char_4_1(link); else if (nbyte%4 == 0) PackInit_DumbType_char_4_0(link); 898eb02082bSJunchao Zhang else if (nbyte == 2) PackInit_DumbType_char_2_1(link); else if (nbyte%2 == 0) PackInit_DumbType_char_2_0(link); 899eb02082bSJunchao Zhang else if (nbyte == 1) PackInit_DumbType_char_1_1(link); else if (nbyte%1 == 0) PackInit_DumbType_char_1_0(link); 900eb02082bSJunchao Zhang link->bs = nbyte; 901b23bfdefSJunchao Zhang link->unitbytes = nbyte; 902b23bfdefSJunchao Zhang link->basicunit = MPI_BYTE; 90340e23c03SJunchao Zhang } else { 904eb02082bSJunchao Zhang nInt = nbyte / sizeof(int); 905eb02082bSJunchao Zhang if (nInt == 8) PackInit_DumbType_DumbInt_8_1(link); else if (nInt%8 == 0) PackInit_DumbType_DumbInt_8_0(link); 906eb02082bSJunchao Zhang else if (nInt == 4) PackInit_DumbType_DumbInt_4_1(link); else if (nInt%4 == 0) PackInit_DumbType_DumbInt_4_0(link); 907eb02082bSJunchao Zhang else if (nInt == 2) PackInit_DumbType_DumbInt_2_1(link); else if (nInt%2 == 0) PackInit_DumbType_DumbInt_2_0(link); 908eb02082bSJunchao Zhang else if (nInt == 1) PackInit_DumbType_DumbInt_1_1(link); else if (nInt%1 == 0) PackInit_DumbType_DumbInt_1_0(link); 909eb02082bSJunchao Zhang link->bs = nInt; 910b23bfdefSJunchao Zhang link->unitbytes = nbyte; 91140e23c03SJunchao Zhang link->basicunit = MPI_INT; 91240e23c03SJunchao Zhang } 9135ad15460SJunchao Zhang if (link->isbuiltin) link->unit = unit; 91440e23c03SJunchao Zhang } 915b23bfdefSJunchao Zhang 9165ad15460SJunchao Zhang if (!link->isbuiltin) {ierr = MPI_Type_dup(unit,&link->unit);CHKERRQ(ierr);} 91740e23c03SJunchao Zhang PetscFunctionReturn(0); 91840e23c03SJunchao Zhang } 91940e23c03SJunchao Zhang 920fcc7397dSJunchao Zhang PetscErrorCode PetscSFLinkGetUnpackAndOp(PetscSFLink link,PetscMemType mtype,MPI_Op op,PetscBool atomic,PetscErrorCode (**UnpackAndOp)(PetscSFLink,PetscInt,PetscInt,PetscSFPackOpt,const PetscInt*,void*,const void*)) 92140e23c03SJunchao Zhang { 92240e23c03SJunchao Zhang PetscFunctionBegin; 92340e23c03SJunchao Zhang *UnpackAndOp = NULL; 924eb02082bSJunchao Zhang if (mtype == PETSC_MEMTYPE_HOST) { 925eb02082bSJunchao Zhang if (op == MPIU_REPLACE) *UnpackAndOp = link->h_UnpackAndInsert; 926eb02082bSJunchao Zhang else if (op == MPI_SUM || op == MPIU_SUM) *UnpackAndOp = link->h_UnpackAndAdd; 927eb02082bSJunchao Zhang else if (op == MPI_PROD) *UnpackAndOp = link->h_UnpackAndMult; 928eb02082bSJunchao Zhang else if (op == MPI_MAX || op == MPIU_MAX) *UnpackAndOp = link->h_UnpackAndMax; 929eb02082bSJunchao Zhang else if (op == MPI_MIN || op == MPIU_MIN) *UnpackAndOp = link->h_UnpackAndMin; 930eb02082bSJunchao Zhang else if (op == MPI_LAND) *UnpackAndOp = link->h_UnpackAndLAND; 931eb02082bSJunchao Zhang else if (op == MPI_BAND) *UnpackAndOp = link->h_UnpackAndBAND; 932eb02082bSJunchao Zhang else if (op == MPI_LOR) *UnpackAndOp = link->h_UnpackAndLOR; 933eb02082bSJunchao Zhang else if (op == MPI_BOR) *UnpackAndOp = link->h_UnpackAndBOR; 934eb02082bSJunchao Zhang else if (op == MPI_LXOR) *UnpackAndOp = link->h_UnpackAndLXOR; 935eb02082bSJunchao Zhang else if (op == MPI_BXOR) *UnpackAndOp = link->h_UnpackAndBXOR; 936eb02082bSJunchao Zhang else if (op == MPI_MAXLOC) *UnpackAndOp = link->h_UnpackAndMaxloc; 937eb02082bSJunchao Zhang else if (op == MPI_MINLOC) *UnpackAndOp = link->h_UnpackAndMinloc; 938eb02082bSJunchao Zhang } 939eb02082bSJunchao Zhang #if defined(PETSC_HAVE_CUDA) 940eb02082bSJunchao Zhang else if (mtype == PETSC_MEMTYPE_DEVICE && !atomic) { 941eb02082bSJunchao Zhang if (op == MPIU_REPLACE) *UnpackAndOp = link->d_UnpackAndInsert; 942eb02082bSJunchao Zhang else if (op == MPI_SUM || op == MPIU_SUM) *UnpackAndOp = link->d_UnpackAndAdd; 943eb02082bSJunchao Zhang else if (op == MPI_PROD) *UnpackAndOp = link->d_UnpackAndMult; 944eb02082bSJunchao Zhang else if (op == MPI_MAX || op == MPIU_MAX) *UnpackAndOp = link->d_UnpackAndMax; 945eb02082bSJunchao Zhang else if (op == MPI_MIN || op == MPIU_MIN) *UnpackAndOp = link->d_UnpackAndMin; 946eb02082bSJunchao Zhang else if (op == MPI_LAND) *UnpackAndOp = link->d_UnpackAndLAND; 947eb02082bSJunchao Zhang else if (op == MPI_BAND) *UnpackAndOp = link->d_UnpackAndBAND; 948eb02082bSJunchao Zhang else if (op == MPI_LOR) *UnpackAndOp = link->d_UnpackAndLOR; 949eb02082bSJunchao Zhang else if (op == MPI_BOR) *UnpackAndOp = link->d_UnpackAndBOR; 950eb02082bSJunchao Zhang else if (op == MPI_LXOR) *UnpackAndOp = link->d_UnpackAndLXOR; 951eb02082bSJunchao Zhang else if (op == MPI_BXOR) *UnpackAndOp = link->d_UnpackAndBXOR; 952eb02082bSJunchao Zhang else if (op == MPI_MAXLOC) *UnpackAndOp = link->d_UnpackAndMaxloc; 953eb02082bSJunchao Zhang else if (op == MPI_MINLOC) *UnpackAndOp = link->d_UnpackAndMinloc; 954eb02082bSJunchao Zhang } else if (mtype == PETSC_MEMTYPE_DEVICE && atomic) { 955eb02082bSJunchao Zhang if (op == MPIU_REPLACE) *UnpackAndOp = link->da_UnpackAndInsert; 956eb02082bSJunchao Zhang else if (op == MPI_SUM || op == MPIU_SUM) *UnpackAndOp = link->da_UnpackAndAdd; 957eb02082bSJunchao Zhang else if (op == MPI_PROD) *UnpackAndOp = link->da_UnpackAndMult; 958eb02082bSJunchao Zhang else if (op == MPI_MAX || op == MPIU_MAX) *UnpackAndOp = link->da_UnpackAndMax; 959eb02082bSJunchao Zhang else if (op == MPI_MIN || op == MPIU_MIN) *UnpackAndOp = link->da_UnpackAndMin; 960eb02082bSJunchao Zhang else if (op == MPI_LAND) *UnpackAndOp = link->da_UnpackAndLAND; 961eb02082bSJunchao Zhang else if (op == MPI_BAND) *UnpackAndOp = link->da_UnpackAndBAND; 962eb02082bSJunchao Zhang else if (op == MPI_LOR) *UnpackAndOp = link->da_UnpackAndLOR; 963eb02082bSJunchao Zhang else if (op == MPI_BOR) *UnpackAndOp = link->da_UnpackAndBOR; 964eb02082bSJunchao Zhang else if (op == MPI_LXOR) *UnpackAndOp = link->da_UnpackAndLXOR; 965eb02082bSJunchao Zhang else if (op == MPI_BXOR) *UnpackAndOp = link->da_UnpackAndBXOR; 966eb02082bSJunchao Zhang else if (op == MPI_MAXLOC) *UnpackAndOp = link->da_UnpackAndMaxloc; 967eb02082bSJunchao Zhang else if (op == MPI_MINLOC) *UnpackAndOp = link->da_UnpackAndMinloc; 968eb02082bSJunchao Zhang } 969eb02082bSJunchao Zhang #endif 97040e23c03SJunchao Zhang PetscFunctionReturn(0); 97140e23c03SJunchao Zhang } 97240e23c03SJunchao Zhang 973fcc7397dSJunchao Zhang PetscErrorCode PetscSFLinkGetScatterAndOp(PetscSFLink link,PetscMemType mtype,MPI_Op op,PetscBool atomic,PetscErrorCode (**ScatterAndOp)(PetscSFLink,PetscInt,PetscInt,PetscSFPackOpt,const PetscInt*,const void*,PetscInt,PetscSFPackOpt,const PetscInt*,void*)) 97440e23c03SJunchao Zhang { 97540e23c03SJunchao Zhang PetscFunctionBegin; 976cd620004SJunchao Zhang *ScatterAndOp = NULL; 977eb02082bSJunchao Zhang if (mtype == PETSC_MEMTYPE_HOST) { 978cd620004SJunchao Zhang if (op == MPIU_REPLACE) *ScatterAndOp = link->h_ScatterAndInsert; 979cd620004SJunchao Zhang else if (op == MPI_SUM || op == MPIU_SUM) *ScatterAndOp = link->h_ScatterAndAdd; 980cd620004SJunchao Zhang else if (op == MPI_PROD) *ScatterAndOp = link->h_ScatterAndMult; 981cd620004SJunchao Zhang else if (op == MPI_MAX || op == MPIU_MAX) *ScatterAndOp = link->h_ScatterAndMax; 982cd620004SJunchao Zhang else if (op == MPI_MIN || op == MPIU_MIN) *ScatterAndOp = link->h_ScatterAndMin; 983cd620004SJunchao Zhang else if (op == MPI_LAND) *ScatterAndOp = link->h_ScatterAndLAND; 984cd620004SJunchao Zhang else if (op == MPI_BAND) *ScatterAndOp = link->h_ScatterAndBAND; 985cd620004SJunchao Zhang else if (op == MPI_LOR) *ScatterAndOp = link->h_ScatterAndLOR; 986cd620004SJunchao Zhang else if (op == MPI_BOR) *ScatterAndOp = link->h_ScatterAndBOR; 987cd620004SJunchao Zhang else if (op == MPI_LXOR) *ScatterAndOp = link->h_ScatterAndLXOR; 988cd620004SJunchao Zhang else if (op == MPI_BXOR) *ScatterAndOp = link->h_ScatterAndBXOR; 989cd620004SJunchao Zhang else if (op == MPI_MAXLOC) *ScatterAndOp = link->h_ScatterAndMaxloc; 990cd620004SJunchao Zhang else if (op == MPI_MINLOC) *ScatterAndOp = link->h_ScatterAndMinloc; 991eb02082bSJunchao Zhang } 992eb02082bSJunchao Zhang #if defined(PETSC_HAVE_CUDA) 993eb02082bSJunchao Zhang else if (mtype == PETSC_MEMTYPE_DEVICE && !atomic) { 994cd620004SJunchao Zhang if (op == MPIU_REPLACE) *ScatterAndOp = link->d_ScatterAndInsert; 995cd620004SJunchao Zhang else if (op == MPI_SUM || op == MPIU_SUM) *ScatterAndOp = link->d_ScatterAndAdd; 996cd620004SJunchao Zhang else if (op == MPI_PROD) *ScatterAndOp = link->d_ScatterAndMult; 997cd620004SJunchao Zhang else if (op == MPI_MAX || op == MPIU_MAX) *ScatterAndOp = link->d_ScatterAndMax; 998cd620004SJunchao Zhang else if (op == MPI_MIN || op == MPIU_MIN) *ScatterAndOp = link->d_ScatterAndMin; 999cd620004SJunchao Zhang else if (op == MPI_LAND) *ScatterAndOp = link->d_ScatterAndLAND; 1000cd620004SJunchao Zhang else if (op == MPI_BAND) *ScatterAndOp = link->d_ScatterAndBAND; 1001cd620004SJunchao Zhang else if (op == MPI_LOR) *ScatterAndOp = link->d_ScatterAndLOR; 1002cd620004SJunchao Zhang else if (op == MPI_BOR) *ScatterAndOp = link->d_ScatterAndBOR; 1003cd620004SJunchao Zhang else if (op == MPI_LXOR) *ScatterAndOp = link->d_ScatterAndLXOR; 1004cd620004SJunchao Zhang else if (op == MPI_BXOR) *ScatterAndOp = link->d_ScatterAndBXOR; 1005cd620004SJunchao Zhang else if (op == MPI_MAXLOC) *ScatterAndOp = link->d_ScatterAndMaxloc; 1006cd620004SJunchao Zhang else if (op == MPI_MINLOC) *ScatterAndOp = link->d_ScatterAndMinloc; 1007eb02082bSJunchao Zhang } else if (mtype == PETSC_MEMTYPE_DEVICE && atomic) { 1008cd620004SJunchao Zhang if (op == MPIU_REPLACE) *ScatterAndOp = link->da_ScatterAndInsert; 1009cd620004SJunchao Zhang else if (op == MPI_SUM || op == MPIU_SUM) *ScatterAndOp = link->da_ScatterAndAdd; 1010cd620004SJunchao Zhang else if (op == MPI_PROD) *ScatterAndOp = link->da_ScatterAndMult; 1011cd620004SJunchao Zhang else if (op == MPI_MAX || op == MPIU_MAX) *ScatterAndOp = link->da_ScatterAndMax; 1012cd620004SJunchao Zhang else if (op == MPI_MIN || op == MPIU_MIN) *ScatterAndOp = link->da_ScatterAndMin; 1013cd620004SJunchao Zhang else if (op == MPI_LAND) *ScatterAndOp = link->da_ScatterAndLAND; 1014cd620004SJunchao Zhang else if (op == MPI_BAND) *ScatterAndOp = link->da_ScatterAndBAND; 1015cd620004SJunchao Zhang else if (op == MPI_LOR) *ScatterAndOp = link->da_ScatterAndLOR; 1016cd620004SJunchao Zhang else if (op == MPI_BOR) *ScatterAndOp = link->da_ScatterAndBOR; 1017cd620004SJunchao Zhang else if (op == MPI_LXOR) *ScatterAndOp = link->da_ScatterAndLXOR; 1018cd620004SJunchao Zhang else if (op == MPI_BXOR) *ScatterAndOp = link->da_ScatterAndBXOR; 1019cd620004SJunchao Zhang else if (op == MPI_MAXLOC) *ScatterAndOp = link->da_ScatterAndMaxloc; 1020cd620004SJunchao Zhang else if (op == MPI_MINLOC) *ScatterAndOp = link->da_ScatterAndMinloc; 1021eb02082bSJunchao Zhang } 1022eb02082bSJunchao Zhang #endif 1023cd620004SJunchao Zhang PetscFunctionReturn(0); 1024cd620004SJunchao Zhang } 1025cd620004SJunchao Zhang 1026fcc7397dSJunchao Zhang PetscErrorCode PetscSFLinkGetFetchAndOp(PetscSFLink link,PetscMemType mtype,MPI_Op op,PetscBool atomic,PetscErrorCode (**FetchAndOp)(PetscSFLink,PetscInt,PetscInt,PetscSFPackOpt,const PetscInt*,void*,void*)) 1027cd620004SJunchao Zhang { 1028cd620004SJunchao Zhang PetscFunctionBegin; 1029cd620004SJunchao Zhang *FetchAndOp = NULL; 1030cd620004SJunchao Zhang if (op != MPI_SUM && op != MPIU_SUM) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"No support for MPI_Op in FetchAndOp"); 1031cd620004SJunchao Zhang if (mtype == PETSC_MEMTYPE_HOST) *FetchAndOp = link->h_FetchAndAdd; 1032cd620004SJunchao Zhang #if defined(PETSC_HAVE_CUDA) 1033cd620004SJunchao Zhang else if (mtype == PETSC_MEMTYPE_DEVICE && !atomic) *FetchAndOp = link->d_FetchAndAdd; 1034cd620004SJunchao Zhang else if (mtype == PETSC_MEMTYPE_DEVICE && atomic) *FetchAndOp = link->da_FetchAndAdd; 1035cd620004SJunchao Zhang #endif 1036cd620004SJunchao Zhang PetscFunctionReturn(0); 1037cd620004SJunchao Zhang } 1038cd620004SJunchao Zhang 1039fcc7397dSJunchao Zhang PetscErrorCode PetscSFLinkGetFetchAndOpLocal(PetscSFLink link,PetscMemType mtype,MPI_Op op,PetscBool atomic,PetscErrorCode (**FetchAndOpLocal)(PetscSFLink,PetscInt,PetscInt,PetscSFPackOpt,const PetscInt*,void*,PetscInt,PetscSFPackOpt,const PetscInt*,const void*,void*)) 1040cd620004SJunchao Zhang { 1041cd620004SJunchao Zhang PetscFunctionBegin; 1042cd620004SJunchao Zhang *FetchAndOpLocal = NULL; 1043cd620004SJunchao Zhang if (op != MPI_SUM && op != MPIU_SUM) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"No support for MPI_Op in FetchAndOp"); 1044cd620004SJunchao Zhang if (mtype == PETSC_MEMTYPE_HOST) *FetchAndOpLocal = link->h_FetchAndAddLocal; 1045cd620004SJunchao Zhang #if defined(PETSC_HAVE_CUDA) 1046cd620004SJunchao Zhang else if (mtype == PETSC_MEMTYPE_DEVICE && !atomic) *FetchAndOpLocal = link->d_FetchAndAddLocal; 1047cd620004SJunchao Zhang else if (mtype == PETSC_MEMTYPE_DEVICE && atomic) *FetchAndOpLocal = link->da_FetchAndAddLocal; 1048cd620004SJunchao Zhang #endif 1049cd620004SJunchao Zhang PetscFunctionReturn(0); 1050cd620004SJunchao Zhang } 1051cd620004SJunchao Zhang 1052cd620004SJunchao Zhang /*============================================================================= 1053cd620004SJunchao Zhang A set of helper routines for Pack/Unpack/Scatter on GPUs 1054cd620004SJunchao Zhang ============================================================================*/ 1055cd620004SJunchao Zhang #if defined(PETSC_HAVE_CUDA) 1056cd620004SJunchao Zhang /* If SF does not know which stream root/leafdata is being computed on, it has to sync the device to 1057cd620004SJunchao Zhang make sure the data is ready for packing. 1058cd620004SJunchao Zhang */ 1059cd620004SJunchao Zhang PETSC_STATIC_INLINE PetscErrorCode PetscSFLinkSyncDeviceBeforePackData(PetscSF sf,PetscSFLink link) 1060cd620004SJunchao Zhang { 1061cd620004SJunchao Zhang PetscFunctionBegin; 1062cd620004SJunchao Zhang if (sf->use_default_stream) PetscFunctionReturn(0); 1063cd620004SJunchao Zhang if (link->rootmtype == PETSC_MEMTYPE_DEVICE || link->leafmtype == PETSC_MEMTYPE_DEVICE) { 1064cd620004SJunchao Zhang cudaError_t cerr = cudaDeviceSynchronize();CHKERRCUDA(cerr); 1065cd620004SJunchao Zhang } 1066cd620004SJunchao Zhang PetscFunctionReturn(0); 1067cd620004SJunchao Zhang } 1068cd620004SJunchao Zhang 1069cd620004SJunchao Zhang /* PetscSFLinkSyncStreamAfterPackXxxData routines make sure root/leafbuf for the remote is ready for MPI */ 1070cd620004SJunchao Zhang PETSC_STATIC_INLINE PetscErrorCode PetscSFLinkSyncStreamAfterPackRootData(PetscSF sf,PetscSFLink link) 1071cd620004SJunchao Zhang { 1072cd620004SJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 1073cd620004SJunchao Zhang 1074cd620004SJunchao Zhang PetscFunctionBegin; 1075b85e67b7SJunchao Zhang /* Do nothing if we use stream aware mpi || has nothing for remote */ 1076b85e67b7SJunchao Zhang if (sf->use_stream_aware_mpi || link->rootmtype != PETSC_MEMTYPE_DEVICE || !bas->rootbuflen[PETSCSF_REMOTE]) PetscFunctionReturn(0); 1077b85e67b7SJunchao Zhang /* If we called a packing kernel || we async-copied rootdata from device to host || No cudaDeviceSynchronize was called (since default stream is assumed) */ 1078c2a741eeSJunchao Zhang if (!link->rootdirect[PETSCSF_REMOTE] || !sf->use_gpu_aware_mpi || sf->use_default_stream) { 1079cd620004SJunchao Zhang cudaError_t cerr = cudaStreamSynchronize(link->stream);CHKERRCUDA(cerr); 1080cd620004SJunchao Zhang } 1081cd620004SJunchao Zhang PetscFunctionReturn(0); 1082cd620004SJunchao Zhang } 1083cd620004SJunchao Zhang PETSC_STATIC_INLINE PetscErrorCode PetscSFLinkSyncStreamAfterPackLeafData(PetscSF sf,PetscSFLink link) 1084cd620004SJunchao Zhang { 1085cd620004SJunchao Zhang PetscFunctionBegin; 1086b85e67b7SJunchao Zhang /* See comments above */ 1087b85e67b7SJunchao Zhang if (sf->use_stream_aware_mpi || link->leafmtype != PETSC_MEMTYPE_DEVICE || !sf->leafbuflen[PETSCSF_REMOTE]) PetscFunctionReturn(0); 1088c2a741eeSJunchao Zhang if (!link->leafdirect[PETSCSF_REMOTE] || !sf->use_gpu_aware_mpi || sf->use_default_stream) { 1089cd620004SJunchao Zhang cudaError_t cerr = cudaStreamSynchronize(link->stream);CHKERRCUDA(cerr); 1090cd620004SJunchao Zhang } 1091cd620004SJunchao Zhang PetscFunctionReturn(0); 1092cd620004SJunchao Zhang } 1093cd620004SJunchao Zhang 1094cd620004SJunchao Zhang /* PetscSFLinkSyncStreamAfterUnpackXxx routines make sure root/leafdata (local & remote) is ready to use for SF callers, when SF 1095cd620004SJunchao Zhang does not know which stream the callers will use. 1096cd620004SJunchao Zhang */ 1097cd620004SJunchao Zhang PETSC_STATIC_INLINE PetscErrorCode PetscSFLinkSyncStreamAfterUnpackRootData(PetscSF sf,PetscSFLink link) 1098cd620004SJunchao Zhang { 1099cd620004SJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 11004a314419SJunchao Zhang PetscBool host2host = (link->rootmtype == PETSC_MEMTYPE_HOST) && (link->leafmtype == PETSC_MEMTYPE_HOST) ? PETSC_TRUE : PETSC_FALSE; 1101cd620004SJunchao Zhang 1102cd620004SJunchao Zhang PetscFunctionBegin; 11034a314419SJunchao Zhang /* Do nothing if host2host OR we are allowed to asynchronously put rootdata on device through the default stream */ 11044a314419SJunchao Zhang if (host2host || (link->rootmtype == PETSC_MEMTYPE_DEVICE && sf->use_default_stream)) PetscFunctionReturn(0); 11054a314419SJunchao Zhang 11064a314419SJunchao Zhang /* If rootmtype is HOST or DEVICE: 11074a314419SJunchao Zhang If we have data from local, then we called a scatter kernel (on link->stream), then we must sync it; 11084a314419SJunchao Zhang If we have data from remote && no rootdirect(i.e., we called an unpack kernel), then we must also sycn it (if rootdirect, 11094a314419SJunchao Zhang i.e., no unpack kernel after MPI, MPI guarentees rootbuf is ready to use so that we do not need the sync). 11104a314419SJunchao Zhang 11114a314419SJunchao Zhang Note a tricky case is when leafmtype=DEVICE, rootmtype=HOST on uni-processor, we must sync the stream otherwise 11124a314419SJunchao Zhang CPU thread might use the yet-to-be-updated rootdata pending in the stream. 1113b85e67b7SJunchao Zhang */ 1114b85e67b7SJunchao Zhang if (bas->rootbuflen[PETSCSF_LOCAL] || (bas->rootbuflen[PETSCSF_REMOTE] && !link->rootdirect[PETSCSF_REMOTE])) { 1115cd620004SJunchao Zhang cudaError_t cerr = cudaStreamSynchronize(link->stream);CHKERRCUDA(cerr); 1116cd620004SJunchao Zhang } 1117cd620004SJunchao Zhang PetscFunctionReturn(0); 1118cd620004SJunchao Zhang } 1119cd620004SJunchao Zhang PETSC_STATIC_INLINE PetscErrorCode PetscSFLinkSyncStreamAfterUnpackLeafData(PetscSF sf,PetscSFLink link) 1120cd620004SJunchao Zhang { 11214a314419SJunchao Zhang PetscBool host2host = (link->rootmtype == PETSC_MEMTYPE_HOST) && (link->leafmtype == PETSC_MEMTYPE_HOST) ? PETSC_TRUE : PETSC_FALSE; 11224a314419SJunchao Zhang 1123cd620004SJunchao Zhang PetscFunctionBegin; 11244a314419SJunchao Zhang /* See comments in PetscSFLinkSyncStreamAfterUnpackRootData*/ 11254a314419SJunchao Zhang if (host2host || (link->leafmtype == PETSC_MEMTYPE_DEVICE && sf->use_default_stream)) PetscFunctionReturn(0); 1126b85e67b7SJunchao Zhang if (sf->leafbuflen[PETSCSF_LOCAL] || (sf->leafbuflen[PETSCSF_REMOTE] && !link->leafdirect[PETSCSF_REMOTE])) { 1127cd620004SJunchao Zhang cudaError_t cerr = cudaStreamSynchronize(link->stream);CHKERRCUDA(cerr); 1128cd620004SJunchao Zhang } 1129cd620004SJunchao Zhang PetscFunctionReturn(0); 1130cd620004SJunchao Zhang } 1131cd620004SJunchao Zhang 1132cd620004SJunchao Zhang /* PetscSFLinkCopyXxxxBufferInCaseNotUseGpuAwareMPI routines are simple: if not use_gpu_aware_mpi, we need 1133cd620004SJunchao Zhang to copy the buffer from GPU to CPU before MPI calls, and from CPU to GPU after MPI calls. 1134cd620004SJunchao Zhang */ 1135cd620004SJunchao Zhang PETSC_STATIC_INLINE PetscErrorCode PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(PetscSF sf,PetscSFLink link,PetscBool device2host) 1136cd620004SJunchao Zhang { 1137cd620004SJunchao Zhang PetscErrorCode ierr; 1138cd620004SJunchao Zhang cudaError_t cerr; 1139cd620004SJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 1140cd620004SJunchao Zhang 1141cd620004SJunchao Zhang PetscFunctionBegin; 1142cd620004SJunchao Zhang if (link->rootmtype == PETSC_MEMTYPE_DEVICE && (link->rootmtype_mpi != link->rootmtype) && bas->rootbuflen[PETSCSF_REMOTE]) { 1143cd620004SJunchao Zhang void *h_buf = link->rootbuf[PETSCSF_REMOTE][PETSC_MEMTYPE_HOST]; 1144cd620004SJunchao Zhang void *d_buf = link->rootbuf[PETSCSF_REMOTE][PETSC_MEMTYPE_DEVICE]; 1145cd620004SJunchao Zhang size_t count = bas->rootbuflen[PETSCSF_REMOTE]*link->unitbytes; 1146cd620004SJunchao Zhang if (device2host) { 1147cd620004SJunchao Zhang cerr = cudaMemcpyAsync(h_buf,d_buf,count,cudaMemcpyDeviceToHost,link->stream);CHKERRCUDA(cerr); 1148cd620004SJunchao Zhang ierr = PetscLogGpuToCpu(count);CHKERRQ(ierr); 1149cd620004SJunchao Zhang } else { 1150cd620004SJunchao Zhang cerr = cudaMemcpyAsync(d_buf,h_buf,count,cudaMemcpyHostToDevice,link->stream);CHKERRCUDA(cerr); 1151cd620004SJunchao Zhang ierr = PetscLogCpuToGpu(count);CHKERRQ(ierr); 1152cd620004SJunchao Zhang } 1153cd620004SJunchao Zhang } 1154cd620004SJunchao Zhang PetscFunctionReturn(0); 1155cd620004SJunchao Zhang } 1156cd620004SJunchao Zhang 1157cd620004SJunchao Zhang PETSC_STATIC_INLINE PetscErrorCode PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(PetscSF sf,PetscSFLink link,PetscBool device2host) 1158cd620004SJunchao Zhang { 1159cd620004SJunchao Zhang PetscErrorCode ierr; 1160cd620004SJunchao Zhang cudaError_t cerr; 1161cd620004SJunchao Zhang 1162cd620004SJunchao Zhang PetscFunctionBegin; 1163cd620004SJunchao Zhang if (link->leafmtype == PETSC_MEMTYPE_DEVICE && (link->leafmtype_mpi != link->leafmtype) && sf->leafbuflen[PETSCSF_REMOTE]) { 1164cd620004SJunchao Zhang void *h_buf = link->leafbuf[PETSCSF_REMOTE][PETSC_MEMTYPE_HOST]; 1165cd620004SJunchao Zhang void *d_buf = link->leafbuf[PETSCSF_REMOTE][PETSC_MEMTYPE_DEVICE]; 1166cd620004SJunchao Zhang size_t count = sf->leafbuflen[PETSCSF_REMOTE]*link->unitbytes; 1167cd620004SJunchao Zhang if (device2host) { 1168cd620004SJunchao Zhang cerr = cudaMemcpyAsync(h_buf,d_buf,count,cudaMemcpyDeviceToHost,link->stream);CHKERRCUDA(cerr); 1169cd620004SJunchao Zhang ierr = PetscLogGpuToCpu(count);CHKERRQ(ierr); 1170cd620004SJunchao Zhang } else { 1171cd620004SJunchao Zhang cerr = cudaMemcpyAsync(d_buf,h_buf,count,cudaMemcpyHostToDevice,link->stream);CHKERRCUDA(cerr); 1172cd620004SJunchao Zhang ierr = PetscLogCpuToGpu(count);CHKERRQ(ierr); 1173cd620004SJunchao Zhang } 1174cd620004SJunchao Zhang } 1175cd620004SJunchao Zhang PetscFunctionReturn(0); 1176cd620004SJunchao Zhang } 1177cd620004SJunchao Zhang #else 1178cd620004SJunchao Zhang 1179cd620004SJunchao Zhang #define PetscSFLinkSyncDeviceBeforePackData(a,b) 0 1180cd620004SJunchao Zhang #define PetscSFLinkSyncStreamAfterPackRootData(a,b) 0 1181cd620004SJunchao Zhang #define PetscSFLinkSyncStreamAfterPackLeafData(a,b) 0 1182cd620004SJunchao Zhang #define PetscSFLinkSyncStreamAfterUnpackRootData(a,b) 0 1183cd620004SJunchao Zhang #define PetscSFLinkSyncStreamAfterUnpackLeafData(a,b) 0 1184cd620004SJunchao Zhang #define PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(a,b,c) 0 1185cd620004SJunchao Zhang #define PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(a,b,c) 0 1186cd620004SJunchao Zhang 1187cd620004SJunchao Zhang #endif 1188cd620004SJunchao Zhang 1189cd620004SJunchao Zhang PETSC_STATIC_INLINE PetscErrorCode PetscSFLinkLogFlopsAfterUnpackRootData(PetscSF sf,PetscSFLink link,PetscSFScope scope,MPI_Op op) 1190cd620004SJunchao Zhang { 1191cd620004SJunchao Zhang PetscErrorCode ierr; 1192cd620004SJunchao Zhang PetscLogDouble flops; 1193cd620004SJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 1194cd620004SJunchao Zhang 1195cd620004SJunchao Zhang 1196cd620004SJunchao Zhang PetscFunctionBegin; 1197cd620004SJunchao Zhang if (op != MPIU_REPLACE && link->basicunit == MPIU_SCALAR) { /* op is a reduction on PetscScalars */ 1198cd620004SJunchao Zhang flops = bas->rootbuflen[scope]*link->bs; /* # of roots in buffer x # of scalars in unit */ 1199cd620004SJunchao Zhang #if defined(PETSC_HAVE_CUDA) 1200cd620004SJunchao Zhang if (link->rootmtype == PETSC_MEMTYPE_DEVICE) {ierr = PetscLogGpuFlops(flops);CHKERRQ(ierr);} else 1201cd620004SJunchao Zhang #endif 1202cd620004SJunchao Zhang {ierr = PetscLogFlops(flops);CHKERRQ(ierr);} 1203cd620004SJunchao Zhang } 1204cd620004SJunchao Zhang PetscFunctionReturn(0); 1205cd620004SJunchao Zhang } 1206cd620004SJunchao Zhang 1207cd620004SJunchao Zhang PETSC_STATIC_INLINE PetscErrorCode PetscSFLinkLogFlopsAfterUnpackLeafData(PetscSF sf,PetscSFLink link,PetscSFScope scope,MPI_Op op) 1208cd620004SJunchao Zhang { 1209cd620004SJunchao Zhang PetscLogDouble flops; 1210cd620004SJunchao Zhang PetscErrorCode ierr; 1211cd620004SJunchao Zhang 1212cd620004SJunchao Zhang PetscFunctionBegin; 1213cd620004SJunchao Zhang if (op != MPIU_REPLACE && link->basicunit == MPIU_SCALAR) { /* op is a reduction on PetscScalars */ 1214cd620004SJunchao Zhang flops = sf->leafbuflen[scope]*link->bs; /* # of roots in buffer x # of scalars in unit */ 1215cd620004SJunchao Zhang #if defined(PETSC_HAVE_CUDA) 1216cd620004SJunchao Zhang if (link->leafmtype == PETSC_MEMTYPE_DEVICE) {ierr = PetscLogGpuFlops(flops);CHKERRQ(ierr);} else 1217cd620004SJunchao Zhang #endif 1218cd620004SJunchao Zhang {ierr = PetscLogFlops(flops);CHKERRQ(ierr);} 1219cd620004SJunchao Zhang } 1220cd620004SJunchao Zhang PetscFunctionReturn(0); 1221cd620004SJunchao Zhang } 1222cd620004SJunchao Zhang 1223cd620004SJunchao Zhang /* When SF could not find a proper UnpackAndOp() from link, it falls back to MPI_Reduce_local. 1224cd620004SJunchao Zhang Input Arguments: 1225cd620004SJunchao Zhang +sf - The StarForest 1226cd620004SJunchao Zhang .link - The link 1227cd620004SJunchao Zhang .count - Number of entries to unpack 1228cd620004SJunchao Zhang .start - The first index, significent when indices=NULL 1229cd620004SJunchao Zhang .indices - Indices of entries in <data>. If NULL, it means indices are contiguous and the first is given in <start> 1230cd620004SJunchao Zhang .buf - A contiguous buffer to unpack from 1231cd620004SJunchao Zhang -op - Operation after unpack 1232cd620004SJunchao Zhang 1233cd620004SJunchao Zhang Output Arguments: 1234cd620004SJunchao Zhang .data - The data to unpack to 1235cd620004SJunchao Zhang */ 1236cd620004SJunchao Zhang PETSC_STATIC_INLINE PetscErrorCode PetscSFLinkUnpackDataWithMPIReduceLocal(PetscSF sf,PetscSFLink link,PetscInt count,PetscInt start,const PetscInt *indices,void *data,const void *buf,MPI_Op op) 1237cd620004SJunchao Zhang { 1238cd620004SJunchao Zhang PetscFunctionBegin; 1239cd620004SJunchao Zhang #if defined(PETSC_HAVE_MPI_REDUCE_LOCAL) 1240cd620004SJunchao Zhang { 1241cd620004SJunchao Zhang PetscErrorCode ierr; 1242cd620004SJunchao Zhang PetscInt i; 1243cd620004SJunchao Zhang PetscMPIInt n; 1244cd620004SJunchao Zhang if (indices) { 1245cd620004SJunchao Zhang /* Note we use link->unit instead of link->basicunit. When op can be mapped to MPI_SUM etc, it operates on 1246cd620004SJunchao Zhang basic units of a root/leaf element-wisely. Otherwise, it is meant to operate on a whole root/leaf. 1247cd620004SJunchao Zhang */ 1248cd620004SJunchao Zhang for (i=0; i<count; i++) {ierr = MPI_Reduce_local((const char*)buf+i*link->unitbytes,(char*)data+indices[i]*link->unitbytes,1,link->unit,op);CHKERRQ(ierr);} 1249cd620004SJunchao Zhang } else { 1250cd620004SJunchao Zhang ierr = PetscMPIIntCast(count,&n);CHKERRQ(ierr); 1251cd620004SJunchao Zhang ierr = MPI_Reduce_local(buf,(char*)data+start*link->unitbytes,n,link->unit,op);CHKERRQ(ierr); 1252cd620004SJunchao Zhang } 1253cd620004SJunchao Zhang } 1254cd620004SJunchao Zhang #else 1255cd620004SJunchao Zhang SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"No unpacking reduction operation for this MPI_Op"); 1256cd620004SJunchao Zhang #endif 1257cd620004SJunchao Zhang PetscFunctionReturn(0); 1258cd620004SJunchao Zhang } 1259cd620004SJunchao Zhang 1260fcc7397dSJunchao Zhang PETSC_STATIC_INLINE PetscErrorCode PetscSFLinkScatterDataWithMPIReduceLocal(PetscSF sf,PetscSFLink link,PetscInt count,PetscInt srcStart,const PetscInt *srcIdx,const void *src,PetscInt dstStart,const PetscInt *dstIdx,void *dst,MPI_Op op) 1261cd620004SJunchao Zhang { 1262cd620004SJunchao Zhang PetscFunctionBegin; 1263cd620004SJunchao Zhang #if defined(PETSC_HAVE_MPI_REDUCE_LOCAL) 1264cd620004SJunchao Zhang { 1265cd620004SJunchao Zhang PetscErrorCode ierr; 1266cd620004SJunchao Zhang PetscInt i,disp; 1267fcc7397dSJunchao Zhang if (!srcIdx) { 1268fcc7397dSJunchao Zhang ierr = PetscSFLinkUnpackDataWithMPIReduceLocal(sf,link,count,dstStart,dstIdx,dst,(const char*)src+srcStart*link->unitbytes,op);CHKERRQ(ierr); 1269fcc7397dSJunchao Zhang } else { 1270cd620004SJunchao Zhang for (i=0; i<count; i++) { 1271fcc7397dSJunchao Zhang disp = dstIdx? dstIdx[i] : dstStart + i; 1272fcc7397dSJunchao Zhang ierr = MPI_Reduce_local((const char*)src+srcIdx[i]*link->unitbytes,(char*)dst+disp*link->unitbytes,1,link->unit,op);CHKERRQ(ierr); 1273fcc7397dSJunchao Zhang } 1274cd620004SJunchao Zhang } 1275cd620004SJunchao Zhang } 1276cd620004SJunchao Zhang #else 1277cd620004SJunchao Zhang SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"No unpacking reduction operation for this MPI_Op"); 1278cd620004SJunchao Zhang #endif 1279cd620004SJunchao Zhang PetscFunctionReturn(0); 1280cd620004SJunchao Zhang } 1281cd620004SJunchao Zhang 1282cd620004SJunchao Zhang /*============================================================================= 1283cd620004SJunchao Zhang Pack/Unpack/Fetch/Scatter routines 1284cd620004SJunchao Zhang ============================================================================*/ 1285cd620004SJunchao Zhang 1286cd620004SJunchao Zhang /* Pack rootdata to rootbuf 1287cd620004SJunchao Zhang Input Arguments: 1288cd620004SJunchao Zhang + sf - The SF this packing works on. 1289cd620004SJunchao Zhang . link - It gives the memtype of the roots and also provides root buffer. 1290cd620004SJunchao Zhang . scope - PETSCSF_LOCAL or PETSCSF_REMOTE. Note SF has the ability to do local and remote communications separately. 1291cd620004SJunchao Zhang - rootdata - Where to read the roots. 1292cd620004SJunchao Zhang 1293cd620004SJunchao Zhang Notes: 1294cd620004SJunchao Zhang When rootdata can be directly used as root buffer, the routine is almost a no-op. After the call, root data is 1295cd620004SJunchao Zhang in a place where the underlying MPI is ready can access (use_gpu_aware_mpi or not) 1296cd620004SJunchao Zhang */ 1297cd620004SJunchao Zhang PetscErrorCode PetscSFLinkPackRootData(PetscSF sf,PetscSFLink link,PetscSFScope scope,const void *rootdata) 1298cd620004SJunchao Zhang { 1299cd620004SJunchao Zhang PetscErrorCode ierr; 1300cd620004SJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 1301cd620004SJunchao Zhang const PetscInt *rootindices = NULL; 1302cd620004SJunchao Zhang PetscInt count,start; 1303fcc7397dSJunchao Zhang PetscErrorCode (*Pack)(PetscSFLink,PetscInt,PetscInt,PetscSFPackOpt,const PetscInt*,const void*,void*) = NULL; 1304cd620004SJunchao Zhang PetscMemType rootmtype = link->rootmtype; 1305fcc7397dSJunchao Zhang PetscSFPackOpt opt = NULL; 1306fcc7397dSJunchao Zhang 1307cd620004SJunchao Zhang 1308cd620004SJunchao Zhang PetscFunctionBegin; 1309cd620004SJunchao Zhang ierr = PetscLogEventBegin(PETSCSF_Pack,sf,0,0,0);CHKERRQ(ierr); 1310cd620004SJunchao Zhang if (scope == PETSCSF_REMOTE) {ierr = PetscSFLinkSyncDeviceBeforePackData(sf,link);CHKERRQ(ierr);} 1311cd620004SJunchao Zhang if (!link->rootdirect[scope] && bas->rootbuflen[scope]) { /* If rootdata works directly as rootbuf, skip packing */ 1312fcc7397dSJunchao Zhang ierr = PetscSFLinkGetRootPackOptAndIndices(sf,link,rootmtype,scope,&count,&start,&opt,&rootindices);CHKERRQ(ierr); 1313cd620004SJunchao Zhang ierr = PetscSFLinkGetPack(link,rootmtype,&Pack);CHKERRQ(ierr); 1314fcc7397dSJunchao Zhang ierr = (*Pack)(link,count,start,opt,rootindices,rootdata,link->rootbuf[scope][rootmtype]);CHKERRQ(ierr); 1315cd620004SJunchao Zhang } 1316cd620004SJunchao Zhang if (scope == PETSCSF_REMOTE) { 1317cd620004SJunchao Zhang ierr = PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf,link,PETSC_TRUE/*device2host*/);CHKERRQ(ierr); 1318cd620004SJunchao Zhang ierr = PetscSFLinkSyncStreamAfterPackRootData(sf,link);CHKERRQ(ierr); 1319cd620004SJunchao Zhang } 1320cd620004SJunchao Zhang ierr = PetscLogEventEnd(PETSCSF_Pack,sf,0,0,0);CHKERRQ(ierr); 1321cd620004SJunchao Zhang PetscFunctionReturn(0); 1322cd620004SJunchao Zhang } 1323cd620004SJunchao Zhang 1324cd620004SJunchao Zhang /* Pack leafdata to leafbuf */ 1325cd620004SJunchao Zhang PetscErrorCode PetscSFLinkPackLeafData(PetscSF sf,PetscSFLink link,PetscSFScope scope,const void *leafdata) 1326cd620004SJunchao Zhang { 1327cd620004SJunchao Zhang PetscErrorCode ierr; 1328cd620004SJunchao Zhang const PetscInt *leafindices = NULL; 1329cd620004SJunchao Zhang PetscInt count,start; 1330fcc7397dSJunchao Zhang PetscErrorCode (*Pack)(PetscSFLink,PetscInt,PetscInt,PetscSFPackOpt,const PetscInt*,const void*,void*) = NULL; 1331cd620004SJunchao Zhang PetscMemType leafmtype = link->leafmtype; 1332fcc7397dSJunchao Zhang PetscSFPackOpt opt = NULL; 1333cd620004SJunchao Zhang 1334cd620004SJunchao Zhang PetscFunctionBegin; 1335cd620004SJunchao Zhang ierr = PetscLogEventBegin(PETSCSF_Pack,sf,0,0,0);CHKERRQ(ierr); 1336cd620004SJunchao Zhang if (scope == PETSCSF_REMOTE) {ierr = PetscSFLinkSyncDeviceBeforePackData(sf,link);CHKERRQ(ierr);} 1337cd620004SJunchao Zhang if (!link->leafdirect[scope] && sf->leafbuflen[scope]) { /* If leafdata works directly as rootbuf, skip packing */ 1338fcc7397dSJunchao Zhang ierr = PetscSFLinkGetLeafPackOptAndIndices(sf,link,leafmtype,scope,&count,&start,&opt,&leafindices);CHKERRQ(ierr); 1339cd620004SJunchao Zhang ierr = PetscSFLinkGetPack(link,leafmtype,&Pack);CHKERRQ(ierr); 1340fcc7397dSJunchao Zhang ierr = (*Pack)(link,count,start,opt,leafindices,leafdata,link->leafbuf[scope][leafmtype]);CHKERRQ(ierr); 1341cd620004SJunchao Zhang } 1342cd620004SJunchao Zhang if (scope == PETSCSF_REMOTE) { 1343cd620004SJunchao Zhang ierr = PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf,link,PETSC_TRUE/*device2host*/);CHKERRQ(ierr); 1344cd620004SJunchao Zhang ierr = PetscSFLinkSyncStreamAfterPackLeafData(sf,link);CHKERRQ(ierr); 1345cd620004SJunchao Zhang } 1346cd620004SJunchao Zhang ierr = PetscLogEventEnd(PETSCSF_Pack,sf,0,0,0);CHKERRQ(ierr); 1347cd620004SJunchao Zhang PetscFunctionReturn(0); 1348cd620004SJunchao Zhang } 1349cd620004SJunchao Zhang 1350cd620004SJunchao Zhang /* Unpack rootbuf to rootdata */ 1351cd620004SJunchao Zhang PetscErrorCode PetscSFLinkUnpackRootData(PetscSF sf,PetscSFLink link,PetscSFScope scope,void *rootdata,MPI_Op op) 1352cd620004SJunchao Zhang { 1353cd620004SJunchao Zhang PetscErrorCode ierr; 1354cd620004SJunchao Zhang const PetscInt *rootindices = NULL; 1355cd620004SJunchao Zhang PetscInt count,start; 1356cd620004SJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 1357fcc7397dSJunchao Zhang PetscErrorCode (*UnpackAndOp)(PetscSFLink,PetscInt,PetscInt,PetscSFPackOpt,const PetscInt*,void*,const void*) = NULL; 1358cd620004SJunchao Zhang PetscMemType rootmtype = link->rootmtype; 1359fcc7397dSJunchao Zhang PetscSFPackOpt opt = NULL; 1360cd620004SJunchao Zhang 1361cd620004SJunchao Zhang PetscFunctionBegin; 1362cd620004SJunchao Zhang ierr = PetscLogEventBegin(PETSCSF_Unpack,sf,0,0,0);CHKERRQ(ierr); 1363cd620004SJunchao Zhang if (scope == PETSCSF_REMOTE) {ierr = PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf,link,PETSC_FALSE);CHKERRQ(ierr);} 1364cd620004SJunchao Zhang if (!link->rootdirect[scope] && bas->rootbuflen[scope]) { /* If rootdata works directly as rootbuf, skip unpacking */ 1365cd620004SJunchao Zhang ierr = PetscSFLinkGetUnpackAndOp(link,rootmtype,op,bas->rootdups[scope],&UnpackAndOp);CHKERRQ(ierr); 1366cd620004SJunchao Zhang if (UnpackAndOp) { 1367fcc7397dSJunchao Zhang ierr = PetscSFLinkGetRootPackOptAndIndices(sf,link,rootmtype,scope,&count,&start,&opt,&rootindices);CHKERRQ(ierr); 1368fcc7397dSJunchao Zhang ierr = (*UnpackAndOp)(link,count,start,opt,rootindices,rootdata,link->rootbuf[scope][rootmtype]);CHKERRQ(ierr); 1369cd620004SJunchao Zhang } else { 1370fcc7397dSJunchao Zhang ierr = PetscSFLinkGetRootPackOptAndIndices(sf,link,PETSC_MEMTYPE_HOST,scope,&count,&start,&opt,&rootindices);CHKERRQ(ierr); 1371cd620004SJunchao Zhang ierr = PetscSFLinkUnpackDataWithMPIReduceLocal(sf,link,count,start,rootindices,rootdata,link->rootbuf[scope][rootmtype],op);CHKERRQ(ierr); 1372cd620004SJunchao Zhang } 1373cd620004SJunchao Zhang } 1374cd620004SJunchao Zhang if (scope == PETSCSF_REMOTE) {ierr = PetscSFLinkSyncStreamAfterUnpackRootData(sf,link);CHKERRQ(ierr);} 1375cd620004SJunchao Zhang ierr = PetscSFLinkLogFlopsAfterUnpackRootData(sf,link,scope,op);CHKERRQ(ierr); 1376cd620004SJunchao Zhang ierr = PetscLogEventEnd(PETSCSF_Unpack,sf,0,0,0);CHKERRQ(ierr); 1377cd620004SJunchao Zhang PetscFunctionReturn(0); 1378cd620004SJunchao Zhang } 1379cd620004SJunchao Zhang 1380cd620004SJunchao Zhang /* Unpack leafbuf to leafdata */ 1381cd620004SJunchao Zhang PetscErrorCode PetscSFLinkUnpackLeafData(PetscSF sf,PetscSFLink link,PetscSFScope scope,void *leafdata,MPI_Op op) 1382cd620004SJunchao Zhang { 1383cd620004SJunchao Zhang PetscErrorCode ierr; 1384cd620004SJunchao Zhang const PetscInt *leafindices = NULL; 1385cd620004SJunchao Zhang PetscInt count,start; 1386fcc7397dSJunchao Zhang PetscErrorCode (*UnpackAndOp)(PetscSFLink,PetscInt,PetscInt,PetscSFPackOpt,const PetscInt*,void*,const void*) = NULL; 1387cd620004SJunchao Zhang PetscMemType leafmtype = link->leafmtype; 1388fcc7397dSJunchao Zhang PetscSFPackOpt opt = NULL; 1389cd620004SJunchao Zhang 1390cd620004SJunchao Zhang PetscFunctionBegin; 1391cd620004SJunchao Zhang ierr = PetscLogEventBegin(PETSCSF_Unpack,sf,0,0,0);CHKERRQ(ierr); 1392cd620004SJunchao Zhang if (scope == PETSCSF_REMOTE) {ierr = PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf,link,PETSC_FALSE);CHKERRQ(ierr);} 1393cd620004SJunchao Zhang if (!link->leafdirect[scope] && sf->leafbuflen[scope]) { /* If leafdata works directly as rootbuf, skip unpacking */ 1394cd620004SJunchao Zhang ierr = PetscSFLinkGetUnpackAndOp(link,leafmtype,op,sf->leafdups[scope],&UnpackAndOp);CHKERRQ(ierr); 1395cd620004SJunchao Zhang if (UnpackAndOp) { 1396fcc7397dSJunchao Zhang ierr = PetscSFLinkGetLeafPackOptAndIndices(sf,link,leafmtype,scope,&count,&start,&opt,&leafindices);CHKERRQ(ierr); 1397fcc7397dSJunchao Zhang ierr = (*UnpackAndOp)(link,count,start,opt,leafindices,leafdata,link->leafbuf[scope][leafmtype]);CHKERRQ(ierr); 1398cd620004SJunchao Zhang } else { 1399fcc7397dSJunchao Zhang ierr = PetscSFLinkGetLeafPackOptAndIndices(sf,link,PETSC_MEMTYPE_HOST,scope,&count,&start,&opt,&leafindices);CHKERRQ(ierr); 1400cd620004SJunchao Zhang ierr = PetscSFLinkUnpackDataWithMPIReduceLocal(sf,link,count,start,leafindices,leafdata,link->leafbuf[scope][leafmtype],op);CHKERRQ(ierr); 1401cd620004SJunchao Zhang } 1402cd620004SJunchao Zhang } 1403cd620004SJunchao Zhang if (scope == PETSCSF_REMOTE) {ierr = PetscSFLinkSyncStreamAfterUnpackLeafData(sf,link);CHKERRQ(ierr);} 1404cd620004SJunchao Zhang ierr = PetscSFLinkLogFlopsAfterUnpackLeafData(sf,link,scope,op);CHKERRQ(ierr); 1405cd620004SJunchao Zhang ierr = PetscLogEventEnd(PETSCSF_Unpack,sf,0,0,0);CHKERRQ(ierr); 1406cd620004SJunchao Zhang PetscFunctionReturn(0); 1407cd620004SJunchao Zhang } 1408cd620004SJunchao Zhang 1409cd620004SJunchao Zhang /* FetchAndOp rootdata with rootbuf */ 1410cd620004SJunchao Zhang PetscErrorCode PetscSFLinkFetchRootData(PetscSF sf,PetscSFLink link,PetscSFScope scope,void *rootdata,MPI_Op op) 1411cd620004SJunchao Zhang { 1412cd620004SJunchao Zhang PetscErrorCode ierr; 1413cd620004SJunchao Zhang const PetscInt *rootindices = NULL; 1414cd620004SJunchao Zhang PetscInt count,start; 1415cd620004SJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 1416fcc7397dSJunchao Zhang PetscErrorCode (*FetchAndOp)(PetscSFLink,PetscInt,PetscInt,PetscSFPackOpt,const PetscInt*,void*,void*) = NULL; 1417cd620004SJunchao Zhang PetscMemType rootmtype = link->rootmtype; 1418fcc7397dSJunchao Zhang PetscSFPackOpt opt = NULL; 1419cd620004SJunchao Zhang 1420cd620004SJunchao Zhang PetscFunctionBegin; 1421cd620004SJunchao Zhang ierr = PetscLogEventBegin(PETSCSF_Unpack,sf,0,0,0);CHKERRQ(ierr); 1422cd620004SJunchao Zhang if (scope == PETSCSF_REMOTE) {ierr = PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf,link,PETSC_FALSE);CHKERRQ(ierr);} 1423cd620004SJunchao Zhang if (bas->rootbuflen[scope]) { 1424cd620004SJunchao Zhang /* Do FetchAndOp on rootdata with rootbuf */ 1425cd620004SJunchao Zhang ierr = PetscSFLinkGetFetchAndOp(link,rootmtype,op,bas->rootdups[scope],&FetchAndOp);CHKERRQ(ierr); 1426fcc7397dSJunchao Zhang ierr = PetscSFLinkGetRootPackOptAndIndices(sf,link,rootmtype,scope,&count,&start,&opt,&rootindices);CHKERRQ(ierr); 1427fcc7397dSJunchao Zhang ierr = (*FetchAndOp)(link,count,start,opt,rootindices,rootdata,link->rootbuf[scope][rootmtype]);CHKERRQ(ierr); 1428cd620004SJunchao Zhang } 1429cd620004SJunchao Zhang if (scope == PETSCSF_REMOTE) { 1430cd620004SJunchao Zhang ierr = PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf,link,PETSC_TRUE);CHKERRQ(ierr); 1431cd620004SJunchao Zhang ierr = PetscSFLinkSyncStreamAfterUnpackRootData(sf,link);CHKERRQ(ierr); 1432cd620004SJunchao Zhang } 1433cd620004SJunchao Zhang ierr = PetscSFLinkLogFlopsAfterUnpackRootData(sf,link,scope,op);CHKERRQ(ierr); 1434cd620004SJunchao Zhang ierr = PetscLogEventEnd(PETSCSF_Unpack,sf,0,0,0);CHKERRQ(ierr); 1435cd620004SJunchao Zhang PetscFunctionReturn(0); 1436cd620004SJunchao Zhang } 1437cd620004SJunchao Zhang 1438cd620004SJunchao Zhang /* Bcast rootdata to leafdata locally (i.e., only for local communication - PETSCSF_LOCAL) */ 1439cd620004SJunchao Zhang PetscErrorCode PetscSFLinkBcastAndOpLocal(PetscSF sf,PetscSFLink link,const void *rootdata,void *leafdata,MPI_Op op) 1440cd620004SJunchao Zhang { 1441cd620004SJunchao Zhang PetscErrorCode ierr; 1442cd620004SJunchao Zhang const PetscInt *rootindices = NULL,*leafindices = NULL; 1443cd620004SJunchao Zhang PetscInt count,rootstart,leafstart; 1444cd620004SJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 1445fcc7397dSJunchao Zhang PetscErrorCode (*ScatterAndOp)(PetscSFLink,PetscInt,PetscInt,PetscSFPackOpt,const PetscInt*,const void*,PetscInt,PetscSFPackOpt,const PetscInt*,void*) = NULL; 1446cd620004SJunchao Zhang const PetscMemType rootmtype = link->rootmtype,leafmtype = link->leafmtype; 1447fcc7397dSJunchao Zhang PetscSFPackOpt leafopt = NULL,rootopt = NULL; 1448cd620004SJunchao Zhang 1449cd620004SJunchao Zhang PetscFunctionBegin; 1450cd620004SJunchao Zhang if (!bas->rootbuflen[PETSCSF_LOCAL]) PetscFunctionReturn(0); 1451cd620004SJunchao Zhang if (rootmtype != leafmtype) { /* Uncommon case */ 1452cd620004SJunchao Zhang /* The local communication has to go through pack and unpack */ 1453cd620004SJunchao Zhang ierr = PetscSFLinkPackRootData(sf,link,PETSCSF_LOCAL,rootdata);CHKERRQ(ierr); 1454f01131f0SJunchao Zhang ierr = PetscSFLinkMemcpy(sf,link,leafmtype,link->leafbuf[PETSCSF_LOCAL][leafmtype],rootmtype,link->rootbuf[PETSCSF_LOCAL][rootmtype],sf->leafbuflen[PETSCSF_LOCAL]*link->unitbytes);CHKERRQ(ierr); 1455cd620004SJunchao Zhang ierr = PetscSFLinkUnpackLeafData(sf,link,PETSCSF_LOCAL,leafdata,op);CHKERRQ(ierr); 1456cd620004SJunchao Zhang } else { 1457cd620004SJunchao Zhang ierr = PetscSFLinkGetScatterAndOp(link,leafmtype,op,sf->leafdups[PETSCSF_LOCAL],&ScatterAndOp);CHKERRQ(ierr); 1458cd620004SJunchao Zhang if (ScatterAndOp) { 1459fcc7397dSJunchao Zhang ierr = PetscSFLinkGetRootPackOptAndIndices(sf,link,rootmtype,PETSCSF_LOCAL,&count,&rootstart,&rootopt,&rootindices);CHKERRQ(ierr); 1460fcc7397dSJunchao Zhang ierr = PetscSFLinkGetLeafPackOptAndIndices(sf,link,leafmtype,PETSCSF_LOCAL,&count,&leafstart,&leafopt,&leafindices);CHKERRQ(ierr); 1461fcc7397dSJunchao Zhang ierr = (*ScatterAndOp)(link,count,rootstart,rootopt,rootindices,rootdata,leafstart,leafopt,leafindices,leafdata);CHKERRQ(ierr); 1462cd620004SJunchao Zhang } else { 1463fcc7397dSJunchao Zhang ierr = PetscSFLinkGetRootPackOptAndIndices(sf,link,PETSC_MEMTYPE_HOST,PETSCSF_LOCAL,&count,&rootstart,&rootopt,&rootindices);CHKERRQ(ierr); 1464fcc7397dSJunchao Zhang ierr = PetscSFLinkGetLeafPackOptAndIndices(sf,link,PETSC_MEMTYPE_HOST,PETSCSF_LOCAL,&count,&leafstart,&leafopt,&leafindices);CHKERRQ(ierr); 1465fcc7397dSJunchao Zhang ierr = PetscSFLinkScatterDataWithMPIReduceLocal(sf,link,count,rootstart,rootindices,rootdata,leafstart,leafindices,leafdata,op);CHKERRQ(ierr); 1466cd620004SJunchao Zhang } 1467cd620004SJunchao Zhang } 1468cd620004SJunchao Zhang PetscFunctionReturn(0); 1469cd620004SJunchao Zhang } 1470cd620004SJunchao Zhang 1471cd620004SJunchao Zhang /* Reduce leafdata to rootdata locally */ 1472cd620004SJunchao Zhang PetscErrorCode PetscSFLinkReduceLocal(PetscSF sf,PetscSFLink link,const void *leafdata,void *rootdata,MPI_Op op) 1473cd620004SJunchao Zhang { 1474cd620004SJunchao Zhang PetscErrorCode ierr; 1475cd620004SJunchao Zhang const PetscInt *rootindices = NULL,*leafindices = NULL; 1476cd620004SJunchao Zhang PetscInt count,rootstart,leafstart; 1477cd620004SJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 1478fcc7397dSJunchao Zhang PetscErrorCode (*ScatterAndOp)(PetscSFLink,PetscInt,PetscInt,PetscSFPackOpt,const PetscInt*,const void*,PetscInt,PetscSFPackOpt,const PetscInt*,void*) = NULL; 1479cd620004SJunchao Zhang const PetscMemType rootmtype = link->rootmtype,leafmtype = link->leafmtype; 1480fcc7397dSJunchao Zhang PetscSFPackOpt leafopt = NULL,rootopt = NULL; 1481cd620004SJunchao Zhang 1482cd620004SJunchao Zhang PetscFunctionBegin; 1483cd620004SJunchao Zhang if (!sf->leafbuflen[PETSCSF_LOCAL]) PetscFunctionReturn(0); 1484cd620004SJunchao Zhang if (rootmtype != leafmtype) { 1485cd620004SJunchao Zhang /* The local communication has to go through pack and unpack */ 1486cd620004SJunchao Zhang ierr = PetscSFLinkPackLeafData(sf,link,PETSCSF_LOCAL,leafdata);CHKERRQ(ierr); 1487f01131f0SJunchao Zhang ierr = PetscSFLinkMemcpy(sf,link,rootmtype,link->rootbuf[PETSCSF_LOCAL][rootmtype],leafmtype,link->leafbuf[PETSCSF_LOCAL][leafmtype],bas->rootbuflen[PETSCSF_LOCAL]*link->unitbytes);CHKERRQ(ierr); 1488cd620004SJunchao Zhang ierr = PetscSFLinkUnpackRootData(sf,link,PETSCSF_LOCAL,rootdata,op);CHKERRQ(ierr); 1489cd620004SJunchao Zhang } else { 1490cd620004SJunchao Zhang ierr = PetscSFLinkGetScatterAndOp(link,rootmtype,op,bas->rootdups[PETSCSF_LOCAL],&ScatterAndOp);CHKERRQ(ierr); 1491cd620004SJunchao Zhang if (ScatterAndOp) { 1492fcc7397dSJunchao Zhang ierr = PetscSFLinkGetRootPackOptAndIndices(sf,link,rootmtype,PETSCSF_LOCAL,&count,&rootstart,&rootopt,&rootindices);CHKERRQ(ierr); 1493fcc7397dSJunchao Zhang ierr = PetscSFLinkGetLeafPackOptAndIndices(sf,link,leafmtype,PETSCSF_LOCAL,&count,&leafstart,&leafopt,&leafindices);CHKERRQ(ierr); 1494fcc7397dSJunchao Zhang ierr = (*ScatterAndOp)(link,count,leafstart,leafopt,leafindices,leafdata,rootstart,rootopt,rootindices,rootdata);CHKERRQ(ierr); 1495cd620004SJunchao Zhang } else { 1496fcc7397dSJunchao Zhang ierr = PetscSFLinkGetRootPackOptAndIndices(sf,link,PETSC_MEMTYPE_HOST,PETSCSF_LOCAL,&count,&rootstart,&rootopt,&rootindices);CHKERRQ(ierr); 1497fcc7397dSJunchao Zhang ierr = PetscSFLinkGetLeafPackOptAndIndices(sf,link,PETSC_MEMTYPE_HOST,PETSCSF_LOCAL,&count,&leafstart,&leafopt,&leafindices);CHKERRQ(ierr); 1498fcc7397dSJunchao Zhang ierr = PetscSFLinkScatterDataWithMPIReduceLocal(sf,link,count,leafstart,leafindices,leafdata,rootstart,rootindices,rootdata,op);CHKERRQ(ierr); 1499cd620004SJunchao Zhang } 1500cd620004SJunchao Zhang } 1501cd620004SJunchao Zhang PetscFunctionReturn(0); 1502cd620004SJunchao Zhang } 1503cd620004SJunchao Zhang 1504cd620004SJunchao Zhang /* Fetch rootdata to leafdata and leafupdate locally */ 1505cd620004SJunchao Zhang PetscErrorCode PetscSFLinkFetchAndOpLocal(PetscSF sf,PetscSFLink link,void *rootdata,const void *leafdata,void *leafupdate,MPI_Op op) 1506cd620004SJunchao Zhang { 1507cd620004SJunchao Zhang PetscErrorCode ierr; 1508cd620004SJunchao Zhang const PetscInt *rootindices = NULL,*leafindices = NULL; 1509cd620004SJunchao Zhang PetscInt count,rootstart,leafstart; 1510cd620004SJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 1511fcc7397dSJunchao Zhang PetscErrorCode (*FetchAndOpLocal)(PetscSFLink,PetscInt,PetscInt,PetscSFPackOpt,const PetscInt*,void*,PetscInt,PetscSFPackOpt,const PetscInt*,const void*,void*) = NULL; 1512cd620004SJunchao Zhang const PetscMemType rootmtype = link->rootmtype,leafmtype = link->leafmtype; 1513fcc7397dSJunchao Zhang PetscSFPackOpt leafopt = NULL,rootopt = NULL; 1514cd620004SJunchao Zhang 1515cd620004SJunchao Zhang PetscFunctionBegin; 1516cd620004SJunchao Zhang if (!bas->rootbuflen[PETSCSF_LOCAL]) PetscFunctionReturn(0); 1517cd620004SJunchao Zhang if (rootmtype != leafmtype) { 1518cd620004SJunchao Zhang /* The local communication has to go through pack and unpack */ 1519cd620004SJunchao Zhang SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Doing PetscSFFetchAndOp with rootdata and leafdata on opposite side of CPU and GPU"); 1520cd620004SJunchao Zhang } else { 1521fcc7397dSJunchao Zhang ierr = PetscSFLinkGetRootPackOptAndIndices(sf,link,rootmtype,PETSCSF_LOCAL,&count,&rootstart,&rootopt,&rootindices);CHKERRQ(ierr); 1522fcc7397dSJunchao Zhang ierr = PetscSFLinkGetLeafPackOptAndIndices(sf,link,leafmtype,PETSCSF_LOCAL,&count,&leafstart,&leafopt,&leafindices);CHKERRQ(ierr); 1523cd620004SJunchao Zhang ierr = PetscSFLinkGetFetchAndOpLocal(link,rootmtype,op,bas->rootdups[PETSCSF_LOCAL],&FetchAndOpLocal);CHKERRQ(ierr); 1524fcc7397dSJunchao Zhang ierr = (*FetchAndOpLocal)(link,count,rootstart,rootopt,rootindices,rootdata,leafstart,leafopt,leafindices,leafdata,leafupdate);CHKERRQ(ierr); 1525cd620004SJunchao Zhang } 152640e23c03SJunchao Zhang PetscFunctionReturn(0); 152740e23c03SJunchao Zhang } 152840e23c03SJunchao Zhang 152940e23c03SJunchao Zhang /* 1530cd620004SJunchao Zhang Create per-rank pack/unpack optimizations based on indice patterns 153140e23c03SJunchao Zhang 153240e23c03SJunchao Zhang Input Parameters: 1533fcc7397dSJunchao Zhang + n - Number of destination ranks 1534eb02082bSJunchao Zhang . offset - [n+1] For the i-th rank, its associated indices are idx[offset[i], offset[i+1]). offset[0] needs not to be 0. 1535b23bfdefSJunchao Zhang - idx - [*] Array storing indices 153640e23c03SJunchao Zhang 153740e23c03SJunchao Zhang Output Parameters: 1538cd620004SJunchao Zhang + opt - Pack optimizations. NULL if no optimizations. 153940e23c03SJunchao Zhang */ 1540cd620004SJunchao Zhang PetscErrorCode PetscSFCreatePackOpt(PetscInt n,const PetscInt *offset,const PetscInt *idx,PetscSFPackOpt *out) 154140e23c03SJunchao Zhang { 154240e23c03SJunchao Zhang PetscErrorCode ierr; 1543fcc7397dSJunchao Zhang PetscInt r,p,start,i,j,k,dx,dy,dz,dydz,m,X,Y; 1544fcc7397dSJunchao Zhang PetscBool optimizable = PETSC_TRUE; 154540e23c03SJunchao Zhang PetscSFPackOpt opt; 154640e23c03SJunchao Zhang 154740e23c03SJunchao Zhang PetscFunctionBegin; 1548fcc7397dSJunchao Zhang ierr = PetscMalloc1(1,&opt);CHKERRQ(ierr); 1549fcc7397dSJunchao Zhang ierr = PetscMalloc1(7*n+2,&opt->array);CHKERRQ(ierr); 1550fcc7397dSJunchao Zhang opt->n = opt->array[0] = n; 1551fcc7397dSJunchao Zhang opt->offset = opt->array + 1; 1552fcc7397dSJunchao Zhang opt->start = opt->array + n + 2; 1553fcc7397dSJunchao Zhang opt->dx = opt->array + 2*n + 2; 1554fcc7397dSJunchao Zhang opt->dy = opt->array + 3*n + 2; 1555fcc7397dSJunchao Zhang opt->dz = opt->array + 4*n + 2; 1556fcc7397dSJunchao Zhang opt->X = opt->array + 5*n + 2; 1557fcc7397dSJunchao Zhang opt->Y = opt->array + 6*n + 2; 1558fcc7397dSJunchao Zhang 1559fcc7397dSJunchao Zhang for (r=0; r<n; r++) { /* For each destination rank */ 1560fcc7397dSJunchao Zhang m = offset[r+1] - offset[r]; /* Total number of indices for this rank. We want to see if m can be factored into dx*dy*dz */ 1561fcc7397dSJunchao Zhang p = offset[r]; 1562fcc7397dSJunchao Zhang start = idx[p]; /* First index for this rank */ 1563fcc7397dSJunchao Zhang p++; 1564fcc7397dSJunchao Zhang 1565fcc7397dSJunchao Zhang /* Search in X dimension */ 1566fcc7397dSJunchao Zhang for (dx=1; dx<m; dx++,p++) { 1567fcc7397dSJunchao Zhang if (start+dx != idx[p]) break; 1568b23bfdefSJunchao Zhang } 1569b23bfdefSJunchao Zhang 1570fcc7397dSJunchao Zhang dydz = m/dx; 1571fcc7397dSJunchao Zhang X = dydz > 1 ? (idx[p]-start) : dx; 1572fcc7397dSJunchao Zhang /* Not optimizable if m is not a multiple of dx, or some unrecognized pattern is found */ 1573fcc7397dSJunchao Zhang if (m%dx || X <= 0) {optimizable = PETSC_FALSE; goto finish;} 1574fcc7397dSJunchao Zhang for (dy=1; dy<dydz; dy++) { /* Search in Y dimension */ 1575fcc7397dSJunchao Zhang for (i=0; i<dx; i++,p++) { 1576fcc7397dSJunchao Zhang if (start+X*dy+i != idx[p]) { 1577fcc7397dSJunchao Zhang if (i) {optimizable = PETSC_FALSE; goto finish;} /* The pattern is violated in the middle of an x-walk */ 1578fcc7397dSJunchao Zhang else goto Z_dimension; 157940e23c03SJunchao Zhang } 158040e23c03SJunchao Zhang } 158140e23c03SJunchao Zhang } 158240e23c03SJunchao Zhang 1583fcc7397dSJunchao Zhang Z_dimension: 1584fcc7397dSJunchao Zhang dz = m/(dx*dy); 1585fcc7397dSJunchao Zhang Y = dz > 1 ? (idx[p]-start)/X : dy; 1586fcc7397dSJunchao Zhang /* Not optimizable if m is not a multiple of dx*dy, or some unrecognized pattern is found */ 1587fcc7397dSJunchao Zhang if (m%(dx*dy) || Y <= 0) {optimizable = PETSC_FALSE; goto finish;} 1588fcc7397dSJunchao Zhang for (k=1; k<dz; k++) { /* Go through Z dimension to see if remaining indices follow the pattern */ 1589fcc7397dSJunchao Zhang for (j=0; j<dy; j++) { 1590fcc7397dSJunchao Zhang for (i=0; i<dx; i++,p++) { 1591fcc7397dSJunchao Zhang if (start+X*Y*k+X*j+i != idx[p]) {optimizable = PETSC_FALSE; goto finish;} 159240e23c03SJunchao Zhang } 159340e23c03SJunchao Zhang } 159440e23c03SJunchao Zhang } 1595fcc7397dSJunchao Zhang opt->start[r] = start; 1596fcc7397dSJunchao Zhang opt->dx[r] = dx; 1597fcc7397dSJunchao Zhang opt->dy[r] = dy; 1598fcc7397dSJunchao Zhang opt->dz[r] = dz; 1599fcc7397dSJunchao Zhang opt->X[r] = X; 1600fcc7397dSJunchao Zhang opt->Y[r] = Y; 160140e23c03SJunchao Zhang } 160240e23c03SJunchao Zhang 1603fcc7397dSJunchao Zhang finish: 1604fcc7397dSJunchao Zhang /* If not optimizable, free arrays to save memory */ 1605fcc7397dSJunchao Zhang if (!n || !optimizable) { 1606fcc7397dSJunchao Zhang ierr = PetscFree(opt->array);CHKERRQ(ierr); 160740e23c03SJunchao Zhang ierr = PetscFree(opt);CHKERRQ(ierr); 160840e23c03SJunchao Zhang *out = NULL; 1609fcc7397dSJunchao Zhang } else { 1610fcc7397dSJunchao Zhang opt->offset[0] = 0; 1611fcc7397dSJunchao Zhang for (r=0; r<n; r++) opt->offset[r+1] = opt->offset[r] + opt->dx[r]*opt->dy[r]*opt->dz[r]; 1612fcc7397dSJunchao Zhang *out = opt; 1613fcc7397dSJunchao Zhang } 161440e23c03SJunchao Zhang PetscFunctionReturn(0); 161540e23c03SJunchao Zhang } 161640e23c03SJunchao Zhang 1617fcc7397dSJunchao Zhang PETSC_STATIC_INLINE PetscErrorCode PetscSFDestroyPackOpt(PetscMemType mtype,PetscSFPackOpt *out) 161840e23c03SJunchao Zhang { 161940e23c03SJunchao Zhang PetscErrorCode ierr; 162040e23c03SJunchao Zhang PetscSFPackOpt opt = *out; 162140e23c03SJunchao Zhang 162240e23c03SJunchao Zhang PetscFunctionBegin; 162340e23c03SJunchao Zhang if (opt) { 1624fcc7397dSJunchao Zhang if (mtype == PETSC_MEMTYPE_HOST) {ierr = PetscFree(opt->array);CHKERRQ(ierr);} 1625fcc7397dSJunchao Zhang #if defined(PETSC_HAVE_CUDA) 1626fcc7397dSJunchao Zhang else {cudaError_t cerr = cudaFree(opt->array);CHKERRCUDA(cerr);opt->array=NULL;} 1627fcc7397dSJunchao Zhang #endif 162840e23c03SJunchao Zhang ierr = PetscFree(opt);CHKERRQ(ierr); 162940e23c03SJunchao Zhang *out = NULL; 163040e23c03SJunchao Zhang } 163140e23c03SJunchao Zhang PetscFunctionReturn(0); 163240e23c03SJunchao Zhang } 1633cd620004SJunchao Zhang 1634cd620004SJunchao Zhang PetscErrorCode PetscSFSetUpPackFields(PetscSF sf) 1635cd620004SJunchao Zhang { 1636cd620004SJunchao Zhang PetscErrorCode ierr; 1637cd620004SJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 1638cd620004SJunchao Zhang PetscInt i,j; 1639cd620004SJunchao Zhang 1640cd620004SJunchao Zhang PetscFunctionBegin; 1641cd620004SJunchao Zhang /* [0] for PETSCSF_LOCAL and [1] for PETSCSF_REMOTE in the following */ 1642cd620004SJunchao Zhang for (i=0; i<2; i++) { /* Set defaults */ 1643cd620004SJunchao Zhang sf->leafstart[i] = 0; 1644cd620004SJunchao Zhang sf->leafcontig[i] = PETSC_TRUE; 1645cd620004SJunchao Zhang sf->leafdups[i] = PETSC_FALSE; 1646cd620004SJunchao Zhang bas->rootstart[i] = 0; 1647cd620004SJunchao Zhang bas->rootcontig[i] = PETSC_TRUE; 1648cd620004SJunchao Zhang bas->rootdups[i] = PETSC_FALSE; 1649cd620004SJunchao Zhang } 1650cd620004SJunchao Zhang 1651cd620004SJunchao Zhang sf->leafbuflen[0] = sf->roffset[sf->ndranks]; 1652cd620004SJunchao Zhang sf->leafbuflen[1] = sf->roffset[sf->nranks] - sf->roffset[sf->ndranks]; 1653cd620004SJunchao Zhang 1654cd620004SJunchao Zhang if (sf->leafbuflen[0]) sf->leafstart[0] = sf->rmine[0]; 1655cd620004SJunchao Zhang if (sf->leafbuflen[1]) sf->leafstart[1] = sf->rmine[sf->roffset[sf->ndranks]]; 1656cd620004SJunchao Zhang 1657cd620004SJunchao Zhang /* Are leaf indices for self and remote contiguous? If yes, it is best for pack/unpack */ 1658cd620004SJunchao Zhang for (i=0; i<sf->roffset[sf->ndranks]; i++) { /* self */ 1659cd620004SJunchao Zhang if (sf->rmine[i] != sf->leafstart[0]+i) {sf->leafcontig[0] = PETSC_FALSE; break;} 1660cd620004SJunchao Zhang } 1661cd620004SJunchao Zhang for (i=sf->roffset[sf->ndranks],j=0; i<sf->roffset[sf->nranks]; i++,j++) { /* remote */ 1662cd620004SJunchao Zhang if (sf->rmine[i] != sf->leafstart[1]+j) {sf->leafcontig[1] = PETSC_FALSE; break;} 1663cd620004SJunchao Zhang } 1664cd620004SJunchao Zhang 1665cd620004SJunchao Zhang /* If not, see if we can have per-rank optimizations by doing index analysis */ 1666cd620004SJunchao Zhang if (!sf->leafcontig[0]) {ierr = PetscSFCreatePackOpt(sf->ndranks, sf->roffset, sf->rmine, &sf->leafpackopt[0]);CHKERRQ(ierr);} 1667cd620004SJunchao Zhang if (!sf->leafcontig[1]) {ierr = PetscSFCreatePackOpt(sf->nranks-sf->ndranks, sf->roffset+sf->ndranks, sf->rmine, &sf->leafpackopt[1]);CHKERRQ(ierr);} 1668cd620004SJunchao Zhang 1669cd620004SJunchao Zhang /* Are root indices for self and remote contiguous? */ 1670cd620004SJunchao Zhang bas->rootbuflen[0] = bas->ioffset[bas->ndiranks]; 1671cd620004SJunchao Zhang bas->rootbuflen[1] = bas->ioffset[bas->niranks] - bas->ioffset[bas->ndiranks]; 1672cd620004SJunchao Zhang 1673cd620004SJunchao Zhang if (bas->rootbuflen[0]) bas->rootstart[0] = bas->irootloc[0]; 1674cd620004SJunchao Zhang if (bas->rootbuflen[1]) bas->rootstart[1] = bas->irootloc[bas->ioffset[bas->ndiranks]]; 1675cd620004SJunchao Zhang 1676cd620004SJunchao Zhang for (i=0; i<bas->ioffset[bas->ndiranks]; i++) { 1677cd620004SJunchao Zhang if (bas->irootloc[i] != bas->rootstart[0]+i) {bas->rootcontig[0] = PETSC_FALSE; break;} 1678cd620004SJunchao Zhang } 1679cd620004SJunchao Zhang for (i=bas->ioffset[bas->ndiranks],j=0; i<bas->ioffset[bas->niranks]; i++,j++) { 1680cd620004SJunchao Zhang if (bas->irootloc[i] != bas->rootstart[1]+j) {bas->rootcontig[1] = PETSC_FALSE; break;} 1681cd620004SJunchao Zhang } 1682cd620004SJunchao Zhang 1683cd620004SJunchao Zhang if (!bas->rootcontig[0]) {ierr = PetscSFCreatePackOpt(bas->ndiranks, bas->ioffset, bas->irootloc, &bas->rootpackopt[0]);CHKERRQ(ierr);} 1684cd620004SJunchao Zhang if (!bas->rootcontig[1]) {ierr = PetscSFCreatePackOpt(bas->niranks-bas->ndiranks, bas->ioffset+bas->ndiranks, bas->irootloc, &bas->rootpackopt[1]);CHKERRQ(ierr);} 1685cd620004SJunchao Zhang 1686cd620004SJunchao Zhang #if defined(PETSC_HAVE_CUDA) 1687cd620004SJunchao Zhang /* Check dups in indices so that CUDA unpacking kernels can use cheaper regular instructions instead of atomics when they know there are no data race chances */ 1688cd620004SJunchao Zhang if (!sf->leafcontig[0]) {ierr = PetscCheckDupsInt(sf->leafbuflen[0], sf->rmine, &sf->leafdups[0]);CHKERRQ(ierr);} 1689cd620004SJunchao Zhang if (!sf->leafcontig[1]) {ierr = PetscCheckDupsInt(sf->leafbuflen[1], sf->rmine+sf->roffset[sf->ndranks], &sf->leafdups[1]);CHKERRQ(ierr);} 1690cd620004SJunchao Zhang if (!bas->rootcontig[0]) {ierr = PetscCheckDupsInt(bas->rootbuflen[0], bas->irootloc, &bas->rootdups[0]);CHKERRQ(ierr);} 1691cd620004SJunchao Zhang if (!bas->rootcontig[1]) {ierr = PetscCheckDupsInt(bas->rootbuflen[1], bas->irootloc+bas->ioffset[bas->ndiranks], &bas->rootdups[1]);CHKERRQ(ierr);} 1692cd620004SJunchao Zhang #endif 1693cd620004SJunchao Zhang 1694cd620004SJunchao Zhang PetscFunctionReturn(0); 1695cd620004SJunchao Zhang } 1696cd620004SJunchao Zhang 1697cd620004SJunchao Zhang PetscErrorCode PetscSFResetPackFields(PetscSF sf) 1698cd620004SJunchao Zhang { 1699cd620004SJunchao Zhang PetscErrorCode ierr; 1700cd620004SJunchao Zhang PetscSF_Basic *bas = (PetscSF_Basic*)sf->data; 1701cd620004SJunchao Zhang PetscInt i; 1702cd620004SJunchao Zhang 1703cd620004SJunchao Zhang PetscFunctionBegin; 1704cd620004SJunchao Zhang for (i=PETSCSF_LOCAL; i<=PETSCSF_REMOTE; i++) { 1705fcc7397dSJunchao Zhang ierr = PetscSFDestroyPackOpt(PETSC_MEMTYPE_HOST,&sf->leafpackopt[i]);CHKERRQ(ierr); 1706fcc7397dSJunchao Zhang ierr = PetscSFDestroyPackOpt(PETSC_MEMTYPE_DEVICE,&sf->leafpackopt_d[i]);CHKERRQ(ierr); 1707fcc7397dSJunchao Zhang ierr = PetscSFDestroyPackOpt(PETSC_MEMTYPE_HOST,&bas->rootpackopt[i]);CHKERRQ(ierr); 1708fcc7397dSJunchao Zhang ierr = PetscSFDestroyPackOpt(PETSC_MEMTYPE_DEVICE,&bas->rootpackopt_d[i]);CHKERRQ(ierr); 1709cd620004SJunchao Zhang } 1710cd620004SJunchao Zhang PetscFunctionReturn(0); 1711cd620004SJunchao Zhang } 1712