Actual source code: veccupmimpl.h
1: #pragma once
3: #include <petsc/private/vecimpl.h>
4: #include <../src/vec/vec/impls/dvecimpl.h>
6: #if PetscDefined(HAVE_NVSHMEM)
7: PETSC_INTERN PetscErrorCode PetscNvshmemInitializeCheck(void);
8: PETSC_INTERN PetscErrorCode PetscNvshmemMalloc(size_t, void **);
9: PETSC_INTERN PetscErrorCode PetscNvshmemCalloc(size_t, void **);
10: PETSC_INTERN PetscErrorCode PetscNvshmemFree_Private(void *);
11: #define PetscNvshmemFree(ptr) ((PetscErrorCode)((ptr) && (PetscNvshmemFree_Private(ptr) || ((ptr) = PETSC_NULLPTR, PETSC_SUCCESS))))
12: PETSC_INTERN PetscErrorCode PetscNvshmemSum(PetscInt, PetscScalar *, const PetscScalar *);
13: PETSC_INTERN PetscErrorCode PetscNvshmemMax(PetscInt, PetscReal *, const PetscReal *);
14: #else
15: #define PetscNvshmemFree(ptr) PETSC_SUCCESS
16: #endif
18: #if defined(__cplusplus) && PetscDefined(HAVE_DEVICE)
19: #include <petsc/private/deviceimpl.h>
20: #include <petsc/private/cupmobject.hpp>
21: #include <petsc/private/cupmblasinterface.hpp>
23: #include <petsc/private/cpp/functional.hpp>
25: #include <limits> // std::numeric_limits
27: namespace Petsc
28: {
30: namespace vec
31: {
33: namespace cupm
34: {
36: namespace impl
37: {
39: namespace
40: {
42: struct no_op {
43: template <typename... T>
44: constexpr PetscErrorCode operator()(T &&...) const noexcept
45: {
46: return PETSC_SUCCESS;
47: }
48: };
50: template <typename T>
51: struct CooPair {
52: using value_type = T;
53: using size_type = PetscCount;
55: value_type *&device;
56: value_type *&host;
57: size_type size;
58: };
60: template <typename U>
61: static constexpr CooPair<U> make_coo_pair(U *&device, U *&host, PetscCount size) noexcept
62: {
63: return {device, host, size};
64: }
66: } // anonymous namespace
68: // forward declarations
69: template <device::cupm::DeviceType>
70: class PETSC_SINGLE_LIBRARY_VISIBILITY_INTERNAL VecSeq_CUPM;
71: template <device::cupm::DeviceType>
72: class PETSC_SINGLE_LIBRARY_VISIBILITY_INTERNAL VecMPI_CUPM;
74: // ==========================================================================================
75: // Vec_CUPMBase
76: //
77: // Base class for the VecSeq and VecMPI CUPM implementations. On top of the usual DeviceType
78: // template parameter it also uses CRTP to be able to use values/calls specific to either
79: // VecSeq or VecMPI. This is in effect "inside-out" polymorphism.
80: // ==========================================================================================
81: template <device::cupm::DeviceType T, typename Derived>
82: class PETSC_SINGLE_LIBRARY_VISIBILITY_INTERNAL Vec_CUPMBase : protected device::cupm::impl::CUPMObject<T> {
83: public:
84: PETSC_CUPMOBJECT_HEADER(T);
86: // ==========================================================================================
87: // Vec_CUPMBase::VectorArray
88: //
89: // RAII versions of the get/restore array routines. Determines constness of the pointer type,
90: // holds the pointer itself provides the implicit conversion operator
91: // ==========================================================================================
92: template <PetscMemType, PetscMemoryAccessMode>
93: class VectorArray;
95: protected:
96: static PetscErrorCode VecView_Debug(Vec v, const char *message = "") noexcept
97: {
98: const auto pobj = PetscObjectCast(v);
99: const auto vimpl = VecIMPLCast(v);
100: const auto vcu = VecCUPMCast(v);
101: PetscMemType mtype;
102: MPI_Comm comm;
104: PetscFunctionBegin;
105: PetscAssertPointer(vimpl, 1);
106: PetscAssertPointer(vcu, 1);
107: PetscCall(PetscObjectGetComm(pobj, &comm));
108: PetscCall(PetscPrintf(comm, "---------- %s ----------\n", message));
109: PetscCall(PetscObjectPrintClassNamePrefixType(pobj, PETSC_VIEWER_STDOUT_(comm)));
110: PetscCall(PetscPrintf(comm, "Address: %p\n", v));
111: PetscCall(PetscPrintf(comm, "Size: %" PetscInt_FMT "\n", v->map->n));
112: PetscCall(PetscPrintf(comm, "Offload mask: %s\n", PetscOffloadMaskToString(v->offloadmask)));
113: PetscCall(PetscPrintf(comm, "Host ptr: %p\n", vimpl->array));
114: PetscCall(PetscPrintf(comm, "Device ptr: %p\n", vcu->array_d));
115: PetscCall(PetscPrintf(comm, "Device alloced ptr: %p\n", vcu->array_allocated_d));
116: PetscCall(PetscCUPMGetMemType(vcu->array_d, &mtype));
117: PetscCall(PetscPrintf(comm, "dptr is device mem? %s\n", PetscBools[static_cast<PetscBool>(PetscMemTypeDevice(mtype))]));
118: PetscFunctionReturn(PETSC_SUCCESS);
119: }
121: // Delete the allocated device array if required and replace it with the given array
122: static PetscErrorCode ResetAllocatedDevicePtr_(PetscDeviceContext, Vec, PetscScalar * = nullptr) noexcept;
123: // Check either the host or device impl pointer is allocated and allocate it if
124: // isn't. CastFunctionType casts the Vec to the required type and returns the pointer
125: template <typename CastFunctionType>
126: static PetscErrorCode VecAllocateCheck_(Vec, void *&, CastFunctionType &&) noexcept;
127: // Check the CUPM part (v->spptr) is allocated, otherwise allocate it
128: static PetscErrorCode VecCUPMAllocateCheck_(Vec) noexcept;
129: // Check the Host part (v->data) is allocated, otherwise allocate it
130: static PetscErrorCode VecIMPLAllocateCheck_(Vec) noexcept;
131: // Check the Host array is allocated, otherwise allocate it
132: static PetscErrorCode HostAllocateCheck_(PetscDeviceContext, Vec) noexcept;
133: // Check the CUPM array is allocated, otherwise allocate it
134: static PetscErrorCode DeviceAllocateCheck_(PetscDeviceContext, Vec) noexcept;
135: // Copy HTOD, allocating device if necessary
136: static PetscErrorCode CopyToDevice_(PetscDeviceContext, Vec, bool = false) noexcept;
137: // Copy DTOH, allocating host if necessary
138: static PetscErrorCode CopyToHost_(PetscDeviceContext, Vec, bool = false) noexcept;
139: static PetscErrorCode DestroyDevice_(Vec) noexcept;
140: static PetscErrorCode DestroyHost_(Vec) noexcept;
142: public:
143: struct Vec_CUPM {
144: PetscScalar *array_d; // gpu data
145: PetscScalar *array_allocated_d; // does PETSc own the array ptr?
146: PetscBool nvshmem; // is array allocated in nvshmem? It is used to allocate
147: // Mvctx->lvec in nvshmem
149: // COO stuff
150: PetscCount *jmap1_d; // [m+1]: i-th entry of the vector has jmap1[i+1]-jmap1[i] repeats
151: // in COO arrays
152: PetscCount *perm1_d; // [tot1]: permutation array for local entries
153: PetscCount *imap2_d; // [nnz2]: i-th unique entry in recvbuf is imap2[i]-th entry in
154: // the vector
155: PetscCount *jmap2_d; // [nnz2+1]
156: PetscCount *perm2_d; // [recvlen]
157: PetscCount *Cperm_d; // [sendlen]: permutation array to fill sendbuf[]. 'C' for
158: // communication
160: // Buffers for remote values in VecSetValuesCOO()
161: PetscScalar *sendbuf_d;
162: PetscScalar *recvbuf_d;
163: };
165: // Cast the Vec to its Vec_CUPM struct, i.e. return the result of (Vec_CUPM *)v->spptr
166: PETSC_NODISCARD static Vec_CUPM *VecCUPMCast(Vec) noexcept;
167: // Cast the Vec to its host struct, i.e. return the result of (Vec_Seq *)v->data
168: template <typename U = Derived>
169: PETSC_NODISCARD static constexpr auto VecIMPLCast(Vec v) noexcept -> decltype(U::VecIMPLCast_(v));
170: // Get the PetscLogEvents for HTOD and DTOH
171: PETSC_NODISCARD static constexpr PetscLogEvent VEC_CUPMCopyToGPU() noexcept;
172: PETSC_NODISCARD static constexpr PetscLogEvent VEC_CUPMCopyFromGPU() noexcept;
173: // Get the VecTypes
174: PETSC_NODISCARD static constexpr VecType VECSEQCUPM() noexcept;
175: PETSC_NODISCARD static constexpr VecType VECMPICUPM() noexcept;
176: PETSC_NODISCARD static constexpr VecType VECCUPM() noexcept;
178: // Get the device VecType of the calling vector
179: template <typename U = Derived>
180: PETSC_NODISCARD static constexpr VecType VECIMPLCUPM() noexcept;
181: // Get the host VecType of the calling vector
182: template <typename U = Derived>
183: PETSC_NODISCARD static constexpr VecType VECIMPL() noexcept;
185: // Call the host destroy function, i.e. VecDestroy_Seq()
186: static PetscErrorCode VecDestroy_IMPL(Vec) noexcept;
187: // Call the host reset function, i.e. VecResetArray_Seq()
188: static PetscErrorCode VecResetArray_IMPL(Vec) noexcept;
189: // ... you get the idea
190: static PetscErrorCode VecPlaceArray_IMPL(Vec, const PetscScalar *) noexcept;
191: // Call the host creation function, i.e. VecCreate_Seq(), and also initialize the CUPM part
192: // along with it if needed
193: static PetscErrorCode VecCreate_IMPL_Private(Vec, PetscBool *, PetscInt = 0, PetscScalar * = nullptr) noexcept;
195: // Shorthand for creating VectorArray's. Need functions to create them, otherwise using them
196: // as an unnamed temporary leads to most vexing parse
197: PETSC_NODISCARD static auto DeviceArrayRead(PetscDeviceContext dctx, Vec v) noexcept PETSC_DECLTYPE_AUTO_RETURNS(VectorArray<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ>{dctx, v});
198: PETSC_NODISCARD static auto DeviceArrayWrite(PetscDeviceContext dctx, Vec v) noexcept PETSC_DECLTYPE_AUTO_RETURNS(VectorArray<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_WRITE>{dctx, v});
199: PETSC_NODISCARD static auto DeviceArrayReadWrite(PetscDeviceContext dctx, Vec v) noexcept PETSC_DECLTYPE_AUTO_RETURNS(VectorArray<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ_WRITE>{dctx, v});
200: PETSC_NODISCARD static auto HostArrayRead(PetscDeviceContext dctx, Vec v) noexcept PETSC_DECLTYPE_AUTO_RETURNS(VectorArray<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ>{dctx, v});
201: PETSC_NODISCARD static auto HostArrayWrite(PetscDeviceContext dctx, Vec v) noexcept PETSC_DECLTYPE_AUTO_RETURNS(VectorArray<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_WRITE>{dctx, v});
202: PETSC_NODISCARD static auto HostArrayReadWrite(PetscDeviceContext dctx, Vec v) noexcept PETSC_DECLTYPE_AUTO_RETURNS(VectorArray<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ_WRITE>{dctx, v});
204: // ops-table functions
205: static PetscErrorCode Create(Vec) noexcept;
206: static PetscErrorCode Destroy(Vec) noexcept;
207: template <PetscMemType, PetscMemoryAccessMode, bool = false>
208: static PetscErrorCode GetArray(Vec, PetscScalar **, PetscDeviceContext) noexcept;
209: template <PetscMemType, PetscMemoryAccessMode, bool = false>
210: static PetscErrorCode GetArray(Vec, PetscScalar **) noexcept;
211: template <PetscMemType, PetscMemoryAccessMode>
212: static PetscErrorCode RestoreArray(Vec, PetscScalar **, PetscDeviceContext) noexcept;
213: template <PetscMemType, PetscMemoryAccessMode>
214: static PetscErrorCode RestoreArray(Vec, PetscScalar **) noexcept;
215: template <PetscMemoryAccessMode>
216: static PetscErrorCode GetArrayAndMemtype(Vec, PetscScalar **, PetscMemType *, PetscDeviceContext) noexcept;
217: template <PetscMemoryAccessMode>
218: static PetscErrorCode GetArrayAndMemtype(Vec, PetscScalar **, PetscMemType *) noexcept;
219: template <PetscMemoryAccessMode>
220: static PetscErrorCode RestoreArrayAndMemtype(Vec, PetscScalar **, PetscDeviceContext) noexcept;
221: template <PetscMemoryAccessMode>
222: static PetscErrorCode RestoreArrayAndMemtype(Vec, PetscScalar **) noexcept;
223: template <PetscMemType>
224: static PetscErrorCode ReplaceArray(Vec, const PetscScalar *) noexcept;
225: template <PetscMemType>
226: static PetscErrorCode ResetArray(Vec) noexcept;
227: template <PetscMemType>
228: static PetscErrorCode PlaceArray(Vec, const PetscScalar *) noexcept;
230: // common ops shared between Seq and MPI
231: static PetscErrorCode Create_CUPM(Vec) noexcept;
232: static PetscErrorCode Create_CUPMBase(MPI_Comm, PetscInt, PetscInt, PetscInt, Vec *, PetscBool, PetscLayout /*reference*/ = nullptr) noexcept;
233: static PetscErrorCode Initialize_CUPMBase(Vec, PetscBool, PetscScalar *, PetscScalar *, PetscDeviceContext) noexcept;
234: template <typename SetupFunctionT = no_op>
235: static PetscErrorCode Duplicate_CUPMBase(Vec, Vec *, PetscDeviceContext, SetupFunctionT && = SetupFunctionT{}) noexcept;
236: static PetscErrorCode BindToCPU_CUPMBase(Vec, PetscBool, PetscDeviceContext) noexcept;
237: static PetscErrorCode GetArrays_CUPMBase(Vec, const PetscScalar **, const PetscScalar **, PetscOffloadMask *, PetscDeviceContext) noexcept;
238: static PetscErrorCode ResetPreallocationCOO_CUPMBase(Vec, PetscDeviceContext) noexcept;
239: template <std::size_t NCount = 0, std::size_t NScal = 0>
240: static PetscErrorCode SetPreallocationCOO_CUPMBase(Vec, PetscCount, const PetscInt[], PetscDeviceContext, const std::array<CooPair<PetscCount>, NCount> & = {}, const std::array<CooPair<PetscScalar>, NScal> & = {}) noexcept;
242: static PetscErrorCode Convert_IMPL_IMPLCUPM(Vec) noexcept;
243: };
245: // ==========================================================================================
246: // Vec_CUPMBase::VectorArray
247: //
248: // RAII versions of the get/restore array routines. Determines constness of the pointer type,
249: // holds the pointer itself and provides the implicit conversion operator.
250: //
251: // On construction this calls the moral equivalent of Vec[CUPM]GetArray[Read|Write]()
252: // (depending on PetscMemoryAccessMode) and on destruction automatically restores the array
253: // for you
254: // ==========================================================================================
255: template <device::cupm::DeviceType T, typename D>
256: template <PetscMemType MT, PetscMemoryAccessMode MA>
257: class Vec_CUPMBase<T, D>::VectorArray : public device::cupm::impl::RestoreableArray<T, MT, MA> {
258: using base_type = device::cupm::impl::RestoreableArray<T, MT, MA>;
260: public:
261: VectorArray(PetscDeviceContext, Vec) noexcept;
262: ~VectorArray() noexcept;
264: private:
265: Vec v_ = nullptr;
266: };
268: // ==========================================================================================
269: // Vec_CUPMBase::VectorArray - Public API
270: // ==========================================================================================
272: template <device::cupm::DeviceType T, typename D>
273: template <PetscMemType MT, PetscMemoryAccessMode MA>
274: inline Vec_CUPMBase<T, D>::VectorArray<MT, MA>::VectorArray(PetscDeviceContext dctx, Vec v) noexcept : base_type{dctx}, v_{v}
275: {
276: PetscFunctionBegin;
277: PetscCallAbort(PETSC_COMM_SELF, Vec_CUPMBase<T, D>::template GetArray<MT, MA, true>(v, &this->ptr_, dctx));
278: PetscFunctionReturnVoid();
279: }
281: template <device::cupm::DeviceType T, typename D>
282: template <PetscMemType MT, PetscMemoryAccessMode MA>
283: inline Vec_CUPMBase<T, D>::VectorArray<MT, MA>::~VectorArray() noexcept
284: {
285: PetscFunctionBegin;
286: PetscCallAbort(PETSC_COMM_SELF, Vec_CUPMBase<T, D>::template RestoreArray<MT, MA>(v_, &this->ptr_, this->dctx_));
287: PetscFunctionReturnVoid();
288: }
290: // ==========================================================================================
291: // Vec_CUPMBase - Protected API
292: // ==========================================================================================
294: template <device::cupm::DeviceType T, typename D>
295: inline PetscErrorCode Vec_CUPMBase<T, D>::ResetAllocatedDevicePtr_(PetscDeviceContext dctx, Vec v, PetscScalar *new_value) noexcept
296: {
297: auto &device_array = VecCUPMCast(v)->array_allocated_d;
299: PetscFunctionBegin;
300: if (device_array) {
301: if (PetscDefined(HAVE_NVSHMEM) && VecCUPMCast(v)->nvshmem) {
302: PetscCall(PetscNvshmemFree(device_array));
303: } else {
304: cupmStream_t stream;
306: PetscCall(GetHandlesFrom_(dctx, &stream));
307: PetscCallCUPM(cupmFreeAsync(device_array, stream));
308: }
309: }
310: device_array = new_value;
311: PetscFunctionReturn(PETSC_SUCCESS);
312: }
314: namespace
315: {
317: inline PetscErrorCode VecCUPMCheckMinimumPinnedMemory_Internal(Vec v, PetscBool *set = nullptr) noexcept
318: {
319: auto mem = static_cast<PetscInt>(v->minimum_bytes_pinned_memory);
320: PetscBool flg;
322: PetscFunctionBegin;
323: PetscObjectOptionsBegin(PetscObjectCast(v));
324: PetscCall(PetscOptionsRangeInt("-vec_pinned_memory_min", "Minimum size (in bytes) for an allocation to use pinned memory on host", "VecSetPinnedMemoryMin", mem, &mem, &flg, 0, std::numeric_limits<decltype(mem)>::max()));
325: if (flg) v->minimum_bytes_pinned_memory = mem;
326: PetscOptionsEnd();
327: if (set) *set = flg;
328: PetscFunctionReturn(PETSC_SUCCESS);
329: }
331: } // anonymous namespace
333: template <device::cupm::DeviceType T, typename D>
334: template <typename CastFunctionType>
335: inline PetscErrorCode Vec_CUPMBase<T, D>::VecAllocateCheck_(Vec v, void *&dest, CastFunctionType &&cast) noexcept
336: {
337: PetscFunctionBegin;
338: if (PetscLikely(dest)) PetscFunctionReturn(PETSC_SUCCESS);
339: // do the check here so we don't have to do it in every function
340: PetscCall(checkCupmBlasIntCast(v->map->n));
341: {
342: auto impl = cast(v);
344: PetscCall(PetscNew(&impl));
345: dest = impl;
346: }
347: PetscFunctionReturn(PETSC_SUCCESS);
348: }
350: template <device::cupm::DeviceType T, typename D>
351: inline PetscErrorCode Vec_CUPMBase<T, D>::VecIMPLAllocateCheck_(Vec v) noexcept
352: {
353: PetscFunctionBegin;
354: PetscCall(VecAllocateCheck_(v, v->data, VecIMPLCast<D>));
355: PetscFunctionReturn(PETSC_SUCCESS);
356: }
358: // allocate the Vec_CUPM struct. this is normally done through DeviceAllocateCheck_(), but in
359: // certain circumstances (such as when the user places the device array) we do not want to do
360: // the full DeviceAllocateCheck_() as it also allocates the array
361: template <device::cupm::DeviceType T, typename D>
362: inline PetscErrorCode Vec_CUPMBase<T, D>::VecCUPMAllocateCheck_(Vec v) noexcept
363: {
364: PetscFunctionBegin;
365: if (!v->spptr) {
366: /*
367: This is where v first gets a CUPM device array, which only the CUPM vector types use. If v is
368: any other type, such as VECSEQ, its operations ignore that array, so any result a device
369: routine writes there is silently lost. PetscCheckTypeNames() catches this only in debug
370: builds, so check the type here too. The check runs only once per vector, on the first
371: allocation.
372: */
373: PetscBool iscupm;
375: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(v), &iscupm, VECSEQCUPM(), VECMPICUPM(), ""));
376: PetscCheck(iscupm, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Vec type %s cannot hold a device array, it must be %s or %s. Obtain vectors for a device operation from MatCreateVecs(), or set their type with VecSetType()", PetscObjectCast(v)->type_name, VECSEQCUPM(), VECMPICUPM());
377: }
378: PetscCall(VecAllocateCheck_(v, v->spptr, VecCUPMCast));
379: PetscFunctionReturn(PETSC_SUCCESS);
380: }
382: template <device::cupm::DeviceType T, typename D>
383: inline PetscErrorCode Vec_CUPMBase<T, D>::HostAllocateCheck_(PetscDeviceContext, Vec v) noexcept
384: {
385: PetscFunctionBegin;
386: PetscCall(VecIMPLAllocateCheck_(v));
387: if (auto &alloc = VecIMPLCast(v)->array_allocated) PetscFunctionReturn(PETSC_SUCCESS);
388: else {
389: PetscCall(VecCUPMCheckMinimumPinnedMemory_Internal(v));
390: {
391: const auto n = v->map->n;
392: const auto useit = UseCUPMHostAlloc((n * sizeof(*alloc)) > v->minimum_bytes_pinned_memory);
394: v->pinned_memory = static_cast<decltype(v->pinned_memory)>(useit.value());
395: PetscCall(PetscMalloc1(n, &alloc));
396: }
397: if (!VecIMPLCast(v)->array) VecIMPLCast(v)->array = alloc;
398: if (v->offloadmask == PETSC_OFFLOAD_UNALLOCATED) v->offloadmask = PETSC_OFFLOAD_CPU;
399: }
400: PetscFunctionReturn(PETSC_SUCCESS);
401: }
403: template <device::cupm::DeviceType T, typename D>
404: inline PetscErrorCode Vec_CUPMBase<T, D>::DeviceAllocateCheck_(PetscDeviceContext dctx, Vec v) noexcept
405: {
406: PetscFunctionBegin;
407: PetscCall(VecCUPMAllocateCheck_(v));
408: if (auto &alloc = VecCUPMCast(v)->array_d) PetscFunctionReturn(PETSC_SUCCESS);
409: else {
410: const auto n = v->map->n;
411: auto &array_allocated_d = VecCUPMCast(v)->array_allocated_d;
412: cupmStream_t stream;
414: PetscCall(GetHandlesFrom_(dctx, &stream));
415: PetscCall(PetscCUPMMallocAsync(&array_allocated_d, n, stream));
416: alloc = array_allocated_d;
417: if (v->offloadmask == PETSC_OFFLOAD_UNALLOCATED) {
418: const auto vimp = VecIMPLCast(v);
419: v->offloadmask = (vimp && vimp->array) ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
420: }
421: }
422: PetscFunctionReturn(PETSC_SUCCESS);
423: }
425: template <device::cupm::DeviceType T, typename D>
426: inline PetscErrorCode Vec_CUPMBase<T, D>::CopyToDevice_(PetscDeviceContext dctx, Vec v, bool forceasync) noexcept
427: {
428: PetscFunctionBegin;
429: PetscCall(DeviceAllocateCheck_(dctx, v));
430: if (v->offloadmask == PETSC_OFFLOAD_CPU) {
431: cupmStream_t stream;
433: v->offloadmask = PETSC_OFFLOAD_BOTH;
434: PetscCall(GetHandlesFrom_(dctx, &stream));
435: PetscCall(PetscLogEventBegin(VEC_CUPMCopyToGPU(), v, 0, 0, 0));
436: PetscCall(PetscCUPMMemcpyAsync(VecCUPMCast(v)->array_d, VecIMPLCast(v)->array, v->map->n, cupmMemcpyHostToDevice, stream, forceasync));
437: PetscCall(PetscLogEventEnd(VEC_CUPMCopyToGPU(), v, 0, 0, 0));
438: }
439: PetscFunctionReturn(PETSC_SUCCESS);
440: }
442: template <device::cupm::DeviceType T, typename D>
443: inline PetscErrorCode Vec_CUPMBase<T, D>::CopyToHost_(PetscDeviceContext dctx, Vec v, bool forceasync) noexcept
444: {
445: PetscFunctionBegin;
446: PetscCall(HostAllocateCheck_(dctx, v));
447: if (v->offloadmask == PETSC_OFFLOAD_GPU) {
448: cupmStream_t stream;
450: v->offloadmask = PETSC_OFFLOAD_BOTH;
451: PetscCall(GetHandlesFrom_(dctx, &stream));
452: PetscCall(PetscLogEventBegin(VEC_CUPMCopyFromGPU(), v, 0, 0, 0));
453: PetscCall(PetscCUPMMemcpyAsync(VecIMPLCast(v)->array, VecCUPMCast(v)->array_d, v->map->n, cupmMemcpyDeviceToHost, stream, forceasync));
454: PetscCall(PetscLogEventEnd(VEC_CUPMCopyFromGPU(), v, 0, 0, 0));
455: }
456: PetscFunctionReturn(PETSC_SUCCESS);
457: }
459: template <device::cupm::DeviceType T, typename D>
460: inline PetscErrorCode Vec_CUPMBase<T, D>::DestroyDevice_(Vec v) noexcept
461: {
462: PetscFunctionBegin;
463: if (const auto vcu = VecCUPMCast(v)) {
464: PetscDeviceContext dctx;
466: PetscCall(GetHandles_(&dctx));
467: PetscCall(ResetAllocatedDevicePtr_(dctx, v));
468: PetscCall(ResetPreallocationCOO_CUPMBase(v, dctx));
469: PetscCall(PetscFree(v->spptr));
470: }
471: PetscFunctionReturn(PETSC_SUCCESS);
472: }
474: template <device::cupm::DeviceType T, typename D>
475: inline PetscErrorCode Vec_CUPMBase<T, D>::DestroyHost_(Vec v) noexcept
476: {
477: PetscFunctionBegin;
478: PetscCall(PetscObjectSAWsViewOff(PetscObjectCast(v)));
479: if (const auto vimpl = VecIMPLCast(v)) {
480: if (auto &array_allocated = vimpl->array_allocated) {
481: const auto useit = UseCUPMHostAlloc(v->pinned_memory);
483: // do this ourselves since we may want to use the cupm functions
484: PetscCall(PetscFree(array_allocated));
485: }
486: }
487: v->pinned_memory = PETSC_FALSE;
488: PetscCall(VecDestroy_IMPL(v));
489: PetscFunctionReturn(PETSC_SUCCESS);
490: }
492: // ==========================================================================================
493: // Vec_CUPMBase - Public API
494: // ==========================================================================================
496: template <device::cupm::DeviceType T, typename D>
497: inline typename Vec_CUPMBase<T, D>::Vec_CUPM *Vec_CUPMBase<T, D>::VecCUPMCast(Vec v) noexcept
498: {
499: return static_cast<Vec_CUPM *>(v->spptr);
500: }
502: // This is a trick to get around the fact that in CRTP the derived class is not yet fully
503: // defined because Base must necessarily be instantiated before Derived is
504: // complete. By using a dummy template parameter we make the type "dependent" and so will
505: // only be determined when the derived class is instantiated (and therefore fully defined)
506: template <device::cupm::DeviceType T, typename D>
507: template <typename U>
508: inline constexpr auto Vec_CUPMBase<T, D>::VecIMPLCast(Vec v) noexcept -> decltype(U::VecIMPLCast_(v))
509: {
510: return U::VecIMPLCast_(v);
511: }
513: template <device::cupm::DeviceType T, typename D>
514: inline PetscErrorCode Vec_CUPMBase<T, D>::VecDestroy_IMPL(Vec v) noexcept
515: {
516: return D::VecDestroy_IMPL_(v);
517: }
519: template <device::cupm::DeviceType T, typename D>
520: inline PetscErrorCode Vec_CUPMBase<T, D>::VecResetArray_IMPL(Vec v) noexcept
521: {
522: return D::VecResetArray_IMPL_(v);
523: }
525: template <device::cupm::DeviceType T, typename D>
526: inline PetscErrorCode Vec_CUPMBase<T, D>::VecPlaceArray_IMPL(Vec v, const PetscScalar *a) noexcept
527: {
528: return D::VecPlaceArray_IMPL_(v, a);
529: }
531: template <device::cupm::DeviceType T, typename D>
532: inline PetscErrorCode Vec_CUPMBase<T, D>::VecCreate_IMPL_Private(Vec v, PetscBool *alloc_missing, PetscInt nghost, PetscScalar *host_array) noexcept
533: {
534: return D::VecCreate_IMPL_Private_(v, alloc_missing, nghost, host_array);
535: }
537: template <device::cupm::DeviceType T, typename D>
538: inline constexpr PetscLogEvent Vec_CUPMBase<T, D>::VEC_CUPMCopyToGPU() noexcept
539: {
540: return T == device::cupm::DeviceType::CUDA ? VEC_CUDACopyToGPU : VEC_HIPCopyToGPU;
541: }
543: template <device::cupm::DeviceType T, typename D>
544: inline constexpr PetscLogEvent Vec_CUPMBase<T, D>::VEC_CUPMCopyFromGPU() noexcept
545: {
546: return T == device::cupm::DeviceType::CUDA ? VEC_CUDACopyFromGPU : VEC_HIPCopyFromGPU;
547: }
549: template <device::cupm::DeviceType T, typename D>
550: inline constexpr VecType Vec_CUPMBase<T, D>::VECSEQCUPM() noexcept
551: {
552: return T == device::cupm::DeviceType::CUDA ? VECSEQCUDA : VECSEQHIP;
553: }
555: template <device::cupm::DeviceType T, typename D>
556: inline constexpr VecType Vec_CUPMBase<T, D>::VECMPICUPM() noexcept
557: {
558: return T == device::cupm::DeviceType::CUDA ? VECMPICUDA : VECMPIHIP;
559: }
561: template <device::cupm::DeviceType T, typename D>
562: inline constexpr VecType Vec_CUPMBase<T, D>::VECCUPM() noexcept
563: {
564: return T == device::cupm::DeviceType::CUDA ? VECCUDA : VECHIP;
565: }
567: template <device::cupm::DeviceType T, typename D>
568: template <typename U>
569: inline constexpr VecType Vec_CUPMBase<T, D>::VECIMPLCUPM() noexcept
570: {
571: return U::VECIMPLCUPM_();
572: }
574: template <device::cupm::DeviceType T, typename D>
575: template <typename U>
576: inline constexpr VecType Vec_CUPMBase<T, D>::VECIMPL() noexcept
577: {
578: return U::VECIMPL_();
579: }
581: // private version that takes a PetscDeviceContext, called by the public variant
582: template <device::cupm::DeviceType T, typename D>
583: template <PetscMemType mtype, PetscMemoryAccessMode access, bool force>
584: inline PetscErrorCode Vec_CUPMBase<T, D>::GetArray(Vec v, PetscScalar **a, PetscDeviceContext dctx) noexcept
585: {
586: constexpr auto hostmem = PetscMemTypeHost(mtype);
587: const auto oldmask = v->offloadmask;
588: auto &mask = v->offloadmask;
589: auto should_sync = false;
591: PetscFunctionBegin;
592: static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), "");
593: PetscCheckTypeNames(v, VECSEQCUPM(), VECMPICUPM());
594: if (PetscMemoryAccessRead(access)) {
595: // READ or READ_WRITE
596: if ((oldmask == PETSC_OFFLOAD_GPU && hostmem) || (oldmask == PETSC_OFFLOAD_CPU && !hostmem)) {
597: // if we move the data we should set the flag to synchronize later on
598: should_sync = true;
599: }
600: PetscCall((hostmem ? CopyToHost_ : CopyToDevice_)(dctx, v, force));
601: } else {
602: // WRITE only
603: PetscCall((hostmem ? HostAllocateCheck_ : DeviceAllocateCheck_)(dctx, v));
604: }
605: *a = hostmem ? VecIMPLCast(v)->array : VecCUPMCast(v)->array_d;
606: // if unallocated previously we should zero things out if we intend to read
607: if (PetscMemoryAccessRead(access) && (oldmask == PETSC_OFFLOAD_UNALLOCATED)) {
608: const auto n = v->map->n;
610: if (hostmem) {
611: PetscCall(PetscArrayzero(*a, n));
612: } else {
613: cupmStream_t stream;
615: PetscCall(GetHandlesFrom_(dctx, &stream));
616: PetscCall(PetscCUPMMemsetAsync(*a, 0, n, stream, force));
617: should_sync = true;
618: }
619: }
620: // update the offloadmask if we intend to write, since we assume immediately modified
621: if (PetscMemoryAccessWrite(access)) {
622: PetscCall(VecSetErrorIfLocked(v, 1));
623: // REVIEW ME: this should probably also call PetscObjectStateIncrease() since we assume it
624: // is immediately modified
625: mask = hostmem ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
626: }
627: // if we are a globally blocking stream and we have MOVED data then we should synchronize,
628: // since even doing async calls on the NULL stream is not synchronous
629: if (!force && should_sync) PetscCall(PetscDeviceContextSynchronize(dctx));
630: PetscFunctionReturn(PETSC_SUCCESS);
631: }
633: // v->ops->getarray[read|write] or VecCUPMGetArray[Read|Write]()
634: template <device::cupm::DeviceType T, typename D>
635: template <PetscMemType mtype, PetscMemoryAccessMode access, bool force>
636: inline PetscErrorCode Vec_CUPMBase<T, D>::GetArray(Vec v, PetscScalar **a) noexcept
637: {
638: PetscDeviceContext dctx;
640: PetscFunctionBegin;
641: PetscCall(GetHandles_(&dctx));
642: PetscCall(D::template GetArray<mtype, access, force>(v, a, dctx));
643: PetscFunctionReturn(PETSC_SUCCESS);
644: }
646: // private version that takes a PetscDeviceContext, called by the public variant
647: template <device::cupm::DeviceType T, typename D>
648: template <PetscMemType mtype, PetscMemoryAccessMode access>
649: inline PetscErrorCode Vec_CUPMBase<T, D>::RestoreArray(Vec v, PetscScalar **a, PetscDeviceContext) noexcept
650: {
651: PetscFunctionBegin;
652: static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), "");
653: PetscCheckTypeNames(v, VECSEQCUPM(), VECMPICUPM());
654: if (PetscMemoryAccessWrite(access)) {
655: // WRITE or READ_WRITE
656: PetscCall(PetscObjectStateIncrease(PetscObjectCast(v)));
657: v->offloadmask = PetscMemTypeHost(mtype) ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
658: }
659: if (a) {
660: PetscCall(CheckPointerMatchesMemType_(*a, mtype));
661: *a = nullptr;
662: }
663: PetscFunctionReturn(PETSC_SUCCESS);
664: }
666: // v->ops->restorearray[read|write] or VecCUPMRestoreArray[Read|Write]()
667: template <device::cupm::DeviceType T, typename D>
668: template <PetscMemType mtype, PetscMemoryAccessMode access>
669: inline PetscErrorCode Vec_CUPMBase<T, D>::RestoreArray(Vec v, PetscScalar **a) noexcept
670: {
671: PetscDeviceContext dctx;
673: PetscFunctionBegin;
674: PetscCall(GetHandles_(&dctx));
675: PetscCall(D::template RestoreArray<mtype, access>(v, a, dctx));
676: PetscFunctionReturn(PETSC_SUCCESS);
677: }
679: template <device::cupm::DeviceType T, typename D>
680: template <PetscMemoryAccessMode access>
681: inline PetscErrorCode Vec_CUPMBase<T, D>::GetArrayAndMemtype(Vec v, PetscScalar **a, PetscMemType *mtype, PetscDeviceContext dctx) noexcept
682: {
683: PetscFunctionBegin;
684: if (a) PetscCall(D::template GetArray<PETSC_MEMTYPE_DEVICE, access>(v, a, dctx));
685: if (mtype) *mtype = (PetscDefined(HAVE_NVSHMEM) && VecCUPMCast(v)->nvshmem) ? PETSC_MEMTYPE_NVSHMEM : PETSC_MEMTYPE_CUPM();
686: PetscFunctionReturn(PETSC_SUCCESS);
687: }
689: // v->ops->getarrayandmemtype
690: template <device::cupm::DeviceType T, typename D>
691: template <PetscMemoryAccessMode access>
692: inline PetscErrorCode Vec_CUPMBase<T, D>::GetArrayAndMemtype(Vec v, PetscScalar **a, PetscMemType *mtype) noexcept
693: {
694: PetscDeviceContext dctx;
696: PetscFunctionBegin;
697: PetscCall(GetHandles_(&dctx));
698: PetscCall(D::template GetArrayAndMemtype<access>(v, a, mtype, dctx));
699: PetscFunctionReturn(PETSC_SUCCESS);
700: }
702: template <device::cupm::DeviceType T, typename D>
703: template <PetscMemoryAccessMode access>
704: inline PetscErrorCode Vec_CUPMBase<T, D>::RestoreArrayAndMemtype(Vec v, PetscScalar **a, PetscDeviceContext dctx) noexcept
705: {
706: PetscFunctionBegin;
707: PetscCall(D::template RestoreArray<PETSC_MEMTYPE_DEVICE, access>(v, a, dctx));
708: PetscFunctionReturn(PETSC_SUCCESS);
709: }
711: // v->ops->restorearrayandmemtype
712: template <device::cupm::DeviceType T, typename D>
713: template <PetscMemoryAccessMode access>
714: inline PetscErrorCode Vec_CUPMBase<T, D>::RestoreArrayAndMemtype(Vec v, PetscScalar **a) noexcept
715: {
716: PetscDeviceContext dctx;
718: PetscFunctionBegin;
719: PetscCall(GetHandles_(&dctx));
720: PetscCall(D::template RestoreArrayAndMemtype<access>(v, a, dctx));
721: PetscFunctionReturn(PETSC_SUCCESS);
722: }
724: // v->ops->placearray or VecCUPMPlaceArray()
725: template <device::cupm::DeviceType T, typename D>
726: template <PetscMemType mtype>
727: inline PetscErrorCode Vec_CUPMBase<T, D>::PlaceArray(Vec v, const PetscScalar *a) noexcept
728: {
729: PetscDeviceContext dctx;
731: PetscFunctionBegin;
732: static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), "");
733: PetscCheckTypeNames(v, VECSEQCUPM(), VECMPICUPM());
734: PetscCall(CheckPointerMatchesMemType_(a, mtype));
735: PetscCall(GetHandles_(&dctx));
736: if (PetscMemTypeHost(mtype)) {
737: PetscCall(CopyToHost_(dctx, v));
738: PetscCall(VecPlaceArray_IMPL(v, a));
739: v->offloadmask = PETSC_OFFLOAD_CPU;
740: } else {
741: PetscCall(VecIMPLAllocateCheck_(v));
742: {
743: auto &backup_array = VecIMPLCast(v)->unplacedarray;
745: PetscCheck(!backup_array, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "VecPlaceArray() was already called on this vector, without a call to VecResetArray()");
746: PetscCall(CopyToDevice_(dctx, v));
747: PetscCall(PetscObjectStateIncrease(PetscObjectCast(v)));
748: backup_array = util::exchange(VecCUPMCast(v)->array_d, const_cast<PetscScalar *>(a));
749: // only update the offload mask if we actually assign a pointer
750: if (a) v->offloadmask = PETSC_OFFLOAD_GPU;
751: }
752: }
753: PetscFunctionReturn(PETSC_SUCCESS);
754: }
756: // v->ops->replacearray or VecCUPMReplaceArray()
757: template <device::cupm::DeviceType T, typename D>
758: template <PetscMemType mtype>
759: inline PetscErrorCode Vec_CUPMBase<T, D>::ReplaceArray(Vec v, const PetscScalar *a) noexcept
760: {
761: const auto aptr = const_cast<PetscScalar *>(a);
762: PetscDeviceContext dctx;
764: PetscFunctionBegin;
765: static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), "");
766: PetscCheckTypeNames(v, VECSEQCUPM(), VECMPICUPM());
767: PetscCall(CheckPointerMatchesMemType_(a, mtype));
768: PetscCall(GetHandles_(&dctx));
769: if (PetscMemTypeHost(mtype)) {
770: PetscCall(VecIMPLAllocateCheck_(v));
771: {
772: const auto vimpl = VecIMPLCast(v);
773: auto &host_array = vimpl->array_allocated;
775: // Preserve caller-owned storage; PETSc-owned storage below is discarded.
776: if (host_array != vimpl->array) {
777: PetscCall(CopyToHost_(dctx, v));
778: PetscCall(PetscDeviceContextSynchronize(dctx));
779: }
780: if (host_array) {
781: const auto useit = UseCUPMHostAlloc(v->pinned_memory);
783: PetscCall(PetscFree(host_array));
784: }
785: host_array = aptr;
786: vimpl->array = host_array;
787: v->pinned_memory = PETSC_FALSE; // REVIEW ME: we can determine this
788: v->offloadmask = PETSC_OFFLOAD_CPU;
789: }
790: } else {
791: PetscCall(VecCUPMAllocateCheck_(v));
792: {
793: const auto vcu = VecCUPMCast(v);
795: PetscCall(ResetAllocatedDevicePtr_(dctx, v, aptr));
796: // don't update the offloadmask if placed pointer is NULL
797: vcu->array_d = vcu->array_allocated_d /* = aptr */;
798: if (aptr) v->offloadmask = PETSC_OFFLOAD_GPU;
799: }
800: }
801: PetscCall(PetscObjectStateIncrease(PetscObjectCast(v)));
802: PetscFunctionReturn(PETSC_SUCCESS);
803: }
805: // v->ops->resetarray or VecCUPMResetArray()
806: template <device::cupm::DeviceType T, typename D>
807: template <PetscMemType mtype>
808: inline PetscErrorCode Vec_CUPMBase<T, D>::ResetArray(Vec v) noexcept
809: {
810: PetscDeviceContext dctx;
812: PetscFunctionBegin;
813: static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), "");
814: PetscCheckTypeNames(v, VECSEQCUPM(), VECMPICUPM());
815: PetscCall(GetHandles_(&dctx));
816: // REVIEW ME:
817: // this is wildly inefficient but must be done if we assume that the placed array must have
818: // correct values
819: if (PetscMemTypeHost(mtype)) {
820: PetscCall(CopyToHost_(dctx, v));
821: PetscCall(VecResetArray_IMPL(v));
822: v->offloadmask = PETSC_OFFLOAD_CPU;
823: // A previous upload may still read the array even when no copy-back was needed.
824: PetscCall(PetscDeviceContextSynchronize(dctx));
825: } else {
826: PetscCall(VecIMPLAllocateCheck_(v));
827: PetscCall(VecCUPMAllocateCheck_(v));
828: {
829: const auto vcu = VecCUPMCast(v);
830: const auto vimpl = VecIMPLCast(v);
831: auto &host_array = vimpl->unplacedarray;
833: PetscCall(CheckPointerMatchesMemType_(host_array, PETSC_MEMTYPE_DEVICE));
834: if (v->offloadmask == PETSC_OFFLOAD_CPU) {
835: PetscCall(CopyToDevice_(dctx, v));
836: PetscCall(PetscDeviceContextSynchronize(dctx)); // Above H2D might be async, so we must sync dctx, otherwise if later user writes v's host array, it could ruin the H2D
837: }
838: PetscCall(PetscObjectStateIncrease(PetscObjectCast(v)));
839: // Need to reset the offloadmask. If we had a stashed pointer we are on the GPU,
840: // otherwise check if the host has a valid pointer. If neither, then we are not
841: // allocated.
842: vcu->array_d = host_array;
843: if (host_array) {
844: host_array = nullptr;
845: v->offloadmask = PETSC_OFFLOAD_GPU;
846: } else if (vimpl->array) {
847: v->offloadmask = PETSC_OFFLOAD_CPU;
848: } else {
849: v->offloadmask = PETSC_OFFLOAD_UNALLOCATED;
850: }
851: }
852: }
853: PetscFunctionReturn(PETSC_SUCCESS);
854: }
856: // v->ops->create
857: template <device::cupm::DeviceType T, typename D>
858: inline PetscErrorCode Vec_CUPMBase<T, D>::Create(Vec v) noexcept
859: {
860: PetscBool alloc_missing;
861: PetscDeviceContext dctx;
863: PetscFunctionBegin;
864: PetscCall(VecCreate_IMPL_Private(v, &alloc_missing));
865: PetscCall(GetHandles_(&dctx));
866: PetscCall(Initialize_CUPMBase(v, alloc_missing, nullptr, nullptr, dctx));
867: PetscFunctionReturn(PETSC_SUCCESS);
868: }
870: // v->ops->destroy
871: template <device::cupm::DeviceType T, typename D>
872: inline PetscErrorCode Vec_CUPMBase<T, D>::Destroy(Vec v) noexcept
873: {
874: PetscFunctionBegin;
875: PetscCall(DestroyDevice_(v));
876: PetscCall(DestroyHost_(v));
877: PetscFunctionReturn(PETSC_SUCCESS);
878: }
880: // ================================================================================== //
881: // Common core between Seq and MPI //
883: // VecCreate_CUPM()
884: template <device::cupm::DeviceType T, typename D>
885: inline PetscErrorCode Vec_CUPMBase<T, D>::Create_CUPM(Vec v) noexcept
886: {
887: PetscMPIInt size;
889: PetscFunctionBegin;
890: PetscCallMPI(MPI_Comm_size(PetscObjectComm(PetscObjectCast(v)), &size));
891: PetscCall(VecSetType(v, size > 1 ? VECMPICUPM() : VECSEQCUPM()));
892: PetscFunctionReturn(PETSC_SUCCESS);
893: }
895: // VecCreateCUPM()
896: template <device::cupm::DeviceType T, typename D>
897: inline PetscErrorCode Vec_CUPMBase<T, D>::Create_CUPMBase(MPI_Comm comm, PetscInt bs, PetscInt n, PetscInt N, Vec *v, PetscBool call_set_type, PetscLayout reference) noexcept
898: {
899: PetscFunctionBegin;
900: PetscCall(VecCreate(comm, v));
901: if (reference) PetscCall(PetscLayoutReference(reference, &(*v)->map));
902: PetscCall(VecSetSizes(*v, n, N));
903: if (bs) PetscCall(VecSetBlockSize(*v, bs));
904: if (call_set_type) PetscCall(VecSetType(*v, VECIMPLCUPM()));
905: PetscFunctionReturn(PETSC_SUCCESS);
906: }
908: // VecCreateIMPL_CUPM(), called through v->ops->create
909: template <device::cupm::DeviceType T, typename D>
910: inline PetscErrorCode Vec_CUPMBase<T, D>::Initialize_CUPMBase(Vec v, PetscBool allocate_missing, PetscScalar *host_array, PetscScalar *device_array, PetscDeviceContext dctx) noexcept
911: {
912: PetscFunctionBegin;
913: // REVIEW ME: perhaps not needed
914: PetscCall(PetscDeviceInitialize(PETSC_DEVICE_CUPM()));
915: PetscCall(PetscObjectChangeTypeName(PetscObjectCast(v), VECIMPLCUPM()));
916: PetscCall(D::BindToCPU(v, PETSC_FALSE));
917: if (device_array) {
918: PetscCall(CheckPointerMatchesMemType_(device_array, PETSC_MEMTYPE_CUPM()));
919: PetscCall(VecCUPMAllocateCheck_(v));
920: VecCUPMCast(v)->array_d = device_array;
921: }
922: if (host_array) {
923: PetscCall(CheckPointerMatchesMemType_(host_array, PETSC_MEMTYPE_HOST));
924: VecIMPLCast(v)->array = host_array;
925: }
926: if (allocate_missing) {
927: PetscCall(DeviceAllocateCheck_(dctx, v));
928: PetscCall(HostAllocateCheck_(dctx, v));
929: // REVIEW ME: junchao, is this needed with new calloc() branch? VecSet() will call
930: // set() for reference
931: // calls device-version
932: PetscCall(VecSet(v, 0));
933: // zero the host while device is underway
934: PetscCall(PetscArrayzero(VecIMPLCast(v)->array, v->map->n));
935: v->offloadmask = PETSC_OFFLOAD_BOTH;
936: } else {
937: if (host_array) {
938: v->offloadmask = device_array ? PETSC_OFFLOAD_BOTH : PETSC_OFFLOAD_CPU;
939: } else {
940: v->offloadmask = device_array ? PETSC_OFFLOAD_GPU : PETSC_OFFLOAD_UNALLOCATED;
941: }
942: }
943: PetscFunctionReturn(PETSC_SUCCESS);
944: }
946: // v->ops->duplicate
947: template <device::cupm::DeviceType T, typename D>
948: template <typename SetupFunctionT>
949: inline PetscErrorCode Vec_CUPMBase<T, D>::Duplicate_CUPMBase(Vec v, Vec *y, PetscDeviceContext dctx, SetupFunctionT &&DerivedCreateIMPLCUPM_Async) noexcept
950: {
951: // if the derived setup is the default no_op then we should call VecSetType()
952: constexpr auto call_set_type = static_cast<PetscBool>(std::is_same<SetupFunctionT, no_op>::value);
953: const auto vobj = PetscObjectCast(v);
954: const auto map = v->map;
955: PetscInt bs;
957: PetscFunctionBegin;
958: PetscCall(VecGetBlockSize(v, &bs));
959: PetscCall(Create_CUPMBase(PetscObjectComm(vobj), bs, map->n, map->N, y, call_set_type, map));
960: // Derived class can set up the remainder of the data structures here
961: PetscCall(DerivedCreateIMPLCUPM_Async(*y));
962: // If the other vector is bound to CPU then the memcpy of the ops struct will give the
963: // duplicated vector the host "getarray" function which does not lazily allocate the array
964: // (as it is assumed to always exist). So we force allocation here, before we overwrite the
965: // ops
966: if (v->boundtocpu) PetscCall(HostAllocateCheck_(dctx, *y));
967: // in case the user has done some VecSetOps() tomfoolery
968: (*y)->ops[0] = v->ops[0];
969: {
970: const auto yobj = PetscObjectCast(*y);
972: PetscCall(PetscObjectListDuplicate(vobj->olist, &yobj->olist));
973: PetscCall(PetscFunctionListDuplicate(vobj->qlist, &yobj->qlist));
974: }
975: (*y)->stash.donotstash = v->stash.donotstash;
976: (*y)->stash.ignorenegidx = v->stash.ignorenegidx;
977: (*y)->map->bs = std::abs(v->map->bs);
978: (*y)->bstash.bs = v->bstash.bs;
979: PetscFunctionReturn(PETSC_SUCCESS);
980: }
982: #define VecSetOp_CUPM(op_name, op_host, ...) \
983: do { \
984: if (usehost) { \
985: v->ops->op_name = op_host; \
986: } else { \
987: v->ops->op_name = __VA_ARGS__; \
988: } \
989: } while (0)
991: // v->ops->bindtocpu
992: template <device::cupm::DeviceType T, typename D>
993: inline PetscErrorCode Vec_CUPMBase<T, D>::BindToCPU_CUPMBase(Vec v, PetscBool usehost, PetscDeviceContext dctx) noexcept
994: {
995: PetscFunctionBegin;
996: v->boundtocpu = usehost;
997: if (usehost) PetscCall(CopyToHost_(dctx, v));
998: PetscCall(PetscStrFreeAllocpy(usehost ? PETSCRANDER48 : PETSCDEVICERAND(), &v->defaultrandtype));
1000: // set the base functions that are guaranteed to be the same for both
1001: v->ops->duplicate = D::Duplicate;
1002: v->ops->create = D::Create;
1003: v->ops->destroy = D::Destroy;
1004: v->ops->bindtocpu = D::BindToCPU;
1005: // Note that setting these to NULL on host breaks convergence in certain areas. I don't know
1006: // why, and I don't know how, but it is IMPERATIVE these are set as such!
1007: v->ops->replacearray = D::template ReplaceArray<PETSC_MEMTYPE_HOST>;
1008: v->ops->restorearray = D::template RestoreArray<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ_WRITE>;
1010: // set device-only common functions
1011: VecSetOp_CUPM(getarray, nullptr, D::template GetArray<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ_WRITE>);
1012: VecSetOp_CUPM(getarraywrite, nullptr, D::template GetArray<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_WRITE>);
1013: VecSetOp_CUPM(restorearraywrite, nullptr, D::template RestoreArray<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_WRITE>);
1015: VecSetOp_CUPM(getarrayread, nullptr, [](Vec v, const PetscScalar **a) { return D::template GetArray<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ>(v, const_cast<PetscScalar **>(a)); });
1016: VecSetOp_CUPM(restorearrayread, nullptr, [](Vec v, const PetscScalar **a) { return D::template RestoreArray<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ>(v, const_cast<PetscScalar **>(a)); });
1018: VecSetOp_CUPM(getarrayandmemtype, nullptr, D::template GetArrayAndMemtype<PETSC_MEMORY_ACCESS_READ_WRITE>);
1019: VecSetOp_CUPM(restorearrayandmemtype, nullptr, D::template RestoreArrayAndMemtype<PETSC_MEMORY_ACCESS_READ_WRITE>);
1021: VecSetOp_CUPM(getarraywriteandmemtype, nullptr, D::template GetArrayAndMemtype<PETSC_MEMORY_ACCESS_WRITE>);
1022: VecSetOp_CUPM(restorearraywriteandmemtype, nullptr, [](Vec v, PetscScalar **a, PetscMemType *) { return D::template RestoreArrayAndMemtype<PETSC_MEMORY_ACCESS_WRITE>(v, a); });
1024: VecSetOp_CUPM(getarrayreadandmemtype, nullptr, [](Vec v, const PetscScalar **a, PetscMemType *m) { return D::template GetArrayAndMemtype<PETSC_MEMORY_ACCESS_READ>(v, const_cast<PetscScalar **>(a), m); });
1025: VecSetOp_CUPM(restorearrayreadandmemtype, nullptr, [](Vec v, const PetscScalar **a) { return D::template RestoreArrayAndMemtype<PETSC_MEMORY_ACCESS_READ>(v, const_cast<PetscScalar **>(a)); });
1027: // set the functions that are always sequential
1028: using VecSeq_T = VecSeq_CUPM<T>;
1029: VecSetOp_CUPM(scale, VecScale_Seq, VecSeq_T::Scale);
1030: VecSetOp_CUPM(copy, VecCopy_Seq, VecSeq_T::Copy);
1031: VecSetOp_CUPM(set, VecSet_Seq, VecSeq_T::Set);
1032: VecSetOp_CUPM(swap, VecSwap_Seq, VecSeq_T::Swap);
1033: VecSetOp_CUPM(axpy, VecAXPY_Seq, VecSeq_T::AXPY);
1034: VecSetOp_CUPM(axpby, VecAXPBY_Seq, VecSeq_T::AXPBY);
1035: VecSetOp_CUPM(maxpy, VecMAXPY_Seq, VecSeq_T::MAXPY);
1036: VecSetOp_CUPM(aypx, VecAYPX_Seq, VecSeq_T::AYPX);
1037: VecSetOp_CUPM(waxpy, VecWAXPY_Seq, VecSeq_T::WAXPY);
1038: VecSetOp_CUPM(axpbypcz, VecAXPBYPCZ_Seq, VecSeq_T::AXPBYPCZ);
1039: VecSetOp_CUPM(pointwisemult, VecPointwiseMult_Seq, VecSeq_T::PointwiseMult);
1040: VecSetOp_CUPM(pointwisedivide, VecPointwiseDivide_Seq, VecSeq_T::PointwiseDivide);
1041: VecSetOp_CUPM(pointwisemax, VecPointwiseMax_Seq, VecSeq_T::PointwiseMax);
1042: VecSetOp_CUPM(pointwisemaxabs, VecPointwiseMaxAbs_Seq, VecSeq_T::PointwiseMaxAbs);
1043: VecSetOp_CUPM(pointwisemin, VecPointwiseMin_Seq, VecSeq_T::PointwiseMin);
1044: VecSetOp_CUPM(setrandom, VecSetRandom_Seq, VecSeq_T::SetRandom);
1045: VecSetOp_CUPM(dot_local, VecDot_Seq, VecSeq_T::Dot);
1046: VecSetOp_CUPM(tdot_local, VecTDot_Seq, VecSeq_T::TDot);
1047: VecSetOp_CUPM(norm_local, VecNorm_Seq, VecSeq_T::Norm);
1048: VecSetOp_CUPM(mdot_local, VecMDot_Seq, VecSeq_T::MDot);
1049: VecSetOp_CUPM(reciprocal, VecReciprocal_Default, VecSeq_T::Reciprocal);
1050: VecSetOp_CUPM(conjugate, VecConjugate_Seq, VecSeq_T::Conjugate);
1051: VecSetOp_CUPM(abs, nullptr, VecSeq_T::Abs);
1052: VecSetOp_CUPM(sqrt, nullptr, VecSeq_T::SqrtAbs);
1053: VecSetOp_CUPM(exp, nullptr, VecSeq_T::Exp);
1054: VecSetOp_CUPM(log, nullptr, VecSeq_T::Log);
1055: VecSetOp_CUPM(shift, nullptr, VecSeq_T::Shift);
1056: VecSetOp_CUPM(dotnorm2, nullptr, D::DotNorm2);
1057: VecSetOp_CUPM(getlocalvector, nullptr, VecSeq_T::template GetLocalVector<PETSC_MEMORY_ACCESS_READ_WRITE>);
1058: VecSetOp_CUPM(restorelocalvector, nullptr, VecSeq_T::template RestoreLocalVector<PETSC_MEMORY_ACCESS_READ_WRITE>);
1059: VecSetOp_CUPM(getlocalvectorread, nullptr, VecSeq_T::template GetLocalVector<PETSC_MEMORY_ACCESS_READ>);
1060: VecSetOp_CUPM(restorelocalvectorread, nullptr, VecSeq_T::template RestoreLocalVector<PETSC_MEMORY_ACCESS_READ>);
1061: VecSetOp_CUPM(sum, nullptr, VecSeq_T::Sum);
1062: VecSetOp_CUPM(errorwnorm, nullptr, D::ErrorWnorm);
1063: VecSetOp_CUPM(duplicatevecs, VecDuplicateVecs_Default, VecDuplicateVecs_Default);
1064: VecSetOp_CUPM(setstdbasis, nullptr, VecSeq_T::SetStdBasis);
1065: // Finish pending transfers before CPU operations can access the host array.
1066: if (usehost) PetscCall(PetscDeviceContextSynchronize(dctx));
1067: PetscFunctionReturn(PETSC_SUCCESS);
1068: }
1070: // Called from VecGetSubVector()
1071: template <device::cupm::DeviceType T, typename D>
1072: inline PetscErrorCode Vec_CUPMBase<T, D>::GetArrays_CUPMBase(Vec v, const PetscScalar **host_array, const PetscScalar **device_array, PetscOffloadMask *mask, PetscDeviceContext dctx) noexcept
1073: {
1074: PetscFunctionBegin;
1075: PetscCheckTypeNames(v, VECSEQCUPM(), VECMPICUPM());
1076: if (host_array) {
1077: PetscCall(HostAllocateCheck_(dctx, v));
1078: *host_array = VecIMPLCast(v)->array;
1079: }
1080: if (device_array) {
1081: PetscCall(DeviceAllocateCheck_(dctx, v));
1082: *device_array = VecCUPMCast(v)->array_d;
1083: }
1084: if (mask) *mask = v->offloadmask;
1085: PetscFunctionReturn(PETSC_SUCCESS);
1086: }
1088: template <device::cupm::DeviceType T, typename D>
1089: inline PetscErrorCode Vec_CUPMBase<T, D>::ResetPreallocationCOO_CUPMBase(Vec v, PetscDeviceContext dctx) noexcept
1090: {
1091: PetscFunctionBegin;
1092: if (const auto vcu = VecCUPMCast(v)) {
1093: cupmStream_t stream;
1094: // clang-format off
1095: const auto cntptrs = util::make_array(
1096: std::ref(vcu->jmap1_d),
1097: std::ref(vcu->perm1_d),
1098: std::ref(vcu->imap2_d),
1099: std::ref(vcu->jmap2_d),
1100: std::ref(vcu->perm2_d),
1101: std::ref(vcu->Cperm_d)
1102: );
1103: // clang-format on
1105: PetscCall(GetHandlesFrom_(dctx, &stream));
1106: for (auto &&ptr : cntptrs) PetscCallCUPM(cupmFreeAsync(ptr.get(), stream));
1107: for (auto &&ptr : util::make_array(std::ref(vcu->sendbuf_d), std::ref(vcu->recvbuf_d))) PetscCallCUPM(cupmFreeAsync(ptr.get(), stream));
1108: }
1109: PetscFunctionReturn(PETSC_SUCCESS);
1110: }
1112: template <device::cupm::DeviceType T, typename D>
1113: template <std::size_t NCount, std::size_t NScal>
1114: inline PetscErrorCode Vec_CUPMBase<T, D>::SetPreallocationCOO_CUPMBase(Vec v, PetscCount, const PetscInt[], PetscDeviceContext dctx, const std::array<CooPair<PetscCount>, NCount> &extra_cntptrs, const std::array<CooPair<PetscScalar>, NScal> &bufptrs) noexcept
1115: {
1116: PetscFunctionBegin;
1117: PetscCall(ResetPreallocationCOO_CUPMBase(v, dctx));
1118: // need to instantiate the private pointer if not already
1119: PetscCall(VecCUPMAllocateCheck_(v));
1120: {
1121: const auto vimpl = VecIMPLCast(v);
1122: const auto vcu = VecCUPMCast(v);
1123: // clang-format off
1124: const auto cntptrs = util::concat_array(
1125: util::make_array(
1126: make_coo_pair(vcu->jmap1_d, vimpl->jmap1, v->map->n + 1),
1127: make_coo_pair(vcu->perm1_d, vimpl->perm1, vimpl->tot1)
1128: ),
1129: extra_cntptrs
1130: );
1131: // clang-format on
1132: cupmStream_t stream;
1134: PetscCall(GetHandlesFrom_(dctx, &stream));
1135: // allocate
1136: for (auto &elem : cntptrs) PetscCall(PetscCUPMMallocAsync(&elem.device, elem.size, stream));
1137: for (auto &elem : bufptrs) PetscCall(PetscCUPMMallocAsync(&elem.device, elem.size, stream));
1138: // copy
1139: for (const auto &elem : cntptrs) PetscCall(PetscCUPMMemcpyAsync(elem.device, elem.host, elem.size, cupmMemcpyHostToDevice, stream, true));
1140: for (const auto &elem : bufptrs) PetscCall(PetscCUPMMemcpyAsync(elem.device, elem.host, elem.size, cupmMemcpyHostToDevice, stream, true));
1141: }
1142: PetscFunctionReturn(PETSC_SUCCESS);
1143: }
1145: template <device::cupm::DeviceType T, typename D>
1146: inline PetscErrorCode Vec_CUPMBase<T, D>::Convert_IMPL_IMPLCUPM(Vec v) noexcept
1147: {
1148: const auto n = v->map->n;
1149: const auto vimpl = VecIMPLCast(v);
1150: auto &impl_arr = vimpl->array;
1151: PetscBool set = PETSC_FALSE;
1152: PetscDeviceContext dctx;
1154: PetscFunctionBegin;
1155: // If users do not explicitly require pinned memory, we prefer keeping the vector's regular
1156: // host array
1157: PetscCall(VecCUPMCheckMinimumPinnedMemory_Internal(v, &set));
1158: if (set && impl_arr && ((n * sizeof(*impl_arr)) > v->minimum_bytes_pinned_memory)) {
1159: auto &impl_alloc = vimpl->array_allocated;
1160: PetscScalar *new_arr;
1162: // users require pinned memory
1163: {
1164: // Allocate pinned memory and copy over the old array
1165: const auto useit = UseCUPMHostAlloc(PETSC_TRUE);
1167: PetscCall(PetscMalloc1(n, &new_arr));
1168: PetscCall(PetscArraycpy(new_arr, impl_arr, n));
1169: }
1170: PetscCall(PetscFree(impl_alloc));
1171: impl_arr = new_arr;
1172: impl_alloc = new_arr;
1173: v->offloadmask = PETSC_OFFLOAD_CPU;
1174: v->pinned_memory = PETSC_TRUE;
1175: }
1176: PetscCall(GetHandles_(&dctx));
1177: PetscCall(Initialize_CUPMBase(v, PETSC_FALSE, impl_arr, nullptr, dctx));
1178: PetscFunctionReturn(PETSC_SUCCESS);
1179: }
1181: #define PETSC_VEC_CUPM_BASE_CLASS_HEADER(name, Tp, ...) \
1182: PETSC_CUPMOBJECT_HEADER(Tp); \
1183: using name = ::Petsc::vec::cupm::impl::Vec_CUPMBase<Tp, __VA_ARGS__>; \
1184: friend name; \
1185: /* introspection */ \
1186: using name::VecCUPMCast; \
1187: using name::VecIMPLCast; \
1188: using name::VECIMPLCUPM; \
1189: using name::VECIMPL; \
1190: using name::VECSEQCUPM; \
1191: using name::VECMPICUPM; \
1192: using name::VECCUPM; \
1193: using name::VecView_Debug; \
1194: /* utility */ \
1195: using typename name::Vec_CUPM; \
1196: using name::VecCUPMAllocateCheck_; \
1197: using name::VecIMPLAllocateCheck_; \
1198: using name::HostAllocateCheck_; \
1199: using name::DeviceAllocateCheck_; \
1200: using name::CopyToDevice_; \
1201: using name::CopyToHost_; \
1202: using name::Create; \
1203: using name::Destroy; \
1204: using name::GetArray; \
1205: using name::RestoreArray; \
1206: using name::GetArrayAndMemtype; \
1207: using name::RestoreArrayAndMemtype; \
1208: using name::PlaceArray; \
1209: using name::ReplaceArray; \
1210: using name::ResetArray; \
1211: /* base functions */ \
1212: using name::Create_CUPMBase; \
1213: using name::Initialize_CUPMBase; \
1214: using name::Duplicate_CUPMBase; \
1215: using name::BindToCPU_CUPMBase; \
1216: using name::Create_CUPM; \
1217: using name::DeviceArrayRead; \
1218: using name::DeviceArrayWrite; \
1219: using name::DeviceArrayReadWrite; \
1220: using name::HostArrayRead; \
1221: using name::HostArrayWrite; \
1222: using name::HostArrayReadWrite; \
1223: using name::ResetPreallocationCOO_CUPMBase; \
1224: using name::SetPreallocationCOO_CUPMBase; \
1225: using name::Convert_IMPL_IMPLCUPM;
1227: } // namespace impl
1229: } // namespace cupm
1231: } // namespace vec
1233: } // namespace Petsc
1235: #endif // __cplusplus && PetscDefined(HAVE_DEVICE)