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; see `VecType`

 41:   Level: intermediate

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

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

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

 64:   PetscFunctionBegin;

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

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

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

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

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

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

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

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

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

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

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

171:   Not Collective

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

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

179:   Level: intermediate

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

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

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

218: /*@C
219:   VecRegister -  Adds a new vector component implementation

221:   Not Collective, No Fortran Support

223:   Input Parameters:
224: + sname    - The name of a new user-defined creation routine
225: - function - The creation routine

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

230:   Example Usage:
231: .vb
232:     VecRegister("my_vec",MyVectorCreate);
233: .ve

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

245:   Level: advanced

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