Actual source code: matdensecupmimpl.h
1: #pragma once
3: #define PETSC_SKIP_IMMINTRIN_H_CUDAWORKAROUND 1
4: #include <petsc/private/matimpl.h>
6: #if defined(__cplusplus)
7: #include <petsc/private/deviceimpl.h>
8: #include <petsc/private/cupmsolverinterface.hpp>
9: #include <petsc/private/cupmobject.hpp>
11: #include "../src/sys/objects/device/impls/cupm/cupmthrustutility.hpp"
12: #include "../src/sys/objects/device/impls/cupm/kernels.hpp"
14: #include <thrust/device_vector.h>
15: #include <thrust/device_ptr.h>
16: #include <thrust/iterator/counting_iterator.h>
17: #include <thrust/iterator/transform_iterator.h>
18: #include <thrust/iterator/permutation_iterator.h>
19: #include <thrust/transform.h>
20: #include <thrust/copy.h>
22: namespace Petsc
23: {
25: namespace vec
26: {
28: namespace cupm
29: {
31: namespace impl
32: {
34: template <device::cupm::DeviceType>
35: class VecSeq_CUPM;
36: template <device::cupm::DeviceType>
37: class VecMPI_CUPM;
39: } // namespace impl
41: } // namespace cupm
43: } // namespace vec
45: namespace mat
46: {
48: namespace cupm
49: {
51: namespace impl
52: {
54: // ==========================================================================================
55: // MatDense_CUPM_Base
56: //
57: // A base class to separate out the CRTP code from the common CUPM stuff (like the composed
58: // function names).
59: // ==========================================================================================
61: template <device::cupm::DeviceType T>
62: class PETSC_SINGLE_LIBRARY_VISIBILITY_INTERNAL MatDense_CUPM_Base : protected device::cupm::impl::CUPMObject<T> {
63: public:
64: PETSC_CUPMOBJECT_HEADER(T);
66: #define MatDenseCUPMComposedOpDecl(OP_NAME) \
67: PETSC_NODISCARD static constexpr const char *PetscConcat(MatDenseCUPM, OP_NAME)() noexcept \
68: { \
69: return T == device::cupm::DeviceType::CUDA ? PetscStringize(PetscConcat(MatDenseCUDA, OP_NAME)) : PetscStringize(PetscConcat(MatDenseHIP, OP_NAME)); \
70: }
72: // clang-format off
73: MatDenseCUPMComposedOpDecl(GetArray_C)
74: MatDenseCUPMComposedOpDecl(GetArrayRead_C)
75: MatDenseCUPMComposedOpDecl(GetArrayWrite_C)
76: MatDenseCUPMComposedOpDecl(RestoreArray_C)
77: MatDenseCUPMComposedOpDecl(RestoreArrayRead_C)
78: MatDenseCUPMComposedOpDecl(RestoreArrayWrite_C)
79: MatDenseCUPMComposedOpDecl(PlaceArray_C)
80: MatDenseCUPMComposedOpDecl(ReplaceArray_C)
81: MatDenseCUPMComposedOpDecl(ResetArray_C)
82: MatDenseCUPMComposedOpDecl(SetPreallocation_C)
83: // clang-format on
85: #undef MatDenseCUPMComposedOpDecl
87: PETSC_NODISCARD static constexpr MatType MATSEQDENSECUPM() noexcept;
88: PETSC_NODISCARD static constexpr MatType MATMPIDENSECUPM() noexcept;
89: PETSC_NODISCARD static constexpr MatType MATDENSECUPM() noexcept;
90: PETSC_NODISCARD static constexpr MatSolverType MATSOLVERCUPM() noexcept;
92: static PetscErrorCode GetCurrentMemType(Mat A, PetscMemType *mtype) noexcept
93: {
94: PetscFunctionBegin;
95: *mtype = A->boundtocpu ? PETSC_MEMTYPE_HOST : PETSC_MEMTYPE_CUPM();
96: PetscFunctionReturn(PETSC_SUCCESS);
97: }
98: };
100: // ==========================================================================================
101: // MatDense_CUPM_Base -- Public API
102: // ==========================================================================================
104: template <device::cupm::DeviceType T>
105: inline constexpr MatType MatDense_CUPM_Base<T>::MATSEQDENSECUPM() noexcept
106: {
107: return T == device::cupm::DeviceType::CUDA ? MATSEQDENSECUDA : MATSEQDENSEHIP;
108: }
110: template <device::cupm::DeviceType T>
111: inline constexpr MatType MatDense_CUPM_Base<T>::MATMPIDENSECUPM() noexcept
112: {
113: return T == device::cupm::DeviceType::CUDA ? MATMPIDENSECUDA : MATMPIDENSEHIP;
114: }
116: template <device::cupm::DeviceType T>
117: inline constexpr MatType MatDense_CUPM_Base<T>::MATDENSECUPM() noexcept
118: {
119: return T == device::cupm::DeviceType::CUDA ? MATDENSECUDA : MATDENSEHIP;
120: }
122: template <device::cupm::DeviceType T>
123: inline constexpr MatSolverType MatDense_CUPM_Base<T>::MATSOLVERCUPM() noexcept
124: {
125: return T == device::cupm::DeviceType::CUDA ? MATSOLVERCUDA : MATSOLVERHIP;
126: }
128: #define MATDENSECUPM_BASE_HEADER(T) \
129: PETSC_CUPMOBJECT_HEADER(T); \
130: using VecSeq_CUPM = ::Petsc::vec::cupm::impl::VecSeq_CUPM<T>; \
131: using VecMPI_CUPM = ::Petsc::vec::cupm::impl::VecMPI_CUPM<T>; \
132: using ::Petsc::mat::cupm::impl::MatDense_CUPM_Base<T>::MATSEQDENSECUPM; \
133: using ::Petsc::mat::cupm::impl::MatDense_CUPM_Base<T>::MATMPIDENSECUPM; \
134: using ::Petsc::mat::cupm::impl::MatDense_CUPM_Base<T>::MATDENSECUPM; \
135: using ::Petsc::mat::cupm::impl::MatDense_CUPM_Base<T>::MATSOLVERCUPM; \
136: using ::Petsc::mat::cupm::impl::MatDense_CUPM_Base<T>::GetCurrentMemType; \
137: using ::Petsc::mat::cupm::impl::MatDense_CUPM_Base<T>::MatDenseCUPMGetArray_C; \
138: using ::Petsc::mat::cupm::impl::MatDense_CUPM_Base<T>::MatDenseCUPMGetArrayRead_C; \
139: using ::Petsc::mat::cupm::impl::MatDense_CUPM_Base<T>::MatDenseCUPMGetArrayWrite_C; \
140: using ::Petsc::mat::cupm::impl::MatDense_CUPM_Base<T>::MatDenseCUPMRestoreArray_C; \
141: using ::Petsc::mat::cupm::impl::MatDense_CUPM_Base<T>::MatDenseCUPMRestoreArrayRead_C; \
142: using ::Petsc::mat::cupm::impl::MatDense_CUPM_Base<T>::MatDenseCUPMRestoreArrayWrite_C; \
143: using ::Petsc::mat::cupm::impl::MatDense_CUPM_Base<T>::MatDenseCUPMPlaceArray_C; \
144: using ::Petsc::mat::cupm::impl::MatDense_CUPM_Base<T>::MatDenseCUPMReplaceArray_C; \
145: using ::Petsc::mat::cupm::impl::MatDense_CUPM_Base<T>::MatDenseCUPMResetArray_C; \
146: using ::Petsc::mat::cupm::impl::MatDense_CUPM_Base<T>::MatDenseCUPMSetPreallocation_C
148: // forward declare
149: template <device::cupm::DeviceType>
150: class MatDense_Seq_CUPM;
151: template <device::cupm::DeviceType>
152: class MatDense_MPI_CUPM;
154: // ==========================================================================================
155: // MatDense_CUPM
156: //
157: // The true "base" class for MatDenseCUPM. The reason MatDense_CUPM and MatDense_CUPM_Base
158: // exist is to separate out the CRTP code from the non-crtp code so that the generic functions
159: // can be called via templates below.
160: // ==========================================================================================
162: template <device::cupm::DeviceType T, typename Derived>
163: class PETSC_SINGLE_LIBRARY_VISIBILITY_INTERNAL MatDense_CUPM : protected MatDense_CUPM_Base<T> {
164: private:
165: static PetscErrorCode CheckSaneSequentialMatSizes_(Mat) noexcept;
167: protected:
168: MATDENSECUPM_BASE_HEADER(T);
170: template <PetscMemType, PetscMemoryAccessMode>
171: class MatrixArray;
173: // Cast the Mat to its host struct, i.e. return the result of (Mat_SeqDense *)m->data
174: template <typename U = Derived>
175: PETSC_NODISCARD static constexpr auto MatIMPLCast(Mat m) noexcept PETSC_DECLTYPE_AUTO_RETURNS(U::MatIMPLCast_(m))
176: PETSC_NODISCARD static constexpr MatType MATIMPLCUPM() noexcept;
178: static PetscErrorCode CreateIMPLDenseCUPM(MPI_Comm, PetscInt, PetscInt, PetscInt, PetscInt, PetscScalar *, Mat *, PetscDeviceContext, bool) noexcept;
179: static PetscErrorCode SetPreallocation(Mat, PetscDeviceContext, PetscScalar *) noexcept;
181: template <typename F>
182: static PetscErrorCode DiagonalUnaryTransform(Mat, PetscDeviceContext, F &&) noexcept;
184: static PetscErrorCode Shift(Mat, PetscScalar) noexcept;
185: static PetscErrorCode GetDiagonal(Mat, Vec) noexcept;
187: PETSC_NODISCARD static auto DeviceArrayRead(PetscDeviceContext dctx, Mat m) noexcept PETSC_DECLTYPE_AUTO_RETURNS(MatrixArray<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ>{dctx, m})
188: PETSC_NODISCARD static auto DeviceArrayWrite(PetscDeviceContext dctx, Mat m) noexcept PETSC_DECLTYPE_AUTO_RETURNS(MatrixArray<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_WRITE>{dctx, m})
189: PETSC_NODISCARD static auto DeviceArrayReadWrite(PetscDeviceContext dctx, Mat m) noexcept PETSC_DECLTYPE_AUTO_RETURNS(MatrixArray<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ_WRITE>{dctx, m})
190: PETSC_NODISCARD static auto HostArrayRead(PetscDeviceContext dctx, Mat m) noexcept PETSC_DECLTYPE_AUTO_RETURNS(MatrixArray<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ>{dctx, m})
191: PETSC_NODISCARD static auto HostArrayWrite(PetscDeviceContext dctx, Mat m) noexcept PETSC_DECLTYPE_AUTO_RETURNS(MatrixArray<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_WRITE>{dctx, m})
192: PETSC_NODISCARD static auto HostArrayReadWrite(PetscDeviceContext dctx, Mat m) noexcept PETSC_DECLTYPE_AUTO_RETURNS(MatrixArray<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ_WRITE>{dctx, m})
193: };
195: // ==========================================================================================
196: // MatDense_CUPM::MatrixArray
197: // ==========================================================================================
199: template <device::cupm::DeviceType T, typename D>
200: template <PetscMemType MT, PetscMemoryAccessMode MA>
201: class PETSC_SINGLE_LIBRARY_VISIBILITY_INTERNAL MatDense_CUPM<T, D>::MatrixArray : public device::cupm::impl::RestoreableArray<T, MT, MA> {
202: using base_type = device::cupm::impl::RestoreableArray<T, MT, MA>;
204: public:
205: MatrixArray(PetscDeviceContext, Mat) noexcept;
206: ~MatrixArray() noexcept;
208: // must declare move constructor since we declare a destructor
209: constexpr MatrixArray(MatrixArray &&) noexcept;
211: private:
212: Mat m_ = nullptr;
213: };
215: // ==========================================================================================
216: // MatDense_CUPM::MatrixArray -- Public API
217: // ==========================================================================================
219: template <device::cupm::DeviceType T, typename D>
220: template <PetscMemType MT, PetscMemoryAccessMode MA>
221: inline MatDense_CUPM<T, D>::MatrixArray<MT, MA>::MatrixArray(PetscDeviceContext dctx, Mat m) noexcept : base_type{dctx}, m_{m}
222: {
223: PetscFunctionBegin;
224: PetscCallAbort(PETSC_COMM_SELF, D::template GetArray<MT, MA>(m, &this->ptr_, dctx));
225: PetscFunctionReturnVoid();
226: }
228: template <device::cupm::DeviceType T, typename D>
229: template <PetscMemType MT, PetscMemoryAccessMode MA>
230: inline MatDense_CUPM<T, D>::MatrixArray<MT, MA>::~MatrixArray() noexcept
231: {
232: PetscFunctionBegin;
233: PetscCallAbort(PETSC_COMM_SELF, D::template RestoreArray<MT, MA>(m_, &this->ptr_, this->dctx_));
234: PetscFunctionReturnVoid();
235: }
237: template <device::cupm::DeviceType T, typename D>
238: template <PetscMemType MT, PetscMemoryAccessMode MA>
239: inline constexpr MatDense_CUPM<T, D>::MatrixArray<MT, MA>::MatrixArray(MatrixArray &&other) noexcept : base_type{std::move(other)}, m_{util::exchange(other.m_, nullptr)}
240: {
241: }
243: // ==========================================================================================
244: // MatDense_CUPM -- Private API
245: // ==========================================================================================
247: template <device::cupm::DeviceType T, typename D>
248: inline PetscErrorCode MatDense_CUPM<T, D>::CheckSaneSequentialMatSizes_(Mat A) noexcept
249: {
250: PetscFunctionBegin;
251: if (PetscDefined(USE_DEBUG)) {
252: PetscBool isseq;
254: PetscCall(PetscObjectTypeCompare(PetscObjectCast(A), D::MATSEQDENSECUPM(), &isseq));
255: if (isseq) {
256: // doing this check allows both sequential and parallel implementations to just pass in
257: // A, otherwise they would need to specify rstart, rend, and cols separately!
258: PetscCheck(A->rmap->rstart == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Sequential matrix row start %" PetscInt_FMT " != 0?", A->rmap->rstart);
259: PetscCheck(A->rmap->rend == A->rmap->n, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Sequential matrix row end %" PetscInt_FMT " != number of rows %" PetscInt_FMT, A->rmap->rend, A->rmap->n);
260: PetscCheck(A->cmap->n == A->cmap->N, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Sequential matrix number of local columns %" PetscInt_FMT " != number of global columns %" PetscInt_FMT, A->cmap->n, A->cmap->n);
261: }
262: }
263: PetscFunctionReturn(PETSC_SUCCESS);
264: }
266: // ==========================================================================================
267: // MatDense_CUPM -- Protected API
268: // ==========================================================================================
270: template <device::cupm::DeviceType T, typename D>
271: inline constexpr MatType MatDense_CUPM<T, D>::MATIMPLCUPM() noexcept
272: {
273: return D::MATIMPLCUPM_();
274: }
276: // Common core for MatCreateSeqDenseCUPM() and MatCreateMPIDenseCUPM()
277: template <device::cupm::DeviceType T, typename D>
278: inline PetscErrorCode MatDense_CUPM<T, D>::CreateIMPLDenseCUPM(MPI_Comm comm, PetscInt m, PetscInt n, PetscInt M, PetscInt N, PetscScalar *data, Mat *A, PetscDeviceContext dctx, bool preallocate) noexcept
279: {
280: Mat mat;
282: PetscFunctionBegin;
283: PetscAssertPointer(A, 7);
284: PetscCall(MatCreate(comm, &mat));
285: PetscCall(MatSetSizes(mat, m, n, M, N));
286: PetscCall(MatSetType(mat, D::MATIMPLCUPM()));
287: if (preallocate) {
288: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
289: PetscCall(D::SetPreallocation(mat, dctx, data));
290: }
291: *A = mat;
292: PetscFunctionReturn(PETSC_SUCCESS);
293: }
295: template <device::cupm::DeviceType T, typename D>
296: inline PetscErrorCode MatDense_CUPM<T, D>::SetPreallocation(Mat A, PetscDeviceContext dctx, PetscScalar *device_array) noexcept
297: {
298: PetscFunctionBegin;
300: // might be the local (sequential) matrix of a MatMPIDense_CUPM. Since this would be called
301: // from the MPI matrix'es impl MATIMPLCUPM() would return MATMPIDENSECUPM().
303: PetscCheckTypeNames(A, D::MATSEQDENSECUPM(), D::MATMPIDENSECUPM());
304: PetscCall(PetscLayoutSetUp(A->rmap));
305: PetscCall(PetscLayoutSetUp(A->cmap));
306: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
307: PetscCall(D::SetPreallocation_(A, dctx, device_array));
308: A->preallocated = PETSC_TRUE;
309: A->assembled = PETSC_TRUE;
310: PetscFunctionReturn(PETSC_SUCCESS);
311: }
313: namespace detail
314: {
316: #if CCCL_VERSION >= 3000000
317: template <class T>
318: using iter_difference_t = cuda::std::iter_difference_t<T>;
319: #else
320: template <class T>
321: using iter_difference_t = typename thrust::iterator_difference<T>::type;
322: #endif
324: // ==========================================================================================
325: // MatrixIteratorBase
326: //
327: // A base class for creating thrust iterators over the local sub-matrix. This will set up the
328: // proper iterator definitions so thrust knows how to handle things properly. Template
329: // parameters are as follows:
330: //
331: // - Iterator:
332: // The type of the primary array iterator. Usually this is
333: // thrust::device_pointer::iterator.
334: //
335: // - IndexFunctor:
336: // This should be a functor which contains an operator() that when called with an index `i`,
337: // returns the i'th permuted index into the array. For example, it could return the i'th
338: // diagonal entry.
339: // ==========================================================================================
340: template <typename Iterator, typename IndexFunctor>
341: class MatrixIteratorBase {
342: public:
343: using array_iterator_type = Iterator;
344: using index_functor_type = IndexFunctor;
346: using difference_type = iter_difference_t<array_iterator_type>;
347: using CountingIterator = thrust::counting_iterator<difference_type>;
348: using TransformIterator = thrust::transform_iterator<index_functor_type, CountingIterator>;
349: using PermutationIterator = thrust::permutation_iterator<array_iterator_type, TransformIterator>;
350: using iterator = PermutationIterator; // type of the begin/end iterator
352: constexpr MatrixIteratorBase(array_iterator_type first, array_iterator_type last, index_functor_type idx_func) noexcept : first{std::move(first)}, last{std::move(last)}, func{std::move(idx_func)} { }
354: PETSC_NODISCARD iterator begin() const noexcept
355: {
356: return PermutationIterator{
357: first, TransformIterator{CountingIterator{0}, func}
358: };
359: }
361: protected:
362: array_iterator_type first;
363: array_iterator_type last;
364: index_functor_type func;
365: };
367: // ==========================================================================================
368: // StridedIndexFunctor
369: //
370: // Iterator which permutes a linear index range into strided matrix indices. Usually used to
371: // get the diagonal.
372: // ==========================================================================================
373: template <typename T>
374: struct StridedIndexFunctor {
375: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL constexpr T operator()(const T &i) const noexcept { return stride * i; }
377: T stride;
378: };
380: template <typename Iterator>
381: class DiagonalIterator : public MatrixIteratorBase<Iterator, StridedIndexFunctor<iter_difference_t<Iterator>>> {
382: public:
383: using base_type = MatrixIteratorBase<Iterator, StridedIndexFunctor<iter_difference_t<Iterator>>>;
385: using difference_type = typename base_type::difference_type;
386: using iterator = typename base_type::iterator;
388: constexpr DiagonalIterator(Iterator first, Iterator last, difference_type stride) noexcept : base_type{std::move(first), std::move(last), {stride}} { }
390: PETSC_NODISCARD iterator end() const noexcept { return this->begin() + (this->last - this->first + this->func.stride - 1) / this->func.stride; }
391: };
393: template <typename T>
394: inline DiagonalIterator<typename thrust::device_vector<T>::iterator> MakeDiagonalIterator(T *data, PetscInt rstart, PetscInt rend, PetscInt cols, PetscInt lda) noexcept
395: {
396: const auto rend2 = std::min(rend, cols);
397: const std::size_t begin = rstart * lda;
398: const std::size_t end = rend2 - rstart + rend2 * lda;
399: const auto dptr = thrust::device_pointer_cast(data);
401: return {dptr + begin, dptr + end, lda + 1};
402: }
404: } // namespace detail
406: template <device::cupm::DeviceType T, typename D>
407: template <typename F>
408: inline PetscErrorCode MatDense_CUPM<T, D>::DiagonalUnaryTransform(Mat A, PetscDeviceContext dctx, F &&functor) noexcept
409: {
410: const auto rstart = A->rmap->rstart;
411: const auto rend = A->rmap->rend;
412: const auto gcols = A->cmap->N;
413: const auto rend2 = std::min(rend, gcols);
415: PetscFunctionBegin;
416: PetscCall(CheckSaneSequentialMatSizes_(A));
417: if (rend2 > rstart) {
418: const auto da = D::DeviceArrayReadWrite(dctx, A);
419: cupmStream_t stream;
420: PetscInt lda;
422: PetscCall(MatDenseGetLDA(A, &lda));
423: PetscCall(D::GetHandlesFrom_(dctx, &stream));
424: {
425: auto diagonal = detail::MakeDiagonalIterator(da.data(), rstart, rend, gcols, lda);
427: // clang-format off
428: PetscCallThrust(
429: THRUST_CALL(
430: thrust::transform,
431: stream,
432: diagonal.begin(), diagonal.end(), diagonal.begin(),
433: std::forward<F>(functor)
434: )
435: );
436: // clang-format on
437: }
438: PetscCall(PetscLogGpuFlops(rend2 - rstart));
439: }
440: PetscFunctionReturn(PETSC_SUCCESS);
441: }
443: template <device::cupm::DeviceType T, typename D>
444: inline PetscErrorCode MatDense_CUPM<T, D>::Shift(Mat A, PetscScalar alpha) noexcept
445: {
446: PetscDeviceContext dctx;
448: PetscFunctionBegin;
449: PetscCall(GetHandles_(&dctx));
450: PetscCall(DiagonalUnaryTransform(A, dctx, device::cupm::functors::make_plus_equals(alpha)));
451: PetscFunctionReturn(PETSC_SUCCESS);
452: }
454: template <device::cupm::DeviceType T, typename D>
455: inline PetscErrorCode MatDense_CUPM<T, D>::GetDiagonal(Mat A, Vec v) noexcept
456: {
457: const auto rstart = A->rmap->rstart;
458: const auto rend = A->rmap->rend;
459: const auto gcols = A->cmap->N;
460: PetscInt lda;
461: PetscDeviceContext dctx;
463: PetscFunctionBegin;
464: PetscCall(CheckSaneSequentialMatSizes_(A));
465: PetscCall(GetHandles_(&dctx));
466: PetscCall(MatDenseGetLDA(A, &lda));
467: {
468: auto dv = VecSeq_CUPM::DeviceArrayWrite(dctx, v);
469: auto da = D::DeviceArrayRead(dctx, A);
470: auto diagonal = detail::MakeDiagonalIterator(da.data(), rstart, rend, gcols, lda);
471: // We must have this cast outside of THRUST_CALL(). Without it, GCC 6.4 - 7.5, and 11.3.0
472: // throw spurious warnings:
473: //
474: // warning: 'MatDense_CUPM<...>::GetDiagonal(Mat, Vec)::' declared with greater
475: // visibility than the type of its field 'MatDense_CUPM<...>::GetDiagonal(Mat,
476: // Vec)::::' [-Wattributes]
477: // 460 | PetscCallThrust(
478: // | ^~~~~~~~~~~~~~~~
479: auto dvp = thrust::device_pointer_cast(dv.data());
480: cupmStream_t stream;
482: PetscCall(GetHandlesFrom_(dctx, &stream));
483: PetscCallThrust(THRUST_CALL(thrust::copy, stream, diagonal.begin(), diagonal.end(), dvp));
484: }
485: PetscFunctionReturn(PETSC_SUCCESS);
486: }
488: #define MatComposeOp_CUPM(use_host, pobj, op_str, op_host, ...) \
489: do { \
490: if (use_host) { \
491: PetscCall(PetscObjectComposeFunction(pobj, op_str, op_host)); \
492: } else { \
493: PetscCall(PetscObjectComposeFunction(pobj, op_str, __VA_ARGS__)); \
494: } \
495: } while (0)
497: #define MatSetOp_CUPM(use_host, mat, op_name, op_host, ...) \
498: do { \
499: if (use_host) { \
500: (mat)->ops->op_name = op_host; \
501: } else { \
502: (mat)->ops->op_name = __VA_ARGS__; \
503: } \
504: } while (0)
506: #define MATDENSECUPM_HEADER(T, ...) \
507: MATDENSECUPM_BASE_HEADER(T); \
508: friend class ::Petsc::mat::cupm::impl::MatDense_CUPM<T, __VA_ARGS__>; \
509: using ::Petsc::mat::cupm::impl::MatDense_CUPM<T, __VA_ARGS__>::MatIMPLCast; \
510: using ::Petsc::mat::cupm::impl::MatDense_CUPM<T, __VA_ARGS__>::MATIMPLCUPM; \
511: using ::Petsc::mat::cupm::impl::MatDense_CUPM<T, __VA_ARGS__>::CreateIMPLDenseCUPM; \
512: using ::Petsc::mat::cupm::impl::MatDense_CUPM<T, __VA_ARGS__>::SetPreallocation; \
513: using ::Petsc::mat::cupm::impl::MatDense_CUPM<T, __VA_ARGS__>::DeviceArrayRead; \
514: using ::Petsc::mat::cupm::impl::MatDense_CUPM<T, __VA_ARGS__>::DeviceArrayWrite; \
515: using ::Petsc::mat::cupm::impl::MatDense_CUPM<T, __VA_ARGS__>::DeviceArrayReadWrite; \
516: using ::Petsc::mat::cupm::impl::MatDense_CUPM<T, __VA_ARGS__>::HostArrayRead; \
517: using ::Petsc::mat::cupm::impl::MatDense_CUPM<T, __VA_ARGS__>::HostArrayWrite; \
518: using ::Petsc::mat::cupm::impl::MatDense_CUPM<T, __VA_ARGS__>::HostArrayReadWrite; \
519: using ::Petsc::mat::cupm::impl::MatDense_CUPM<T, __VA_ARGS__>::DiagonalUnaryTransform; \
520: using ::Petsc::mat::cupm::impl::MatDense_CUPM<T, __VA_ARGS__>::Shift; \
521: using ::Petsc::mat::cupm::impl::MatDense_CUPM<T, __VA_ARGS__>::GetDiagonal
523: } // namespace impl
525: // ==========================================================================================
526: // MatDense_CUPM -- Implementations
527: // ==========================================================================================
529: template <device::cupm::DeviceType T, PetscMemoryAccessMode access>
530: inline PetscErrorCode MatDenseCUPMGetArray_Private(Mat A, PetscScalar **array) noexcept
531: {
532: PetscFunctionBegin;
534: PetscAssertPointer(array, 2);
535: switch (access) {
536: case PETSC_MEMORY_ACCESS_READ:
537: PetscUseMethod(A, impl::MatDense_CUPM_Base<T>::MatDenseCUPMGetArrayRead_C(), (Mat, PetscScalar **), (A, array));
538: break;
539: case PETSC_MEMORY_ACCESS_WRITE:
540: PetscUseMethod(A, impl::MatDense_CUPM_Base<T>::MatDenseCUPMGetArrayWrite_C(), (Mat, PetscScalar **), (A, array));
541: break;
542: case PETSC_MEMORY_ACCESS_READ_WRITE:
543: PetscUseMethod(A, impl::MatDense_CUPM_Base<T>::MatDenseCUPMGetArray_C(), (Mat, PetscScalar **), (A, array));
544: break;
545: }
546: if (PetscMemoryAccessWrite(access)) PetscCall(PetscObjectStateIncrease(PetscObjectCast(A)));
547: PetscFunctionReturn(PETSC_SUCCESS);
548: }
550: template <device::cupm::DeviceType T, PetscMemoryAccessMode access>
551: inline PetscErrorCode MatDenseCUPMRestoreArray_Private(Mat A, PetscScalar **array) noexcept
552: {
553: PetscFunctionBegin;
555: if (array) PetscAssertPointer(array, 2);
556: switch (access) {
557: case PETSC_MEMORY_ACCESS_READ:
558: PetscUseMethod(A, impl::MatDense_CUPM_Base<T>::MatDenseCUPMRestoreArrayRead_C(), (Mat, PetscScalar **), (A, array));
559: break;
560: case PETSC_MEMORY_ACCESS_WRITE:
561: PetscUseMethod(A, impl::MatDense_CUPM_Base<T>::MatDenseCUPMRestoreArrayWrite_C(), (Mat, PetscScalar **), (A, array));
562: break;
563: case PETSC_MEMORY_ACCESS_READ_WRITE:
564: PetscUseMethod(A, impl::MatDense_CUPM_Base<T>::MatDenseCUPMRestoreArray_C(), (Mat, PetscScalar **), (A, array));
565: break;
566: }
567: if (PetscMemoryAccessWrite(access)) {
568: PetscCall(PetscObjectStateIncrease(PetscObjectCast(A)));
569: A->offloadmask = PETSC_OFFLOAD_GPU;
570: }
571: if (array) *array = nullptr;
572: PetscFunctionReturn(PETSC_SUCCESS);
573: }
575: template <device::cupm::DeviceType T>
576: inline PetscErrorCode MatDenseCUPMGetArray(Mat A, PetscScalar **array) noexcept
577: {
578: PetscFunctionBegin;
579: PetscCall(MatDenseCUPMGetArray_Private<T, PETSC_MEMORY_ACCESS_READ_WRITE>(A, array));
580: PetscFunctionReturn(PETSC_SUCCESS);
581: }
583: template <device::cupm::DeviceType T>
584: inline PetscErrorCode MatDenseCUPMGetArrayRead(Mat A, const PetscScalar **array) noexcept
585: {
586: PetscFunctionBegin;
587: PetscCall(MatDenseCUPMGetArray_Private<T, PETSC_MEMORY_ACCESS_READ>(A, const_cast<PetscScalar **>(array)));
588: PetscFunctionReturn(PETSC_SUCCESS);
589: }
591: template <device::cupm::DeviceType T>
592: inline PetscErrorCode MatDenseCUPMGetArrayWrite(Mat A, PetscScalar **array) noexcept
593: {
594: PetscFunctionBegin;
595: PetscCall(MatDenseCUPMGetArray_Private<T, PETSC_MEMORY_ACCESS_WRITE>(A, array));
596: PetscFunctionReturn(PETSC_SUCCESS);
597: }
599: template <device::cupm::DeviceType T>
600: inline PetscErrorCode MatDenseCUPMRestoreArray(Mat A, PetscScalar **array) noexcept
601: {
602: PetscFunctionBegin;
603: PetscCall(MatDenseCUPMRestoreArray_Private<T, PETSC_MEMORY_ACCESS_READ_WRITE>(A, array));
604: PetscFunctionReturn(PETSC_SUCCESS);
605: }
607: template <device::cupm::DeviceType T>
608: inline PetscErrorCode MatDenseCUPMRestoreArrayRead(Mat A, const PetscScalar **array) noexcept
609: {
610: PetscFunctionBegin;
611: PetscCall(MatDenseCUPMRestoreArray_Private<T, PETSC_MEMORY_ACCESS_READ>(A, const_cast<PetscScalar **>(array)));
612: PetscFunctionReturn(PETSC_SUCCESS);
613: }
615: template <device::cupm::DeviceType T>
616: inline PetscErrorCode MatDenseCUPMRestoreArrayWrite(Mat A, PetscScalar **array) noexcept
617: {
618: PetscFunctionBegin;
619: PetscCall(MatDenseCUPMRestoreArray_Private<T, PETSC_MEMORY_ACCESS_WRITE>(A, array));
620: PetscFunctionReturn(PETSC_SUCCESS);
621: }
623: template <device::cupm::DeviceType T>
624: inline PetscErrorCode MatDenseCUPMPlaceArray(Mat A, const PetscScalar *array) noexcept
625: {
626: PetscFunctionBegin;
628: PetscUseMethod(A, impl::MatDense_CUPM_Base<T>::MatDenseCUPMPlaceArray_C(), (Mat, const PetscScalar *), (A, array));
629: PetscCall(PetscObjectStateIncrease(PetscObjectCast(A)));
630: A->offloadmask = PETSC_OFFLOAD_GPU;
631: PetscFunctionReturn(PETSC_SUCCESS);
632: }
634: template <device::cupm::DeviceType T>
635: inline PetscErrorCode MatDenseCUPMReplaceArray(Mat A, const PetscScalar *array) noexcept
636: {
637: PetscFunctionBegin;
639: PetscUseMethod(A, impl::MatDense_CUPM_Base<T>::MatDenseCUPMReplaceArray_C(), (Mat, const PetscScalar *), (A, array));
640: PetscCall(PetscObjectStateIncrease(PetscObjectCast(A)));
641: A->offloadmask = PETSC_OFFLOAD_GPU;
642: PetscFunctionReturn(PETSC_SUCCESS);
643: }
645: template <device::cupm::DeviceType T>
646: inline PetscErrorCode MatDenseCUPMResetArray(Mat A) noexcept
647: {
648: PetscFunctionBegin;
650: PetscUseMethod(A, impl::MatDense_CUPM_Base<T>::MatDenseCUPMResetArray_C(), (Mat), (A));
651: PetscCall(PetscObjectStateIncrease(PetscObjectCast(A)));
652: PetscFunctionReturn(PETSC_SUCCESS);
653: }
655: template <device::cupm::DeviceType T>
656: inline PetscErrorCode MatDenseCUPMSetPreallocation(Mat A, PetscScalar *device_data, PetscDeviceContext dctx = nullptr) noexcept
657: {
658: PetscFunctionBegin;
660: PetscUseMethod(A, impl::MatDense_CUPM_Base<T>::MatDenseCUPMSetPreallocation_C(), (Mat, PetscDeviceContext, PetscScalar *), (A, dctx, device_data));
661: PetscFunctionReturn(PETSC_SUCCESS);
662: }
664: } // namespace cupm
666: } // namespace mat
668: } // namespace Petsc
670: #endif // __cplusplus