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:   PetscCallHYPRE(HYPRE_IJVectorCreate(map->comm, map->rstart, map->rend - 1, &nij->ij));
     17:   PetscCallHYPRE(HYPRE_IJVectorSetObjectType(nij->ij, HYPRE_PARCSR));
     18: #if defined(PETSC_HAVE_HYPRE_DEVICE)
     19:   {
     20:     HYPRE_MemoryLocation memloc;
     21:     PetscCall(PetscHYPREInitialize());
     22:     PetscCallHYPRE(HYPRE_GetMemoryLocation(&memloc));
     23:     PetscCallHYPRE(HYPRE_IJVectorInitialize_v2(nij->ij, memloc));
     24:   }
     25: #else
     26:   PetscCallHYPRE(HYPRE_IJVectorInitialize(nij->ij));
     27: #endif
     28:   PetscCallHYPRE(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:   PetscCallHYPRE(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:     PetscCall(PetscHYPREInitialize());
     52:     PetscCallHYPRE(HYPRE_GetMemoryLocation(&memloc));
     53:     PetscCallHYPRE(HYPRE_IJVectorInitialize_v2(ij->ij, memloc));
     54:   }
     55: #else
     56:   PetscCallHYPRE(HYPRE_IJVectorInitialize(ij->ij));
     57: #endif
     58:   PetscCall(VecGetArrayRead(v, &array));
     59:   PetscCallHYPRE(HYPRE_IJVectorSetValues(ij->ij, (HYPRE_Int)v->map->n, NULL, (HYPRE_Complex *)array));
     60:   PetscCall(VecRestoreArrayRead(v, &array));
     61:   PetscCallHYPRE(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 || !PetscDefined(HAVE_HYPRE_DEVICE) ? 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 PETSC_PKG_HYPRE_VERSION_GT(2, 19, 0)
    212:   if (hmem != VecHYPRE_IJVectorMemoryLocation(ij->ij)) {
    213:     PetscCallHYPRE(HYPRE_IJVectorGetObject(ij->ij, (void **)&hij));
    214:     PetscCallHYPRE(hypre_ParVectorMigrate(hij, hmem));
    215:   }
    216: #endif
    217:   PetscFunctionReturn(PETSC_SUCCESS);
    218: }