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:   PetscCallExternal(HYPRE_IJVectorInitialize_v2, nij->ij, HYPRE_MEMORY_DEVICE);
     20: #else
     21:   PetscCallExternal(HYPRE_IJVectorInitialize, nij->ij);
     22: #endif
     23:   PetscCallExternal(HYPRE_IJVectorAssemble, nij->ij);
     24:   *ij = nij;
     25:   PetscFunctionReturn(PETSC_SUCCESS);
     26: }

     28: PetscErrorCode VecHYPRE_IJVectorDestroy(VecHYPRE_IJVector *ij)
     29: {
     30:   PetscFunctionBegin;
     31:   if (!*ij) PetscFunctionReturn(PETSC_SUCCESS);
     32:   PetscCheck(!(*ij)->pvec, PetscObjectComm((PetscObject)((*ij)->pvec)), PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
     33:   PetscCallExternal(HYPRE_IJVectorDestroy, (*ij)->ij);
     34:   PetscCall(PetscFree(*ij));
     35:   PetscFunctionReturn(PETSC_SUCCESS);
     36: }

     38: PetscErrorCode VecHYPRE_IJVectorCopy(Vec v, VecHYPRE_IJVector ij)
     39: {
     40:   const PetscScalar *array;

     42:   PetscFunctionBegin;
     43: #if defined(PETSC_HAVE_HYPRE_DEVICE)
     44:   PetscCallExternal(HYPRE_IJVectorInitialize_v2, ij->ij, HYPRE_MEMORY_DEVICE);
     45: #else
     46:   PetscCallExternal(HYPRE_IJVectorInitialize, ij->ij);
     47: #endif
     48:   PetscCall(VecGetArrayRead(v, &array));
     49:   PetscCallExternal(HYPRE_IJVectorSetValues, ij->ij, v->map->n, NULL, (HYPRE_Complex *)array);
     50:   PetscCall(VecRestoreArrayRead(v, &array));
     51:   PetscCallExternal(HYPRE_IJVectorAssemble, ij->ij);
     52:   PetscFunctionReturn(PETSC_SUCCESS);
     53: }

     55: /*
     56:     Replaces the address where the HYPRE vector points to its data with the address of
     57:   PETSc's data. Saves the old address so it can be reset when we are finished with it.
     58:   Allows use to get the data into a HYPRE vector without the cost of memcopies
     59: */
     60: #define VecHYPRE_ParVectorReplacePointer(b, newvalue, savedvalue) \
     61:   do { \
     62:     hypre_ParVector *par_vector   = (hypre_ParVector *)hypre_IJVectorObject((hypre_IJVector *)b); \
     63:     hypre_Vector    *local_vector = hypre_ParVectorLocalVector(par_vector); \
     64:     savedvalue                    = local_vector->data; \
     65:     local_vector->data            = newvalue; \
     66:   } while (0)

     68: /*
     69:   This routine access the pointer to the raw data of the "v" to be passed to HYPRE
     70:    - rw values indicate the type of access : 0 -> read, 1 -> write, 2 -> read-write
     71:    - hmem is the location HYPRE is expecting
     72:    - the function returns a pointer to the data (ptr) and the corresponding restore
     73:   Could be extended to VECKOKKOS if we had a way to access the raw pointer to device data.
     74: */
     75: static PetscErrorCode VecGetArrayForHYPRE(Vec v, int rw, HYPRE_MemoryLocation hmem, PetscScalar **ptr, PetscErrorCode (**res)(Vec, PetscScalar **))
     76: {
     77:   PetscMemType mtype;
     78:   MPI_Comm     comm;

     80:   PetscFunctionBegin;
     81: #if !defined(PETSC_HAVE_HYPRE_DEVICE)
     82:   hmem = HYPRE_MEMORY_HOST; /* this is just a convenience because HYPRE_MEMORY_HOST and HYPRE_MEMORY_DEVICE are the same in this case */
     83: #endif
     84:   *ptr = NULL;
     85:   *res = NULL;
     86:   PetscCall(PetscObjectGetComm((PetscObject)v, &comm));
     87:   switch (rw) {
     88:   case 0: /* read */
     89:     if (hmem == HYPRE_MEMORY_HOST) {
     90:       PetscCall(VecGetArrayRead(v, (const PetscScalar **)ptr));
     91:       *res = (PetscErrorCode (*)(Vec, PetscScalar **))VecRestoreArrayRead;
     92:     } else {
     93:       PetscCall(VecGetArrayReadAndMemType(v, (const PetscScalar **)ptr, &mtype));
     94:       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");
     95:       *res = (PetscErrorCode (*)(Vec, PetscScalar **))VecRestoreArrayReadAndMemType;
     96:     }
     97:     break;
     98:   case 1: /* write */
     99:     if (hmem == HYPRE_MEMORY_HOST) {
    100:       PetscCall(VecGetArrayWrite(v, ptr));
    101:       *res = VecRestoreArrayWrite;
    102:     } else {
    103:       PetscCall(VecGetArrayWriteAndMemType(v, (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 = VecRestoreArrayWriteAndMemType;
    106:     }
    107:     break;
    108:   case 2: /* read/write */
    109:     if (hmem == HYPRE_MEMORY_HOST) {
    110:       PetscCall(VecGetArray(v, ptr));
    111:       *res = VecRestoreArray;
    112:     } else {
    113:       PetscCall(VecGetArrayAndMemType(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 = VecRestoreArrayAndMemType;
    116:     }
    117:     break;
    118:   default:
    119:     SETERRQ(comm, PETSC_ERR_SUP, "Unhandled case %d", rw);
    120:   }
    121:   PetscFunctionReturn(PETSC_SUCCESS);
    122: }

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

    126: /* Temporarily pushes the array of the data in v to ij (read access)
    127:    depending on the value of the ij memory location
    128:    Must be completed with a call to VecHYPRE_IJVectorPopVec */
    129: PetscErrorCode VecHYPRE_IJVectorPushVecRead(VecHYPRE_IJVector ij, Vec v)
    130: {
    131:   HYPRE_Complex *pv;

    133:   PetscFunctionBegin;
    135:   PetscCheck(!ij->pvec, PetscObjectComm((PetscObject)v), PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
    136:   PetscCheck(!ij->hv, PetscObjectComm((PetscObject)v), PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
    137:   PetscCall(VecGetArrayForHYPRE(v, 0, VecHYPRE_IJVectorMemoryLocation(ij->ij), (PetscScalar **)&pv, &ij->restore));
    138:   VecHYPRE_ParVectorReplacePointer(ij->ij, pv, ij->hv);
    139:   ij->pvec = v;
    140:   PetscFunctionReturn(PETSC_SUCCESS);
    141: }

    143: /* Temporarily pushes the array of the data in v to ij (write access)
    144:    depending on the value of the ij memory location
    145:    Must be completed with a call to VecHYPRE_IJVectorPopVec */
    146: PetscErrorCode VecHYPRE_IJVectorPushVecWrite(VecHYPRE_IJVector ij, Vec v)
    147: {
    148:   HYPRE_Complex *pv;

    150:   PetscFunctionBegin;
    152:   PetscCheck(!ij->pvec, PetscObjectComm((PetscObject)v), PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
    153:   PetscCheck(!ij->hv, PetscObjectComm((PetscObject)v), PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
    154:   PetscCall(VecGetArrayForHYPRE(v, 1, VecHYPRE_IJVectorMemoryLocation(ij->ij), (PetscScalar **)&pv, &ij->restore));
    155:   VecHYPRE_ParVectorReplacePointer(ij->ij, pv, ij->hv);
    156:   ij->pvec = v;
    157:   PetscFunctionReturn(PETSC_SUCCESS);
    158: }

    160: /* Temporarily pushes the array of the data in v to ij (read/write access)
    161:    depending on the value of the ij memory location
    162:    Must be completed with a call to VecHYPRE_IJVectorPopVec */
    163: PetscErrorCode VecHYPRE_IJVectorPushVec(VecHYPRE_IJVector ij, Vec v)
    164: {
    165:   HYPRE_Complex *pv;

    167:   PetscFunctionBegin;
    169:   PetscCheck(!ij->pvec, PetscObjectComm((PetscObject)v), PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
    170:   PetscCheck(!ij->hv, PetscObjectComm((PetscObject)v), PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
    171:   PetscCall(VecGetArrayForHYPRE(v, 2, VecHYPRE_IJVectorMemoryLocation(ij->ij), (PetscScalar **)&pv, &ij->restore));
    172:   VecHYPRE_ParVectorReplacePointer(ij->ij, pv, ij->hv);
    173:   ij->pvec = v;
    174:   PetscFunctionReturn(PETSC_SUCCESS);
    175: }

    177: /* Restores the pointer data to v */
    178: PetscErrorCode VecHYPRE_IJVectorPopVec(VecHYPRE_IJVector ij)
    179: {
    180:   HYPRE_Complex *pv;

    182:   PetscFunctionBegin;
    183:   PetscCheck(ij->pvec, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPushVec()");
    184:   PetscCheck(ij->restore, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPushVec()");
    185:   VecHYPRE_ParVectorReplacePointer(ij->ij, ij->hv, pv);
    186:   PetscCall((*ij->restore)(ij->pvec, (PetscScalar **)&pv));
    187:   ij->hv      = NULL;
    188:   ij->pvec    = NULL;
    189:   ij->restore = NULL;
    190:   PetscFunctionReturn(PETSC_SUCCESS);
    191: }

    193: PetscErrorCode VecHYPRE_IJBindToCPU(VecHYPRE_IJVector ij, PetscBool bind)
    194: {
    195:   HYPRE_MemoryLocation hmem = bind ? HYPRE_MEMORY_HOST : HYPRE_MEMORY_DEVICE;
    196:   hypre_ParVector     *hij;

    198:   PetscFunctionBegin;
    199:   PetscCheck(!ij->pvec, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
    200:   PetscCheck(!ij->hv, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Forgot to call VecHYPRE_IJVectorPopVec()");
    201: #if !defined(PETSC_HAVE_HYPRE_DEVICE)
    202:   hmem = HYPRE_MEMORY_HOST;
    203: #endif
    204: #if PETSC_PKG_HYPRE_VERSION_GT(2, 19, 0)
    205:   if (hmem != VecHYPRE_IJVectorMemoryLocation(ij->ij)) {
    206:     PetscCallExternal(HYPRE_IJVectorGetObject, ij->ij, (void **)&hij);
    207:     PetscCallExternal(hypre_ParVectorMigrate, hij, hmem);
    208:   }
    209: #endif
    210:   PetscFunctionReturn(PETSC_SUCCESS);
    211: }