Actual source code: vhyp.c

  1: /*
  2:     Creates hypre ijvector from PETSc vector
  3: */

  5: #include <petsc/private/vecimpl.h>
  6: #include <../src/vec/vec/impls/hypre/vhyp.h>
  7: #include <HYPRE.h>

  9: PetscErrorCode VecHYPRE_IJVectorCreate(PetscLayout map, VecHYPRE_IJVector *ij)
 10: {
 11:   VecHYPRE_IJVector nij;

 13:   PetscFunctionBegin;
 14:   PetscCall(PetscNew(&nij));
 15:   PetscCall(PetscLayoutSetUp(map));
 16:   PetscCallExternal(HYPRE_IJVectorCreate, map->comm, map->rstart, map->rend - 1, &nij->ij);
 17:   PetscCallExternal(HYPRE_IJVectorSetObjectType, nij->ij, HYPRE_PARCSR);
 18: #if defined(PETSC_HAVE_HYPRE_DEVICE)
 19:   {
 20:     HYPRE_MemoryLocation memloc;
 21:     PetscHYPREInitialize();
 22:     PetscCallExternal(HYPRE_GetMemoryLocation, &memloc);
 23:     PetscCallExternal(HYPRE_IJVectorInitialize_v2, nij->ij, memloc);
 24:   }
 25: #else
 26:   PetscCallExternal(HYPRE_IJVectorInitialize, nij->ij);
 27: #endif
 28:   PetscCallExternal(HYPRE_IJVectorAssemble, nij->ij);
 29:   *ij = nij;
 30:   PetscFunctionReturn(PETSC_SUCCESS);
 31: }

 33: PetscErrorCode VecHYPRE_IJVectorDestroy(VecHYPRE_IJVector *ij)
 34: {
 35:   PetscFunctionBegin;
 36:   if (!*ij) PetscFunctionReturn(PETSC_SUCCESS);
 37:   PetscCheck(!(*ij)->pvec, PetscObjectComm((PetscObject)((*ij)->pvec)), PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
 38:   PetscCallExternal(HYPRE_IJVectorDestroy, (*ij)->ij);
 39:   PetscCall(PetscFree(*ij));
 40:   PetscFunctionReturn(PETSC_SUCCESS);
 41: }

 43: PetscErrorCode VecHYPRE_IJVectorCopy(Vec v, VecHYPRE_IJVector ij)
 44: {
 45:   const PetscScalar *array;

 47:   PetscFunctionBegin;
 48: #if defined(PETSC_HAVE_HYPRE_DEVICE)
 49:   {
 50:     HYPRE_MemoryLocation memloc;
 51:     PetscHYPREInitialize();
 52:     PetscCallExternal(HYPRE_GetMemoryLocation, &memloc);
 53:     PetscCallExternal(HYPRE_IJVectorInitialize_v2, ij->ij, memloc);
 54:   }
 55: #else
 56:   PetscCallExternal(HYPRE_IJVectorInitialize, ij->ij);
 57: #endif
 58:   PetscCall(VecGetArrayRead(v, &array));
 59:   PetscCallExternal(HYPRE_IJVectorSetValues, ij->ij, v->map->n, NULL, (HYPRE_Complex *)array);
 60:   PetscCall(VecRestoreArrayRead(v, &array));
 61:   PetscCallExternal(HYPRE_IJVectorAssemble, ij->ij);
 62:   PetscFunctionReturn(PETSC_SUCCESS);
 63: }

 65: /*
 66:     Replaces the address where the HYPRE vector points to its data with the address of
 67:   PETSc's data. Saves the old address so it can be reset when we are finished with it.
 68:   Allows use to get the data into a HYPRE vector without the cost of memcopies
 69: */
 70: #define VecHYPRE_ParVectorReplacePointer(b, newvalue, savedvalue) \
 71:   do { \
 72:     hypre_ParVector *par_vector   = (hypre_ParVector *)hypre_IJVectorObject((hypre_IJVector *)b); \
 73:     hypre_Vector    *local_vector = hypre_ParVectorLocalVector(par_vector); \
 74:     savedvalue                    = local_vector->data; \
 75:     local_vector->data            = newvalue; \
 76:   } while (0)

 78: /*
 79:   This routine access the pointer to the raw data of the "v" to be passed to HYPRE
 80:    - rw values indicate the type of access : 0 -> read, 1 -> write, 2 -> read-write
 81:    - hmem is the location HYPRE is expecting
 82:    - the function returns a pointer to the data (ptr) and the corresponding restore
 83:   Could be extended to VECKOKKOS if we had a way to access the raw pointer to device data.
 84: */
 85: static PetscErrorCode VecGetArrayForHYPRE(Vec v, int rw, HYPRE_MemoryLocation hmem, PetscScalar **ptr, PetscErrorCode (**res)(Vec, PetscScalar **))
 86: {
 87:   PetscMemType mtype;
 88:   MPI_Comm     comm;

 90:   PetscFunctionBegin;
 91: #if !defined(PETSC_HAVE_HYPRE_DEVICE)
 92:   hmem = HYPRE_MEMORY_HOST; /* this is just a convenience because HYPRE_MEMORY_HOST and HYPRE_MEMORY_DEVICE are the same in this case */
 93: #endif
 94:   *ptr = NULL;
 95:   *res = NULL;
 96:   PetscCall(PetscObjectGetComm((PetscObject)v, &comm));
 97:   switch (rw) {
 98:   case 0: /* read */
 99:     if (hmem == HYPRE_MEMORY_HOST) {
100:       PetscCall(VecGetArrayRead(v, (const PetscScalar **)ptr));
101:       *res = (PetscErrorCode (*)(Vec, PetscScalar **))VecRestoreArrayRead;
102:     } else {
103:       PetscCall(VecGetArrayReadAndMemType(v, (const PetscScalar **)ptr, &mtype));
104:       PetscCheck(PetscMemTypeDevice(mtype), comm, PETSC_ERR_ARG_WRONG, "HYPRE_MEMORY_DEVICE expects a device vector. You need to enable PETSc device support, for example, in some cases, -vec_type cuda");
105:       *res = (PetscErrorCode (*)(Vec, PetscScalar **))VecRestoreArrayReadAndMemType;
106:     }
107:     break;
108:   case 1: /* write */
109:     if (hmem == HYPRE_MEMORY_HOST) {
110:       PetscCall(VecGetArrayWrite(v, ptr));
111:       *res = VecRestoreArrayWrite;
112:     } else {
113:       PetscCall(VecGetArrayWriteAndMemType(v, (PetscScalar **)ptr, &mtype));
114:       PetscCheck(PetscMemTypeDevice(mtype), comm, PETSC_ERR_ARG_WRONG, "HYPRE_MEMORY_DEVICE expects a device vector. You need to enable PETSc device support, for example, in some cases, -vec_type cuda");
115:       *res = VecRestoreArrayWriteAndMemType;
116:     }
117:     break;
118:   case 2: /* read/write */
119:     if (hmem == HYPRE_MEMORY_HOST) {
120:       PetscCall(VecGetArray(v, ptr));
121:       *res = VecRestoreArray;
122:     } else {
123:       PetscCall(VecGetArrayAndMemType(v, (PetscScalar **)ptr, &mtype));
124:       PetscCheck(PetscMemTypeDevice(mtype), comm, PETSC_ERR_ARG_WRONG, "HYPRE_MEMORY_DEVICE expects a device vector. You need to enable PETSc device support, for example, in some cases, -vec_type cuda");
125:       *res = VecRestoreArrayAndMemType;
126:     }
127:     break;
128:   default:
129:     SETERRQ(comm, PETSC_ERR_SUP, "Unhandled case %d", rw);
130:   }
131:   PetscFunctionReturn(PETSC_SUCCESS);
132: }

134: #define VecHYPRE_IJVectorMemoryLocation(v) hypre_IJVectorMemoryLocation((hypre_IJVector *)(v))

136: /* Temporarily pushes the array of the data in v to ij (read access)
137:    depending on the value of the ij memory location
138:    Must be completed with a call to VecHYPRE_IJVectorPopVec */
139: PetscErrorCode VecHYPRE_IJVectorPushVecRead(VecHYPRE_IJVector ij, Vec v)
140: {
141:   HYPRE_Complex *pv;

143:   PetscFunctionBegin;
145:   PetscCheck(!ij->pvec, PetscObjectComm((PetscObject)v), PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
146:   PetscCheck(!ij->hv, PetscObjectComm((PetscObject)v), PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
147:   PetscCall(VecGetArrayForHYPRE(v, 0, VecHYPRE_IJVectorMemoryLocation(ij->ij), (PetscScalar **)&pv, &ij->restore));
148:   VecHYPRE_ParVectorReplacePointer(ij->ij, pv, ij->hv);
149:   ij->pvec = v;
150:   PetscFunctionReturn(PETSC_SUCCESS);
151: }

153: /* Temporarily pushes the array of the data in v to ij (write access)
154:    depending on the value of the ij memory location
155:    Must be completed with a call to VecHYPRE_IJVectorPopVec */
156: PetscErrorCode VecHYPRE_IJVectorPushVecWrite(VecHYPRE_IJVector ij, Vec v)
157: {
158:   HYPRE_Complex *pv;

160:   PetscFunctionBegin;
162:   PetscCheck(!ij->pvec, PetscObjectComm((PetscObject)v), PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
163:   PetscCheck(!ij->hv, PetscObjectComm((PetscObject)v), PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
164:   PetscCall(VecGetArrayForHYPRE(v, 1, VecHYPRE_IJVectorMemoryLocation(ij->ij), (PetscScalar **)&pv, &ij->restore));
165:   VecHYPRE_ParVectorReplacePointer(ij->ij, pv, ij->hv);
166:   ij->pvec = v;
167:   PetscFunctionReturn(PETSC_SUCCESS);
168: }

170: /* Temporarily pushes the array of the data in v to ij (read/write access)
171:    depending on the value of the ij memory location
172:    Must be completed with a call to VecHYPRE_IJVectorPopVec */
173: PetscErrorCode VecHYPRE_IJVectorPushVec(VecHYPRE_IJVector ij, Vec v)
174: {
175:   HYPRE_Complex *pv;

177:   PetscFunctionBegin;
179:   PetscCheck(!ij->pvec, PetscObjectComm((PetscObject)v), PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
180:   PetscCheck(!ij->hv, PetscObjectComm((PetscObject)v), PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
181:   PetscCall(VecGetArrayForHYPRE(v, 2, VecHYPRE_IJVectorMemoryLocation(ij->ij), (PetscScalar **)&pv, &ij->restore));
182:   VecHYPRE_ParVectorReplacePointer(ij->ij, pv, ij->hv);
183:   ij->pvec = v;
184:   PetscFunctionReturn(PETSC_SUCCESS);
185: }

187: /* Restores the pointer data to v */
188: PetscErrorCode VecHYPRE_IJVectorPopVec(VecHYPRE_IJVector ij)
189: {
190:   HYPRE_Complex *pv;

192:   PetscFunctionBegin;
193:   PetscCheck(ij->pvec, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPushVec()");
194:   PetscCheck(ij->restore, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPushVec()");
195:   VecHYPRE_ParVectorReplacePointer(ij->ij, ij->hv, pv);
196:   PetscCall((*ij->restore)(ij->pvec, (PetscScalar **)&pv));
197:   ij->hv      = NULL;
198:   ij->pvec    = NULL;
199:   ij->restore = NULL;
200:   PetscFunctionReturn(PETSC_SUCCESS);
201: }

203: PetscErrorCode VecHYPRE_IJBindToCPU(VecHYPRE_IJVector ij, PetscBool bind)
204: {
205:   HYPRE_MemoryLocation hmem = bind ? HYPRE_MEMORY_HOST : HYPRE_MEMORY_DEVICE;
206:   hypre_ParVector     *hij;

208:   PetscFunctionBegin;
209:   PetscCheck(!ij->pvec, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
210:   PetscCheck(!ij->hv, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
211: #if !defined(PETSC_HAVE_HYPRE_DEVICE)
212:   hmem = HYPRE_MEMORY_HOST;
213: #endif
214: #if PETSC_PKG_HYPRE_VERSION_GT(2, 19, 0)
215:   if (hmem != VecHYPRE_IJVectorMemoryLocation(ij->ij)) {
216:     PetscCallExternal(HYPRE_IJVectorGetObject, ij->ij, (void **)&hij);
217:     PetscCallExternal(hypre_ParVectorMigrate, hij, hmem);
218:   }
219: #endif
220:   PetscFunctionReturn(PETSC_SUCCESS);
221: }