Actual source code: vecseqcupm_impl.hpp
1: #pragma once
3: #include "vecseqcupm.hpp"
5: #include <petsc/private/randomimpl.h>
7: #include "../src/sys/objects/device/impls/cupm/cupmthrustutility.hpp"
8: #include "../src/sys/objects/device/impls/cupm/kernels.hpp"
10: #if PetscDefined(USE_COMPLEX)
11: #include <cstring>
12: #endif
14: #include <thrust/transform_reduce.h>
15: #include <thrust/transform.h>
16: #include <thrust/reduce.h>
17: #include <thrust/functional.h>
18: #include <thrust/tuple.h>
19: #include <thrust/device_ptr.h>
20: #include <thrust/iterator/zip_iterator.h>
21: #include <thrust/iterator/counting_iterator.h>
22: #include <thrust/iterator/constant_iterator.h>
23: #include <thrust/inner_product.h>
24: #if CCCL_VERSION >= 3004000
25: #include <cuda/iterator>
26: #endif
28: namespace Petsc
29: {
31: namespace vec
32: {
34: namespace cupm
35: {
37: namespace impl
38: {
40: // ==========================================================================================
41: // VecSeq_CUPM - Private API
42: // ==========================================================================================
44: template <device::cupm::DeviceType T>
45: inline Vec_Seq *VecSeq_CUPM<T>::VecIMPLCast_(Vec v) noexcept
46: {
47: return static_cast<Vec_Seq *>(v->data);
48: }
50: template <device::cupm::DeviceType T>
51: inline constexpr VecType VecSeq_CUPM<T>::VECIMPLCUPM_() noexcept
52: {
53: return VECSEQCUPM();
54: }
56: template <device::cupm::DeviceType T>
57: inline constexpr VecType VecSeq_CUPM<T>::VECIMPL_() noexcept
58: {
59: return VECSEQ;
60: }
62: template <device::cupm::DeviceType T>
63: inline PetscErrorCode VecSeq_CUPM<T>::ClearAsyncFunctions(Vec v) noexcept
64: {
65: PetscFunctionBegin;
66: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Abs), nullptr));
67: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(AXPBY), nullptr));
68: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(AXPBYPCZ), nullptr));
69: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(AXPY), nullptr));
70: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(AYPX), nullptr));
71: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Conjugate), nullptr));
72: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Copy), nullptr));
73: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Exp), nullptr));
74: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Log), nullptr));
75: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(MAXPY), nullptr));
76: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(PointwiseDivide), nullptr));
77: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(PointwiseMax), nullptr));
78: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(PointwiseMaxAbs), nullptr));
79: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(PointwiseMin), nullptr));
80: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(PointwiseMult), nullptr));
81: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(PointwiseSign), nullptr));
82: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Reciprocal), nullptr));
83: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Scale), nullptr));
84: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Set), nullptr));
85: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Shift), nullptr));
86: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(SqrtAbs), nullptr));
87: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Swap), nullptr));
88: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(WAXPY), nullptr));
89: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(SetStdBasis), nullptr));
90: PetscFunctionReturn(PETSC_SUCCESS);
91: }
93: template <device::cupm::DeviceType T>
94: inline PetscErrorCode VecSeq_CUPM<T>::InitializeAsyncFunctions(Vec v) noexcept
95: {
96: PetscFunctionBegin;
97: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Abs), VecSeq_CUPM<T>::AbsAsync));
98: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(AXPBY), VecSeq_CUPM<T>::AXPBYAsync));
99: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(AXPBYPCZ), VecSeq_CUPM<T>::AXPBYPCZAsync));
100: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(AXPY), VecSeq_CUPM<T>::AXPYAsync));
101: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(AYPX), VecSeq_CUPM<T>::AYPXAsync));
102: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Conjugate), VecSeq_CUPM<T>::ConjugateAsync));
103: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Copy), VecSeq_CUPM<T>::CopyAsync));
104: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Exp), VecSeq_CUPM<T>::ExpAsync));
105: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Log), VecSeq_CUPM<T>::LogAsync));
106: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(MAXPY), VecSeq_CUPM<T>::MAXPYAsync));
107: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(PointwiseDivide), VecSeq_CUPM<T>::PointwiseDivideAsync));
108: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(PointwiseMax), VecSeq_CUPM<T>::PointwiseMaxAsync));
109: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(PointwiseMaxAbs), VecSeq_CUPM<T>::PointwiseMaxAbsAsync));
110: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(PointwiseMin), VecSeq_CUPM<T>::PointwiseMinAsync));
111: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(PointwiseMult), VecSeq_CUPM<T>::PointwiseMultAsync));
112: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(PointwiseSign), VecSeq_CUPM<T>::PointwiseSignAsync));
113: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Reciprocal), VecSeq_CUPM<T>::ReciprocalAsync));
114: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Scale), VecSeq_CUPM<T>::ScaleAsync));
115: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Set), VecSeq_CUPM<T>::SetAsync));
116: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Shift), VecSeq_CUPM<T>::ShiftAsync));
117: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(SqrtAbs), VecSeq_CUPM<T>::SqrtAbsAsync));
118: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(Swap), VecSeq_CUPM<T>::SwapAsync));
119: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(WAXPY), VecSeq_CUPM<T>::WAXPYAsync));
120: PetscCall(PetscObjectComposeFunction(PetscObjectCast(v), VecAsyncFnName(SetStdBasis), VecSeq_CUPM<T>::SetStdBasisAsync));
121: PetscFunctionReturn(PETSC_SUCCESS);
122: }
124: template <device::cupm::DeviceType T>
125: inline PetscErrorCode VecSeq_CUPM<T>::VecDestroy_IMPL_(Vec v) noexcept
126: {
127: PetscFunctionBegin;
128: PetscCall(ClearAsyncFunctions(v));
129: PetscCall(VecDestroy_Seq(v));
130: PetscFunctionReturn(PETSC_SUCCESS);
131: }
133: template <device::cupm::DeviceType T>
134: inline PetscErrorCode VecSeq_CUPM<T>::VecResetArray_IMPL_(Vec v) noexcept
135: {
136: return VecResetArray_Seq(v);
137: }
139: template <device::cupm::DeviceType T>
140: inline PetscErrorCode VecSeq_CUPM<T>::VecPlaceArray_IMPL_(Vec v, const PetscScalar *a) noexcept
141: {
142: return VecPlaceArray_Seq(v, a);
143: }
145: template <device::cupm::DeviceType T>
146: inline PetscErrorCode VecSeq_CUPM<T>::VecCreate_IMPL_Private_(Vec v, PetscBool *alloc_missing, PetscInt, PetscScalar *host_array) noexcept
147: {
148: PetscMPIInt size;
150: PetscFunctionBegin;
151: if (alloc_missing) *alloc_missing = PETSC_FALSE;
152: PetscCallMPI(MPI_Comm_size(PetscObjectComm(PetscObjectCast(v)), &size));
153: PetscCheck(size <= 1, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Must create VecSeq on communicator of size 1, have size %d", size);
154: PetscCall(VecCreate_Seq_Private(v, host_array));
155: PetscCall(InitializeAsyncFunctions(v));
156: PetscFunctionReturn(PETSC_SUCCESS);
157: }
159: // for functions with an early return based one vec size we still need to artificially bump the
160: // object state. This is to prevent the following:
161: //
162: // 0. Suppose you have a Vec {
163: // rank 0: [0],
164: // rank 1: []
165: // }
166: // 1. both ranks have Vec with PetscObjectState = 0, stashed norm of 0
167: // 2. Vec enters e.g. VecSet(10)
168: // 3. rank 1 has local size 0 and bails immediately
169: // 4. rank 0 has local size 1 and enters function, eventually calls DeviceArrayWrite()
170: // 5. DeviceArrayWrite() calls PetscObjectStateIncrease(), now state = 1
171: // 6. Vec enters VecNorm(), and calls VecNormAvailable()
172: // 7. rank 1 has object state = 0, equal to stash and returns early with norm = 0
173: // 8. rank 0 has object state = 1, not equal to stash, continues to impl function
174: // 9. rank 0 deadlocks on MPI_Allreduce() because rank 1 bailed early
175: template <device::cupm::DeviceType T>
176: inline PetscErrorCode VecSeq_CUPM<T>::MaybeIncrementEmptyLocalVec(Vec v) noexcept
177: {
178: PetscFunctionBegin;
179: if (PetscUnlikely((v->map->n == 0) && (v->map->N != 0))) PetscCall(PetscObjectStateIncrease(PetscObjectCast(v)));
180: PetscFunctionReturn(PETSC_SUCCESS);
181: }
183: template <device::cupm::DeviceType T>
184: inline PetscErrorCode VecSeq_CUPM<T>::CreateSeqCUPM_(Vec v, PetscDeviceContext dctx, PetscScalar *host_array, PetscScalar *device_array) noexcept
185: {
186: PetscFunctionBegin;
187: PetscCall(base_type::VecCreate_IMPL_Private(v, nullptr, 0, host_array));
188: PetscCall(Initialize_CUPMBase(v, PETSC_FALSE, host_array, device_array, dctx));
189: PetscFunctionReturn(PETSC_SUCCESS);
190: }
192: template <device::cupm::DeviceType T>
193: template <typename BinaryFuncT>
194: inline PetscErrorCode VecSeq_CUPM<T>::PointwiseBinary_(BinaryFuncT &&binary, Vec xin, Vec yin, Vec zout, PetscDeviceContext dctx) noexcept
195: {
196: PetscFunctionBegin;
197: if (const auto n = zout->map->n) {
198: cupmStream_t stream;
200: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
201: PetscCall(GetHandlesFrom_(dctx, &stream));
202: // clang-format off
203: PetscCallThrust(
204: const auto dxptr = thrust::device_pointer_cast(DeviceArrayRead(dctx, xin).data());
206: THRUST_CALL(
207: thrust::transform,
208: stream,
209: dxptr, dxptr + n,
210: thrust::device_pointer_cast(DeviceArrayRead(dctx, yin).data()),
211: thrust::device_pointer_cast(DeviceArrayWrite(dctx, zout).data()),
212: std::forward<BinaryFuncT>(binary)
213: )
214: );
215: // clang-format on
216: PetscCall(PetscLogGpuFlops(n));
217: PetscCall(PetscDeviceContextSynchronizeIfWithBarrier_Internal(dctx));
218: } else {
219: PetscCall(MaybeIncrementEmptyLocalVec(zout));
220: }
221: PetscFunctionReturn(PETSC_SUCCESS);
222: }
224: template <device::cupm::DeviceType T>
225: template <typename BinaryFuncT>
226: inline PetscErrorCode VecSeq_CUPM<T>::PointwiseBinaryDispatch_(PetscErrorCode (*VecSeqFunction)(Vec, Vec, Vec), BinaryFuncT &&binary, Vec wout, Vec xin, Vec yin, PetscDeviceContext dctx) noexcept
227: {
228: PetscFunctionBegin;
229: if (xin->boundtocpu || yin->boundtocpu) PetscCall((*VecSeqFunction)(wout, xin, yin));
230: else PetscCall(PointwiseBinary_(std::forward<BinaryFuncT>(binary), xin, yin, wout, dctx)); // note order of arguments! xin and yin are read, wout is written!
231: PetscFunctionReturn(PETSC_SUCCESS);
232: }
234: template <device::cupm::DeviceType T>
235: template <typename UnaryFuncT>
236: inline PetscErrorCode VecSeq_CUPM<T>::PointwiseUnary_(UnaryFuncT &&unary, Vec xinout, Vec yout, PetscDeviceContext dctx) noexcept
237: {
238: const auto inplace = !yout || (xinout == yout);
240: PetscFunctionBegin;
241: if (const auto n = xinout->map->n) {
242: cupmStream_t stream;
243: const auto apply = [&](PetscScalar *xinout, PetscScalar *yout = nullptr) {
244: PetscFunctionBegin;
245: // clang-format off
246: PetscCallThrust(
247: const auto xptr = thrust::device_pointer_cast(xinout);
249: THRUST_CALL(
250: thrust::transform,
251: stream,
252: xptr, xptr + n,
253: (yout && (yout != xinout)) ? thrust::device_pointer_cast(yout) : xptr,
254: std::forward<UnaryFuncT>(unary)
255: )
256: );
257: // clang-format on
258: PetscFunctionReturn(PETSC_SUCCESS);
259: };
261: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
262: PetscCall(GetHandlesFrom_(dctx, &stream));
263: if (inplace) {
264: PetscCall(apply(DeviceArrayReadWrite(dctx, xinout).data()));
265: } else {
266: PetscCall(apply(DeviceArrayRead(dctx, xinout).data(), DeviceArrayWrite(dctx, yout).data()));
267: }
268: PetscCall(PetscLogGpuFlops(n));
269: PetscCall(PetscDeviceContextSynchronizeIfWithBarrier_Internal(dctx));
270: } else {
271: if (inplace) {
272: PetscCall(MaybeIncrementEmptyLocalVec(xinout));
273: } else {
274: PetscCall(MaybeIncrementEmptyLocalVec(yout));
275: }
276: }
277: PetscFunctionReturn(PETSC_SUCCESS);
278: }
280: // ==========================================================================================
281: // VecSeq_CUPM - Public API - Constructors
282: // ==========================================================================================
284: // VecCreateSeqCUPM()
285: template <device::cupm::DeviceType T>
286: inline PetscErrorCode VecSeq_CUPM<T>::CreateSeqCUPM(MPI_Comm comm, PetscInt bs, PetscInt n, Vec *v, PetscBool call_set_type) noexcept
287: {
288: PetscFunctionBegin;
289: PetscCall(Create_CUPMBase(comm, bs, n, n, v, call_set_type));
290: PetscFunctionReturn(PETSC_SUCCESS);
291: }
293: // VecCreateSeqCUPMWithArrays()
294: template <device::cupm::DeviceType T>
295: inline PetscErrorCode VecSeq_CUPM<T>::CreateSeqCUPMWithBothArrays(MPI_Comm comm, PetscInt bs, PetscInt n, const PetscScalar host_array[], const PetscScalar device_array[], Vec *v) noexcept
296: {
297: PetscDeviceContext dctx;
299: PetscFunctionBegin;
300: PetscCall(GetHandles_(&dctx));
301: // do NOT call VecSetType(), otherwise ops->create() -> create() ->
302: // CreateSeqCUPM_() is called!
303: PetscCall(CreateSeqCUPM(comm, bs, n, v, PETSC_FALSE));
304: PetscCall(CreateSeqCUPM_(*v, dctx, PetscRemoveConstCast(host_array), PetscRemoveConstCast(device_array)));
305: PetscFunctionReturn(PETSC_SUCCESS);
306: }
308: // v->ops->duplicate
309: template <device::cupm::DeviceType T>
310: inline PetscErrorCode VecSeq_CUPM<T>::Duplicate(Vec v, Vec *y) noexcept
311: {
312: PetscDeviceContext dctx;
314: PetscFunctionBegin;
315: PetscCall(GetHandles_(&dctx));
316: PetscCall(Duplicate_CUPMBase(v, y, dctx));
317: PetscFunctionReturn(PETSC_SUCCESS);
318: }
320: // ==========================================================================================
321: // VecSeq_CUPM - Public API - Utility
322: // ==========================================================================================
324: // v->ops->bindtocpu
325: template <device::cupm::DeviceType T>
326: inline PetscErrorCode VecSeq_CUPM<T>::BindToCPU(Vec v, PetscBool usehost) noexcept
327: {
328: PetscDeviceContext dctx;
330: PetscFunctionBegin;
331: PetscCall(GetHandles_(&dctx));
332: PetscCall(BindToCPU_CUPMBase(v, usehost, dctx));
334: // REVIEW ME: this absolutely should be some sort of bulk mempcy rather than this mess
335: VecSetOp_CUPM(dot, VecDot_Seq, Dot);
336: VecSetOp_CUPM(norm, VecNorm_Seq, Norm);
337: VecSetOp_CUPM(tdot, VecTDot_Seq, TDot);
338: VecSetOp_CUPM(mdot, VecMDot_Seq, MDot);
339: VecSetOp_CUPM(resetarray, VecResetArray_Seq, base_type::template ResetArray<PETSC_MEMTYPE_HOST>);
340: VecSetOp_CUPM(placearray, VecPlaceArray_Seq, base_type::template PlaceArray<PETSC_MEMTYPE_HOST>);
341: v->ops->mtdot = v->ops->mtdot_local = VecMTDot_Seq;
342: VecSetOp_CUPM(max, VecMax_Seq, Max);
343: VecSetOp_CUPM(min, VecMin_Seq, Min);
344: VecSetOp_CUPM(setpreallocationcoo, VecSetPreallocationCOO_Seq, SetPreallocationCOO);
345: VecSetOp_CUPM(setvaluescoo, VecSetValuesCOO_Seq, SetValuesCOO);
346: PetscFunctionReturn(PETSC_SUCCESS);
347: }
349: // ==========================================================================================
350: // VecSeq_CUPM - Public API - Mutators
351: // ==========================================================================================
353: // v->ops->getlocalvector or v->ops->getlocalvectorread
354: template <device::cupm::DeviceType T>
355: template <PetscMemoryAccessMode access>
356: inline PetscErrorCode VecSeq_CUPM<T>::GetLocalVector(Vec v, Vec w) noexcept
357: {
358: PetscBool wisseqcupm;
360: PetscFunctionBegin;
361: PetscCheckTypeNames(v, VECSEQCUPM(), VECMPICUPM());
362: PetscCall(PetscObjectTypeCompare(PetscObjectCast(w), VECSEQCUPM(), &wisseqcupm));
363: if (wisseqcupm) {
364: if (const auto wseq = VecIMPLCast(w)) {
365: if (auto &alloced = wseq->array_allocated) {
366: const auto useit = UseCUPMHostAlloc(util::exchange(w->pinned_memory, PETSC_FALSE));
368: PetscCall(PetscFree(alloced));
369: }
370: wseq->array = nullptr;
371: wseq->unplacedarray = nullptr;
372: }
373: if (const auto wcu = VecCUPMCast(w)) {
374: if (auto &device_array = wcu->array_d) {
375: cupmStream_t stream;
377: PetscCall(GetHandles_(&stream));
378: PetscCallCUPM(cupmFreeAsync(device_array, stream));
379: }
380: PetscCall(PetscFree(w->spptr /* wcu */));
381: }
382: }
383: if (v->petscnative && wisseqcupm) {
384: PetscCall(PetscFree(w->data));
385: w->data = v->data;
386: w->offloadmask = v->offloadmask;
387: w->pinned_memory = v->pinned_memory;
388: w->spptr = v->spptr;
389: PetscCall(PetscObjectStateIncrease(PetscObjectCast(w)));
390: } else {
391: const auto array = &VecIMPLCast(w)->array;
393: if (access == PETSC_MEMORY_ACCESS_READ) {
394: PetscCall(VecGetArrayRead(v, const_cast<const PetscScalar **>(array)));
395: } else {
396: PetscCall(VecGetArray(v, array));
397: }
398: w->offloadmask = PETSC_OFFLOAD_CPU;
399: if (wisseqcupm) {
400: PetscDeviceContext dctx;
402: PetscCall(GetHandles_(&dctx));
403: PetscCall(DeviceAllocateCheck_(dctx, w));
404: }
405: }
406: PetscFunctionReturn(PETSC_SUCCESS);
407: }
409: // v->ops->restorelocalvector or v->ops->restorelocalvectorread
410: template <device::cupm::DeviceType T>
411: template <PetscMemoryAccessMode access>
412: inline PetscErrorCode VecSeq_CUPM<T>::RestoreLocalVector(Vec v, Vec w) noexcept
413: {
414: PetscBool wisseqcupm;
416: PetscFunctionBegin;
417: PetscCheckTypeNames(v, VECSEQCUPM(), VECMPICUPM());
418: PetscCall(PetscObjectTypeCompare(PetscObjectCast(w), VECSEQCUPM(), &wisseqcupm));
419: if (v->petscnative && wisseqcupm) {
420: // the assignments to nullptr are __critical__, as w may persist after this call returns
421: // and shouldn't share data with v!
422: v->pinned_memory = w->pinned_memory;
423: v->offloadmask = util::exchange(w->offloadmask, PETSC_OFFLOAD_UNALLOCATED);
424: v->data = util::exchange(w->data, nullptr);
425: v->spptr = util::exchange(w->spptr, nullptr);
426: } else {
427: const auto array = &VecIMPLCast(w)->array;
429: if (access == PETSC_MEMORY_ACCESS_READ) {
430: PetscCall(VecRestoreArrayRead(v, const_cast<const PetscScalar **>(array)));
431: } else {
432: PetscCall(VecRestoreArray(v, array));
433: }
434: if (w->spptr && wisseqcupm) {
435: cupmStream_t stream;
437: PetscCall(GetHandles_(&stream));
438: PetscCallCUPM(cupmFreeAsync(VecCUPMCast(w)->array_d, stream));
439: PetscCall(PetscFree(w->spptr));
440: }
441: }
442: PetscFunctionReturn(PETSC_SUCCESS);
443: }
445: // ==========================================================================================
446: // VecSeq_CUPM - Public API - Compute Methods
447: // ==========================================================================================
449: // VecAYPXAsync_Private
450: template <device::cupm::DeviceType T>
451: inline PetscErrorCode VecSeq_CUPM<T>::AYPXAsync(Vec yin, PetscScalar alpha, Vec xin, PetscDeviceContext dctx) noexcept
452: {
453: const auto n = static_cast<cupmBlasInt_t>(yin->map->n);
454: PetscBool xiscupm;
456: PetscFunctionBegin;
457: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(xin), &xiscupm, VECSEQCUPM(), VECMPICUPM(), ""));
458: if (!xiscupm) {
459: PetscCall(VecAYPX_Seq(yin, alpha, xin));
460: PetscFunctionReturn(PETSC_SUCCESS);
461: }
462: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
463: if (alpha == PetscScalar(0.0)) {
464: cupmStream_t stream;
466: PetscCall(GetHandlesFrom_(dctx, &stream));
467: PetscCall(PetscLogGpuTimeBegin());
468: PetscCall(PetscCUPMMemcpyAsync(DeviceArrayWrite(dctx, yin).data(), DeviceArrayRead(dctx, xin).data(), n, cupmMemcpyDeviceToDevice, stream));
469: PetscCall(PetscLogGpuTimeEnd());
470: } else if (n) {
471: const auto alphaIsOne = alpha == PetscScalar(1.0);
472: const auto calpha = cupmScalarPtrCast(&alpha);
473: cupmBlasHandle_t cupmBlasHandle;
475: PetscCall(GetHandlesFrom_(dctx, &cupmBlasHandle));
476: {
477: const auto yptr = DeviceArrayReadWrite(dctx, yin);
478: const auto xptr = DeviceArrayRead(dctx, xin);
480: PetscCall(PetscLogGpuTimeBegin());
481: if (alphaIsOne) {
482: PetscCallCUPMBLAS(cupmBlasXaxpy(cupmBlasHandle, n, calpha, xptr.cupmdata(), 1, yptr.cupmdata(), 1));
483: } else {
484: const auto one = cupmScalarCast(1.0);
486: PetscCallCUPMBLAS(cupmBlasXscal(cupmBlasHandle, n, calpha, yptr.cupmdata(), 1));
487: PetscCallCUPMBLAS(cupmBlasXaxpy(cupmBlasHandle, n, &one, xptr.cupmdata(), 1, yptr.cupmdata(), 1));
488: }
489: PetscCall(PetscLogGpuTimeEnd());
490: }
491: PetscCall(PetscLogGpuFlops((alphaIsOne ? 1 : 2) * n));
492: }
493: if (n > 0) PetscCall(PetscDeviceContextSynchronizeIfWithBarrier_Internal(dctx));
494: PetscFunctionReturn(PETSC_SUCCESS);
495: }
497: // v->ops->aypx
498: template <device::cupm::DeviceType T>
499: inline PetscErrorCode VecSeq_CUPM<T>::AYPX(Vec yin, PetscScalar alpha, Vec xin) noexcept
500: {
501: PetscFunctionBegin;
502: PetscCall(AYPXAsync(yin, alpha, xin, nullptr));
503: PetscFunctionReturn(PETSC_SUCCESS);
504: }
506: // VecAXPYAsync_Private
507: template <device::cupm::DeviceType T>
508: inline PetscErrorCode VecSeq_CUPM<T>::AXPYAsync(Vec yin, PetscScalar alpha, Vec xin, PetscDeviceContext dctx) noexcept
509: {
510: PetscBool xiscupm;
512: PetscFunctionBegin;
513: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(xin), &xiscupm, VECSEQCUPM(), VECMPICUPM(), ""));
514: if (xiscupm) {
515: const auto n = static_cast<cupmBlasInt_t>(yin->map->n);
516: cupmBlasHandle_t cupmBlasHandle;
518: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
519: PetscCall(GetHandlesFrom_(dctx, &cupmBlasHandle));
520: PetscCall(PetscLogGpuTimeBegin());
521: PetscCallCUPMBLAS(cupmBlasXaxpy(cupmBlasHandle, n, cupmScalarPtrCast(&alpha), DeviceArrayRead(dctx, xin), 1, DeviceArrayReadWrite(dctx, yin), 1));
522: PetscCall(PetscLogGpuTimeEnd());
523: PetscCall(PetscLogGpuFlops(2 * n));
524: if (n > 0) PetscCall(PetscDeviceContextSynchronizeIfWithBarrier_Internal(dctx));
525: } else {
526: PetscCall(VecAXPY_Seq(yin, alpha, xin));
527: }
528: PetscFunctionReturn(PETSC_SUCCESS);
529: }
531: // v->ops->axpy
532: template <device::cupm::DeviceType T>
533: inline PetscErrorCode VecSeq_CUPM<T>::AXPY(Vec yin, PetscScalar alpha, Vec xin) noexcept
534: {
535: PetscFunctionBegin;
536: PetscCall(AXPYAsync(yin, alpha, xin, nullptr));
537: PetscFunctionReturn(PETSC_SUCCESS);
538: }
540: namespace detail
541: {
543: struct divides {
544: PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &lhs, const PetscScalar &rhs) const noexcept { return rhs == PetscScalar{0.0} ? (lhs == PetscScalar{0.0} ? PetscScalar{1.0} : rhs) : lhs / rhs; }
545: };
547: } // namespace detail
549: // VecPointwiseDivideAsync_Private
550: template <device::cupm::DeviceType T>
551: inline PetscErrorCode VecSeq_CUPM<T>::PointwiseDivideAsync(Vec wout, Vec xin, Vec yin, PetscDeviceContext dctx) noexcept
552: {
553: PetscFunctionBegin;
554: PetscCall(PointwiseBinaryDispatch_(VecPointwiseDivide_Seq, detail::divides{}, wout, xin, yin, dctx));
555: PetscFunctionReturn(PETSC_SUCCESS);
556: }
558: // v->ops->pointwisedivide
559: template <device::cupm::DeviceType T>
560: inline PetscErrorCode VecSeq_CUPM<T>::PointwiseDivide(Vec wout, Vec xin, Vec yin) noexcept
561: {
562: PetscFunctionBegin;
563: PetscCall(PointwiseDivideAsync(wout, xin, yin, nullptr));
564: PetscFunctionReturn(PETSC_SUCCESS);
565: }
567: namespace detail
568: {
570: struct multiplies {
571: PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &lhs, const PetscScalar &rhs) const noexcept { return lhs * rhs; }
572: };
574: } // namespace detail
576: // VecPointwiseMultAsync_Private
577: template <device::cupm::DeviceType T>
578: inline PetscErrorCode VecSeq_CUPM<T>::PointwiseMultAsync(Vec wout, Vec xin, Vec yin, PetscDeviceContext dctx) noexcept
579: {
580: PetscFunctionBegin;
581: PetscCall(PointwiseBinaryDispatch_(VecPointwiseMult_Seq, detail::multiplies{}, wout, xin, yin, dctx));
582: PetscFunctionReturn(PETSC_SUCCESS);
583: }
585: // v->ops->pointwisemult
586: template <device::cupm::DeviceType T>
587: inline PetscErrorCode VecSeq_CUPM<T>::PointwiseMult(Vec wout, Vec xin, Vec yin) noexcept
588: {
589: PetscFunctionBegin;
590: PetscCall(PointwiseMultAsync(wout, xin, yin, nullptr));
591: PetscFunctionReturn(PETSC_SUCCESS);
592: }
594: namespace detail
595: {
597: struct MaximumRealPart {
598: PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &lhs, const PetscScalar &rhs) const noexcept { return thrust::maximum<PetscReal>{}(PetscRealPart(lhs), PetscRealPart(rhs)); }
599: };
601: } // namespace detail
603: // VecPointwiseMaxAsync_Private
604: template <device::cupm::DeviceType T>
605: inline PetscErrorCode VecSeq_CUPM<T>::PointwiseMaxAsync(Vec wout, Vec xin, Vec yin, PetscDeviceContext dctx) noexcept
606: {
607: PetscFunctionBegin;
608: PetscCall(PointwiseBinaryDispatch_(VecPointwiseMax_Seq, detail::MaximumRealPart{}, wout, xin, yin, dctx));
609: PetscFunctionReturn(PETSC_SUCCESS);
610: }
612: // v->ops->pointwisemax
613: template <device::cupm::DeviceType T>
614: inline PetscErrorCode VecSeq_CUPM<T>::PointwiseMax(Vec wout, Vec xin, Vec yin) noexcept
615: {
616: PetscFunctionBegin;
617: PetscCall(PointwiseMaxAsync(wout, xin, yin, nullptr));
618: PetscFunctionReturn(PETSC_SUCCESS);
619: }
621: namespace detail
622: {
624: struct MaximumAbsoluteValue {
625: PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &lhs, const PetscScalar &rhs) const noexcept { return thrust::maximum<PetscReal>{}(PetscAbsScalar(lhs), PetscAbsScalar(rhs)); }
626: };
628: } // namespace detail
630: // VecPointwiseMaxAbsAsync_Private
631: template <device::cupm::DeviceType T>
632: inline PetscErrorCode VecSeq_CUPM<T>::PointwiseMaxAbsAsync(Vec wout, Vec xin, Vec yin, PetscDeviceContext dctx) noexcept
633: {
634: PetscFunctionBegin;
635: PetscCall(PointwiseBinaryDispatch_(VecPointwiseMaxAbs_Seq, detail::MaximumAbsoluteValue{}, wout, xin, yin, dctx));
636: PetscFunctionReturn(PETSC_SUCCESS);
637: }
639: // v->ops->pointwisemaxabs
640: template <device::cupm::DeviceType T>
641: inline PetscErrorCode VecSeq_CUPM<T>::PointwiseMaxAbs(Vec wout, Vec xin, Vec yin) noexcept
642: {
643: PetscFunctionBegin;
644: PetscCall(PointwiseMaxAbsAsync(wout, xin, yin, nullptr));
645: PetscFunctionReturn(PETSC_SUCCESS);
646: }
648: namespace detail
649: {
651: struct MinimumRealPart {
652: PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &lhs, const PetscScalar &rhs) const noexcept { return thrust::minimum<PetscReal>{}(PetscRealPart(lhs), PetscRealPart(rhs)); }
653: };
655: } // namespace detail
657: // VecPointwiseMinAsync_Private
658: template <device::cupm::DeviceType T>
659: inline PetscErrorCode VecSeq_CUPM<T>::PointwiseMinAsync(Vec wout, Vec xin, Vec yin, PetscDeviceContext dctx) noexcept
660: {
661: PetscFunctionBegin;
662: PetscCall(PointwiseBinaryDispatch_(VecPointwiseMin_Seq, detail::MinimumRealPart{}, wout, xin, yin, dctx));
663: PetscFunctionReturn(PETSC_SUCCESS);
664: }
666: // v->ops->pointwisemin
667: template <device::cupm::DeviceType T>
668: inline PetscErrorCode VecSeq_CUPM<T>::PointwiseMin(Vec wout, Vec xin, Vec yin) noexcept
669: {
670: PetscFunctionBegin;
671: PetscCall(PointwiseMinAsync(wout, xin, yin, nullptr));
672: PetscFunctionReturn(PETSC_SUCCESS);
673: }
675: namespace detail
676: {
678: struct Reciprocal {
679: PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &s) const noexcept
680: {
681: // yes all of this verbosity is needed because sometimes PetscScalar is a thrust::complex
682: // and then it matters whether we do s ? true : false vs s == 0, as well as whether we wrap
683: // everything in PetscScalar...
684: return s == PetscScalar{0.0} ? s : PetscScalar{1.0} / s;
685: }
686: };
688: } // namespace detail
690: // VecReciprocalAsync_Private
691: template <device::cupm::DeviceType T>
692: inline PetscErrorCode VecSeq_CUPM<T>::ReciprocalAsync(Vec xin, PetscDeviceContext dctx) noexcept
693: {
694: PetscFunctionBegin;
695: PetscCall(PointwiseUnary_(detail::Reciprocal{}, xin, nullptr, dctx));
696: PetscFunctionReturn(PETSC_SUCCESS);
697: }
699: // v->ops->reciprocal
700: template <device::cupm::DeviceType T>
701: inline PetscErrorCode VecSeq_CUPM<T>::Reciprocal(Vec xin) noexcept
702: {
703: PetscFunctionBegin;
704: PetscCall(ReciprocalAsync(xin, nullptr));
705: PetscFunctionReturn(PETSC_SUCCESS);
706: }
708: namespace detail
709: {
711: struct AbsoluteValue {
712: PETSC_HOSTDEVICE_INLINE_DECL PetscReal operator()(const PetscScalar &s) const noexcept { return PetscAbsScalar(s); }
713: };
715: struct NormMax {
716: // Propagate NaNs as in VecNorm_Seq().
717: PETSC_HOSTDEVICE_INLINE_DECL PetscReal operator()(PetscReal lhs, PetscReal rhs) const noexcept { return lhs != lhs || lhs > rhs ? lhs : rhs; }
718: };
720: } // namespace detail
722: // VecAbsAsync_Private
723: template <device::cupm::DeviceType T>
724: inline PetscErrorCode VecSeq_CUPM<T>::AbsAsync(Vec xin, PetscDeviceContext dctx) noexcept
725: {
726: PetscFunctionBegin;
727: PetscCall(PointwiseUnary_(detail::AbsoluteValue{}, xin, nullptr, dctx));
728: PetscFunctionReturn(PETSC_SUCCESS);
729: }
731: // v->ops->abs
732: template <device::cupm::DeviceType T>
733: inline PetscErrorCode VecSeq_CUPM<T>::Abs(Vec xin) noexcept
734: {
735: PetscFunctionBegin;
736: PetscCall(AbsAsync(xin, nullptr));
737: PetscFunctionReturn(PETSC_SUCCESS);
738: }
740: namespace detail
741: {
743: struct SignZeroToSignedUnit {
744: PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &s) const noexcept { return VecSignZeroToSignedUnit_Private(PetscRealPart(s)); }
745: };
747: struct SignZeroToZero {
748: PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &s) const noexcept { return VecSignZeroToZero_Private(PetscRealPart(s)); }
749: };
751: struct SignZeroToSignedZero {
752: PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &s) const noexcept { return VecSignZeroToSignedZero_Private(PetscRealPart(s)); }
753: };
755: } // namespace detail
757: // VecPointwiseSignAsync_Private
758: template <device::cupm::DeviceType T>
759: inline PetscErrorCode VecSeq_CUPM<T>::PointwiseSignAsync(Vec yout, Vec xin, VecSignMode sign_type, PetscDeviceContext dctx) noexcept
760: {
761: PetscFunctionBegin;
762: switch (sign_type) {
763: case VEC_SIGN_ZERO_TO_ZERO:
764: PetscCall(PointwiseUnary_(detail::SignZeroToZero{}, xin, yout, dctx));
765: break;
766: case VEC_SIGN_ZERO_TO_SIGNED_ZERO:
767: PetscCall(PointwiseUnary_(detail::SignZeroToSignedZero{}, xin, yout, dctx));
768: break;
769: case VEC_SIGN_ZERO_TO_SIGNED_UNIT:
770: PetscCall(PointwiseUnary_(detail::SignZeroToSignedUnit{}, xin, yout, dctx));
771: break;
772: }
773: PetscFunctionReturn(PETSC_SUCCESS);
774: }
776: namespace detail
777: {
779: struct SquareRootAbsoluteValue {
780: PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &s) const noexcept { return PetscSqrtReal(PetscAbsScalar(s)); }
781: };
783: } // namespace detail
785: // VecSqrtAbsAsync_Private
786: template <device::cupm::DeviceType T>
787: inline PetscErrorCode VecSeq_CUPM<T>::SqrtAbsAsync(Vec xin, PetscDeviceContext dctx) noexcept
788: {
789: PetscFunctionBegin;
790: PetscCall(PointwiseUnary_(detail::SquareRootAbsoluteValue{}, xin, nullptr, dctx));
791: PetscFunctionReturn(PETSC_SUCCESS);
792: }
794: // v->ops->sqrt
795: template <device::cupm::DeviceType T>
796: inline PetscErrorCode VecSeq_CUPM<T>::SqrtAbs(Vec xin) noexcept
797: {
798: PetscFunctionBegin;
799: PetscCall(SqrtAbsAsync(xin, nullptr));
800: PetscFunctionReturn(PETSC_SUCCESS);
801: }
803: namespace detail
804: {
806: struct Exponent {
807: PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &s) const noexcept { return PetscExpScalar(s); }
808: };
810: } // namespace detail
812: // VecExpAsync_Private
813: template <device::cupm::DeviceType T>
814: inline PetscErrorCode VecSeq_CUPM<T>::ExpAsync(Vec xin, PetscDeviceContext dctx) noexcept
815: {
816: PetscFunctionBegin;
817: PetscCall(PointwiseUnary_(detail::Exponent{}, xin, nullptr, dctx));
818: PetscFunctionReturn(PETSC_SUCCESS);
819: }
821: // v->ops->exp
822: template <device::cupm::DeviceType T>
823: inline PetscErrorCode VecSeq_CUPM<T>::Exp(Vec xin) noexcept
824: {
825: PetscFunctionBegin;
826: PetscCall(ExpAsync(xin, nullptr));
827: PetscFunctionReturn(PETSC_SUCCESS);
828: }
830: namespace detail
831: {
833: struct Logarithm {
834: PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &s) const noexcept { return PetscLogScalar(s); }
835: };
837: } // namespace detail
839: // VecLogAsync_Private
840: template <device::cupm::DeviceType T>
841: inline PetscErrorCode VecSeq_CUPM<T>::LogAsync(Vec xin, PetscDeviceContext dctx) noexcept
842: {
843: PetscFunctionBegin;
844: PetscCall(PointwiseUnary_(detail::Logarithm{}, xin, nullptr, dctx));
845: PetscFunctionReturn(PETSC_SUCCESS);
846: }
848: // v->ops->log
849: template <device::cupm::DeviceType T>
850: inline PetscErrorCode VecSeq_CUPM<T>::Log(Vec xin) noexcept
851: {
852: PetscFunctionBegin;
853: PetscCall(LogAsync(xin, nullptr));
854: PetscFunctionReturn(PETSC_SUCCESS);
855: }
857: // v->ops->waxpy
858: template <device::cupm::DeviceType T>
859: inline PetscErrorCode VecSeq_CUPM<T>::WAXPYAsync(Vec win, PetscScalar alpha, Vec xin, Vec yin, PetscDeviceContext dctx) noexcept
860: {
861: PetscBool xiscupm, yiscupm;
863: PetscFunctionBegin;
864: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(xin), &xiscupm, VECSEQCUPM(), VECMPICUPM(), ""));
865: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(yin), &yiscupm, VECSEQCUPM(), VECMPICUPM(), ""));
866: if (!xiscupm || !yiscupm) {
867: PetscCall(VecWAXPY_Seq(win, alpha, xin, yin));
868: PetscFunctionReturn(PETSC_SUCCESS);
869: }
870: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
871: if (alpha == PetscScalar(0.0)) {
872: PetscCall(CopyAsync(yin, win, dctx));
873: } else if (const auto n = static_cast<cupmBlasInt_t>(win->map->n)) {
874: cupmBlasHandle_t cupmBlasHandle;
875: cupmStream_t stream;
876: PetscBool xiscupm, yiscupm;
878: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(xin), &xiscupm, VECSEQCUPM(), VECMPICUPM(), ""));
879: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(yin), &yiscupm, VECSEQCUPM(), VECMPICUPM(), ""));
880: if (!xiscupm || !yiscupm) {
881: PetscCall(VecWAXPY_Seq(win, alpha, xin, yin));
882: PetscFunctionReturn(PETSC_SUCCESS);
883: }
884: PetscCall(GetHandlesFrom_(dctx, &cupmBlasHandle, NULL, &stream));
885: {
886: const auto wptr = DeviceArrayWrite(dctx, win);
888: PetscCall(PetscLogGpuTimeBegin());
889: PetscCall(PetscCUPMMemcpyAsync(wptr.data(), DeviceArrayRead(dctx, yin).data(), n, cupmMemcpyDeviceToDevice, stream, true));
890: PetscCallCUPMBLAS(cupmBlasXaxpy(cupmBlasHandle, n, cupmScalarPtrCast(&alpha), DeviceArrayRead(dctx, xin), 1, wptr.cupmdata(), 1));
891: PetscCall(PetscLogGpuTimeEnd());
892: }
893: PetscCall(PetscLogGpuFlops(2 * n));
894: PetscCall(PetscDeviceContextSynchronizeIfWithBarrier_Internal(dctx));
895: }
896: PetscFunctionReturn(PETSC_SUCCESS);
897: }
899: // v->ops->waxpy
900: template <device::cupm::DeviceType T>
901: inline PetscErrorCode VecSeq_CUPM<T>::WAXPY(Vec win, PetscScalar alpha, Vec xin, Vec yin) noexcept
902: {
903: PetscFunctionBegin;
904: PetscCall(WAXPYAsync(win, alpha, xin, yin, nullptr));
905: PetscFunctionReturn(PETSC_SUCCESS);
906: }
908: namespace detail
909: {
910: struct stdbasis_functor {
911: PetscInt _i;
913: stdbasis_functor(PetscInt i) : _i(i) { }
915: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &, const PetscScalar &i) const noexcept { return i == _i ? 1.0 : 0.0; }
916: };
917: } // namespace detail
919: // v->ops->setstdbasis
920: template <device::cupm::DeviceType T>
921: inline PetscErrorCode VecSeq_CUPM<T>::SetStdBasisAsync(Vec s, PetscInt i, PetscDeviceContext dctx) noexcept
922: {
923: cupmStream_t stream;
925: PetscFunctionBegin;
926: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
927: PetscCall(GetHandlesFrom_(dctx, &stream));
928: {
929: const auto nl = s->map->n;
930: const auto st = s->map->rstart;
931: auto nit = thrust::make_counting_iterator(PetscInt{st});
932: const auto sptr = thrust::device_pointer_cast(DeviceArrayWrite(dctx, s).data());
933: PetscCall(PetscLogGpuTimeBegin());
934: // clang-format off
935: PetscCallThrust(
936: THRUST_CALL(
937: thrust::transform,
938: stream,
939: sptr,
940: sptr + nl,
941: nit,
942: sptr,
943: detail::stdbasis_functor(i)
944: )
945: );
946: // clang-format on
947: PetscCall(PetscLogGpuFlops(nl));
948: PetscCall(PetscDeviceContextSynchronizeIfWithBarrier_Internal(dctx));
949: }
950: PetscFunctionReturn(PETSC_SUCCESS);
951: }
953: // v->ops->setstdbasis
954: template <device::cupm::DeviceType T>
955: inline PetscErrorCode VecSeq_CUPM<T>::SetStdBasis(Vec s, PetscInt i) noexcept
956: {
957: PetscFunctionBegin;
958: PetscCall(SetStdBasisAsync(s, i, nullptr));
959: PetscFunctionReturn(PETSC_SUCCESS);
960: }
962: namespace kernels
963: {
965: template <std::size_t N>
966: struct MAXPYCoefficients {
967: PetscScalar values[N];
968: };
970: template <typename... Args>
971: PETSC_KERNEL_DECL static void MAXPY_kernel(const PetscInt size, PetscScalar *PETSC_RESTRICT xptr, const MAXPYCoefficients<sizeof...(Args)> alpha, Args... yptr)
972: {
973: constexpr int N = sizeof...(Args);
974: const PetscScalar *yptr_p[] = {yptr...};
976: ::Petsc::device::cupm::kernels::util::grid_stride_1D(size, [&](PetscInt i) {
977: auto sum = xptr[i];
979: #pragma unroll
980: for (auto j = 0; j < N; ++j) sum += alpha.values[j] * yptr_p[j][i];
981: xptr[i] = sum;
982: });
983: return;
984: }
986: } // namespace kernels
988: namespace detail
989: {
991: // a helper-struct to gobble the size_t input, it is used with template parameter pack
992: // expansion such that
993: // typename repeat_type...
994: // expands to
995: // MyType, MyType, MyType, ... [repeated sizeof...(IdxParamPack) times]
996: template <typename T, std::size_t>
997: struct repeat_type {
998: using type = T;
999: };
1001: } // namespace detail
1003: template <device::cupm::DeviceType T>
1004: template <std::size_t... Idx>
1005: inline PetscErrorCode VecSeq_CUPM<T>::MAXPY_kernel_dispatch_(PetscDeviceContext dctx, cupmStream_t stream, PetscScalar *xptr, const PetscScalar *aptr, const Vec *yin, PetscInt size, util::index_sequence<Idx...>) noexcept
1006: {
1007: #if PetscDefined(USE_COMPLEX)
1008: kernels::MAXPYCoefficients<sizeof...(Idx)> alpha;
1009: #else
1010: const kernels::MAXPYCoefficients<sizeof...(Idx)> alpha{{aptr[Idx]...}};
1011: #endif
1013: PetscFunctionBegin;
1014: #if PetscDefined(USE_COMPLEX)
1015: /* C complex coefficients may have weaker alignment than the CUDA/HIP complex type.
1016: Use a byte copy: PetscArraycpy() can select a typed scalar copy at configure time. */
1017: std::memcpy(static_cast<void *>(alpha.values), static_cast<const void *>(aptr), sizeof(alpha.values));
1018: #endif
1019: // clang-format off
1020: PetscCall(
1021: PetscCUPMLaunchKernel1D(
1022: size, 0, stream,
1023: kernels::MAXPY_kernel<typename detail::repeat_type<const PetscScalar *, Idx>::type...>,
1024: size, xptr, alpha, DeviceArrayRead(dctx, yin[Idx]).data()...
1025: )
1026: );
1027: // clang-format on
1028: PetscFunctionReturn(PETSC_SUCCESS);
1029: }
1031: template <device::cupm::DeviceType T>
1032: template <int N>
1033: inline PetscErrorCode VecSeq_CUPM<T>::MAXPY_kernel_dispatch_(PetscDeviceContext dctx, cupmStream_t stream, PetscScalar *xptr, const PetscScalar *aptr, const Vec *yin, PetscInt size, PetscInt &yidx) noexcept
1034: {
1035: PetscFunctionBegin;
1036: PetscCall(MAXPY_kernel_dispatch_(dctx, stream, xptr, aptr + yidx, yin + yidx, size, util::make_index_sequence<N>{}));
1037: yidx += N;
1038: PetscFunctionReturn(PETSC_SUCCESS);
1039: }
1041: // VecMAXPYAsync_Private
1042: template <device::cupm::DeviceType T>
1043: inline PetscErrorCode VecSeq_CUPM<T>::MAXPYAsync(Vec xin, PetscInt nv, const PetscScalar *alpha, Vec *yin, PetscDeviceContext dctx) noexcept
1044: {
1045: const auto n = xin->map->n;
1046: cupmStream_t stream;
1047: PetscBool yiscupm = PETSC_TRUE;
1049: PetscFunctionBegin;
1050: for (PetscInt i = 0; i < nv && yiscupm; i++) PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(yin[i]), &yiscupm, VECSEQCUPM(), VECMPICUPM(), ""));
1051: if (!yiscupm) {
1052: PetscCall(VecMAXPY_Seq(xin, nv, alpha, yin));
1053: PetscFunctionReturn(PETSC_SUCCESS);
1054: }
1055: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
1056: PetscCall(GetHandlesFrom_(dctx, &stream));
1057: {
1058: const auto xptr = DeviceArrayReadWrite(dctx, xin);
1059: PetscInt yidx = 0;
1061: // placement of early-return is deliberate, we would like to capture the
1062: // DeviceArrayReadWrite() call (which calls PetscObjectStateIncreate()) before we bail
1063: if (!n || !nv) PetscFunctionReturn(PETSC_SUCCESS);
1064: PetscCall(PetscLogGpuTimeBegin());
1065: do {
1066: switch (nv - yidx) {
1067: case 7:
1068: PetscCall(MAXPY_kernel_dispatch_<7>(dctx, stream, xptr.data(), alpha, yin, n, yidx));
1069: break;
1070: case 6:
1071: PetscCall(MAXPY_kernel_dispatch_<6>(dctx, stream, xptr.data(), alpha, yin, n, yidx));
1072: break;
1073: case 5:
1074: PetscCall(MAXPY_kernel_dispatch_<5>(dctx, stream, xptr.data(), alpha, yin, n, yidx));
1075: break;
1076: case 4:
1077: PetscCall(MAXPY_kernel_dispatch_<4>(dctx, stream, xptr.data(), alpha, yin, n, yidx));
1078: break;
1079: case 3:
1080: PetscCall(MAXPY_kernel_dispatch_<3>(dctx, stream, xptr.data(), alpha, yin, n, yidx));
1081: break;
1082: case 2:
1083: PetscCall(MAXPY_kernel_dispatch_<2>(dctx, stream, xptr.data(), alpha, yin, n, yidx));
1084: break;
1085: case 1:
1086: PetscCall(MAXPY_kernel_dispatch_<1>(dctx, stream, xptr.data(), alpha, yin, n, yidx));
1087: break;
1088: default: // 8 or more
1089: PetscCall(MAXPY_kernel_dispatch_<8>(dctx, stream, xptr.data(), alpha, yin, n, yidx));
1090: break;
1091: }
1092: } while (yidx < nv);
1093: PetscCall(PetscLogGpuTimeEnd());
1094: PetscCall(PetscDeviceContextSynchronizeIfWithBarrier_Internal(dctx));
1095: }
1096: PetscCall(PetscLogGpuFlops(nv * 2 * n));
1097: PetscFunctionReturn(PETSC_SUCCESS);
1098: }
1100: // v->ops->maxpy
1101: template <device::cupm::DeviceType T>
1102: inline PetscErrorCode VecSeq_CUPM<T>::MAXPY(Vec xin, PetscInt nv, const PetscScalar *alpha, Vec *yin) noexcept
1103: {
1104: PetscFunctionBegin;
1105: PetscCall(MAXPYAsync(xin, nv, alpha, yin, nullptr));
1106: PetscFunctionReturn(PETSC_SUCCESS);
1107: }
1109: template <device::cupm::DeviceType T>
1110: inline PetscErrorCode VecSeq_CUPM<T>::Dot(Vec xin, Vec yin, PetscScalar *z) noexcept
1111: {
1112: PetscBool yiscupm;
1114: PetscFunctionBegin;
1115: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(yin), &yiscupm, VECSEQCUPM(), VECMPICUPM(), ""));
1116: if (!yiscupm) {
1117: PetscCall(VecDot_Seq(xin, yin, z));
1118: PetscFunctionReturn(PETSC_SUCCESS);
1119: }
1120: if (const auto n = static_cast<cupmBlasInt_t>(xin->map->n)) {
1121: PetscDeviceContext dctx;
1122: cupmBlasHandle_t cupmBlasHandle;
1124: PetscCall(GetHandles_(&dctx, &cupmBlasHandle));
1125: // arguments y, x are reversed because BLAS complex conjugates the first argument, PETSc the
1126: // second
1127: PetscCall(PetscLogGpuTimeBegin());
1128: PetscCallCUPMBLAS(cupmBlasXdot(cupmBlasHandle, n, DeviceArrayRead(dctx, yin), 1, DeviceArrayRead(dctx, xin), 1, cupmScalarPtrCast(z)));
1129: PetscCall(PetscLogGpuTimeEnd());
1130: PetscCall(PetscLogGpuFlops(2 * n - 1));
1131: } else {
1132: *z = 0.0;
1133: }
1134: PetscFunctionReturn(PETSC_SUCCESS);
1135: }
1137: #define MDOT_WORKGROUP_NUM 128
1138: #define MDOT_WORKGROUP_SIZE MDOT_WORKGROUP_NUM
1140: namespace kernels
1141: {
1143: PETSC_DEVICE_INLINE_DECL static PetscInt EntriesPerGroup(const PetscInt size) noexcept
1144: {
1145: const auto group_entries = (size - 1) / gridDim.x + 1;
1146: // for very small vectors, a group should still do some work
1147: return group_entries ? group_entries : 1;
1148: }
1150: template <typename... ConstPetscScalarPointer>
1151: PETSC_KERNEL_DECL static void MDot_kernel(const PetscScalar *PETSC_RESTRICT x, const PetscInt size, PetscScalar *PETSC_RESTRICT results, ConstPetscScalarPointer... y)
1152: {
1153: constexpr int N = sizeof...(ConstPetscScalarPointer);
1154: const PetscScalar *ylocal[] = {y...};
1155: PetscScalar sumlocal[N];
1157: PETSC_SHAREDMEM_DECL PetscScalar shmem[N * MDOT_WORKGROUP_SIZE];
1159: // HIP -- for whatever reason -- has threadIdx, blockIdx, blockDim, and gridDim as separate
1160: // types, so each of these go on separate lines...
1161: const auto tx = threadIdx.x;
1162: const auto bx = blockIdx.x;
1163: const auto bdx = blockDim.x;
1164: const auto gdx = gridDim.x;
1165: const auto worksize = EntriesPerGroup(size);
1166: const auto begin = tx + bx * worksize;
1167: const auto end = min((bx + 1) * worksize, size);
1169: #pragma unroll
1170: for (auto i = 0; i < N; ++i) sumlocal[i] = 0;
1172: for (auto i = begin; i < end; i += bdx) {
1173: const auto xi = x[i]; // load only once from global memory!
1175: #pragma unroll
1176: for (auto j = 0; j < N; ++j) sumlocal[j] += ylocal[j][i] * xi;
1177: }
1179: #pragma unroll
1180: for (auto i = 0; i < N; ++i) shmem[tx + i * MDOT_WORKGROUP_SIZE] = sumlocal[i];
1182: // parallel reduction
1183: for (auto stride = bdx / 2; stride > 0; stride /= 2) {
1184: __syncthreads();
1185: if (tx < stride) {
1186: #pragma unroll
1187: for (auto i = 0; i < N; ++i) shmem[tx + i * MDOT_WORKGROUP_SIZE] += shmem[tx + stride + i * MDOT_WORKGROUP_SIZE];
1188: }
1189: }
1190: // bottom N threads per block write to global memory
1191: // REVIEW ME: I am ~pretty~ sure we don't need another __syncthreads() here since each thread
1192: // writes to the same sections in the above loop that it is about to read from below, but
1193: // running this under the racecheck tool of compute-sanitizer reports a write-after-write hazard.
1194: __syncthreads();
1195: if (tx < N) results[bx + tx * gdx] = shmem[tx * MDOT_WORKGROUP_SIZE];
1196: return;
1197: }
1199: namespace
1200: {
1202: PETSC_KERNEL_DECL void sum_kernel(const PetscInt size, PetscScalar *PETSC_RESTRICT results)
1203: {
1204: int local_i = 0;
1205: PetscScalar local_results[8];
1207: // each thread sums up MDOT_WORKGROUP_NUM entries of the result, storing it in a local buffer
1208: //
1209: // *-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*
1210: // | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | ...
1211: // *-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*
1212: // | ______________________________________________________/
1213: // | / <- MDOT_WORKGROUP_NUM ->
1214: // |/
1215: // +
1216: // v
1217: // *-*-*
1218: // | | | ...
1219: // *-*-*
1220: //
1221: ::Petsc::device::cupm::kernels::util::grid_stride_1D(size, [&](PetscInt i) {
1222: PetscScalar z_sum = 0;
1224: for (auto j = i * MDOT_WORKGROUP_SIZE; j < (i + 1) * MDOT_WORKGROUP_SIZE; ++j) z_sum += results[j];
1225: local_results[local_i++] = z_sum;
1226: });
1227: // if we needed more than 1 workgroup to handle the vector we should sync since other threads
1228: // may currently be reading from results
1229: if (size >= MDOT_WORKGROUP_SIZE) __syncthreads();
1230: // Local buffer is now written to global memory
1231: ::Petsc::device::cupm::kernels::util::grid_stride_1D(size, [&](PetscInt i) {
1232: const auto j = --local_i;
1234: if (j >= 0) results[i] = local_results[j];
1235: });
1236: return;
1237: }
1239: } // namespace
1241: #if PetscDefined(USING_HCC)
1242: namespace do_not_use
1243: {
1245: inline void silence_warning_function_sum_kernel_is_not_needed_and_will_not_be_emitted()
1246: {
1247: (void)sum_kernel;
1248: }
1250: } // namespace do_not_use
1251: #endif
1253: } // namespace kernels
1255: template <device::cupm::DeviceType T>
1256: template <std::size_t... Idx>
1257: inline PetscErrorCode VecSeq_CUPM<T>::MDot_kernel_dispatch_(PetscDeviceContext dctx, cupmStream_t stream, const PetscScalar *xarr, const Vec yin[], PetscInt size, PetscScalar *results, util::index_sequence<Idx...>) noexcept
1258: {
1259: PetscFunctionBegin;
1260: // REVIEW ME: convert this kernel launch to PetscCUPMLaunchKernel1D(), it currently launches
1261: // 128 blocks of 128 threads every time which may be wasteful
1262: // clang-format off
1263: PetscCallCUPM(
1264: cupmLaunchKernel(
1265: kernels::MDot_kernel<typename detail::repeat_type<const PetscScalar *, Idx>::type...>,
1266: MDOT_WORKGROUP_NUM, MDOT_WORKGROUP_SIZE, 0, stream,
1267: xarr, size, results, DeviceArrayRead(dctx, yin[Idx]).data()...
1268: )
1269: );
1270: // clang-format on
1271: PetscFunctionReturn(PETSC_SUCCESS);
1272: }
1274: template <device::cupm::DeviceType T>
1275: template <int N>
1276: inline PetscErrorCode VecSeq_CUPM<T>::MDot_kernel_dispatch_(PetscDeviceContext dctx, cupmStream_t stream, const PetscScalar *xarr, const Vec yin[], PetscInt size, PetscScalar *results, PetscInt &yidx) noexcept
1277: {
1278: PetscFunctionBegin;
1279: PetscCall(MDot_kernel_dispatch_(dctx, stream, xarr, yin + yidx, size, results + yidx * MDOT_WORKGROUP_NUM, util::make_index_sequence<N>{}));
1280: yidx += N;
1281: PetscFunctionReturn(PETSC_SUCCESS);
1282: }
1284: template <device::cupm::DeviceType T>
1285: inline PetscErrorCode VecSeq_CUPM<T>::MDot_(std::false_type, Vec xin, PetscInt nv, const Vec yin[], PetscScalar *z, PetscDeviceContext dctx) noexcept
1286: {
1287: // the largest possible size of a batch
1288: constexpr PetscInt batchsize = 8;
1289: // how many sub streams to create, if nv <= batchsize we can do this without looping, so we
1290: // do not create substreams. Note we don't create more than 8 streams, in practice we could
1291: // not get more parallelism with higher numbers.
1292: const auto num_sub_streams = nv > batchsize ? std::min((nv + batchsize) / batchsize, batchsize) : 0;
1293: const auto n = xin->map->n;
1294: const auto nwork = nv * MDOT_WORKGROUP_NUM;
1295: PetscScalar *d_results;
1296: cupmStream_t stream;
1298: PetscFunctionBegin;
1299: PetscCall(GetHandlesFrom_(dctx, &stream));
1300: // allocate scratchpad memory for the results of individual work groups
1301: PetscCall(PetscDeviceMalloc(dctx, PETSC_MEMTYPE_CUPM(), nwork, &d_results));
1302: {
1303: const auto xptr = DeviceArrayRead(dctx, xin);
1304: PetscInt yidx = 0;
1305: auto subidx = 0;
1306: auto cur_stream = stream;
1307: auto cur_ctx = dctx;
1308: PetscDeviceContext *sub = nullptr;
1309: PetscStreamType stype;
1311: // REVIEW ME: maybe PetscDeviceContextFork() should insert dctx into the first entry of
1312: // sub. Ideally the parent context should also join in on the fork, but it is extremely
1313: // fiddly to do so presently
1314: PetscCall(PetscDeviceContextGetStreamType(dctx, &stype));
1315: if (stype == PETSC_STREAM_DEFAULT || stype == PETSC_STREAM_DEFAULT_WITH_BARRIER) stype = PETSC_STREAM_NONBLOCKING;
1316: if (num_sub_streams) {
1317: // Upload inputs before forking so child contexts cannot race on repeated vectors.
1318: for (PetscInt i = 0; i < nv; ++i) PetscCall(CopyToDevice_(dctx, yin[i]));
1319: }
1320: // If we have a default stream create nonblocking streams instead (as we can
1321: // locally exploit the parallelism). Otherwise use the prescribed stream type.
1322: PetscCall(PetscDeviceContextForkWithStreamType(dctx, stype, num_sub_streams, &sub));
1323: PetscCall(PetscLogGpuTimeBegin());
1324: do {
1325: if (num_sub_streams) {
1326: cur_ctx = sub[subidx++ % num_sub_streams];
1327: PetscCall(GetHandlesFrom_(cur_ctx, &cur_stream));
1328: }
1329: // REVIEW ME: Should probably try and load-balance these. Consider the case where nv = 9;
1330: // it is very likely better to do 4+5 rather than 8+1
1331: switch (nv - yidx) {
1332: case 7:
1333: PetscCall(MDot_kernel_dispatch_<7>(cur_ctx, cur_stream, xptr.data(), yin, n, d_results, yidx));
1334: break;
1335: case 6:
1336: PetscCall(MDot_kernel_dispatch_<6>(cur_ctx, cur_stream, xptr.data(), yin, n, d_results, yidx));
1337: break;
1338: case 5:
1339: PetscCall(MDot_kernel_dispatch_<5>(cur_ctx, cur_stream, xptr.data(), yin, n, d_results, yidx));
1340: break;
1341: case 4:
1342: PetscCall(MDot_kernel_dispatch_<4>(cur_ctx, cur_stream, xptr.data(), yin, n, d_results, yidx));
1343: break;
1344: case 3:
1345: PetscCall(MDot_kernel_dispatch_<3>(cur_ctx, cur_stream, xptr.data(), yin, n, d_results, yidx));
1346: break;
1347: case 2:
1348: PetscCall(MDot_kernel_dispatch_<2>(cur_ctx, cur_stream, xptr.data(), yin, n, d_results, yidx));
1349: break;
1350: case 1:
1351: PetscCall(MDot_kernel_dispatch_<1>(cur_ctx, cur_stream, xptr.data(), yin, n, d_results, yidx));
1352: break;
1353: default: // 8 or more
1354: PetscCall(MDot_kernel_dispatch_<8>(cur_ctx, cur_stream, xptr.data(), yin, n, d_results, yidx));
1355: break;
1356: }
1357: } while (yidx < nv);
1358: PetscCall(PetscLogGpuTimeEnd());
1359: PetscCall(PetscDeviceContextJoin(dctx, num_sub_streams, PETSC_DEVICE_CONTEXT_JOIN_DESTROY, &sub));
1360: }
1362: PetscCall(PetscCUPMLaunchKernel1D(nv, 0, stream, kernels::sum_kernel, nv, d_results));
1363: // copy result of device reduction to host
1364: PetscCall(PetscCUPMMemcpyAsync(z, d_results, nv, cupmMemcpyDeviceToHost, stream));
1365: // do these now while final reduction is in flight
1366: PetscCall(PetscLogGpuFlops(nwork));
1367: PetscCall(PetscDeviceFree(dctx, d_results));
1368: PetscFunctionReturn(PETSC_SUCCESS);
1369: }
1371: #undef MDOT_WORKGROUP_NUM
1372: #undef MDOT_WORKGROUP_SIZE
1374: template <device::cupm::DeviceType T>
1375: inline PetscErrorCode VecSeq_CUPM<T>::MDot_(std::true_type, Vec xin, PetscInt nv, const Vec yin[], PetscScalar *z, PetscDeviceContext dctx) noexcept
1376: {
1377: // probably not worth it to run more than 8 of these at a time?
1378: const auto n_sub = PetscMin(nv, 8);
1379: const auto n = static_cast<cupmBlasInt_t>(xin->map->n);
1380: const auto xptr = DeviceArrayRead(dctx, xin);
1381: PetscScalar *d_z;
1382: PetscDeviceContext *subctx;
1383: cupmStream_t stream;
1385: PetscFunctionBegin;
1386: PetscCall(GetHandlesFrom_(dctx, &stream));
1387: PetscCall(PetscDeviceMalloc(dctx, PETSC_MEMTYPE_CUPM(), nv, &d_z));
1388: // Upload inputs before forking so child contexts cannot race on repeated vectors.
1389: for (PetscInt i = 0; i < nv; ++i) PetscCall(CopyToDevice_(dctx, yin[i]));
1390: PetscCall(PetscDeviceContextFork(dctx, n_sub, &subctx));
1391: PetscCall(PetscLogGpuTimeBegin());
1392: for (PetscInt i = 0; i < nv; ++i) {
1393: const auto sub = subctx[i % n_sub];
1394: cupmBlasHandle_t handle;
1396: PetscCall(GetHandlesFrom_(sub, &handle));
1397: {
1398: const auto pointer_mode = CUPMBlasPointerModeGuard{handle, CUPMBLAS_POINTER_MODE_DEVICE};
1400: PetscCallCUPMBLAS(cupmBlasXdot(handle, n, DeviceArrayRead(sub, yin[i]), 1, xptr.cupmdata(), 1, cupmScalarPtrCast(d_z + i)));
1401: }
1402: }
1403: PetscCall(PetscLogGpuTimeEnd());
1404: PetscCall(PetscDeviceContextJoin(dctx, n_sub, PETSC_DEVICE_CONTEXT_JOIN_DESTROY, &subctx));
1405: PetscCall(PetscCUPMMemcpyAsync(z, d_z, nv, cupmMemcpyDeviceToHost, stream));
1406: PetscCall(PetscDeviceFree(dctx, d_z));
1407: // REVIEW ME: flops?????
1408: PetscFunctionReturn(PETSC_SUCCESS);
1409: }
1411: // v->ops->mdot
1412: template <device::cupm::DeviceType T>
1413: inline PetscErrorCode VecSeq_CUPM<T>::MDot(Vec xin, PetscInt nv, const Vec yin[], PetscScalar *z) noexcept
1414: {
1415: PetscFunctionBegin;
1416: if (PetscUnlikely(nv == 1)) {
1417: // dot handles nv = 0 correctly
1418: PetscCall(Dot(xin, const_cast<Vec>(yin[0]), z));
1419: } else if (const auto n = xin->map->n) {
1420: PetscDeviceContext dctx;
1422: PetscCheck(nv > 0, PETSC_COMM_SELF, PETSC_ERR_LIB, "Number of vectors provided to %s %" PetscInt_FMT " not positive", PETSC_FUNCTION_NAME, nv);
1423: PetscCall(GetHandles_(&dctx));
1424: PetscCall(MDot_(std::integral_constant<bool, PetscDefined(USE_COMPLEX)>{}, xin, nv, yin, z, dctx));
1425: // REVIEW ME: double count of flops??
1426: PetscCall(PetscLogGpuFlops(nv * (2 * n - 1)));
1427: PetscCall(PetscDeviceContextSynchronize(dctx));
1428: } else {
1429: PetscCall(PetscArrayzero(z, nv));
1430: }
1431: PetscFunctionReturn(PETSC_SUCCESS);
1432: }
1434: // VecSetAsync_Private
1435: template <device::cupm::DeviceType T>
1436: inline PetscErrorCode VecSeq_CUPM<T>::SetAsync(Vec xin, PetscScalar alpha, PetscDeviceContext dctx) noexcept
1437: {
1438: const auto n = xin->map->n;
1439: cupmStream_t stream;
1441: PetscFunctionBegin;
1442: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
1443: PetscCall(GetHandlesFrom_(dctx, &stream));
1444: {
1445: const auto xptr = DeviceArrayWrite(dctx, xin);
1447: if (alpha == PetscScalar(0.0)) {
1448: PetscCall(PetscCUPMMemsetAsync(xptr.data(), 0, n, stream));
1449: } else {
1450: const auto dptr = thrust::device_pointer_cast(xptr.data());
1452: PetscCallThrust(THRUST_CALL(thrust::fill, stream, dptr, dptr + n, alpha));
1453: }
1454: }
1455: if (n > 0) PetscCall(PetscDeviceContextSynchronizeIfWithBarrier_Internal(dctx));
1456: PetscFunctionReturn(PETSC_SUCCESS);
1457: }
1459: // v->ops->set
1460: template <device::cupm::DeviceType T>
1461: inline PetscErrorCode VecSeq_CUPM<T>::Set(Vec xin, PetscScalar alpha) noexcept
1462: {
1463: PetscFunctionBegin;
1464: PetscCall(SetAsync(xin, alpha, nullptr));
1465: PetscFunctionReturn(PETSC_SUCCESS);
1466: }
1468: // VecScaleAsync_Private
1469: template <device::cupm::DeviceType T>
1470: inline PetscErrorCode VecSeq_CUPM<T>::ScaleAsync(Vec xin, PetscScalar alpha, PetscDeviceContext dctx) noexcept
1471: {
1472: PetscFunctionBegin;
1473: if (PetscUnlikely(alpha == PetscScalar(1.0))) PetscFunctionReturn(PETSC_SUCCESS);
1474: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
1475: if (PetscUnlikely(alpha == PetscScalar(0.0))) {
1476: PetscCall(SetAsync(xin, alpha, dctx));
1477: } else if (const auto n = static_cast<cupmBlasInt_t>(xin->map->n)) {
1478: cupmBlasHandle_t cupmBlasHandle;
1480: PetscCall(GetHandlesFrom_(dctx, &cupmBlasHandle));
1481: PetscCall(PetscLogGpuTimeBegin());
1482: PetscCallCUPMBLAS(cupmBlasXscal(cupmBlasHandle, n, cupmScalarPtrCast(&alpha), DeviceArrayReadWrite(dctx, xin), 1));
1483: PetscCall(PetscLogGpuTimeEnd());
1484: PetscCall(PetscLogGpuFlops(n));
1485: PetscCall(PetscDeviceContextSynchronizeIfWithBarrier_Internal(dctx));
1486: } else {
1487: PetscCall(MaybeIncrementEmptyLocalVec(xin));
1488: }
1489: PetscFunctionReturn(PETSC_SUCCESS);
1490: }
1492: // v->ops->scale
1493: template <device::cupm::DeviceType T>
1494: inline PetscErrorCode VecSeq_CUPM<T>::Scale(Vec xin, PetscScalar alpha) noexcept
1495: {
1496: PetscFunctionBegin;
1497: PetscCall(ScaleAsync(xin, alpha, nullptr));
1498: PetscFunctionReturn(PETSC_SUCCESS);
1499: }
1501: // v->ops->tdot
1502: template <device::cupm::DeviceType T>
1503: inline PetscErrorCode VecSeq_CUPM<T>::TDot(Vec xin, Vec yin, PetscScalar *z) noexcept
1504: {
1505: PetscBool yiscupm;
1507: PetscFunctionBegin;
1508: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(yin), &yiscupm, VECSEQCUPM(), VECMPICUPM(), ""));
1509: if (!yiscupm) {
1510: PetscCall(VecTDot_Seq(xin, yin, z));
1511: PetscFunctionReturn(PETSC_SUCCESS);
1512: }
1513: if (const auto n = static_cast<cupmBlasInt_t>(xin->map->n)) {
1514: PetscDeviceContext dctx;
1515: cupmBlasHandle_t cupmBlasHandle;
1517: PetscCall(GetHandles_(&dctx, &cupmBlasHandle));
1518: PetscCall(PetscLogGpuTimeBegin());
1519: PetscCallCUPMBLAS(cupmBlasXdotu(cupmBlasHandle, n, DeviceArrayRead(dctx, xin), 1, DeviceArrayRead(dctx, yin), 1, cupmScalarPtrCast(z)));
1520: PetscCall(PetscLogGpuTimeEnd());
1521: PetscCall(PetscLogGpuFlops(2 * n - 1));
1522: } else {
1523: *z = 0.0;
1524: }
1525: PetscFunctionReturn(PETSC_SUCCESS);
1526: }
1528: // VecCopyAsync_Private
1529: template <device::cupm::DeviceType T>
1530: inline PetscErrorCode VecSeq_CUPM<T>::CopyAsync(Vec xin, Vec yout, PetscDeviceContext dctx) noexcept
1531: {
1532: PetscFunctionBegin;
1533: if (xin == yout) PetscFunctionReturn(PETSC_SUCCESS);
1534: if (const auto n = xin->map->n) {
1535: const auto xmask = xin->offloadmask;
1536: // silence buggy gcc warning: mode may be used uninitialized in this function
1537: auto mode = cupmMemcpyDeviceToDevice;
1538: cupmStream_t stream;
1540: // translate from PetscOffloadMask to cupmMemcpyKind
1541: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
1542: switch (const auto ymask = yout->offloadmask) {
1543: case PETSC_OFFLOAD_CPU:
1544: case PETSC_OFFLOAD_UNALLOCATED: {
1545: PetscBool yiscupm;
1547: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(yout), &yiscupm, VECSEQCUPM(), VECMPICUPM(), ""));
1548: if (yiscupm && !yout->boundtocpu) {
1549: /* If GPU vector, also ensure output is on GPU unless explicitly bound to CPU */
1550: mode = PetscOffloadDevice(xmask) ? cupmMemcpyDeviceToDevice : cupmMemcpyHostToDevice;
1551: } else {
1552: mode = PetscOffloadDevice(xmask) ? cupmMemcpyDeviceToHost : cupmMemcpyHostToHost;
1553: }
1554: break;
1555: }
1556: case PETSC_OFFLOAD_BOTH:
1557: case PETSC_OFFLOAD_GPU:
1558: mode = PetscOffloadDevice(xmask) ? cupmMemcpyDeviceToDevice : cupmMemcpyHostToDevice;
1559: break;
1560: default:
1561: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Incompatible offload mask %s", PetscOffloadMaskToString(ymask));
1562: }
1564: PetscCall(GetHandlesFrom_(dctx, &stream));
1565: switch (mode) {
1566: case cupmMemcpyDeviceToDevice: // the best case
1567: case cupmMemcpyHostToDevice: { // not terrible
1568: const auto yptr = DeviceArrayWrite(dctx, yout);
1569: const auto xptr = mode == cupmMemcpyDeviceToDevice ? DeviceArrayRead(dctx, xin).data() : HostArrayRead(dctx, xin).data();
1571: PetscCall(PetscLogGpuTimeBegin());
1572: PetscCall(PetscCUPMMemcpyAsync(yptr.data(), xptr, n, mode, stream));
1573: PetscCall(PetscLogGpuTimeEnd());
1574: } break;
1575: case cupmMemcpyDeviceToHost: // not great
1576: case cupmMemcpyHostToHost: { // worst case
1577: const auto xptr = mode == cupmMemcpyDeviceToHost ? DeviceArrayRead(dctx, xin).data() : HostArrayRead(dctx, xin).data();
1578: PetscScalar *yptr;
1580: PetscCall(VecGetArrayWrite(yout, &yptr));
1581: if (mode == cupmMemcpyDeviceToHost) PetscCall(PetscLogGpuTimeBegin());
1582: PetscCall(PetscCUPMMemcpyAsync(yptr, xptr, n, mode, stream, /* force async */ true));
1583: if (mode == cupmMemcpyDeviceToHost) PetscCall(PetscLogGpuTimeEnd());
1584: PetscCall(VecRestoreArrayWrite(yout, &yptr));
1585: } break;
1586: default:
1587: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_GPU, "Unknown cupmMemcpyKind %d", static_cast<int>(mode));
1588: }
1589: // Host-to-host copies are synchronous, even through the Async copy API.
1590: // Complete device-to-host copies before CPU operations can access yout.
1591: if (mode == cupmMemcpyDeviceToHost) PetscCall(PetscDeviceContextSynchronize(dctx));
1592: else PetscCall(PetscDeviceContextSynchronizeIfWithBarrier_Internal(dctx));
1593: } else {
1594: PetscCall(MaybeIncrementEmptyLocalVec(yout));
1595: }
1596: PetscFunctionReturn(PETSC_SUCCESS);
1597: }
1599: // v->ops->copy
1600: template <device::cupm::DeviceType T>
1601: inline PetscErrorCode VecSeq_CUPM<T>::Copy(Vec xin, Vec yout) noexcept
1602: {
1603: PetscFunctionBegin;
1604: PetscCall(CopyAsync(xin, yout, nullptr));
1605: PetscFunctionReturn(PETSC_SUCCESS);
1606: }
1608: // VecSwapAsync_Private
1609: template <device::cupm::DeviceType T>
1610: inline PetscErrorCode VecSeq_CUPM<T>::SwapAsync(Vec xin, Vec yin, PetscDeviceContext dctx) noexcept
1611: {
1612: PetscBool yiscupm;
1614: PetscFunctionBegin;
1615: if (xin == yin) PetscFunctionReturn(PETSC_SUCCESS);
1616: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(yin), &yiscupm, VECSEQCUPM(), VECMPICUPM(), ""));
1617: PetscCheck(yiscupm, PetscObjectComm(PetscObjectCast(yin)), PETSC_ERR_SUP, "Cannot swap with Y of type %s", PetscObjectCast(yin)->type_name);
1618: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
1619: if (const auto n = static_cast<cupmBlasInt_t>(xin->map->n)) {
1620: cupmBlasHandle_t cupmBlasHandle;
1622: PetscCall(GetHandlesFrom_(dctx, &cupmBlasHandle));
1623: PetscCall(PetscLogGpuTimeBegin());
1624: PetscCallCUPMBLAS(cupmBlasXswap(cupmBlasHandle, n, DeviceArrayReadWrite(dctx, xin), 1, DeviceArrayReadWrite(dctx, yin), 1));
1625: PetscCall(PetscLogGpuTimeEnd());
1626: PetscCall(PetscDeviceContextSynchronizeIfWithBarrier_Internal(dctx));
1627: } else {
1628: PetscCall(MaybeIncrementEmptyLocalVec(xin));
1629: PetscCall(MaybeIncrementEmptyLocalVec(yin));
1630: }
1631: PetscFunctionReturn(PETSC_SUCCESS);
1632: }
1634: // v->ops->swap
1635: template <device::cupm::DeviceType T>
1636: inline PetscErrorCode VecSeq_CUPM<T>::Swap(Vec xin, Vec yin) noexcept
1637: {
1638: PetscFunctionBegin;
1639: PetscCall(SwapAsync(xin, yin, nullptr));
1640: PetscFunctionReturn(PETSC_SUCCESS);
1641: }
1643: // VecAXPYBYAsync_Private
1644: template <device::cupm::DeviceType T>
1645: inline PetscErrorCode VecSeq_CUPM<T>::AXPBYAsync(Vec yin, PetscScalar alpha, PetscScalar beta, Vec xin, PetscDeviceContext dctx) noexcept
1646: {
1647: PetscBool xiscupm;
1649: PetscFunctionBegin;
1650: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(xin), &xiscupm, VECSEQCUPM(), VECMPICUPM(), ""));
1651: if (!xiscupm) {
1652: PetscCall(VecAXPBY_Seq(yin, alpha, beta, xin));
1653: PetscFunctionReturn(PETSC_SUCCESS);
1654: }
1655: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
1656: if (alpha == PetscScalar(0.0)) {
1657: PetscCall(ScaleAsync(yin, beta, dctx));
1658: } else if (beta == PetscScalar(1.0)) {
1659: PetscCall(AXPYAsync(yin, alpha, xin, dctx));
1660: } else if (alpha == PetscScalar(1.0)) {
1661: PetscCall(AYPXAsync(yin, beta, xin, dctx));
1662: } else if (const auto n = static_cast<cupmBlasInt_t>(yin->map->n)) {
1663: PetscBool xiscupm;
1665: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(xin), &xiscupm, VECSEQCUPM(), VECMPICUPM(), ""));
1666: if (!xiscupm) {
1667: PetscCall(VecAXPBY_Seq(yin, alpha, beta, xin));
1668: PetscFunctionReturn(PETSC_SUCCESS);
1669: }
1671: const auto betaIsZero = beta == PetscScalar(0.0);
1672: const auto aptr = cupmScalarPtrCast(&alpha);
1673: cupmBlasHandle_t cupmBlasHandle;
1675: PetscCall(GetHandlesFrom_(dctx, &cupmBlasHandle));
1676: {
1677: const auto xptr = DeviceArrayRead(dctx, xin);
1679: if (betaIsZero /* beta = 0 */) {
1680: // here we can get away with purely write-only as we memcpy into it first
1681: const auto yptr = DeviceArrayWrite(dctx, yin);
1682: cupmStream_t stream;
1684: PetscCall(GetHandlesFrom_(dctx, &stream));
1685: PetscCall(PetscLogGpuTimeBegin());
1686: PetscCall(PetscCUPMMemcpyAsync(yptr.data(), xptr.data(), n, cupmMemcpyDeviceToDevice, stream));
1687: PetscCallCUPMBLAS(cupmBlasXscal(cupmBlasHandle, n, aptr, yptr.cupmdata(), 1));
1688: } else {
1689: const auto yptr = DeviceArrayReadWrite(dctx, yin);
1691: PetscCall(PetscLogGpuTimeBegin());
1692: PetscCallCUPMBLAS(cupmBlasXscal(cupmBlasHandle, n, cupmScalarPtrCast(&beta), yptr.cupmdata(), 1));
1693: PetscCallCUPMBLAS(cupmBlasXaxpy(cupmBlasHandle, n, aptr, xptr.cupmdata(), 1, yptr.cupmdata(), 1));
1694: }
1695: }
1696: PetscCall(PetscLogGpuTimeEnd());
1697: PetscCall(PetscLogGpuFlops((betaIsZero ? 1 : 3) * n));
1698: PetscCall(PetscDeviceContextSynchronizeIfWithBarrier_Internal(dctx));
1699: } else {
1700: PetscCall(MaybeIncrementEmptyLocalVec(yin));
1701: }
1702: PetscFunctionReturn(PETSC_SUCCESS);
1703: }
1705: // v->ops->axpby
1706: template <device::cupm::DeviceType T>
1707: inline PetscErrorCode VecSeq_CUPM<T>::AXPBY(Vec yin, PetscScalar alpha, PetscScalar beta, Vec xin) noexcept
1708: {
1709: PetscFunctionBegin;
1710: PetscCall(AXPBYAsync(yin, alpha, beta, xin, nullptr));
1711: PetscFunctionReturn(PETSC_SUCCESS);
1712: }
1714: // VecAXPBYPCZAsync_Private
1715: template <device::cupm::DeviceType T>
1716: inline PetscErrorCode VecSeq_CUPM<T>::AXPBYPCZAsync(Vec zin, PetscScalar alpha, PetscScalar beta, PetscScalar gamma, Vec xin, Vec yin, PetscDeviceContext dctx) noexcept
1717: {
1718: PetscFunctionBegin;
1719: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
1720: if (gamma != PetscScalar(1.0)) PetscCall(ScaleAsync(zin, gamma, dctx));
1721: PetscCall(AXPYAsync(zin, alpha, xin, dctx));
1722: PetscCall(AXPYAsync(zin, beta, yin, dctx));
1723: PetscFunctionReturn(PETSC_SUCCESS);
1724: }
1726: // v->ops->axpbypcz
1727: template <device::cupm::DeviceType T>
1728: inline PetscErrorCode VecSeq_CUPM<T>::AXPBYPCZ(Vec zin, PetscScalar alpha, PetscScalar beta, PetscScalar gamma, Vec xin, Vec yin) noexcept
1729: {
1730: PetscFunctionBegin;
1731: PetscCall(AXPBYPCZAsync(zin, alpha, beta, gamma, xin, yin, nullptr));
1732: PetscFunctionReturn(PETSC_SUCCESS);
1733: }
1735: // v->ops->norm
1736: template <device::cupm::DeviceType T>
1737: inline PetscErrorCode VecSeq_CUPM<T>::Norm(Vec xin, NormType type, PetscReal *z) noexcept
1738: {
1739: PetscDeviceContext dctx;
1740: cupmBlasHandle_t cupmBlasHandle;
1741: cupmStream_t stream;
1743: PetscFunctionBegin;
1744: PetscCall(GetHandles_(&dctx, &cupmBlasHandle, &stream));
1745: if (const auto n = static_cast<cupmBlasInt_t>(xin->map->n)) {
1746: const auto xptr = DeviceArrayRead(dctx, xin);
1747: PetscInt flopCount = 0;
1749: PetscCall(PetscLogGpuTimeBegin());
1750: switch (type) {
1751: case NORM_1_AND_2:
1752: case NORM_1:
1753: if (PetscDefined(USE_COMPLEX)) {
1754: const auto dptr = thrust::device_pointer_cast(xptr.data());
1756: // Complex BLAS reductions use |Re(x)| + |Im(x)|, but VecNorm() uses |x|.
1757: PetscCallThrust(*z = THRUST_CALL(thrust::transform_reduce, stream, dptr, dptr + n, detail::AbsoluteValue{}, PetscReal{0.0}, thrust::plus<PetscReal>{}));
1758: } else PetscCallCUPMBLAS(cupmBlasXasum(cupmBlasHandle, n, xptr.cupmdata(), 1, cupmRealPtrCast(z)));
1759: flopCount = std::max(n - 1, 0);
1760: if (type == NORM_1) break;
1761: ++z; // fall-through
1762: #if PETSC_CPP_VERSION >= 17
1763: [[fallthrough]];
1764: #endif
1765: case NORM_2:
1766: case NORM_FROBENIUS:
1767: PetscCallCUPMBLAS(cupmBlasXnrm2(cupmBlasHandle, n, xptr.cupmdata(), 1, cupmRealPtrCast(z)));
1768: flopCount += std::max(2 * n - 1, 0); // += in case we've fallen through from NORM_1_AND_2
1769: break;
1770: case NORM_INFINITY:
1771: if (PetscDefined(USE_COMPLEX)) {
1772: const auto dptr = thrust::device_pointer_cast(xptr.data());
1774: PetscCallThrust(*z = THRUST_CALL(thrust::transform_reduce, stream, dptr, dptr + n, detail::AbsoluteValue{}, PetscReal{0.0}, detail::NormMax{}));
1775: } else {
1776: cupmBlasInt_t max_loc = 0;
1777: PetscScalar xv = 0.;
1779: PetscCallCUPMBLAS(cupmBlasXamax(cupmBlasHandle, n, xptr.cupmdata(), 1, &max_loc));
1780: PetscCall(PetscCUPMMemcpyAsync(&xv, xptr.data() + max_loc - 1, 1, cupmMemcpyDeviceToHost, stream));
1781: PetscCall(PetscDeviceContextSynchronize(dctx));
1782: *z = PetscAbsScalar(xv);
1783: }
1784: // REVIEW ME: flopCount = ???
1785: break;
1786: }
1787: PetscCall(PetscLogGpuTimeEnd());
1788: PetscCall(PetscLogGpuFlops(flopCount));
1789: } else {
1790: z[0] = 0.0;
1791: z[type == NORM_1_AND_2] = 0.0;
1792: }
1793: PetscFunctionReturn(PETSC_SUCCESS);
1794: }
1796: namespace detail
1797: {
1799: template <NormType wnormtype>
1800: class ErrorWNormTransformBase {
1801: public:
1802: using result_type = thrust::tuple<PetscReal, PetscReal, PetscReal, PetscInt, PetscInt, PetscInt>;
1804: constexpr explicit ErrorWNormTransformBase(PetscReal v) noexcept : ignore_max_{v} { }
1806: protected:
1807: struct NormTuple {
1808: PetscReal norm;
1809: PetscInt loc;
1810: };
1812: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL static NormTuple compute_norm_(PetscReal err, PetscReal tol) noexcept
1813: {
1814: if (tol > 0.) {
1815: const auto val = err / tol;
1817: return {wnormtype == NORM_INFINITY ? val : PetscSqr(val), 1};
1818: } else {
1819: return {0.0, 0};
1820: }
1821: }
1823: PetscReal ignore_max_;
1824: };
1826: template <NormType wnormtype>
1827: struct ErrorWNormTransform : ErrorWNormTransformBase<wnormtype> {
1828: using base_type = ErrorWNormTransformBase<wnormtype>;
1829: using result_type = typename base_type::result_type;
1830: using argument_type = thrust::tuple<PetscScalar, PetscScalar, PetscScalar, PetscScalar>;
1832: using base_type::base_type;
1834: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL result_type operator()(const argument_type &x) const noexcept
1835: {
1836: const auto u = thrust::get<0>(x); // with x.get<0>(), cuda-12.4.0 gives error: class "cuda::std::__4::tuple" has no member "get"
1837: const auto y = thrust::get<1>(x);
1838: const auto au = PetscAbsScalar(u);
1839: const auto ay = PetscAbsScalar(y);
1840: const auto skip = au < this->ignore_max_ || ay < this->ignore_max_;
1841: const auto tola = skip ? 0.0 : PetscRealPart(thrust::get<2>(x));
1842: const auto tolr = skip ? 0.0 : PetscRealPart(thrust::get<3>(x)) * PetscMax(au, ay);
1843: const auto tol = tola + tolr;
1844: const auto err = PetscAbsScalar(u - y);
1845: const auto tup_a = this->compute_norm_(err, tola);
1846: const auto tup_r = this->compute_norm_(err, tolr);
1847: const auto tup_n = this->compute_norm_(err, tol);
1849: return {tup_n.norm, tup_a.norm, tup_r.norm, tup_n.loc, tup_a.loc, tup_r.loc};
1850: }
1851: };
1853: template <NormType wnormtype>
1854: struct ErrorWNormETransform : ErrorWNormTransformBase<wnormtype> {
1855: using base_type = ErrorWNormTransformBase<wnormtype>;
1856: using result_type = typename base_type::result_type;
1857: using argument_type = thrust::tuple<PetscScalar, PetscScalar, PetscScalar, PetscScalar, PetscScalar>;
1859: using base_type::base_type;
1861: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL result_type operator()(const argument_type &x) const noexcept
1862: {
1863: const auto au = PetscAbsScalar(thrust::get<0>(x));
1864: const auto ay = PetscAbsScalar(thrust::get<1>(x));
1865: const auto skip = au < this->ignore_max_ || ay < this->ignore_max_;
1866: const auto tola = skip ? 0.0 : PetscRealPart(thrust::get<3>(x));
1867: const auto tolr = skip ? 0.0 : PetscRealPart(thrust::get<4>(x)) * PetscMax(au, ay);
1868: const auto tol = tola + tolr;
1869: const auto err = PetscAbsScalar(thrust::get<2>(x));
1870: const auto tup_a = this->compute_norm_(err, tola);
1871: const auto tup_r = this->compute_norm_(err, tolr);
1872: const auto tup_n = this->compute_norm_(err, tol);
1874: return {tup_n.norm, tup_a.norm, tup_r.norm, tup_n.loc, tup_a.loc, tup_r.loc};
1875: }
1876: };
1878: template <NormType wnormtype>
1879: struct ErrorWNormReduce {
1880: using value_type = typename ErrorWNormTransformBase<wnormtype>::result_type;
1882: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL value_type operator()(const value_type &lhs, const value_type &rhs) const noexcept
1883: {
1884: // cannot use lhs.get<0>() etc since the using decl above ambiguates the fact that
1885: // result_type is a template, so in order to fix this we would need to write:
1886: //
1887: // lhs.template get<0>()
1888: //
1889: // which is unseemly.
1890: if (wnormtype == NORM_INFINITY) {
1891: // clang-format off
1892: return {
1893: PetscMax(thrust::get<0>(lhs), thrust::get<0>(rhs)),
1894: PetscMax(thrust::get<1>(lhs), thrust::get<1>(rhs)),
1895: PetscMax(thrust::get<2>(lhs), thrust::get<2>(rhs)),
1896: thrust::get<3>(lhs) + thrust::get<3>(rhs),
1897: thrust::get<4>(lhs) + thrust::get<4>(rhs),
1898: thrust::get<5>(lhs) + thrust::get<5>(rhs)
1899: };
1900: // clang-format on
1901: } else {
1902: // clang-format off
1903: return {
1904: thrust::get<0>(lhs) + thrust::get<0>(rhs),
1905: thrust::get<1>(lhs) + thrust::get<1>(rhs),
1906: thrust::get<2>(lhs) + thrust::get<2>(rhs),
1907: thrust::get<3>(lhs) + thrust::get<3>(rhs),
1908: thrust::get<4>(lhs) + thrust::get<4>(rhs),
1909: thrust::get<5>(lhs) + thrust::get<5>(rhs)
1910: };
1911: // clang-format on
1912: }
1913: }
1914: };
1916: template <template <NormType> class WNormTransformType, typename Tuple, typename cupmStream_t>
1917: inline PetscErrorCode ExecuteWNorm(Tuple &&first, Tuple &&last, NormType wnormtype, cupmStream_t stream, PetscReal ignore_max, PetscReal *norm, PetscInt *norm_loc, PetscReal *norma, PetscInt *norma_loc, PetscReal *normr, PetscInt *normr_loc) noexcept
1918: {
1919: auto begin = thrust::make_zip_iterator(std::forward<Tuple>(first));
1920: auto end = thrust::make_zip_iterator(std::forward<Tuple>(last));
1921: PetscReal n = 0, na = 0, nr = 0;
1922: PetscInt n_loc = 0, na_loc = 0, nr_loc = 0;
1924: PetscFunctionBegin;
1925: // clang-format off
1926: if (wnormtype == NORM_INFINITY) {
1927: PetscCallThrust(
1928: thrust::tie(*norm, *norma, *normr, *norm_loc, *norma_loc, *normr_loc) = THRUST_CALL(
1929: thrust::transform_reduce,
1930: stream,
1931: std::move(begin),
1932: std::move(end),
1933: WNormTransformType<NORM_INFINITY>{ignore_max},
1934: thrust::make_tuple(n, na, nr, n_loc, na_loc, nr_loc),
1935: ErrorWNormReduce<NORM_INFINITY>{}
1936: )
1937: );
1938: } else {
1939: PetscCallThrust(
1940: thrust::tie(*norm, *norma, *normr, *norm_loc, *norma_loc, *normr_loc) = THRUST_CALL(
1941: thrust::transform_reduce,
1942: stream,
1943: std::move(begin),
1944: std::move(end),
1945: WNormTransformType<NORM_2>{ignore_max},
1946: thrust::make_tuple(n, na, nr, n_loc, na_loc, nr_loc),
1947: ErrorWNormReduce<NORM_2>{}
1948: )
1949: );
1950: }
1951: // clang-format on
1952: if (wnormtype == NORM_2) {
1953: *norm = PetscSqrtReal(*norm);
1954: *norma = PetscSqrtReal(*norma);
1955: *normr = PetscSqrtReal(*normr);
1956: }
1957: PetscFunctionReturn(PETSC_SUCCESS);
1958: }
1960: } // namespace detail
1962: // v->ops->errorwnorm
1963: template <device::cupm::DeviceType T>
1964: inline PetscErrorCode VecSeq_CUPM<T>::ErrorWnorm(Vec U, Vec Y, Vec E, NormType wnormtype, PetscReal atol, Vec vatol, PetscReal rtol, Vec vrtol, PetscReal ignore_max, PetscReal *norm, PetscInt *norm_loc, PetscReal *norma, PetscInt *norma_loc, PetscReal *normr, PetscInt *normr_loc) noexcept
1965: {
1966: const auto nl = U->map->n;
1967: #if CCCL_VERSION >= 3004000
1968: auto ait = cuda::make_constant_iterator(static_cast<PetscScalar>(atol));
1969: auto rit = cuda::make_constant_iterator(static_cast<PetscScalar>(rtol));
1970: #else
1971: auto ait = thrust::make_constant_iterator(static_cast<PetscScalar>(atol));
1972: auto rit = thrust::make_constant_iterator(static_cast<PetscScalar>(rtol));
1973: #endif
1974: PetscDeviceContext dctx;
1975: cupmStream_t stream;
1977: PetscFunctionBegin;
1978: PetscCall(GetHandles_(&dctx, &stream));
1979: {
1980: const auto ConditionalDeviceArrayRead = [&](Vec v) {
1981: if (v) {
1982: return thrust::device_pointer_cast(DeviceArrayRead(dctx, v).data());
1983: } else {
1984: return thrust::device_ptr<PetscScalar>{nullptr};
1985: }
1986: };
1988: const auto uarr = DeviceArrayRead(dctx, U);
1989: const auto yarr = DeviceArrayRead(dctx, Y);
1990: const auto uptr = thrust::device_pointer_cast(uarr.data());
1991: const auto yptr = thrust::device_pointer_cast(yarr.data());
1992: const auto eptr = ConditionalDeviceArrayRead(E);
1993: const auto rptr = ConditionalDeviceArrayRead(vrtol);
1994: const auto aptr = ConditionalDeviceArrayRead(vatol);
1996: if (!vatol && !vrtol) {
1997: if (E) {
1998: // clang-format off
1999: PetscCall(
2000: detail::ExecuteWNorm<detail::ErrorWNormETransform>(
2001: thrust::make_tuple(uptr, yptr, eptr, ait, rit),
2002: thrust::make_tuple(uptr + nl, yptr + nl, eptr + nl, ait, rit),
2003: wnormtype, stream, ignore_max, norm, norm_loc, norma, norma_loc, normr, normr_loc
2004: )
2005: );
2006: // clang-format on
2007: } else {
2008: // clang-format off
2009: PetscCall(
2010: detail::ExecuteWNorm<detail::ErrorWNormTransform>(
2011: thrust::make_tuple(uptr, yptr, ait, rit),
2012: thrust::make_tuple(uptr + nl, yptr + nl, ait, rit),
2013: wnormtype, stream, ignore_max, norm, norm_loc, norma, norma_loc, normr, normr_loc
2014: )
2015: );
2016: // clang-format on
2017: }
2018: } else if (!vatol) {
2019: if (E) {
2020: // clang-format off
2021: PetscCall(
2022: detail::ExecuteWNorm<detail::ErrorWNormETransform>(
2023: thrust::make_tuple(uptr, yptr, eptr, ait, rptr),
2024: thrust::make_tuple(uptr + nl, yptr + nl, eptr + nl, ait, rptr + nl),
2025: wnormtype, stream, ignore_max, norm, norm_loc, norma, norma_loc, normr, normr_loc
2026: )
2027: );
2028: // clang-format on
2029: } else {
2030: // clang-format off
2031: PetscCall(
2032: detail::ExecuteWNorm<detail::ErrorWNormTransform>(
2033: thrust::make_tuple(uptr, yptr, ait, rptr),
2034: thrust::make_tuple(uptr + nl, yptr + nl, ait, rptr + nl),
2035: wnormtype, stream, ignore_max, norm, norm_loc, norma, norma_loc, normr, normr_loc
2036: )
2037: );
2038: // clang-format on
2039: }
2040: } else if (!vrtol) {
2041: if (E) {
2042: // clang-format off
2043: PetscCall(
2044: detail::ExecuteWNorm<detail::ErrorWNormETransform>(
2045: thrust::make_tuple(uptr, yptr, eptr, aptr, rit),
2046: thrust::make_tuple(uptr + nl, yptr + nl, eptr + nl, aptr + nl, rit),
2047: wnormtype, stream, ignore_max, norm, norm_loc, norma, norma_loc, normr, normr_loc
2048: )
2049: );
2050: // clang-format on
2051: } else {
2052: // clang-format off
2053: PetscCall(
2054: detail::ExecuteWNorm<detail::ErrorWNormTransform>(
2055: thrust::make_tuple(uptr, yptr, aptr, rit),
2056: thrust::make_tuple(uptr + nl, yptr + nl, aptr + nl, rit),
2057: wnormtype, stream, ignore_max, norm, norm_loc, norma, norma_loc, normr, normr_loc
2058: )
2059: );
2060: // clang-format on
2061: }
2062: } else {
2063: if (E) {
2064: // clang-format off
2065: PetscCall(
2066: detail::ExecuteWNorm<detail::ErrorWNormETransform>(
2067: thrust::make_tuple(uptr, yptr, eptr, aptr, rptr),
2068: thrust::make_tuple(uptr + nl, yptr + nl, eptr + nl, aptr + nl, rptr + nl),
2069: wnormtype, stream, ignore_max, norm, norm_loc, norma, norma_loc, normr, normr_loc
2070: )
2071: );
2072: // clang-format on
2073: } else {
2074: // clang-format off
2075: PetscCall(
2076: detail::ExecuteWNorm<detail::ErrorWNormTransform>(
2077: thrust::make_tuple(uptr, yptr, aptr, rptr),
2078: thrust::make_tuple(uptr + nl, yptr + nl, aptr + nl, rptr + nl),
2079: wnormtype, stream, ignore_max, norm, norm_loc, norma, norma_loc, normr, normr_loc
2080: )
2081: );
2082: // clang-format on
2083: }
2084: }
2085: }
2086: PetscFunctionReturn(PETSC_SUCCESS);
2087: }
2089: namespace detail
2090: {
2091: struct dotnorm2_mult {
2092: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL thrust::tuple<PetscScalar, PetscScalar> operator()(const PetscScalar &s, const PetscScalar &t) const noexcept
2093: {
2094: const auto conjt = PetscConj(t);
2096: return {s * conjt, t * conjt};
2097: }
2098: };
2100: // it is positively __bananas__ that thrust does not define default operator+ for tuples... I
2101: // would do it myself but now I am worried that they do so on purpose...
2102: struct dotnorm2_tuple_plus {
2103: using value_type = thrust::tuple<PetscScalar, PetscScalar>;
2105: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL value_type operator()(const value_type &lhs, const value_type &rhs) const noexcept { return {thrust::get<0>(lhs) + thrust::get<0>(rhs), thrust::get<1>(lhs) + thrust::get<1>(rhs)}; }
2106: };
2108: } // namespace detail
2110: // v->ops->dotnorm2
2111: template <device::cupm::DeviceType T>
2112: inline PetscErrorCode VecSeq_CUPM<T>::DotNorm2(Vec s, Vec t, PetscScalar *dp, PetscScalar *nm) noexcept
2113: {
2114: PetscDeviceContext dctx;
2115: cupmStream_t stream;
2117: PetscFunctionBegin;
2118: PetscCall(GetHandles_(&dctx, &stream));
2119: {
2120: PetscScalar dpt = 0.0, nmt = 0.0;
2121: const auto sdptr = thrust::device_pointer_cast(DeviceArrayRead(dctx, s).data());
2123: // clang-format off
2124: PetscCallThrust(
2125: thrust::tie(*dp, *nm) = THRUST_CALL(
2126: thrust::inner_product,
2127: stream,
2128: sdptr, sdptr+s->map->n, thrust::device_pointer_cast(DeviceArrayRead(dctx, t).data()),
2129: thrust::make_tuple(dpt, nmt),
2130: detail::dotnorm2_tuple_plus{}, detail::dotnorm2_mult{}
2131: );
2132: );
2133: // clang-format on
2134: }
2135: PetscFunctionReturn(PETSC_SUCCESS);
2136: }
2138: namespace detail
2139: {
2140: struct conjugate {
2141: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &x) const noexcept { return PetscConj(x); }
2142: };
2144: } // namespace detail
2146: // v->ops->conjugate
2147: template <device::cupm::DeviceType T>
2148: inline PetscErrorCode VecSeq_CUPM<T>::ConjugateAsync(Vec xin, PetscDeviceContext dctx) noexcept
2149: {
2150: PetscFunctionBegin;
2151: if (PetscDefined(USE_COMPLEX)) PetscCall(PointwiseUnary_(detail::conjugate{}, xin, nullptr, dctx));
2152: PetscFunctionReturn(PETSC_SUCCESS);
2153: }
2155: // v->ops->conjugate
2156: template <device::cupm::DeviceType T>
2157: inline PetscErrorCode VecSeq_CUPM<T>::Conjugate(Vec xin) noexcept
2158: {
2159: PetscFunctionBegin;
2160: PetscCall(ConjugateAsync(xin, nullptr));
2161: PetscFunctionReturn(PETSC_SUCCESS);
2162: }
2164: namespace detail
2165: {
2167: struct real_part {
2168: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL thrust::tuple<PetscReal, PetscInt> operator()(const thrust::tuple<PetscScalar, PetscInt> &x) const noexcept { return {PetscRealPart(thrust::get<0>(x)), thrust::get<1>(x)}; }
2170: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL PetscReal operator()(const PetscScalar &x) const noexcept { return PetscRealPart(x); }
2171: };
2173: // deriving from Operator allows us to "store" an instance of the operator in the class but
2174: // also take advantage of empty base class optimization if the operator is stateless
2175: template <typename Operator>
2176: class tuple_compare : Operator {
2177: public:
2178: using tuple_type = thrust::tuple<PetscReal, PetscInt>;
2179: using operator_type = Operator;
2181: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL tuple_type operator()(const tuple_type &x, const tuple_type &y) const noexcept
2182: {
2183: if (op_()(thrust::get<0>(y), thrust::get<0>(x))) {
2184: // if y is strictly greater/less than x, return y
2185: return y;
2186: } else if (thrust::get<0>(y) == thrust::get<0>(x)) {
2187: // if equal, prefer lower index
2188: return thrust::get<1>(y) < thrust::get<1>(x) ? y : x;
2189: }
2190: // otherwise return x
2191: return x;
2192: }
2194: private:
2195: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL const operator_type &op_() const noexcept { return *this; }
2196: };
2198: } // namespace detail
2200: template <device::cupm::DeviceType T>
2201: template <typename TupleFuncT, typename UnaryFuncT>
2202: inline PetscErrorCode VecSeq_CUPM<T>::MinMax_(TupleFuncT &&tuple_ftr, UnaryFuncT &&unary_ftr, Vec v, PetscInt *p, PetscReal *m) noexcept
2203: {
2204: PetscFunctionBegin;
2205: PetscCheckTypeNames(v, VECSEQCUPM(), VECMPICUPM());
2206: if (p) *p = -1;
2207: if (const auto n = v->map->n) {
2208: PetscDeviceContext dctx;
2209: cupmStream_t stream;
2211: PetscCall(GetHandles_(&dctx, &stream));
2212: // needed to:
2213: // 1. switch between transform_reduce and reduce
2214: // 2. strip the real_part functor from the arguments
2215: #if PetscDefined(USE_COMPLEX)
2216: #define THRUST_MINMAX_REDUCE(...) THRUST_CALL(thrust::transform_reduce, __VA_ARGS__)
2217: #else
2218: #define THRUST_MINMAX_REDUCE(s, b, e, real_part__, ...) THRUST_CALL(thrust::reduce, s, b, e, __VA_ARGS__)
2219: #endif
2220: {
2221: const auto vptr = thrust::device_pointer_cast(DeviceArrayRead(dctx, v).data());
2223: if (p) {
2224: // clang-format off
2225: const auto zip = thrust::make_zip_iterator(
2226: thrust::make_tuple(std::move(vptr), thrust::make_counting_iterator(PetscInt{0}))
2227: );
2228: // clang-format on
2229: // need to use preprocessor conditionals since otherwise thrust complains about not being
2230: // able to convert a thrust::device_reference to a PetscReal on complex
2231: // builds...
2232: // clang-format off
2233: PetscCallThrust(
2234: thrust::tie(*m, *p) = THRUST_MINMAX_REDUCE(
2235: stream, zip, zip + n, detail::real_part{},
2236: thrust::make_tuple(*m, *p), std::forward<TupleFuncT>(tuple_ftr)
2237: );
2238: );
2239: // clang-format on
2240: } else {
2241: // clang-format off
2242: PetscCallThrust(
2243: *m = THRUST_MINMAX_REDUCE(
2244: stream, vptr, vptr + n, detail::real_part{},
2245: *m, std::forward<UnaryFuncT>(unary_ftr)
2246: );
2247: );
2248: // clang-format on
2249: }
2250: }
2251: #undef THRUST_MINMAX_REDUCE
2252: }
2253: // REVIEW ME: flops?
2254: PetscFunctionReturn(PETSC_SUCCESS);
2255: }
2257: // v->ops->max
2258: template <device::cupm::DeviceType T>
2259: inline PetscErrorCode VecSeq_CUPM<T>::Max(Vec v, PetscInt *p, PetscReal *m) noexcept
2260: {
2261: #if CCCL_VERSION >= 3001000
2262: using tuple_functor = detail::tuple_compare<cuda::std::greater<PetscReal>>;
2263: using unary_functor = cuda::maximum<PetscReal>;
2264: #else
2265: using tuple_functor = detail::tuple_compare<thrust::greater<PetscReal>>;
2266: using unary_functor = thrust::maximum<PetscReal>;
2267: #endif
2269: PetscFunctionBegin;
2270: *m = PETSC_MIN_REAL;
2271: // use {} constructor syntax otherwise most vexing parse
2272: PetscCall(MinMax_(tuple_functor{}, unary_functor{}, v, p, m));
2273: PetscFunctionReturn(PETSC_SUCCESS);
2274: }
2276: // v->ops->min
2277: template <device::cupm::DeviceType T>
2278: inline PetscErrorCode VecSeq_CUPM<T>::Min(Vec v, PetscInt *p, PetscReal *m) noexcept
2279: {
2280: #if CCCL_VERSION >= 3001000
2281: using tuple_functor = detail::tuple_compare<cuda::std::less<PetscReal>>;
2282: using unary_functor = cuda::minimum<PetscReal>;
2283: #else
2284: using tuple_functor = detail::tuple_compare<thrust::less<PetscReal>>;
2285: using unary_functor = thrust::minimum<PetscReal>;
2286: #endif
2288: PetscFunctionBegin;
2289: *m = PETSC_MAX_REAL;
2290: // use {} constructor syntax otherwise most vexing parse
2291: PetscCall(MinMax_(tuple_functor{}, unary_functor{}, v, p, m));
2292: PetscFunctionReturn(PETSC_SUCCESS);
2293: }
2295: // v->ops->sum
2296: template <device::cupm::DeviceType T>
2297: inline PetscErrorCode VecSeq_CUPM<T>::Sum(Vec v, PetscScalar *sum) noexcept
2298: {
2299: PetscFunctionBegin;
2300: if (const auto n = v->map->n) {
2301: PetscDeviceContext dctx;
2302: cupmStream_t stream;
2304: PetscCall(GetHandles_(&dctx, &stream));
2305: const auto dptr = thrust::device_pointer_cast(DeviceArrayRead(dctx, v).data());
2306: // REVIEW ME: why not cupmBlasXasum()?
2307: PetscCallThrust(*sum = THRUST_CALL(thrust::reduce, stream, dptr, dptr + n, PetscScalar{0.0}););
2308: // REVIEW ME: must be at least n additions
2309: PetscCall(PetscLogGpuFlops(n));
2310: } else {
2311: *sum = 0.0;
2312: }
2313: PetscFunctionReturn(PETSC_SUCCESS);
2314: }
2316: template <device::cupm::DeviceType T>
2317: inline PetscErrorCode VecSeq_CUPM<T>::ShiftAsync(Vec v, PetscScalar shift, PetscDeviceContext dctx) noexcept
2318: {
2319: PetscFunctionBegin;
2320: PetscCall(PointwiseUnary_(device::cupm::functors::make_plus_equals(shift), v, nullptr, dctx));
2321: PetscFunctionReturn(PETSC_SUCCESS);
2322: }
2324: template <device::cupm::DeviceType T>
2325: inline PetscErrorCode VecSeq_CUPM<T>::Shift(Vec v, PetscScalar shift) noexcept
2326: {
2327: PetscFunctionBegin;
2328: PetscCall(ShiftAsync(v, shift, nullptr));
2329: PetscFunctionReturn(PETSC_SUCCESS);
2330: }
2332: template <device::cupm::DeviceType T>
2333: inline PetscErrorCode VecSeq_CUPM<T>::SetRandom(Vec v, PetscRandom rand) noexcept
2334: {
2335: PetscFunctionBegin;
2336: if (const auto n = v->map->n) {
2337: PetscBool iscurand;
2338: PetscDeviceContext dctx;
2340: PetscCall(GetHandles_(&dctx));
2341: PetscCall(PetscObjectTypeCompare(PetscObjectCast(rand), PETSCCURAND, &iscurand));
2342: if (iscurand) PetscCall(PetscRandomGetValues(rand, n, DeviceArrayWrite(dctx, v)));
2343: else PetscCall(PetscRandomGetValues(rand, n, HostArrayWrite(dctx, v)));
2344: } else {
2345: PetscCall(MaybeIncrementEmptyLocalVec(v));
2346: }
2347: // REVIEW ME: flops????
2348: // REVIEW ME: Timing???
2349: PetscFunctionReturn(PETSC_SUCCESS);
2350: }
2352: // v->ops->setpreallocation
2353: template <device::cupm::DeviceType T>
2354: inline PetscErrorCode VecSeq_CUPM<T>::SetPreallocationCOO(Vec v, PetscCount ncoo, const PetscInt coo_i[]) noexcept
2355: {
2356: PetscDeviceContext dctx;
2358: PetscFunctionBegin;
2359: PetscCall(GetHandles_(&dctx));
2360: PetscCall(VecSetPreallocationCOO_Seq(v, ncoo, coo_i));
2361: PetscCall(SetPreallocationCOO_CUPMBase(v, ncoo, coo_i, dctx));
2362: PetscFunctionReturn(PETSC_SUCCESS);
2363: }
2365: // v->ops->setvaluescoo
2366: template <device::cupm::DeviceType T>
2367: inline PetscErrorCode VecSeq_CUPM<T>::SetValuesCOO(Vec x, const PetscScalar v[], InsertMode imode) noexcept
2368: {
2369: auto vv = const_cast<PetscScalar *>(v);
2370: PetscMemType memtype;
2371: PetscDeviceContext dctx;
2372: cupmStream_t stream;
2374: PetscFunctionBegin;
2375: PetscCall(GetHandles_(&dctx, &stream));
2376: PetscCall(PetscGetMemType(v, &memtype));
2377: if (PetscMemTypeHost(memtype)) {
2378: const auto size = VecIMPLCast(x)->coo_n;
2380: // If user gave v[] in host, we might need to copy it to device if any
2381: PetscCall(PetscDeviceMalloc(dctx, PETSC_MEMTYPE_CUPM(), size, &vv));
2382: PetscCall(PetscCUPMMemcpyAsync(vv, v, size, cupmMemcpyHostToDevice, stream));
2383: PetscCall(PetscLogCpuToGpu(size * sizeof(PetscScalar)));
2384: }
2386: if (const auto n = x->map->n) {
2387: const auto vcu = VecCUPMCast(x);
2389: PetscCall(PetscCUPMLaunchKernel1D(n, 0, stream, kernels::add_coo_values, vv, n, vcu->jmap1_d, vcu->perm1_d, imode, imode == INSERT_VALUES ? DeviceArrayWrite(dctx, x).data() : DeviceArrayReadWrite(dctx, x).data()));
2390: } else {
2391: PetscCall(MaybeIncrementEmptyLocalVec(x));
2392: }
2394: if (PetscMemTypeHost(memtype)) PetscCall(PetscDeviceFree(dctx, vv));
2395: PetscCall(PetscDeviceContextSynchronize(dctx));
2396: PetscFunctionReturn(PETSC_SUCCESS);
2397: }
2399: } // namespace impl
2401: } // namespace cupm
2403: } // namespace vec
2405: } // namespace Petsc