1 #include <petsc/private/vecimpl_kokkos.hpp> 2 #include <petsc/private/dmdaimpl.h> 3 #include <petscdmda_kokkos.hpp> 4 5 /* Use macro instead of inlined function just to avoid annoying warnings like: 'dof' may be used uninitialized in this function [-Wmaybe-uninitialized] */ 6 #define DMDA_VEC_GET_SHAPE(da,vec,xs,ys,zs,xm,ym,zm,gxs,gys,gzs,gxm,gym,gzm,N,dim,dof) \ 7 do { \ 8 PetscErrorCode ierr; \ 9 ierr = DMDAGetCorners(da,&xs,&ys,&zs,&xm,&ym,&zm);CHKERRQ(ierr); \ 10 ierr = DMDAGetGhostCorners(da,&gxs,&gys,&gzs,&gxm,&gym,&gzm);CHKERRQ(ierr); \ 11 ierr = DMDAGetInfo(da,&dim,NULL,NULL,NULL,NULL,NULL,NULL,&dof,NULL,NULL,NULL,NULL,NULL);CHKERRQ(ierr); \ 12 /* Handle case where user passes in global vector as opposed to local */ \ 13 ierr = VecGetLocalSize(vec,&N);CHKERRQ(ierr); \ 14 if (N == xm*ym*zm*dof) { \ 15 gxm = xm; gym = ym; gzm = zm; \ 16 gxs = xs; gys = ys; gzs = zs; \ 17 } else if (N != gxm*gym*gzm*dof) SETERRQ3(PETSC_COMM_SELF,PETSC_ERR_ARG_INCOMP,"Vector local size %D is not compatible with DMDA local sizes %D %D\n",N,xm*ym*zm*dof,gxm*gym*gzm*dof); \ 18 if (dim != 1) SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_ARG_CORRUPT,"KokkosOffsetView is 1D but DMDA is %DD\n",dim); \ 19 } while (0) 20 21 template<class MemorySpace> 22 PetscErrorCode DMDAVecGetKokkosOffsetView_Private(DM da,Vec vec,PetscScalarKokkosOffsetView1DType<MemorySpace> *ov,PetscBool overwrite) 23 { 24 PetscErrorCode ierr; 25 PetscInt xs,ys,zs,xm,ym,zm,gxs,gys,gzs,gxm,gym,gzm,N,dim,dof; 26 PetscScalarKokkosViewType<MemorySpace> kv; 27 28 PetscFunctionBegin; 29 PetscValidHeaderSpecificType(da,DM_CLASSID,1,DMDA); 30 PetscValidHeaderSpecific(vec,VEC_CLASSID,2); 31 PetscValidPointer(ov,3); 32 DMDA_VEC_GET_SHAPE(da,vec,xs,ys,zs,xm,ym,zm,gxs,gys,gzs,gxm,gym,gzm,N,dim,dof); 33 if (overwrite) {ierr = VecGetKokkosViewWrite(vec,&kv);CHKERRQ(ierr);} 34 else {ierr = VecGetKokkosView(vec,&kv);CHKERRQ(ierr);} 35 *ov = PetscScalarKokkosOffsetView1DType<MemorySpace>(kv,{gxs*dof}); /* View to OffsetView by giving the start. The extent is already known. */ 36 PetscFunctionReturn(0); 37 } 38 39 template<class MemorySpace> 40 PetscErrorCode DMDAVecRestoreKokkosOffsetView_Private(DM da,Vec vec,PetscScalarKokkosOffsetView1DType<MemorySpace> *ov,PetscBool overwrite) 41 { 42 PetscErrorCode ierr; 43 PetscScalarKokkosViewType<MemorySpace> kv; 44 45 PetscFunctionBegin; 46 PetscValidHeaderSpecificType(da,DM_CLASSID,1,DMDA); 47 PetscValidHeaderSpecific(vec,VEC_CLASSID,2); 48 PetscValidPointer(ov,3); 49 kv = ov->view(); /* OffsetView to View */ 50 if (overwrite) {ierr = VecRestoreKokkosViewWrite(vec,&kv);CHKERRQ(ierr);} 51 else {ierr = VecRestoreKokkosView(vec,&kv);CHKERRQ(ierr);} 52 PetscFunctionReturn(0); 53 } 54 55 template<class MemorySpace> 56 PetscErrorCode DMDAVecGetKokkosOffsetView(DM da,Vec vec,ConstPetscScalarKokkosOffsetView1DType<MemorySpace> *ov) 57 { 58 PetscErrorCode ierr; 59 PetscInt xs,ys,zs,xm,ym,zm,gxs,gys,gzs,gxm,gym,gzm,N,dim,dof; 60 ConstPetscScalarKokkosViewType<MemorySpace> kv; 61 62 PetscFunctionBegin; 63 PetscValidHeaderSpecificType(da,DM_CLASSID,1,DMDA); 64 PetscValidHeaderSpecific(vec,VEC_CLASSID,2); 65 PetscValidPointer(ov,3); 66 DMDA_VEC_GET_SHAPE(da,vec,xs,ys,zs,xm,ym,zm,gxs,gys,gzs,gxm,gym,gzm,N,dim,dof); 67 ierr = VecGetKokkosView(vec,&kv);CHKERRQ(ierr); 68 *ov = ConstPetscScalarKokkosOffsetView1DType<MemorySpace>(kv,{gxs*dof}); /* View to OffsetView */ 69 PetscFunctionReturn(0); 70 } 71 72 template<class MemorySpace> 73 PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da,Vec vec,ConstPetscScalarKokkosOffsetView1DType<MemorySpace> *ov) 74 { 75 PetscErrorCode ierr; 76 ConstPetscScalarKokkosViewType<MemorySpace> kv; 77 78 PetscFunctionBegin; 79 PetscValidHeaderSpecificType(da,DM_CLASSID,1,DMDA); 80 PetscValidHeaderSpecific(vec,VEC_CLASSID,2); 81 PetscValidPointer(ov,3); 82 kv = ov->view(); 83 ierr = VecRestoreKokkosView(vec,&kv);CHKERRQ(ierr); 84 PetscFunctionReturn(0); 85 } 86 87 template<class MemorySpace> 88 PetscErrorCode DMDAVecGetKokkosOffsetView_Private(DM da,Vec vec,PetscScalarKokkosOffsetView2DType<MemorySpace> *ov,PetscBool overwrite) 89 { 90 PetscErrorCode ierr; 91 PetscInt xs,ys,zs,xm,ym,zm,gxs,gys,gzs,gxm,gym,gzm,N,dim,dof; 92 PetscScalarKokkosViewType<MemorySpace> kv; 93 94 PetscFunctionBegin; 95 PetscValidHeaderSpecificType(da,DM_CLASSID,1,DMDA); 96 PetscValidHeaderSpecific(vec,VEC_CLASSID,2); 97 PetscValidPointer(ov,3); 98 DMDA_VEC_GET_SHAPE(da,vec,xs,ys,zs,xm,ym,zm,gxs,gys,gzs,gxm,gym,gzm,N,dim,dof); 99 if (overwrite) {ierr = VecGetKokkosViewWrite(vec,&kv);CHKERRQ(ierr);} 100 else {ierr = VecGetKokkosView(vec,&kv);CHKERRQ(ierr);} 101 *ov = PetscScalarKokkosOffsetView2DType<MemorySpace>(kv.data(), {gxs*dof,(gxs+gxm)*dof}, {gys*dof,(gys+gym)*dof}); /* View to OffsetView */ 102 PetscFunctionReturn(0); 103 } 104 105 template<class MemorySpace> 106 PetscErrorCode DMDAVecRestoreKokkosOffsetView_Private(DM da,Vec vec,PetscScalarKokkosOffsetView2DType<MemorySpace> *ov,PetscBool overwrite) 107 { 108 PetscErrorCode ierr; 109 PetscScalarKokkosViewType<MemorySpace> kv; 110 111 PetscFunctionBegin; 112 PetscValidHeaderSpecificType(da,DM_CLASSID,1,DMDA); 113 PetscValidHeaderSpecific(vec,VEC_CLASSID,2); 114 PetscValidPointer(ov,3); 115 // kv = ov->view(); /* 2D OffsetView => 2D View => 1D View. Why does it not work? */ 116 kv = PetscScalarKokkosViewType<MemorySpace>(ov->data(),ov->extent(0)*ov->extent(1)); 117 if (overwrite) {ierr = VecRestoreKokkosViewWrite(vec,&kv);CHKERRQ(ierr);} 118 else {ierr = VecRestoreKokkosView(vec,&kv);CHKERRQ(ierr);} 119 PetscFunctionReturn(0); 120 } 121 122 template<class MemorySpace> 123 PetscErrorCode DMDAVecGetKokkosOffsetView(DM da,Vec vec,ConstPetscScalarKokkosOffsetView2DType<MemorySpace> *ov) 124 { 125 PetscErrorCode ierr; 126 PetscInt xs,ys,zs,xm,ym,zm,gxs,gys,gzs,gxm,gym,gzm,N,dim,dof; 127 ConstPetscScalarKokkosViewType<MemorySpace> kv; 128 129 PetscFunctionBegin; 130 PetscValidHeaderSpecificType(da,DM_CLASSID,1,DMDA); 131 PetscValidHeaderSpecific(vec,VEC_CLASSID,2); 132 PetscValidPointer(ov,3); 133 DMDA_VEC_GET_SHAPE(da,vec,xs,ys,zs,xm,ym,zm,gxs,gys,gzs,gxm,gym,gzm,N,dim,dof); 134 ierr = VecGetKokkosView(vec,&kv);CHKERRQ(ierr); 135 *ov = ConstPetscScalarKokkosOffsetView2DType<MemorySpace>(kv.data(), {gxs*dof,(gxs+gxm)*dof}, {gys*dof,(gys+gym)*dof}); /* View to OffsetView */ 136 PetscFunctionReturn(0); 137 } 138 139 template<class MemorySpace> 140 PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da,Vec vec,ConstPetscScalarKokkosOffsetView2DType<MemorySpace> *ov) 141 { 142 PetscErrorCode ierr; 143 ConstPetscScalarKokkosViewType<MemorySpace> kv; 144 145 PetscFunctionBegin; 146 PetscValidHeaderSpecificType(da,DM_CLASSID,1,DMDA); 147 PetscValidHeaderSpecific(vec,VEC_CLASSID,2); 148 PetscValidPointer(ov,3); 149 kv = ConstPetscScalarKokkosViewType<MemorySpace>(ov->data(),ov->extent(0)*ov->extent(1)); 150 ierr = VecRestoreKokkosView(vec,&kv);CHKERRQ(ierr); 151 PetscFunctionReturn(0); 152 } 153 154 /* Function template explicit instantiation */ 155 template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView (DM,Vec,ConstPetscScalarKokkosOffsetView1D*); 156 template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView (DM,Vec,ConstPetscScalarKokkosOffsetView1D*); 157 template<> PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView (DM da,Vec vec,PetscScalarKokkosOffsetView1D* ov) {return DMDAVecGetKokkosOffsetView_Private(da,vec,ov,PETSC_FALSE);} 158 template<> PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView (DM da,Vec vec,PetscScalarKokkosOffsetView1D* ov) {return DMDAVecRestoreKokkosOffsetView_Private(da,vec,ov,PETSC_FALSE);} 159 template<> PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewWrite (DM da,Vec vec,PetscScalarKokkosOffsetView1D* ov) {return DMDAVecGetKokkosOffsetView_Private(da,vec,ov,PETSC_TRUE);} 160 template<> PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewWrite(DM da,Vec vec,PetscScalarKokkosOffsetView1D* ov) {return DMDAVecRestoreKokkosOffsetView_Private(da,vec,ov,PETSC_TRUE);} 161 162 template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView (DM,Vec,ConstPetscScalarKokkosOffsetView2D*); 163 template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView (DM,Vec,ConstPetscScalarKokkosOffsetView2D*); 164 template<> PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView (DM da,Vec vec,PetscScalarKokkosOffsetView2D* ov) {return DMDAVecGetKokkosOffsetView_Private(da,vec,ov,PETSC_FALSE);} 165 template<> PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView (DM da,Vec vec,PetscScalarKokkosOffsetView2D* ov) {return DMDAVecRestoreKokkosOffsetView_Private(da,vec,ov,PETSC_FALSE);} 166 template<> PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewWrite (DM da,Vec vec,PetscScalarKokkosOffsetView2D* ov) {return DMDAVecGetKokkosOffsetView_Private(da,vec,ov,PETSC_TRUE);} 167 template<> PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewWrite(DM da,Vec vec,PetscScalarKokkosOffsetView2D* ov) {return DMDAVecRestoreKokkosOffsetView_Private(da,vec,ov,PETSC_TRUE);} 168 169