Actual source code: vecreg.c

      1: #include <petsc/private/vecimpl.h>

      3: PetscFunctionList VecList = NULL;

      5: /* compare a vector type against a list of target vector types */
      6: static inline PetscErrorCode VecTypeCompareAny_Private(VecType srcType, PetscBool *match, const char tgtTypes[], ...)
      7: {
      8:   PetscBool flg = PETSC_FALSE;
      9:   va_list   Argp;

     11:   PetscFunctionBegin;
     12:   PetscAssertPointer(match, 2);
     13:   *match = PETSC_FALSE;
     14:   va_start(Argp, tgtTypes);
     15:   while (tgtTypes && tgtTypes[0]) {
     16:     PetscCall(PetscStrcmp(srcType, tgtTypes, &flg));
     17:     if (flg) {
     18:       *match = PETSC_TRUE;
     19:       break;
     20:     }
     21:     tgtTypes = va_arg(Argp, const char *);
     22:   }
     23:   va_end(Argp);
     24:   PetscFunctionReturn(PETSC_SUCCESS);
     25: }

     27: #define PETSC_MAX_VECTYPE_LEN 64

     29: /*@
     30:   VecSetType - Builds a vector, for a particular vector implementation.

     32:   Collective

     34:   Input Parameters:
     35: + vec     - The vector object
     36: - newType - The name of the vector type

     38:   Options Database Key:
     39: . -vec_type <type> - Sets the vector type; use -help for a list
     40:                      of available types

     42:   Level: intermediate

     44:   Notes:
     45:   See `VecType` for available vector types (for instance, `VECSEQ` or `VECMPI`)
     46:   Changing a vector to a new type will retain its old value if any.

     48:   Use `VecDuplicate()` or `VecDuplicateVecs()` to form additional vectors of the same type as an existing vector.

     50: .seealso: [](ch_vectors), `Vec`, `VecType`, `VecGetType()`, `VecCreate()`, `VecDuplicate()`, `VecDuplicateVecs()`
     51: @*/
     52: PetscErrorCode VecSetType(Vec vec, VecType newType)
     53: {
     54:   PetscErrorCode (*r)(Vec);
     55:   VecType      curType;
     56:   PetscBool    match;
     57:   PetscMPIInt  size;
     58:   PetscBool    dstSeq = PETSC_FALSE; // type info of the new type
     59:   MPI_Comm     comm;
     60:   char         seqType[PETSC_MAX_VECTYPE_LEN] = {0};
     61:   char         mpiType[PETSC_MAX_VECTYPE_LEN] = {0};
     62:   PetscScalar *oldValue;
     63:   PetscBool    srcStandard, dstStandard;

     65:   PetscFunctionBegin;

     68:   PetscCall(VecGetType(vec, &curType));
     69:   if (!curType) goto newvec; // vec's type is not set yet

     71:   /* return if exactly the same type */
     72:   PetscCall(PetscObjectTypeCompare((PetscObject)vec, newType, &match));
     73:   if (match) PetscFunctionReturn(PETSC_SUCCESS);

     75:   /* error on illegal mpi to seq conversion */
     76:   PetscCall(PetscObjectGetComm((PetscObject)vec, &comm));
     77:   PetscCallMPI(MPI_Comm_size(comm, &size));

     79:   PetscCall(PetscStrbeginswith(newType, VECSEQ, &dstSeq));
     80:   PetscCheck(!(size > 1 && dstSeq), comm, PETSC_ERR_ARG_WRONG, "Cannot convert MPI vectors to sequential ones");

     82:   /* return if standard => standard */
     83:   if (size == 1) PetscCall(PetscObjectTypeCompare((PetscObject)vec, VECSEQ, &srcStandard));
     84:   else PetscCall(PetscObjectTypeCompare((PetscObject)vec, VECMPI, &srcStandard));
     85:   PetscCall(VecTypeCompareAny_Private(newType, &dstStandard, VECSTANDARD, VECSEQ, VECMPI, ""));
     86:   if (srcStandard && dstStandard) PetscFunctionReturn(PETSC_SUCCESS);

     88:   /* return if curType = "seq" | "mpi" + newType */
     89:   PetscCall(PetscStrncpy(mpiType, "mpi", 4));
     90:   PetscCall(PetscStrlcat(mpiType, newType, PETSC_MAX_VECTYPE_LEN));
     91:   PetscCall(PetscStrncpy(seqType, "seq", 4));
     92:   PetscCall(PetscStrlcat(seqType, newType, PETSC_MAX_VECTYPE_LEN));
     93:   PetscCall(PetscObjectTypeCompareAny((PetscObject)vec, &match, seqType, mpiType, ""));
     94:   if (match) PetscFunctionReturn(PETSC_SUCCESS);

     96:   /* downcast VECSTANDARD to VECCUDA/HIP/KOKKOS in place. We don't do in-place upcasting
     97:   for those vectors. At least, it is not always possible to upcast a VECCUDA to VECSTANDARD
     98:   in place, since the host array might be pinned (i.e., allocated by cudaMallocHost()). If
     99:   we upcast it to VECSTANDARD, we could not free the pinned array with PetscFree(), which
    100:   is assumed for VECSTANDARD. Thus we just create a new vector, though it is expensive.
    101:   Upcasting is rare and users are not recommended to use it.
    102:   */
    103: #if defined(PETSC_HAVE_CUDA)
    104:   {
    105:     PetscBool dstCUDA = PETSC_FALSE;
    106:     if (!dstStandard) PetscCall(VecTypeCompareAny_Private(newType, &dstCUDA, VECCUDA, VECSEQCUDA, VECMPICUDA, ""));
    107:     if (srcStandard && dstCUDA) {
    108:       if (size == 1) PetscCall(VecConvert_Seq_SeqCUDA_inplace(vec));
    109:       else PetscCall(VecConvert_MPI_MPICUDA_inplace(vec));
    110:       PetscFunctionReturn(PETSC_SUCCESS);
    111:     }
    112:   }
    113: #endif
    114: #if defined(PETSC_HAVE_HIP)
    115:   {
    116:     PetscBool dstHIP = PETSC_FALSE;
    117:     if (!dstStandard) PetscCall(VecTypeCompareAny_Private(newType, &dstHIP, VECHIP, VECSEQHIP, VECMPIHIP, ""));
    118:     if (srcStandard && dstHIP) {
    119:       if (size == 1) PetscCall(VecConvert_Seq_SeqHIP_inplace(vec));
    120:       else PetscCall(VecConvert_MPI_MPIHIP_inplace(vec));
    121:       PetscFunctionReturn(PETSC_SUCCESS);
    122:     }
    123:   }
    124: #endif
    125: #if defined(PETSC_HAVE_KOKKOS_KERNELS)
    126:   {
    127:     PetscBool dstKokkos = PETSC_FALSE;
    128:     if (!dstStandard) PetscCall(VecTypeCompareAny_Private(newType, &dstKokkos, VECKOKKOS, VECSEQKOKKOS, VECMPIKOKKOS, ""));
    129:     if (srcStandard && dstKokkos) {
    130:       if (size == 1) PetscCall(VecConvert_Seq_SeqKokkos_inplace(vec));
    131:       else PetscCall(VecConvert_MPI_MPIKokkos_inplace(vec));
    132:       PetscFunctionReturn(PETSC_SUCCESS);
    133:     }
    134:   }
    135: #endif

    137:   /* Other conversion scenarios: create a new vector but retain old value */
    138: newvec:
    139:   PetscCall(PetscFunctionListFind(VecList, newType, &r));
    140:   PetscCheck(r, PetscObjectComm((PetscObject)vec), PETSC_ERR_ARG_UNKNOWN_TYPE, "Unknown vector type: %s", newType);
    141:   if (curType) { /* no need to destroy a vec without type */
    142:     const PetscScalar *array;
    143:     PetscCall(VecGetArrayRead(vec, &array));
    144:     if (array) {                                       /* record the old value if any before destroy */
    145:       PetscCall(PetscMalloc1(vec->map->n, &oldValue)); /* no need to free since we'll drop it into vec */
    146:       PetscCall(PetscArraycpy(oldValue, array, vec->map->n));
    147:     } else {
    148:       oldValue = NULL;
    149:     }
    150:     PetscCall(VecRestoreArrayRead(vec, &array));
    151:     PetscTryTypeMethod(vec, destroy);
    152:     PetscCall(PetscMemzero(vec->ops, sizeof(struct _VecOps)));
    153:     PetscCall(PetscFree(vec->defaultrandtype));
    154:     PetscCall(PetscFree(((PetscObject)vec)->type_name)); /* free type_name to make vec clean to use, as we might call VecSetType() again */
    155:   }

    157:   if (vec->map->n < 0 && vec->map->N < 0) {
    158:     vec->ops->create = r;
    159:     vec->ops->load   = VecLoad_Default;
    160:   } else {
    161:     PetscCall((*r)(vec));
    162:   }

    164:   /* drop in the old value */
    165:   if (curType && vec->map->n) PetscCall(VecReplaceArray(vec, oldValue));
    166:   PetscFunctionReturn(PETSC_SUCCESS);
    167: }

    169: /*@
    170:   VecGetType - Gets the vector type name (as a string) from a `Vec`.

    172:   Not Collective

    174:   Input Parameter:
    175: . vec - The vector

    177:   Output Parameter:
    178: . type - The `VecType` of the vector

    180:   Level: intermediate

    182: .seealso: [](ch_vectors), `Vec`, `VecType`, `VecCreate()`, `VecDuplicate()`, `VecDuplicateVecs()`
    183: @*/
    184: PetscErrorCode VecGetType(Vec vec, VecType *type)
    185: {
    186:   PetscFunctionBegin;
    188:   PetscAssertPointer(type, 2);
    189:   PetscCall(VecRegisterAll());
    190:   *type = ((PetscObject)vec)->type_name;
    191:   PetscFunctionReturn(PETSC_SUCCESS);
    192: }

    194: PetscErrorCode VecGetRootType_Private(Vec vec, VecType *vtype)
    195: {
    196:   PetscBool iscuda, iship, iskokkos, isvcl;

    198:   PetscFunctionBegin;
    200:   PetscAssertPointer(vtype, 2);
    201:   PetscCall(PetscObjectTypeCompareAny((PetscObject)vec, &iscuda, VECCUDA, VECMPICUDA, VECSEQCUDA, ""));
    202:   PetscCall(PetscObjectTypeCompareAny((PetscObject)vec, &iship, VECHIP, VECMPIHIP, VECSEQHIP, ""));
    203:   PetscCall(PetscObjectTypeCompareAny((PetscObject)vec, &iskokkos, VECKOKKOS, VECMPIKOKKOS, VECSEQKOKKOS, ""));
    204:   PetscCall(PetscObjectTypeCompareAny((PetscObject)vec, &isvcl, VECVIENNACL, VECMPIVIENNACL, VECSEQVIENNACL, ""));
    205:   if (iscuda) {
    206:     *vtype = VECCUDA;
    207:   } else if (iship) {
    208:     *vtype = VECHIP;
    209:   } else if (iskokkos) {
    210:     *vtype = VECKOKKOS;
    211:   } else if (isvcl) {
    212:     *vtype = VECVIENNACL;
    213:   } else {
    214:     *vtype = VECSTANDARD;
    215:   }
    216:   PetscFunctionReturn(PETSC_SUCCESS);
    217: }

    219: /*--------------------------------------------------------------------------------------------------------------------*/

    221: /*@C
    222:   VecRegister -  Adds a new vector component implementation

    224:   Not Collective, No Fortran Support

    226:   Input Parameters:
    227: + sname    - The name of a new user-defined creation routine
    228: - function - The creation routine

    230:   Notes:
    231:   `VecRegister()` may be called multiple times to add several user-defined vectors

    233:   Example Usage:
    234: .vb
    235:     VecRegister("my_vec",MyVectorCreate);
    236: .ve

    238:   Then, your vector type can be chosen with the procedural interface via
    239: .vb
    240:     VecCreate(MPI_Comm, Vec *);
    241:     VecSetType(Vec,"my_vector_name");
    242: .ve
    243:   or at runtime via the option
    244: .vb
    245:     -vec_type my_vector_name
    246: .ve

    248:   Level: advanced

    250: .seealso: `VecRegisterAll()`, `VecRegisterDestroy()`
    251: @*/
    252: PetscErrorCode VecRegister(const char sname[], PetscErrorCode (*function)(Vec))
    253: {
    254:   PetscFunctionBegin;
    255:   PetscCall(VecInitializePackage());
    256:   PetscCall(PetscFunctionListAdd(&VecList, sname, function));
    257:   PetscFunctionReturn(PETSC_SUCCESS);
    258: }