Actual source code: matseqdensecupm.hpp
1: #pragma once
3: #include <petsc/private/matdensecupmimpl.h>
4: #include <../src/mat/impls/dense/seq/dense.h>
6: #include <petsc/private/deviceimpl.h>
7: #include <petsc/private/randomimpl.h>
8: #include <petsc/private/vecimpl.h>
9: #include <petsc/private/cupmobject.hpp>
10: #include <petsc/private/cupmsolverinterface.hpp>
12: #include <petsc/private/cpp/type_traits.hpp>
13: #include <petsc/private/cpp/utility.hpp>
15: #include <../src/vec/vec/impls/seq/cupm/vecseqcupm.hpp>
17: #include <thrust/functional.h> // thrust::plus, thrust::maximum
18: #include <thrust/reduce.h> // thrust::reduce, thrust::reduce_by_key
19: #include <thrust/transform_reduce.h> // thrust::transform_reduce
20: #include <thrust/iterator/discard_iterator.h> // thrust::make_discard_iterator
22: namespace Petsc
23: {
25: namespace mat
26: {
28: namespace cupm
29: {
31: namespace impl
32: {
34: template <device::cupm::DeviceType T>
35: class PETSC_SINGLE_LIBRARY_VISIBILITY_INTERNAL MatDense_Seq_CUPM : MatDense_CUPM<T, MatDense_Seq_CUPM<T>> {
36: public:
37: MATDENSECUPM_HEADER(T, MatDense_Seq_CUPM<T>);
39: private:
40: struct Mat_SeqDenseCUPM {
41: PetscScalar *d_v; // pointer to the matrix on the GPU
42: PetscScalar *unplacedarray; // if one called MatCUPMDensePlaceArray(), this is where it stashed the original
43: bool d_user_alloc;
44: bool d_unplaced_user_alloc;
45: // factorization support
46: cupmBlasInt_t *d_fact_ipiv; // device pivots
47: cupmScalar_t *d_fact_tau; // device QR tau vector
48: cupmBlasInt_t *d_fact_info; // device info
49: cupmScalar_t *d_fact_work; // device workspace
50: cupmBlasInt_t d_fact_lwork; // size of device workspace
51: // workspace
52: Vec workvec;
53: };
55: static PetscErrorCode SetPreallocation_(Mat, PetscDeviceContext, PetscScalar *) noexcept;
57: static PetscErrorCode HostToDevice_(Mat, PetscDeviceContext) noexcept;
58: static PetscErrorCode DeviceToHost_(Mat, PetscDeviceContext) noexcept;
60: static PetscErrorCode CheckCUPMSolverInfo_(const cupmBlasInt_t *, cupmStream_t) noexcept;
62: template <typename Derived>
63: struct SolveCommon;
64: struct SolveQR;
65: struct SolveCholesky;
66: struct SolveLU;
68: template <typename Solver, bool transpose>
69: static PetscErrorCode MatSolve_Factored_Dispatch_(Mat, Vec, Vec) noexcept;
70: template <typename Solver, bool transpose>
71: static PetscErrorCode MatMatSolve_Factored_Dispatch_(Mat, Mat, Mat) noexcept;
72: template <bool transpose, bool hermitian>
73: static PetscErrorCode MatMultAddColumnRange_Dispatch_(Mat, Vec, Vec, Vec, PetscInt, PetscInt) noexcept;
74: template <bool transpose, bool hermitian>
75: static PetscErrorCode MatMultColumnRange_Dispatch_(Mat, Vec, Vec, PetscInt, PetscInt) noexcept;
76: template <bool transpose, bool hermitian>
77: static PetscErrorCode MatMultAdd_Dispatch_(Mat, Vec, Vec, Vec) noexcept;
79: template <bool to_host>
80: static PetscErrorCode Convert_Dispatch_(Mat, MatType, MatReuse, Mat *) noexcept;
82: PETSC_NODISCARD static constexpr MatType MATIMPLCUPM_() noexcept;
83: PETSC_NODISCARD static constexpr Mat_SeqDense *MatIMPLCast_(Mat) noexcept;
85: public:
86: PETSC_NODISCARD static constexpr Mat_SeqDenseCUPM *MatCUPMCast(Mat) noexcept;
88: // define these by hand since they don't fit the above mold
89: PETSC_NODISCARD static constexpr const char *MatConvert_seqdensecupm_seqdense_C() noexcept;
90: PETSC_NODISCARD static constexpr const char *MatProductSetFromOptions_seqaij_seqdensecupm_C() noexcept;
92: static PetscErrorCode Create(Mat) noexcept;
93: static PetscErrorCode Destroy(Mat) noexcept;
94: static PetscErrorCode SetUp(Mat) noexcept;
95: static PetscErrorCode Reset(Mat) noexcept;
97: static PetscErrorCode BindToCPU(Mat, PetscBool) noexcept;
98: static PetscErrorCode Convert_SeqDense_SeqDenseCUPM(Mat, MatType, MatReuse, Mat *) noexcept;
99: static PetscErrorCode Convert_SeqDenseCUPM_SeqDense(Mat, MatType, MatReuse, Mat *) noexcept;
101: static PetscErrorCode PlaceColumnVecArray(Vec, PetscScalar *) noexcept;
102: static PetscErrorCode ResetColumnVecArray(Vec) noexcept;
104: template <PetscMemType, PetscMemoryAccessMode>
105: static PetscErrorCode GetArray(Mat, PetscScalar **, PetscDeviceContext) noexcept;
106: template <PetscMemType, PetscMemoryAccessMode>
107: static PetscErrorCode RestoreArray(Mat, PetscScalar **, PetscDeviceContext) noexcept;
108: template <PetscMemoryAccessMode>
109: static PetscErrorCode GetArrayAndMemType(Mat, PetscScalar **, PetscMemType *, PetscDeviceContext) noexcept;
110: template <PetscMemoryAccessMode>
111: static PetscErrorCode RestoreArrayAndMemType(Mat, PetscScalar **, PetscDeviceContext) noexcept;
113: private:
114: template <PetscMemType mtype, PetscMemoryAccessMode mode>
115: static PetscErrorCode GetArrayC_(Mat m, PetscScalar **p) noexcept
116: {
117: PetscDeviceContext dctx;
119: PetscFunctionBegin;
120: PetscCall(GetHandles_(&dctx));
121: PetscCall(GetArray<mtype, mode>(m, p, dctx));
122: PetscFunctionReturn(PETSC_SUCCESS);
123: }
125: template <PetscMemType mtype, PetscMemoryAccessMode mode>
126: static PetscErrorCode RestoreArrayC_(Mat m, PetscScalar **p) noexcept
127: {
128: PetscDeviceContext dctx;
130: PetscFunctionBegin;
131: PetscCall(GetHandles_(&dctx));
132: PetscCall(RestoreArray<mtype, mode>(m, p, dctx));
133: PetscFunctionReturn(PETSC_SUCCESS);
134: }
136: template <PetscMemoryAccessMode mode>
137: static PetscErrorCode GetArrayAndMemTypeC_(Mat m, PetscScalar **p, PetscMemType *tp) noexcept
138: {
139: PetscDeviceContext dctx;
141: PetscFunctionBegin;
142: PetscCall(GetHandles_(&dctx));
143: PetscCall(GetArrayAndMemType<mode>(m, p, tp, dctx));
144: PetscFunctionReturn(PETSC_SUCCESS);
145: }
147: template <PetscMemoryAccessMode mode>
148: static PetscErrorCode RestoreArrayAndMemTypeC_(Mat m, PetscScalar **p) noexcept
149: {
150: PetscDeviceContext dctx;
152: PetscFunctionBegin;
153: PetscCall(GetHandles_(&dctx));
154: PetscCall(RestoreArrayAndMemType<mode>(m, p, dctx));
155: PetscFunctionReturn(PETSC_SUCCESS);
156: }
158: public:
159: static PetscErrorCode PlaceArray(Mat, const PetscScalar *) noexcept;
160: static PetscErrorCode ReplaceArray(Mat, const PetscScalar *) noexcept;
161: static PetscErrorCode ResetArray(Mat) noexcept;
163: template <bool transpose_A, bool transpose_B>
164: static PetscErrorCode MatMatMult_Numeric_Dispatch(Mat, Mat, Mat) noexcept;
165: static PetscErrorCode Copy(Mat, Mat, MatStructure) noexcept;
166: static PetscErrorCode ZeroEntries(Mat) noexcept;
167: static PetscErrorCode Conjugate(Mat) noexcept;
168: static PetscErrorCode Scale(Mat, PetscScalar) noexcept;
169: static PetscErrorCode Norm(Mat, NormType, PetscReal *) noexcept;
170: static PetscErrorCode DiagonalScale(Mat, Vec, Vec) noexcept;
171: static PetscErrorCode AXPY(Mat, PetscScalar, Mat, MatStructure) noexcept;
172: static PetscErrorCode Duplicate(Mat, MatDuplicateOption, Mat *) noexcept;
173: static PetscErrorCode SetRandom(Mat, PetscRandom) noexcept;
175: static PetscErrorCode GetColumnVector(Mat, Vec, PetscInt) noexcept;
176: template <PetscMemoryAccessMode>
177: static PetscErrorCode GetColumnVec(Mat, PetscInt, Vec *) noexcept;
178: template <PetscMemoryAccessMode>
179: static PetscErrorCode RestoreColumnVec(Mat, PetscInt, Vec *) noexcept;
181: static PetscErrorCode GetFactor(Mat, MatFactorType, Mat *) noexcept;
182: static PetscErrorCode InvertFactors(Mat) noexcept;
184: static PetscErrorCode GetSubMatrix(Mat, PetscInt, PetscInt, PetscInt, PetscInt, Mat *) noexcept;
185: static PetscErrorCode RestoreSubMatrix(Mat, Mat *) noexcept;
186: };
188: } // namespace impl
190: namespace
191: {
193: // Declare this here so that the functions below can make use of it
194: template <device::cupm::DeviceType T>
195: inline PetscErrorCode MatCreateSeqDenseCUPM(MPI_Comm comm, PetscInt m, PetscInt n, PetscScalar *data, Mat *A, PetscDeviceContext dctx = nullptr, bool preallocate = true) noexcept
196: {
197: PetscFunctionBegin;
198: PetscCall(impl::MatDense_Seq_CUPM<T>::CreateIMPLDenseCUPM(comm, m, n, m, n, data, A, dctx, preallocate));
199: PetscFunctionReturn(PETSC_SUCCESS);
200: }
202: } // anonymous namespace
204: namespace impl
205: {
207: // ==========================================================================================
208: // MatDense_Seq_CUPM - Private API - Utility
209: // ==========================================================================================
211: template <device::cupm::DeviceType T>
212: inline PetscErrorCode MatDense_Seq_CUPM<T>::SetPreallocation_(Mat m, PetscDeviceContext dctx, PetscScalar *user_device_array) noexcept
213: {
214: const auto mcu = MatCUPMCast(m);
215: const auto nrows = m->rmap->n;
216: const auto ncols = m->cmap->n;
217: auto &lda = MatIMPLCast(m)->lda;
218: cupmStream_t stream;
220: PetscFunctionBegin;
221: PetscCheckTypeName(m, MATSEQDENSECUPM());
223: PetscCall(checkCupmBlasIntCast(nrows));
224: PetscCall(checkCupmBlasIntCast(ncols));
225: PetscCall(GetHandlesFrom_(dctx, &stream));
226: if (lda <= 0) lda = nrows;
227: if (!mcu->d_user_alloc) PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream));
228: if (user_device_array) {
229: mcu->d_user_alloc = PETSC_TRUE;
230: mcu->d_v = user_device_array;
231: } else {
232: std::size_t size;
234: mcu->d_user_alloc = PETSC_FALSE;
235: size = lda * ncols;
236: PetscCall(PetscCUPMMallocAsync(&mcu->d_v, size, stream));
237: PetscCall(PetscCUPMMemsetAsync(mcu->d_v, 0, size, stream));
238: }
239: m->offloadmask = PETSC_OFFLOAD_GPU;
240: PetscFunctionReturn(PETSC_SUCCESS);
241: }
243: template <device::cupm::DeviceType T>
244: inline PetscErrorCode MatDense_Seq_CUPM<T>::HostToDevice_(Mat m, PetscDeviceContext dctx) noexcept
245: {
246: const auto nrows = m->rmap->n;
247: const auto ncols = m->cmap->n;
248: const auto copy = m->offloadmask == PETSC_OFFLOAD_CPU || m->offloadmask == PETSC_OFFLOAD_UNALLOCATED;
250: PetscFunctionBegin;
251: PetscCheckTypeName(m, MATSEQDENSECUPM());
252: if (m->boundtocpu) PetscFunctionReturn(PETSC_SUCCESS);
253: PetscCall(PetscInfo(m, "%s matrix %" PetscInt_FMT " x %" PetscInt_FMT "\n", copy ? "Copy" : "Reusing", nrows, ncols));
254: if (copy) {
255: const auto mcu = MatCUPMCast(m);
256: cupmStream_t stream;
258: // Allocate GPU memory if not present
259: if (!mcu->d_v) PetscCall(SetPreallocation(m, dctx, nullptr));
260: PetscCall(GetHandlesFrom_(dctx, &stream));
261: PetscCall(PetscLogEventBegin(MAT_DenseCopyToGPU, m, 0, 0, 0));
262: {
263: const auto mimpl = MatIMPLCast(m);
264: const auto lda = mimpl->lda;
265: const auto src = mimpl->v;
266: const auto dest = mcu->d_v;
268: if (lda > nrows) {
269: PetscCall(PetscCUPMMemcpy2DAsync(dest, lda, src, lda, nrows, ncols, cupmMemcpyHostToDevice, stream));
270: } else {
271: PetscCall(PetscCUPMMemcpyAsync(dest, src, lda * ncols, cupmMemcpyHostToDevice, stream));
272: }
273: }
274: PetscCall(PetscLogEventEnd(MAT_DenseCopyToGPU, m, 0, 0, 0));
275: // order important, ensure that offloadmask is PETSC_OFFLOAD_BOTH
276: m->offloadmask = PETSC_OFFLOAD_BOTH;
277: }
278: PetscFunctionReturn(PETSC_SUCCESS);
279: }
281: template <device::cupm::DeviceType T>
282: inline PetscErrorCode MatDense_Seq_CUPM<T>::DeviceToHost_(Mat m, PetscDeviceContext dctx) noexcept
283: {
284: const auto nrows = m->rmap->n;
285: const auto ncols = m->cmap->n;
286: const auto copy = m->offloadmask == PETSC_OFFLOAD_GPU;
288: PetscFunctionBegin;
289: PetscCheckTypeName(m, MATSEQDENSECUPM());
290: PetscCall(PetscInfo(m, "%s matrix %" PetscInt_FMT " x %" PetscInt_FMT "\n", copy ? "Copy" : "Reusing", nrows, ncols));
291: if (copy) {
292: const auto mimpl = MatIMPLCast(m);
293: cupmStream_t stream;
295: // MatCreateSeqDenseCUPM may not allocate CPU memory. Allocate if needed
296: if (!mimpl->v) PetscCall(MatSeqDenseSetPreallocation(m, nullptr));
297: PetscCall(GetHandlesFrom_(dctx, &stream));
298: PetscCall(PetscLogEventBegin(MAT_DenseCopyFromGPU, m, 0, 0, 0));
299: {
300: const auto lda = mimpl->lda;
301: const auto dest = mimpl->v;
302: const auto src = MatCUPMCast(m)->d_v;
304: if (lda > nrows) {
305: PetscCall(PetscCUPMMemcpy2DAsync(dest, lda, src, lda, nrows, ncols, cupmMemcpyDeviceToHost, stream));
306: } else {
307: PetscCall(PetscCUPMMemcpyAsync(dest, src, lda * ncols, cupmMemcpyDeviceToHost, stream));
308: }
309: }
310: PetscCall(PetscLogEventEnd(MAT_DenseCopyFromGPU, m, 0, 0, 0));
311: // order is important, MatSeqDenseSetPreallocation() might set offloadmask
312: m->offloadmask = PETSC_OFFLOAD_BOTH;
313: }
314: PetscFunctionReturn(PETSC_SUCCESS);
315: }
317: template <device::cupm::DeviceType T>
318: inline PetscErrorCode MatDense_Seq_CUPM<T>::CheckCUPMSolverInfo_(const cupmBlasInt_t *fact_info, cupmStream_t stream) noexcept
319: {
320: PetscFunctionBegin;
321: if (PetscDefined(USE_DEBUG)) {
322: cupmBlasInt_t info = 0;
324: PetscCall(PetscCUPMMemcpyAsync(&info, fact_info, 1, cupmMemcpyDeviceToHost, stream));
325: if (stream) PetscCallCUPM(cupmStreamSynchronize(stream));
326: static_assert(std::is_same<decltype(info), int>::value, "");
327: PetscCheck(info <= 0, PETSC_COMM_SELF, PETSC_ERR_MAT_CH_ZRPVT, "Bad factorization: zero pivot in row %d", info - 1);
328: PetscCheck(info >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Wrong argument to cupmSolver %d", -info);
329: }
330: PetscFunctionReturn(PETSC_SUCCESS);
331: }
333: // ==========================================================================================
334: // MatDense_Seq_CUPM - Private API - Solver Dispatch
335: // ==========================================================================================
337: // specific solvers called through the dispatch_() family of functions
338: template <device::cupm::DeviceType T>
339: template <typename Derived>
340: struct MatDense_Seq_CUPM<T>::SolveCommon {
341: using derived_type = Derived;
343: template <typename F>
344: static PetscErrorCode ResizeFactLwork(Mat_SeqDenseCUPM *mcu, cupmStream_t stream, F &&cupmSolverComputeFactLwork) noexcept
345: {
346: cupmBlasInt_t lwork;
348: PetscFunctionBegin;
349: PetscCallCUPMSOLVER(cupmSolverComputeFactLwork(&lwork));
350: if (lwork > mcu->d_fact_lwork) {
351: mcu->d_fact_lwork = lwork;
352: PetscCallCUPM(cupmFreeAsync(mcu->d_fact_work, stream));
353: PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_work, lwork, stream));
354: }
355: PetscFunctionReturn(PETSC_SUCCESS);
356: }
358: static PetscErrorCode FactorPrepare(Mat A, cupmStream_t stream) noexcept
359: {
360: const auto mcu = MatCUPMCast(A);
362: PetscFunctionBegin;
363: PetscCall(PetscInfo(A, "%s factor %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", derived_type::NAME(), A->rmap->n, A->cmap->n));
364: A->factortype = derived_type::MATFACTORTYPE();
365: A->ops->solve = MatSolve_Factored_Dispatch_<derived_type, false>;
366: A->ops->solvetranspose = MatSolve_Factored_Dispatch_<derived_type, true>;
367: A->ops->matsolve = MatMatSolve_Factored_Dispatch_<derived_type, false>;
368: A->ops->matsolvetranspose = MatMatSolve_Factored_Dispatch_<derived_type, true>;
370: PetscCall(PetscStrFreeAllocpy(MATSOLVERCUPM(), &A->solvertype));
371: if (!mcu->d_fact_info) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_info, 1, stream));
372: PetscFunctionReturn(PETSC_SUCCESS);
373: }
374: };
376: template <device::cupm::DeviceType T>
377: struct MatDense_Seq_CUPM<T>::SolveLU : SolveCommon<SolveLU> {
378: using base_type = SolveCommon<SolveLU>;
380: static constexpr const char *NAME() noexcept { return "LU"; }
381: static constexpr MatFactorType MATFACTORTYPE() noexcept { return MAT_FACTOR_LU; }
383: static PetscErrorCode Factor(Mat A, IS, IS, const MatFactorInfo *) noexcept
384: {
385: const auto m = static_cast<cupmBlasInt_t>(A->rmap->n);
386: const auto n = static_cast<cupmBlasInt_t>(A->cmap->n);
387: cupmStream_t stream;
388: cupmSolverHandle_t handle;
389: PetscDeviceContext dctx;
391: PetscFunctionBegin;
392: if (!m || !n) PetscFunctionReturn(PETSC_SUCCESS);
393: PetscCall(GetHandles_(&dctx, &handle, &stream));
394: PetscCall(base_type::FactorPrepare(A, stream));
395: {
396: const auto mcu = MatCUPMCast(A);
397: const auto da = DeviceArrayReadWrite(dctx, A);
398: const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);
400: // clang-format off
401: PetscCall(
402: base_type::ResizeFactLwork(
403: mcu, stream,
404: [&](cupmBlasInt_t *fact_lwork)
405: {
406: return cupmSolverXgetrf_bufferSize(handle, m, n, da.cupmdata(), lda, fact_lwork);
407: }
408: )
409: );
410: // clang-format on
411: if (!mcu->d_fact_ipiv) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_ipiv, n, stream));
413: PetscCall(PetscLogGpuTimeBegin());
414: PetscCallCUPMSOLVER(cupmSolverXgetrf(handle, m, n, da.cupmdata(), lda, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_ipiv, mcu->d_fact_info));
415: PetscCall(PetscLogGpuTimeEnd());
416: PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream));
417: }
418: PetscCall(PetscLogGpuFlops(2.0 * n * n * m / 3.0));
419: PetscFunctionReturn(PETSC_SUCCESS);
420: }
422: template <bool transpose>
423: static PetscErrorCode Solve(Mat A, cupmScalar_t *x, cupmBlasInt_t ldx, cupmBlasInt_t m, cupmBlasInt_t nrhs, cupmBlasInt_t k, PetscDeviceContext dctx, cupmStream_t stream) noexcept
424: {
425: const auto mcu = MatCUPMCast(A);
426: const auto fact_info = mcu->d_fact_info;
427: const auto fact_ipiv = mcu->d_fact_ipiv;
428: cupmSolverHandle_t handle;
430: PetscFunctionBegin;
431: PetscCall(GetHandlesFrom_(dctx, &handle));
432: PetscCall(PetscInfo(A, "%s solve %d x %d on backend\n", NAME(), m, k));
433: PetscCall(PetscLogGpuTimeBegin());
434: {
435: constexpr auto op = transpose ? CUPMSOLVER_OP_T : CUPMSOLVER_OP_N;
436: const auto da = DeviceArrayRead(dctx, A);
437: const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);
439: // clang-format off
440: PetscCall(
441: base_type::ResizeFactLwork(
442: mcu, stream,
443: [&](cupmBlasInt_t *lwork)
444: {
445: return cupmSolverXgetrs_bufferSize(
446: handle, op, m, nrhs, da.cupmdata(), lda, fact_ipiv, x, ldx, lwork
447: );
448: }
449: )
450: );
451: // clang-format on
452: PetscCallCUPMSOLVER(cupmSolverXgetrs(handle, op, m, nrhs, da.cupmdata(), lda, fact_ipiv, x, ldx, mcu->d_fact_work, mcu->d_fact_lwork, fact_info));
453: PetscCall(CheckCUPMSolverInfo_(fact_info, stream));
454: }
455: PetscCall(PetscLogGpuTimeEnd());
456: PetscCall(PetscLogGpuFlops(nrhs * (2.0 * m * m - m)));
457: PetscFunctionReturn(PETSC_SUCCESS);
458: }
459: };
461: template <device::cupm::DeviceType T>
462: struct MatDense_Seq_CUPM<T>::SolveCholesky : SolveCommon<SolveCholesky> {
463: using base_type = SolveCommon<SolveCholesky>;
465: static constexpr const char *NAME() noexcept { return "Cholesky"; }
466: static constexpr MatFactorType MATFACTORTYPE() noexcept { return MAT_FACTOR_CHOLESKY; }
468: static PetscErrorCode Factor(Mat A, IS, const MatFactorInfo *) noexcept
469: {
470: const auto n = static_cast<cupmBlasInt_t>(A->rmap->n);
471: PetscDeviceContext dctx;
472: cupmSolverHandle_t handle;
473: cupmStream_t stream;
475: PetscFunctionBegin;
476: if (!n || !A->cmap->n) PetscFunctionReturn(PETSC_SUCCESS);
477: PetscCheck(A->spd == PETSC_BOOL3_TRUE, PETSC_COMM_SELF, PETSC_ERR_SUP, "%ssytrs unavailable. Use MAT_FACTOR_LU", cupmSolverName());
478: PetscCall(GetHandles_(&dctx, &handle, &stream));
479: PetscCall(base_type::FactorPrepare(A, stream));
480: {
481: const auto mcu = MatCUPMCast(A);
482: const auto da = DeviceArrayReadWrite(dctx, A);
483: const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);
485: // clang-format off
486: PetscCall(
487: base_type::ResizeFactLwork(
488: mcu, stream,
489: [&](cupmBlasInt_t *fact_lwork)
490: {
491: return cupmSolverXpotrf_bufferSize(
492: handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, fact_lwork
493: );
494: }
495: )
496: );
497: // clang-format on
498: PetscCall(PetscLogGpuTimeBegin());
499: PetscCallCUPMSOLVER(cupmSolverXpotrf(handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_info));
500: PetscCall(PetscLogGpuTimeEnd());
501: PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream));
502: }
503: PetscCall(PetscLogGpuFlops(1.0 * n * n * n / 3.0));
505: #if 0
506: // At the time of writing this interface (cuda 10.0), cusolverDn does not implement *sytrs
507: // and *hetr* routines. The code below should work, and it can be activated when *sytrs
508: // routines will be available
509: if (!mcu->d_fact_ipiv) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_ipiv, n, stream));
510: if (!mcu->d_fact_lwork) {
511: PetscCallCUPMSOLVER(cupmSolverDnXsytrf_bufferSize(handle, n, da.cupmdata(), lda, &mcu->d_fact_lwork));
512: PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_work, mcu->d_fact_lwork, stream));
513: }
514: if (mcu->d_fact_info) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_info, 1, stream));
515: PetscCall(PetscLogGpuTimeBegin());
516: PetscCallCUPMSOLVER(cupmSolverXsytrf(handle, CUPMSOLVER_FILL_MODE_LOWER, n, da, lda, mcu->d_fact_ipiv, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_info));
517: PetscCall(PetscLogGpuTimeEnd());
518: #endif
519: PetscFunctionReturn(PETSC_SUCCESS);
520: }
522: template <bool transpose>
523: static PetscErrorCode Solve(Mat A, cupmScalar_t *x, cupmBlasInt_t ldx, cupmBlasInt_t m, cupmBlasInt_t nrhs, cupmBlasInt_t k, PetscDeviceContext dctx, cupmStream_t stream) noexcept
524: {
525: const auto mcu = MatCUPMCast(A);
526: const auto fact_info = mcu->d_fact_info;
527: cupmSolverHandle_t handle;
529: PetscFunctionBegin;
530: PetscAssert(!mcu->d_fact_ipiv, PETSC_COMM_SELF, PETSC_ERR_LIB, "%ssytrs not implemented", cupmSolverName());
531: PetscCall(GetHandlesFrom_(dctx, &handle));
532: PetscCall(PetscInfo(A, "%s solve %d x %d on backend\n", NAME(), m, k));
533: PetscCall(PetscLogGpuTimeBegin());
534: {
535: const auto da = DeviceArrayRead(dctx, A);
536: const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);
538: // clang-format off
539: PetscCall(
540: base_type::ResizeFactLwork(
541: mcu, stream,
542: [&](cupmBlasInt_t *lwork)
543: {
544: return cupmSolverXpotrs_bufferSize(
545: handle, CUPMSOLVER_FILL_MODE_LOWER, m, nrhs, da.cupmdata(), lda, x, ldx, lwork
546: );
547: }
548: )
549: );
550: // clang-format on
551: PetscCallCUPMSOLVER(cupmSolverXpotrs(handle, CUPMSOLVER_FILL_MODE_LOWER, m, nrhs, da.cupmdata(), lda, x, ldx, mcu->d_fact_work, mcu->d_fact_lwork, fact_info));
552: PetscCall(CheckCUPMSolverInfo_(fact_info, stream));
553: }
554: PetscCall(PetscLogGpuTimeEnd());
555: PetscCall(PetscLogGpuFlops(nrhs * (2.0 * m * m - m)));
556: PetscFunctionReturn(PETSC_SUCCESS);
557: }
558: };
560: template <device::cupm::DeviceType T>
561: struct MatDense_Seq_CUPM<T>::SolveQR : SolveCommon<SolveQR> {
562: using base_type = SolveCommon<SolveQR>;
564: static constexpr const char *NAME() noexcept { return "QR"; }
565: static constexpr MatFactorType MATFACTORTYPE() noexcept { return MAT_FACTOR_QR; }
567: static PetscErrorCode Factor(Mat A, IS, const MatFactorInfo *) noexcept
568: {
569: const auto m = static_cast<cupmBlasInt_t>(A->rmap->n);
570: const auto n = static_cast<cupmBlasInt_t>(A->cmap->n);
571: const auto min = std::min(m, n);
572: const auto mimpl = MatIMPLCast(A);
573: cupmStream_t stream;
574: cupmSolverHandle_t handle;
575: PetscDeviceContext dctx;
577: PetscFunctionBegin;
578: if (!m || !n) PetscFunctionReturn(PETSC_SUCCESS);
579: PetscCall(GetHandles_(&dctx, &handle, &stream));
580: PetscCall(base_type::FactorPrepare(A, stream));
581: mimpl->rank = min;
582: {
583: const auto mcu = MatCUPMCast(A);
584: const auto da = DeviceArrayReadWrite(dctx, A);
585: const auto lda = static_cast<cupmBlasInt_t>(mimpl->lda);
587: if (!mcu->workvec) PetscCall(vec::cupm::VecCreateSeqCUPMAsync<T>(PetscObjectComm(PetscObjectCast(A)), m, &mcu->workvec));
588: if (!mcu->d_fact_tau) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_tau, min, stream));
589: // clang-format off
590: PetscCall(
591: base_type::ResizeFactLwork(
592: mcu, stream,
593: [&](cupmBlasInt_t *fact_lwork)
594: {
595: return cupmSolverXgeqrf_bufferSize(handle, m, n, da.cupmdata(), lda, fact_lwork);
596: }
597: )
598: );
599: // clang-format on
600: PetscCall(PetscLogGpuTimeBegin());
601: PetscCallCUPMSOLVER(cupmSolverXgeqrf(handle, m, n, da.cupmdata(), lda, mcu->d_fact_tau, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_info));
602: PetscCall(PetscLogGpuTimeEnd());
603: PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream));
604: }
605: PetscCall(PetscLogGpuFlops(2.0 * min * min * (std::max(m, n) - min / 3.0)));
606: PetscFunctionReturn(PETSC_SUCCESS);
607: }
609: template <bool transpose>
610: static PetscErrorCode Solve(Mat A, cupmScalar_t *x, cupmBlasInt_t ldx, cupmBlasInt_t m, cupmBlasInt_t nrhs, cupmBlasInt_t k, PetscDeviceContext dctx, cupmStream_t stream) noexcept
611: {
612: const auto mimpl = MatIMPLCast(A);
613: const auto rank = static_cast<cupmBlasInt_t>(mimpl->rank);
614: const auto mcu = MatCUPMCast(A);
615: const auto fact_info = mcu->d_fact_info;
616: const auto fact_tau = mcu->d_fact_tau;
617: const auto fact_work = mcu->d_fact_work;
618: const auto fact_lwork = mcu->d_fact_lwork;
619: cupmSolverHandle_t solver_handle;
620: cupmBlasHandle_t blas_handle;
622: PetscFunctionBegin;
623: PetscCall(GetHandlesFrom_(dctx, &blas_handle, &solver_handle));
624: PetscCall(PetscInfo(A, "%s solve %d x %d on backend\n", NAME(), m, k));
625: PetscCall(PetscLogGpuTimeBegin());
626: {
627: const auto da = DeviceArrayRead(dctx, A);
628: const auto one = cupmScalarCast(1.0);
629: const auto lda = static_cast<cupmBlasInt_t>(mimpl->lda);
631: if (transpose) {
632: PetscCallCUPMBLAS(cupmBlasXtrsm(blas_handle, CUPMBLAS_SIDE_LEFT, CUPMBLAS_FILL_MODE_UPPER, CUPMBLAS_OP_T, CUPMBLAS_DIAG_NON_UNIT, rank, nrhs, &one, da.cupmdata(), lda, x, ldx));
633: PetscCallCUPMSOLVER(cupmSolverXormqr(solver_handle, CUPMSOLVER_SIDE_LEFT, CUPMSOLVER_OP_N, m, nrhs, rank, da.cupmdata(), lda, fact_tau, x, ldx, fact_work, fact_lwork, fact_info));
634: PetscCall(CheckCUPMSolverInfo_(fact_info, stream));
635: } else {
636: constexpr auto op = PetscDefined(USE_COMPLEX) ? CUPMSOLVER_OP_C : CUPMSOLVER_OP_T;
638: PetscCallCUPMSOLVER(cupmSolverXormqr(solver_handle, CUPMSOLVER_SIDE_LEFT, op, m, nrhs, rank, da.cupmdata(), lda, fact_tau, x, ldx, fact_work, fact_lwork, fact_info));
639: PetscCall(CheckCUPMSolverInfo_(fact_info, stream));
640: PetscCallCUPMBLAS(cupmBlasXtrsm(blas_handle, CUPMBLAS_SIDE_LEFT, CUPMBLAS_FILL_MODE_UPPER, CUPMBLAS_OP_N, CUPMBLAS_DIAG_NON_UNIT, rank, nrhs, &one, da.cupmdata(), lda, x, ldx));
641: }
642: }
643: PetscCall(PetscLogGpuTimeEnd());
644: PetscCall(PetscLogFlops(nrhs * (4.0 * m * rank - (rank * rank))));
645: PetscFunctionReturn(PETSC_SUCCESS);
646: }
647: };
649: template <device::cupm::DeviceType T>
650: template <typename Solver, bool transpose>
651: inline PetscErrorCode MatDense_Seq_CUPM<T>::MatSolve_Factored_Dispatch_(Mat A, Vec x, Vec y) noexcept
652: {
653: using namespace vec::cupm;
654: const auto pobj_A = PetscObjectCast(A);
655: const auto m = static_cast<cupmBlasInt_t>(A->rmap->n);
656: const auto k = static_cast<cupmBlasInt_t>(A->cmap->n);
657: auto &workvec = MatCUPMCast(A)->workvec;
658: PetscScalar *y_array = nullptr;
659: PetscDeviceContext dctx;
660: PetscBool xiscupm, yiscupm, aiscupm;
661: bool use_y_array_directly;
662: cupmStream_t stream;
664: PetscFunctionBegin;
665: PetscCheck(A->factortype != MAT_FACTOR_NONE, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix must be factored to solve");
666: PetscCall(PetscObjectTypeCompare(PetscObjectCast(x), VecSeq_CUPM::VECSEQCUPM(), &xiscupm));
667: PetscCall(PetscObjectTypeCompare(PetscObjectCast(y), VecSeq_CUPM::VECSEQCUPM(), &yiscupm));
668: PetscCall(PetscObjectTypeCompare(pobj_A, MATSEQDENSECUPM(), &aiscupm));
669: PetscAssert(aiscupm, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Matrix A is somehow not CUPM?????????????????????????????");
670: PetscCall(GetHandles_(&dctx, &stream));
671: use_y_array_directly = yiscupm && (k >= m);
672: {
673: const PetscScalar *x_array;
674: const auto xisdevice = xiscupm && PetscOffloadDevice(x->offloadmask);
675: const auto copy_mode = xisdevice ? cupmMemcpyDeviceToDevice : cupmMemcpyHostToDevice;
677: if (!use_y_array_directly && !workvec) PetscCall(VecCreateSeqCUPMAsync<T>(PetscObjectComm(pobj_A), m, &workvec));
678: // The logic here is to try to minimize the amount of memory copying:
679: //
680: // If we call VecCUPMGetArrayRead(X, &x) every time xiscupm and the data is not offloaded
681: // to the GPU yet, then the data is copied to the GPU. But we are only trying to get the
682: // data in order to copy it into the y array. So the array x will be wherever the data
683: // already is so that only one memcpy is performed
684: if (xisdevice) {
685: PetscCall(VecCUPMGetArrayReadAsync<T>(x, &x_array, dctx));
686: } else {
687: PetscCall(VecGetArrayRead(x, &x_array));
688: }
689: PetscCall(VecCUPMGetArrayWriteAsync<T>(use_y_array_directly ? y : workvec, &y_array, dctx));
690: PetscCall(PetscCUPMMemcpyAsync(y_array, x_array, m, copy_mode, stream));
691: if (xisdevice) {
692: PetscCall(VecCUPMRestoreArrayReadAsync<T>(x, &x_array, dctx));
693: } else {
694: PetscCall(VecRestoreArrayRead(x, &x_array));
695: }
696: }
698: if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
699: PetscCall(Solver{}.template Solve<transpose>(A, cupmScalarPtrCast(y_array), m, m, 1, k, dctx, stream));
700: if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A));
702: if (use_y_array_directly) {
703: PetscCall(VecCUPMRestoreArrayWriteAsync<T>(y, &y_array, dctx));
704: } else {
705: const auto copy_mode = yiscupm ? cupmMemcpyDeviceToDevice : cupmMemcpyDeviceToHost;
706: PetscScalar *yv;
708: // The logic here is that the data is not yet in either y's GPU array or its CPU array.
709: // There is nothing in the interface to say where the user would like it to end up. So we
710: // choose the GPU, because it is the faster option
711: if (yiscupm) {
712: PetscCall(VecCUPMGetArrayWriteAsync<T>(y, &yv, dctx));
713: } else {
714: PetscCall(VecGetArray(y, &yv));
715: }
716: PetscCall(PetscCUPMMemcpyAsync(yv, y_array, k, copy_mode, stream));
717: if (yiscupm) {
718: PetscCall(VecCUPMRestoreArrayWriteAsync<T>(y, &yv, dctx));
719: } else {
720: PetscCall(VecRestoreArray(y, &yv));
721: }
722: PetscCall(VecCUPMRestoreArrayWriteAsync<T>(workvec, &y_array));
723: }
724: PetscFunctionReturn(PETSC_SUCCESS);
725: }
727: template <device::cupm::DeviceType T>
728: template <typename Solver, bool transpose>
729: inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMatSolve_Factored_Dispatch_(Mat A, Mat B, Mat X) noexcept
730: {
731: const auto m = static_cast<cupmBlasInt_t>(A->rmap->n);
732: const auto k = static_cast<cupmBlasInt_t>(A->cmap->n);
733: cupmBlasInt_t nrhs, ldb, ldx, ldy;
734: PetscScalar *y;
735: PetscBool biscupm, xiscupm, aiscupm;
736: PetscDeviceContext dctx;
737: cupmStream_t stream;
739: PetscFunctionBegin;
740: PetscCheck(A->factortype != MAT_FACTOR_NONE, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix must be factored to solve");
741: PetscCall(PetscObjectTypeCompare(PetscObjectCast(B), MATSEQDENSECUPM(), &biscupm));
742: PetscCall(PetscObjectTypeCompare(PetscObjectCast(X), MATSEQDENSECUPM(), &xiscupm));
743: PetscCall(PetscObjectTypeCompare(PetscObjectCast(A), MATSEQDENSECUPM(), &aiscupm));
744: PetscCall(GetHandles_(&dctx, &stream));
745: {
746: PetscInt n;
748: PetscCall(MatGetSize(B, nullptr, &n));
749: PetscCall(PetscCUPMBlasIntCast(n, &nrhs));
750: PetscCall(MatDenseGetLDA(B, &n));
751: PetscCall(PetscCUPMBlasIntCast(n, &ldb));
752: PetscCall(MatDenseGetLDA(X, &n));
753: PetscCall(PetscCUPMBlasIntCast(n, &ldx));
754: }
755: {
756: // The logic here is to try to minimize the amount of memory copying:
757: //
758: // If we call MatDenseCUPMGetArrayRead(B, &b) every time biscupm and the data is not
759: // offloaded to the GPU yet, then the data is copied to the GPU. But we are only trying to
760: // get the data in order to copy it into the y array. So the array b will be wherever the
761: // data already is so that only one memcpy is performed
762: const auto bisdevice = biscupm && PetscOffloadDevice(B->offloadmask);
763: const auto copy_mode = bisdevice ? cupmMemcpyDeviceToDevice : cupmMemcpyHostToDevice;
764: const PetscScalar *b;
766: if (bisdevice) {
767: b = DeviceArrayRead(dctx, B);
768: } else if (biscupm) {
769: b = HostArrayRead(dctx, B);
770: } else {
771: PetscCall(MatDenseGetArrayRead(B, &b));
772: }
774: if (ldx < m || !xiscupm) {
775: // X's array cannot serve as the array (too small or not on device), B's array cannot
776: // serve as the array (const), so allocate a new array
777: ldy = m;
778: PetscCall(PetscCUPMMallocAsync(&y, nrhs * m));
779: } else {
780: // X's array should serve as the array
781: ldy = ldx;
782: y = DeviceArrayWrite(dctx, X);
783: }
784: PetscCall(PetscCUPMMemcpy2DAsync(y, ldy, b, ldb, m, nrhs, copy_mode, stream));
785: if (!bisdevice && !biscupm) PetscCall(MatDenseRestoreArrayRead(B, &b));
786: }
788: // convert to CUPM twice??????????????????????????????????
789: // but A should already be CUPM??????????????????????????????????????
790: if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
791: PetscCall(Solver{}.template Solve<transpose>(A, cupmScalarPtrCast(y), ldy, m, nrhs, k, dctx, stream));
792: if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
794: if (ldx < m || !xiscupm) {
795: const auto copy_mode = xiscupm ? cupmMemcpyDeviceToDevice : cupmMemcpyDeviceToHost;
796: PetscScalar *x;
798: // The logic here is that the data is not yet in either X's GPU array or its CPU
799: // array. There is nothing in the interface to say where the user would like it to end up.
800: // So we choose the GPU, because it is the faster option
801: if (xiscupm) {
802: x = DeviceArrayWrite(dctx, X);
803: } else {
804: PetscCall(MatDenseGetArray(X, &x));
805: }
806: PetscCall(PetscCUPMMemcpy2DAsync(x, ldx, y, ldy, k, nrhs, copy_mode, stream));
807: if (!xiscupm) PetscCall(MatDenseRestoreArray(X, &x));
808: PetscCallCUPM(cupmFreeAsync(y, stream));
809: }
810: PetscFunctionReturn(PETSC_SUCCESS);
811: }
813: template <device::cupm::DeviceType T>
814: template <bool transpose, bool hermitian>
815: inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMultAddColumnRange_Dispatch_(Mat A, Vec xx, Vec yy, Vec zz, PetscInt c_start, PetscInt c_end) noexcept
816: {
817: const auto m = static_cast<cupmBlasInt_t>(A->rmap->n);
818: const auto n = static_cast<cupmBlasInt_t>(c_end - c_start);
819: const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);
820: PetscBool xiscupm, yiscupm, ziscupm;
821: cupmBlasHandle_t handle;
822: Vec x = xx, y = yy, z = zz;
823: PetscDeviceContext dctx;
825: PetscFunctionBegin;
826: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(xx), &xiscupm, VecSeq_CUPM::VECSEQCUPM(), VecSeq_CUPM::VECMPICUPM(), VecSeq_CUPM::VECCUPM(), ""));
827: if (!xiscupm || xx->boundtocpu) {
828: PetscCall(VecCreate(PetscObjectComm(PetscObjectCast(xx)), &x));
829: PetscCall(VecSetLayout(x, xx->map));
830: PetscCall(VecSetType(x, VecSeq_CUPM::VECCUPM()));
831: PetscCall(VecCopy(xx, x));
832: }
834: if (yy) {
835: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(yy), &yiscupm, VecSeq_CUPM::VECSEQCUPM(), VecSeq_CUPM::VECMPICUPM(), VecSeq_CUPM::VECCUPM(), ""));
836: if (!yiscupm || yy->boundtocpu) {
837: PetscCall(VecCreate(PetscObjectComm(PetscObjectCast(yy)), &y));
838: PetscCall(VecSetLayout(y, yy->map));
839: PetscCall(VecSetType(y, VecSeq_CUPM::VECCUPM()));
840: PetscCall(VecCopy(yy, y));
841: }
842: }
844: if (zz != yy) {
845: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(zz), &ziscupm, VecSeq_CUPM::VECSEQCUPM(), VecSeq_CUPM::VECMPICUPM(), VecSeq_CUPM::VECCUPM(), ""));
846: if (!ziscupm || zz->boundtocpu) {
847: PetscCall(VecCreate(PetscObjectComm(PetscObjectCast(zz)), &z));
848: PetscCall(VecSetLayout(z, zz->map));
849: PetscCall(VecSetType(z, VecSeq_CUPM::VECCUPM()));
850: }
851: } else {
852: z = y;
853: }
855: if (y && y != z) PetscCall(VecSeq_CUPM::Copy(y, z)); // mult add
856: if (!m || !n) {
857: // mult only
858: if (!y) PetscCall(VecSeq_CUPM::Set(z, 0.0));
859: PetscFunctionReturn(PETSC_SUCCESS);
860: }
861: PetscCall(PetscInfo(A, "Matrix-vector product %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " on backend\n", m, n));
862: PetscCall(GetHandles_(&dctx, &handle));
863: {
864: constexpr auto op = transpose ? (hermitian ? CUPMBLAS_OP_C : CUPMBLAS_OP_T) : CUPMBLAS_OP_N;
865: const auto one = cupmScalarCast(1.0);
866: const auto zero = cupmScalarCast(0.0);
867: const auto da = DeviceArrayRead(dctx, A);
868: const auto dxx = VecSeq_CUPM::DeviceArrayRead(dctx, x);
869: const auto dzz = VecSeq_CUPM::DeviceArrayReadWrite(dctx, z);
871: PetscCall(PetscLogGpuTimeBegin());
872: PetscCallCUPMBLAS(cupmBlasXgemv(handle, op, m, n, &one, da.cupmdata() + c_start * lda, lda, dxx.cupmdata() + (transpose ? 0 : c_start), 1, y ? &one : &zero, dzz.cupmdata() + (transpose ? c_start : 0), 1));
873: PetscCall(PetscLogGpuTimeEnd());
874: }
875: PetscCall(PetscLogGpuFlops(2.0 * m * n - (yy ? 0 : m)));
876: if (z != zz) {
877: PetscCall(VecCopy(z, zz));
878: if (z != y) PetscCall(VecDestroy(&z));
879: }
880: if (y != yy) PetscCall(VecDestroy(&y));
881: if (x != xx) PetscCall(VecDestroy(&x));
882: PetscCall(PetscDeviceContextSynchronizeIfWithBarrier_Internal(dctx));
883: PetscFunctionReturn(PETSC_SUCCESS);
884: }
886: template <device::cupm::DeviceType T>
887: template <bool transpose, bool hermitian>
888: inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMultColumnRange_Dispatch_(Mat A, Vec xx, Vec yy, PetscInt c_start, PetscInt c_end) noexcept
889: {
890: PetscFunctionBegin;
891: PetscCall(MatMultAddColumnRange_Dispatch_<transpose, hermitian>(A, xx, nullptr, yy, c_start, c_end));
892: PetscFunctionReturn(PETSC_SUCCESS);
893: }
895: template <device::cupm::DeviceType T>
896: template <bool transpose, bool hermitian>
897: inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMultAdd_Dispatch_(Mat A, Vec xx, Vec yy, Vec zz) noexcept
898: {
899: PetscFunctionBegin;
900: PetscCall(MatMultAddColumnRange_Dispatch_<transpose, hermitian>(A, xx, yy, zz, 0, A->cmap->n));
901: PetscFunctionReturn(PETSC_SUCCESS);
902: }
904: // ==========================================================================================
905: // MatDense_Seq_CUPM - Private API - Conversion Dispatch
906: // ==========================================================================================
908: template <device::cupm::DeviceType T>
909: template <bool to_host>
910: inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_Dispatch_(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept
911: {
912: PetscFunctionBegin;
913: if (reuse == MAT_REUSE_MATRIX || reuse == MAT_INITIAL_MATRIX) {
914: // TODO these cases should be optimized
915: PetscCall(MatConvert_Basic(M, type, reuse, newmat));
916: } else {
917: const auto B = *newmat;
918: const auto pobj = PetscObjectCast(B);
920: if (to_host) {
921: PetscCall(BindToCPU(B, PETSC_TRUE));
922: PetscCall(Reset(B));
923: } else {
924: PetscCall(PetscDeviceInitialize(PETSC_DEVICE_CUPM()));
925: }
927: PetscCall(PetscStrFreeAllocpy(to_host ? VECSTANDARD : VecSeq_CUPM::VECCUPM(), &B->defaultvectype));
928: PetscCall(PetscObjectChangeTypeName(pobj, to_host ? MATSEQDENSE : MATSEQDENSECUPM()));
929: // cvec might be the wrong VecType, destroy and rebuild it if necessary
930: // REVIEW ME: this is possibly very inefficient
931: PetscCall(VecDestroy(&MatIMPLCast(B)->cvec));
933: MatComposeOp_CUPM(to_host, pobj, MatConvert_seqdensecupm_seqdense_C(), nullptr, Convert_SeqDenseCUPM_SeqDense);
934: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArray_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ_WRITE>);
935: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArrayRead_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ>);
936: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArrayWrite_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_WRITE>);
937: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArray_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ_WRITE>);
938: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArrayRead_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ>);
939: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArrayWrite_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_WRITE>);
940: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMPlaceArray_C(), nullptr, PlaceArray);
941: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMResetArray_C(), nullptr, ResetArray);
942: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMReplaceArray_C(), nullptr, ReplaceArray);
943: MatComposeOp_CUPM(to_host, pobj, MatProductSetFromOptions_seqaij_seqdensecupm_C(), nullptr, MatProductSetFromOptions_SeqAIJ_SeqDense);
944: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMSetPreallocation_C(), nullptr, SetPreallocation);
946: if (to_host) {
947: B->offloadmask = PETSC_OFFLOAD_CPU;
948: } else {
949: Mat_SeqDenseCUPM *mcu;
951: PetscCall(PetscNew(&mcu));
952: B->spptr = mcu;
953: B->offloadmask = PETSC_OFFLOAD_UNALLOCATED; // REVIEW ME: why not offload host??
954: PetscCall(BindToCPU(B, PETSC_FALSE));
955: }
957: MatSetOp_CUPM(to_host, B, bindtocpu, nullptr, BindToCPU);
958: MatSetOp_CUPM(to_host, B, getcurrentmemtype, nullptr, GetCurrentMemType);
959: MatSetOp_CUPM(to_host, B, destroy, MatDestroy_SeqDense, Destroy);
960: }
961: PetscFunctionReturn(PETSC_SUCCESS);
962: }
964: // ==========================================================================================
965: // MatDense_Seq_CUPM - Public API
966: // ==========================================================================================
968: template <device::cupm::DeviceType T>
969: inline constexpr MatType MatDense_Seq_CUPM<T>::MATIMPLCUPM_() noexcept
970: {
971: return MATSEQDENSECUPM();
972: }
974: template <device::cupm::DeviceType T>
975: inline constexpr typename MatDense_Seq_CUPM<T>::Mat_SeqDenseCUPM *MatDense_Seq_CUPM<T>::MatCUPMCast(Mat m) noexcept
976: {
977: return static_cast<Mat_SeqDenseCUPM *>(m->spptr);
978: }
980: template <device::cupm::DeviceType T>
981: inline constexpr Mat_SeqDense *MatDense_Seq_CUPM<T>::MatIMPLCast_(Mat m) noexcept
982: {
983: return static_cast<Mat_SeqDense *>(m->data);
984: }
986: template <device::cupm::DeviceType T>
987: inline constexpr const char *MatDense_Seq_CUPM<T>::MatConvert_seqdensecupm_seqdense_C() noexcept
988: {
989: return T == device::cupm::DeviceType::CUDA ? "MatConvert_seqdensecuda_seqdense_C" : "MatConvert_seqdensehip_seqdense_C";
990: }
992: template <device::cupm::DeviceType T>
993: inline constexpr const char *MatDense_Seq_CUPM<T>::MatProductSetFromOptions_seqaij_seqdensecupm_C() noexcept
994: {
995: return T == device::cupm::DeviceType::CUDA ? "MatProductSetFromOptions_seqaij_seqdensecuda_C" : "MatProductSetFromOptions_seqaij_seqdensehip_C";
996: }
998: // ==========================================================================================
1000: // MatCreate_SeqDenseCUPM()
1001: template <device::cupm::DeviceType T>
1002: inline PetscErrorCode MatDense_Seq_CUPM<T>::Create(Mat A) noexcept
1003: {
1004: PetscFunctionBegin;
1005: PetscCall(PetscDeviceInitialize(PETSC_DEVICE_CUPM()));
1006: PetscCall(MatCreate_SeqDense(A));
1007: PetscCall(Convert_SeqDense_SeqDenseCUPM(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
1008: PetscFunctionReturn(PETSC_SUCCESS);
1009: }
1011: template <device::cupm::DeviceType T>
1012: inline PetscErrorCode MatDense_Seq_CUPM<T>::Destroy(Mat A) noexcept
1013: {
1014: PetscFunctionBegin;
1015: // prevent copying back data if we own the data pointer
1016: if (!MatIMPLCast(A)->user_alloc) A->offloadmask = PETSC_OFFLOAD_CPU;
1017: PetscCall(Convert_SeqDenseCUPM_SeqDense(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A));
1018: PetscCall(MatDestroy_SeqDense(A));
1019: PetscFunctionReturn(PETSC_SUCCESS);
1020: }
1022: // obj->ops->setup()
1023: template <device::cupm::DeviceType T>
1024: inline PetscErrorCode MatDense_Seq_CUPM<T>::SetUp(Mat A) noexcept
1025: {
1026: PetscFunctionBegin;
1027: PetscCall(PetscLayoutSetUp(A->rmap));
1028: PetscCall(PetscLayoutSetUp(A->cmap));
1029: if (!A->preallocated) {
1030: PetscDeviceContext dctx;
1032: PetscCall(GetHandles_(&dctx));
1033: PetscCall(SetPreallocation(A, dctx, nullptr));
1034: }
1035: PetscFunctionReturn(PETSC_SUCCESS);
1036: }
1038: template <device::cupm::DeviceType T>
1039: inline PetscErrorCode MatDense_Seq_CUPM<T>::Reset(Mat A) noexcept
1040: {
1041: PetscFunctionBegin;
1042: if (const auto mcu = MatCUPMCast(A)) {
1043: cupmStream_t stream;
1045: PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME());
1046: PetscCall(GetHandles_(&stream));
1047: if (!mcu->d_user_alloc) PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream));
1048: PetscCallCUPM(cupmFreeAsync(mcu->d_fact_tau, stream));
1049: PetscCallCUPM(cupmFreeAsync(mcu->d_fact_ipiv, stream));
1050: PetscCallCUPM(cupmFreeAsync(mcu->d_fact_info, stream));
1051: PetscCallCUPM(cupmFreeAsync(mcu->d_fact_work, stream));
1052: PetscCall(VecDestroy(&mcu->workvec));
1053: PetscCall(PetscFree(A->spptr /* mcu */));
1054: }
1055: PetscFunctionReturn(PETSC_SUCCESS);
1056: }
1058: // ==========================================================================================
1060: template <device::cupm::DeviceType T>
1061: inline PetscErrorCode MatDense_Seq_CUPM<T>::BindToCPU(Mat A, PetscBool to_host) noexcept
1062: {
1063: const auto mimpl = MatIMPLCast(A);
1064: const auto pobj = PetscObjectCast(A);
1066: PetscFunctionBegin;
1067: PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
1068: PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1069: A->boundtocpu = to_host;
1070: PetscCall(PetscStrFreeAllocpy(to_host ? PETSCRANDER48 : PETSCDEVICERAND(), &A->defaultrandtype));
1071: if (to_host) {
1072: PetscDeviceContext dctx;
1074: // make sure we have an up-to-date copy on the CPU
1075: PetscCall(GetHandles_(&dctx));
1076: PetscCall(DeviceToHost_(A, dctx));
1077: } else {
1078: PetscBool iscupm;
1080: if (auto &cvec = mimpl->cvec) {
1081: PetscCall(PetscObjectTypeCompare(PetscObjectCast(cvec), VecSeq_CUPM::VECSEQCUPM(), &iscupm));
1082: if (!iscupm) PetscCall(VecDestroy(&cvec));
1083: }
1084: if (auto &cmat = mimpl->cmat) {
1085: PetscCall(PetscObjectTypeCompare(PetscObjectCast(cmat), MATSEQDENSECUPM(), &iscupm));
1086: if (!iscupm) PetscCall(MatDestroy(&cmat));
1087: }
1088: }
1090: // ============================================================
1091: // Composed ops
1092: // ============================================================
1093: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArray_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ_WRITE>);
1094: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayRead_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ>);
1095: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayWrite_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_WRITE>);
1096: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ_WRITE>);
1097: MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ_WRITE>);
1098: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayReadAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ>);
1099: MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayReadAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ>);
1100: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayWriteAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_WRITE>);
1101: MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayWriteAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_WRITE>);
1102: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVec_C", MatDenseGetColumnVec_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_READ_WRITE>);
1103: MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVec_C", MatDenseRestoreColumnVec_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_READ_WRITE>);
1104: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVecRead_C", MatDenseGetColumnVecRead_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_READ>);
1105: MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVecRead_C", MatDenseRestoreColumnVecRead_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_READ>);
1106: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVecWrite_C", MatDenseGetColumnVecWrite_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_WRITE>);
1107: MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVecWrite_C", MatDenseRestoreColumnVecWrite_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_WRITE>);
1108: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetSubMatrix_C", MatDenseGetSubMatrix_SeqDense, GetSubMatrix);
1109: MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreSubMatrix_C", MatDenseRestoreSubMatrix_SeqDense, RestoreSubMatrix);
1110: MatComposeOp_CUPM(to_host, pobj, "MatQRFactor_C", MatQRFactor_SeqDense, SolveQR::Factor);
1111: MatComposeOp_CUPM(to_host, pobj, "MatMultColumnRange_C", MatMultColumnRange_SeqDense, MatMultColumnRange_Dispatch_</* transpose */ false, /* hermitian */ false>);
1112: MatComposeOp_CUPM(to_host, pobj, "MatMultAddColumnRange_C", MatMultAddColumnRange_SeqDense, MatMultAddColumnRange_Dispatch_</* transpose */ false, /* hermitian */ false>);
1113: MatComposeOp_CUPM(to_host, pobj, "MatMultHermitianTransposeColumnRange_C", MatMultHermitianTransposeColumnRange_SeqDense, MatMultColumnRange_Dispatch_</* transpose */ true, /* hermitian */ true>);
1114: MatComposeOp_CUPM(to_host, pobj, "MatMultHermitianTransposeAddColumnRange_C", MatMultHermitianTransposeAddColumnRange_SeqDense, MatMultAddColumnRange_Dispatch_</* transpose */ true, /* hermitian */ true>);
1115: // always the same
1116: PetscCall(PetscObjectComposeFunction(pobj, "MatDenseSetLDA_C", MatDenseSetLDA_SeqDense));
1118: // ============================================================
1119: // Function pointer ops
1120: // ============================================================
1121: MatSetOp_CUPM(to_host, A, duplicate, MatDuplicate_SeqDense, Duplicate);
1122: MatSetOp_CUPM(to_host, A, mult, MatMult_SeqDense, [](Mat A, Vec xx, Vec yy) { return MatMultAdd_Dispatch_</* transpose */ false, /* hermitian */ false>(A, xx, nullptr, yy); });
1123: MatSetOp_CUPM(to_host, A, multtranspose, MatMultTranspose_SeqDense, [](Mat A, Vec xx, Vec yy) { return MatMultAdd_Dispatch_</* transpose */ true, /* hermitian */ false>(A, xx, nullptr, yy); });
1124: MatSetOp_CUPM(to_host, A, multhermitiantranspose, MatMultTranspose_SeqDense, [](Mat A, Vec xx, Vec yy) { return MatMultAdd_Dispatch_</* transpose */ true, /* hermitian */ true>(A, xx, nullptr, yy); });
1125: MatSetOp_CUPM(to_host, A, multadd, MatMultAdd_SeqDense, MatMultAdd_Dispatch_</* transpose */ false, /* hermitian */ false>);
1126: MatSetOp_CUPM(to_host, A, multtransposeadd, MatMultTransposeAdd_SeqDense, MatMultAdd_Dispatch_</* transpose */ true, /* hermitian */ false>);
1127: MatSetOp_CUPM(to_host, A, multhermitiantransposeadd, MatMultHermitianTransposeAdd_SeqDense, MatMultAdd_Dispatch_</* transpose */ true, /* hermitian */ true>);
1128: MatSetOp_CUPM(to_host, A, matmultnumeric, MatMatMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ false, /* transpose_B */ false>);
1129: MatSetOp_CUPM(to_host, A, mattransposemultnumeric, MatMatTransposeMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ false, /* transpose_B */ true>);
1130: MatSetOp_CUPM(to_host, A, transposematmultnumeric, MatTransposeMatMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ true, /* transpose_B */ false>);
1131: MatSetOp_CUPM(to_host, A, axpy, MatAXPY_SeqDense, AXPY);
1132: MatSetOp_CUPM(to_host, A, choleskyfactor, MatCholeskyFactor_SeqDense, SolveCholesky::Factor);
1133: MatSetOp_CUPM(to_host, A, lufactor, MatLUFactor_SeqDense, SolveLU::Factor);
1134: MatSetOp_CUPM(to_host, A, getcolumnvector, MatGetColumnVector_SeqDense, GetColumnVector);
1135: MatSetOp_CUPM(to_host, A, conjugate, MatConjugate_SeqDense, Conjugate);
1136: MatSetOp_CUPM(to_host, A, scale, MatScale_SeqDense, Scale);
1137: MatSetOp_CUPM(to_host, A, norm, MatNorm_SeqDense, Norm);
1138: MatSetOp_CUPM(to_host, A, diagonalscale, MatDiagonalScale_SeqDense, DiagonalScale);
1139: MatSetOp_CUPM(to_host, A, shift, MatShift_SeqDense, Shift);
1140: MatSetOp_CUPM(to_host, A, copy, MatCopy_SeqDense, Copy);
1141: MatSetOp_CUPM(to_host, A, zeroentries, MatZeroEntries_SeqDense, ZeroEntries);
1142: MatSetOp_CUPM(to_host, A, setup, MatSetUp_SeqDense, SetUp);
1143: MatSetOp_CUPM(to_host, A, setrandom, MatSetRandom_SeqDense, SetRandom);
1144: MatSetOp_CUPM(to_host, A, getdiagonal, MatGetDiagonal_SeqDense, GetDiagonal);
1145: // seemingly always the same
1146: A->ops->productsetfromoptions = MatProductSetFromOptions_SeqDense;
1148: if (const auto cmat = mimpl->cmat) PetscCall(MatBindToCPU(cmat, to_host));
1149: PetscFunctionReturn(PETSC_SUCCESS);
1150: }
1152: template <device::cupm::DeviceType T>
1153: inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_SeqDenseCUPM_SeqDense(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept
1154: {
1155: PetscFunctionBegin;
1156: PetscCall(Convert_Dispatch_</* to host */ true>(M, type, reuse, newmat));
1157: PetscFunctionReturn(PETSC_SUCCESS);
1158: }
1160: template <device::cupm::DeviceType T>
1161: inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_SeqDense_SeqDenseCUPM(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept
1162: {
1163: PetscFunctionBegin;
1164: PetscCall(Convert_Dispatch_</* to host */ false>(M, type, reuse, newmat));
1165: PetscFunctionReturn(PETSC_SUCCESS);
1166: }
1168: // ==========================================================================================
1170: template <device::cupm::DeviceType T>
1171: template <PetscMemType mtype, PetscMemoryAccessMode access>
1172: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetArray(Mat m, PetscScalar **array, PetscDeviceContext dctx) noexcept
1173: {
1174: constexpr auto hostmem = PetscMemTypeHost(mtype);
1175: constexpr auto read_access = PetscMemoryAccessRead(access);
1177: PetscFunctionBegin;
1178: static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), "");
1179: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
1180: if (hostmem) {
1181: if (read_access) {
1182: PetscCall(DeviceToHost_(m, dctx));
1183: } else if (!MatIMPLCast(m)->v) {
1184: // MatCreateSeqDenseCUPM may not allocate CPU memory. Allocate if needed
1185: PetscCall(MatSeqDenseSetPreallocation(m, nullptr));
1186: }
1187: *array = MatIMPLCast(m)->v;
1188: } else {
1189: if (read_access) {
1190: PetscCall(HostToDevice_(m, dctx));
1191: } else if (!MatCUPMCast(m)->d_v) {
1192: // write-only
1193: PetscCall(SetPreallocation(m, dctx, nullptr));
1194: }
1195: *array = MatCUPMCast(m)->d_v;
1196: }
1197: if (PetscMemoryAccessWrite(access)) {
1198: m->offloadmask = hostmem ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
1199: PetscCall(PetscObjectStateIncrease(PetscObjectCast(m)));
1200: }
1201: PetscFunctionReturn(PETSC_SUCCESS);
1202: }
1204: template <device::cupm::DeviceType T>
1205: template <PetscMemType mtype, PetscMemoryAccessMode access>
1206: inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreArray(Mat m, PetscScalar **array, PetscDeviceContext) noexcept
1207: {
1208: PetscFunctionBegin;
1209: static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), "");
1210: if (PetscMemoryAccessWrite(access)) {
1211: // WRITE or READ_WRITE
1212: m->offloadmask = PetscMemTypeHost(mtype) ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
1213: PetscCall(PetscObjectStateIncrease(PetscObjectCast(m)));
1214: }
1215: if (array) {
1216: PetscCall(CheckPointerMatchesMemType_(*array, mtype));
1217: *array = nullptr;
1218: }
1219: PetscFunctionReturn(PETSC_SUCCESS);
1220: }
1222: template <device::cupm::DeviceType T>
1223: template <PetscMemoryAccessMode access>
1224: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetArrayAndMemType(Mat m, PetscScalar **array, PetscMemType *mtype, PetscDeviceContext dctx) noexcept
1225: {
1226: PetscFunctionBegin;
1227: PetscCall(GetArray<PETSC_MEMTYPE_DEVICE, access>(m, array, dctx));
1228: if (mtype) *mtype = PETSC_MEMTYPE_CUPM();
1229: PetscFunctionReturn(PETSC_SUCCESS);
1230: }
1232: template <device::cupm::DeviceType T>
1233: template <PetscMemoryAccessMode access>
1234: inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreArrayAndMemType(Mat m, PetscScalar **array, PetscDeviceContext dctx) noexcept
1235: {
1236: PetscFunctionBegin;
1237: PetscCall(RestoreArray<PETSC_MEMTYPE_DEVICE, access>(m, array, dctx));
1238: PetscFunctionReturn(PETSC_SUCCESS);
1239: }
1241: // ==========================================================================================
1243: template <device::cupm::DeviceType T>
1244: inline PetscErrorCode MatDense_Seq_CUPM<T>::PlaceArray(Mat A, const PetscScalar *array) noexcept
1245: {
1246: const auto mimpl = MatIMPLCast(A);
1247: const auto mcu = MatCUPMCast(A);
1249: PetscFunctionBegin;
1250: PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
1251: PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1252: PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME());
1253: if (mimpl->v) {
1254: PetscDeviceContext dctx;
1256: PetscCall(GetHandles_(&dctx));
1257: PetscCall(HostToDevice_(A, dctx));
1258: }
1259: mcu->unplacedarray = util::exchange(mcu->d_v, const_cast<PetscScalar *>(array));
1260: mcu->d_unplaced_user_alloc = util::exchange(mcu->d_user_alloc, PETSC_TRUE);
1261: PetscFunctionReturn(PETSC_SUCCESS);
1262: }
1264: template <device::cupm::DeviceType T>
1265: inline PetscErrorCode MatDense_Seq_CUPM<T>::ReplaceArray(Mat A, const PetscScalar *array) noexcept
1266: {
1267: const auto mimpl = MatIMPLCast(A);
1268: const auto mcu = MatCUPMCast(A);
1270: PetscFunctionBegin;
1271: PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
1272: PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1273: PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME());
1274: if (!mcu->d_user_alloc) {
1275: cupmStream_t stream;
1277: PetscCall(GetHandles_(&stream));
1278: PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream));
1279: }
1280: mcu->d_v = const_cast<PetscScalar *>(array);
1281: mcu->d_user_alloc = PETSC_FALSE;
1282: PetscFunctionReturn(PETSC_SUCCESS);
1283: }
1285: template <device::cupm::DeviceType T>
1286: inline PetscErrorCode MatDense_Seq_CUPM<T>::ResetArray(Mat A) noexcept
1287: {
1288: const auto mimpl = MatIMPLCast(A);
1289: const auto mcu = MatCUPMCast(A);
1291: PetscFunctionBegin;
1292: PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
1293: PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1294: if (mimpl->v) {
1295: PetscDeviceContext dctx;
1297: PetscCall(GetHandles_(&dctx));
1298: PetscCall(HostToDevice_(A, dctx));
1299: }
1300: mcu->d_v = util::exchange(mcu->unplacedarray, nullptr);
1301: mcu->d_user_alloc = mcu->d_unplaced_user_alloc;
1302: PetscFunctionReturn(PETSC_SUCCESS);
1303: }
1305: // ==========================================================================================
1307: template <device::cupm::DeviceType T>
1308: template <bool transpose_A, bool transpose_B>
1309: inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMatMult_Numeric_Dispatch(Mat A, Mat B, Mat C) noexcept
1310: {
1311: cupmBlasInt_t m, n, k;
1312: PetscBool Aiscupm, Biscupm;
1313: PetscDeviceContext dctx;
1314: cupmBlasHandle_t handle;
1316: PetscFunctionBegin;
1317: PetscCall(PetscCUPMBlasIntCast(C->rmap->n, &m));
1318: PetscCall(PetscCUPMBlasIntCast(C->cmap->n, &n));
1319: PetscCall(PetscCUPMBlasIntCast(transpose_A ? A->rmap->n : A->cmap->n, &k));
1320: if (!m || !n || !k) {
1321: PetscCall(ZeroEntries(C));
1322: PetscFunctionReturn(PETSC_SUCCESS);
1323: }
1325: // we may end up with SEQDENSE as one of the arguments
1326: // REVIEW ME: how? and why is it not B and C????????
1327: PetscCall(PetscObjectTypeCompare(PetscObjectCast(A), MATSEQDENSECUPM(), &Aiscupm));
1328: PetscCall(PetscObjectTypeCompare(PetscObjectCast(B), MATSEQDENSECUPM(), &Biscupm));
1329: if (!Aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
1330: if (!Biscupm) PetscCall(MatConvert(B, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &B));
1331: PetscCall(PetscInfo(C, "Matrix-Matrix product %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " on backend\n", m, k, n));
1332: PetscCall(GetHandles_(&dctx, &handle));
1334: PetscCall(PetscLogGpuTimeBegin());
1335: {
1336: const auto one = cupmScalarCast(1.0);
1337: const auto zero = cupmScalarCast(0.0);
1338: const auto da = DeviceArrayRead(dctx, A);
1339: const auto db = DeviceArrayRead(dctx, B);
1340: const auto dc = DeviceArrayWrite(dctx, C);
1341: PetscInt alda, blda, clda;
1343: PetscCall(MatDenseGetLDA(A, &alda));
1344: PetscCall(MatDenseGetLDA(B, &blda));
1345: PetscCall(MatDenseGetLDA(C, &clda));
1346: PetscCallCUPMBLAS(cupmBlasXgemm(handle, transpose_A ? CUPMBLAS_OP_T : CUPMBLAS_OP_N, transpose_B ? CUPMBLAS_OP_T : CUPMBLAS_OP_N, m, n, k, &one, da.cupmdata(), alda, db.cupmdata(), blda, &zero, dc.cupmdata(), clda));
1347: }
1348: PetscCall(PetscLogGpuTimeEnd());
1350: PetscCall(PetscLogGpuFlops(1.0 * m * n * k + 1.0 * m * n * (k - 1)));
1351: if (!Aiscupm) PetscCall(MatConvert(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A));
1352: if (!Biscupm) PetscCall(MatConvert(B, MATSEQDENSE, MAT_INPLACE_MATRIX, &B));
1353: PetscFunctionReturn(PETSC_SUCCESS);
1354: }
1356: template <device::cupm::DeviceType T>
1357: inline PetscErrorCode MatDense_Seq_CUPM<T>::Copy(Mat A, Mat B, MatStructure str) noexcept
1358: {
1359: const auto m = A->rmap->n;
1360: const auto n = A->cmap->n;
1362: PetscFunctionBegin;
1363: PetscAssert(m == B->rmap->n && n == B->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "size(B) != size(A)");
1364: // The two matrices must have the same copy implementation to be eligible for fast copy
1365: if (A->ops->copy == B->ops->copy) {
1366: PetscDeviceContext dctx;
1367: cupmStream_t stream;
1369: PetscCall(GetHandles_(&dctx, &stream));
1370: PetscCall(PetscLogGpuTimeBegin());
1371: {
1372: const auto va = DeviceArrayRead(dctx, A);
1373: const auto vb = DeviceArrayWrite(dctx, B);
1374: // order is important, DeviceArrayRead/Write() might call SetPreallocation() which sets
1375: // lda!
1376: const auto lda_a = MatIMPLCast(A)->lda;
1377: const auto lda_b = MatIMPLCast(B)->lda;
1379: if (lda_a > m || lda_b > m) {
1380: PetscAssert(lda_b > 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "B lda (%" PetscBLASInt_FMT ") must be > 0 at this point, this indicates Mat%sSetPreallocation() was not called when it should have been!", lda_b, cupmNAME());
1381: PetscAssert(lda_a > 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "A lda (%" PetscBLASInt_FMT ") must be > 0 at this point, this indicates Mat%sSetPreallocation() was not called when it should have been!", lda_a, cupmNAME());
1382: PetscCall(PetscCUPMMemcpy2DAsync(vb.data(), lda_b, va.data(), lda_a, m, n, cupmMemcpyDeviceToDevice, stream));
1383: } else {
1384: PetscCall(PetscCUPMMemcpyAsync(vb.data(), va.data(), m * n, cupmMemcpyDeviceToDevice, stream));
1385: }
1386: }
1387: PetscCall(PetscLogGpuTimeEnd());
1388: } else {
1389: PetscCall(MatCopy_Basic(A, B, str));
1390: }
1391: PetscFunctionReturn(PETSC_SUCCESS);
1392: }
1394: template <device::cupm::DeviceType T>
1395: inline PetscErrorCode MatDense_Seq_CUPM<T>::ZeroEntries(Mat m) noexcept
1396: {
1397: PetscDeviceContext dctx;
1398: cupmStream_t stream;
1400: PetscFunctionBegin;
1401: PetscCall(GetHandles_(&dctx, &stream));
1402: PetscCall(PetscLogGpuTimeBegin());
1403: {
1404: const auto va = DeviceArrayWrite(dctx, m);
1405: const auto lda = MatIMPLCast(m)->lda;
1406: const auto ma = m->rmap->n;
1407: const auto na = m->cmap->n;
1409: if (lda > ma) {
1410: PetscCall(PetscCUPMMemset2DAsync(va.data(), lda, 0, ma, na, stream));
1411: } else {
1412: PetscCall(PetscCUPMMemsetAsync(va.data(), 0, ma * na, stream));
1413: }
1414: }
1415: PetscCall(PetscLogGpuTimeEnd());
1416: PetscFunctionReturn(PETSC_SUCCESS);
1417: }
1419: namespace detail
1420: {
1422: // ==========================================================================================
1423: // SubMatIndexFunctor
1424: //
1425: // Iterator which permutes a linear index range into matrix indices for an nrows x ncols
1426: // submat with leading dimension lda. Essentially SubMatIndexFunctor(i) returns the index for
1427: // the i'th sequential entry in the matrix.
1428: // ==========================================================================================
1429: template <typename T>
1430: struct SubMatIndexFunctor {
1431: PETSC_HOSTDEVICE_INLINE_DECL T operator()(T x) const noexcept { return ((x / nrows) * lda) + (x % nrows); }
1433: PetscInt nrows;
1434: PetscInt ncols;
1435: PetscInt lda;
1436: };
1438: template <typename Iterator>
1439: struct SubMatrixIterator : MatrixIteratorBase<Iterator, SubMatIndexFunctor<iter_difference_t<Iterator>>> {
1440: using base_type = MatrixIteratorBase<Iterator, SubMatIndexFunctor<iter_difference_t<Iterator>>>;
1442: using iterator = typename base_type::iterator;
1444: constexpr SubMatrixIterator(Iterator first, Iterator last, PetscInt nrows, PetscInt ncols, PetscInt lda) noexcept :
1445: base_type{
1446: std::move(first), std::move(last), {nrows, ncols, lda}
1447: }
1448: {
1449: }
1451: PETSC_NODISCARD iterator end() const noexcept { return this->begin() + (this->func.nrows * this->func.ncols); }
1452: };
1454: namespace
1455: {
1457: template <typename T>
1458: PETSC_NODISCARD inline SubMatrixIterator<typename thrust::device_vector<T>::iterator> make_submat_iterator(PetscInt rstart, PetscInt rend, PetscInt cstart, PetscInt cend, PetscInt lda, T *ptr) noexcept
1459: {
1460: const auto nrows = rend - rstart;
1461: const auto ncols = cend - cstart;
1462: const auto dptr = thrust::device_pointer_cast(ptr);
1464: return {dptr + (rstart * lda) + cstart, dptr + ((rstart + nrows) * lda) + cstart, nrows, ncols, lda};
1465: }
1467: } // namespace
1469: struct conjugate {
1470: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &x) const noexcept { return PetscConj(x); }
1471: };
1473: struct real_abs {
1474: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL PetscReal operator()(const PetscScalar &x) const noexcept { return PetscAbsScalar(x); }
1475: };
1477: struct real_abs_squared {
1478: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL PetscReal operator()(const PetscScalar &x) const noexcept { return PetscRealPart(PetscConj(x) * x); }
1479: };
1481: // ==========================================================================================
1482: // RowMajorIndexFunctor
1483: //
1484: // Iterator which permutes a linear row-major index range into the memory offsets of a matrix
1485: // with ncols columns which is stored column-major with leading dimension lda. Essentially
1486: // RowMajorIndexFunctor(k) returns the index of the k'th entry of the matrix when it is
1487: // traversed one row at a time.
1488: // ==========================================================================================
1489: template <typename T>
1490: struct RowMajorIndexFunctor {
1491: PETSC_HOSTDEVICE_INLINE_DECL T operator()(T x) const noexcept { return ((x % ncols) * lda) + (x / ncols); }
1493: PetscInt ncols;
1494: PetscInt lda;
1495: };
1497: // ==========================================================================================
1498: // GroupIndexFunctor
1499: //
1500: // Maps a linear index range onto the index of the group of group_size consecutive entries
1501: // that each index belongs to. Used to generate the keys of a segmented reduction, where each
1502: // group is one column (or one row) of the matrix.
1503: // ==========================================================================================
1504: template <typename T>
1505: struct GroupIndexFunctor {
1506: PETSC_HOSTDEVICE_INLINE_DECL T operator()(T x) const noexcept { return x / group_size; }
1508: PetscInt group_size;
1509: };
1511: } // namespace detail
1513: template <device::cupm::DeviceType T>
1514: inline PetscErrorCode MatDense_Seq_CUPM<T>::Conjugate(Mat A) noexcept
1515: {
1516: const auto m = A->rmap->n;
1517: const auto n = A->cmap->n;
1518: const auto N = m * n;
1519: PetscDeviceContext dctx;
1520: cupmStream_t stream;
1522: PetscFunctionBegin;
1523: if (PetscDefined(USE_COMPLEX)) {
1524: PetscCall(GetHandles_(&dctx, &stream));
1525: PetscCall(PetscLogGpuTimeBegin());
1526: {
1527: const auto da = DeviceArrayReadWrite(dctx, A);
1528: const auto lda = MatIMPLCast(A)->lda;
1529: cupmStream_t stream;
1530: PetscCall(GetHandlesFrom_(dctx, &stream));
1532: if (lda > m) {
1533: // clang-format off
1534: PetscCallThrust(
1535: const auto sub_mat = detail::make_submat_iterator(0, m, 0, n, lda, da.data());
1537: THRUST_CALL(
1538: thrust::transform,
1539: stream,
1540: sub_mat.begin(), sub_mat.end(), sub_mat.begin(),
1541: detail::conjugate{}
1542: )
1543: );
1544: // clang-format on
1545: } else {
1546: // clang-format off
1547: PetscCallThrust(
1548: const auto aptr = thrust::device_pointer_cast(da.data());
1550: THRUST_CALL(
1551: thrust::transform,
1552: stream,
1553: aptr, aptr + N, aptr,
1554: detail::conjugate{}
1555: )
1556: );
1557: // clang-format on
1558: }
1559: }
1560: PetscCall(PetscLogGpuTimeEnd());
1561: }
1562: PetscFunctionReturn(PETSC_SUCCESS);
1563: }
1565: template <device::cupm::DeviceType T>
1566: inline PetscErrorCode MatDense_Seq_CUPM<T>::Scale(Mat A, PetscScalar alpha) noexcept
1567: {
1568: const auto m = A->rmap->n;
1569: const auto n = A->cmap->n;
1570: const auto N = m * n;
1571: PetscDeviceContext dctx;
1573: PetscFunctionBegin;
1574: PetscCall(PetscInfo(A, "Performing Scale %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m, n));
1575: PetscCall(GetHandles_(&dctx));
1576: {
1577: const auto da = DeviceArrayReadWrite(dctx, A);
1578: const auto lda = MatIMPLCast(A)->lda;
1580: if (lda > m) {
1581: cupmStream_t stream;
1583: PetscCall(GetHandlesFrom_(dctx, &stream));
1584: // clang-format off
1585: PetscCallThrust(
1586: const auto sub_mat = detail::make_submat_iterator(0, m, 0, n, lda, da.data());
1588: THRUST_CALL(
1589: thrust::transform,
1590: stream,
1591: sub_mat.begin(), sub_mat.end(), sub_mat.begin(),
1592: device::cupm::functors::make_times_equals(alpha)
1593: )
1594: );
1595: // clang-format on
1596: } else {
1597: const auto cu_alpha = cupmScalarCast(alpha);
1598: cupmBlasHandle_t handle;
1600: PetscCall(GetHandlesFrom_(dctx, &handle));
1601: PetscCall(PetscLogGpuTimeBegin());
1602: PetscCallCUPMBLAS(cupmBlasXscal(handle, N, &cu_alpha, da.cupmdata(), 1));
1603: PetscCall(PetscLogGpuTimeEnd());
1604: }
1605: }
1606: PetscCall(PetscLogGpuFlops(N));
1607: PetscFunctionReturn(PETSC_SUCCESS);
1608: }
1610: template <device::cupm::DeviceType T>
1611: inline PetscErrorCode MatDense_Seq_CUPM<T>::Norm(Mat A, NormType type, PetscReal *nrm) noexcept
1612: {
1613: #if PetscDefined(USING_NVCC) && CCCL_VERSION >= 3001000
1614: using max_functor = cuda::maximum<PetscReal>;
1615: #else
1616: using max_functor = thrust::maximum<PetscReal>;
1617: #endif
1618: const auto m = A->rmap->n;
1619: const auto n = A->cmap->n;
1620: const auto N = m * n;
1621: PetscDeviceContext dctx;
1622: cupmStream_t stream;
1624: PetscFunctionBegin;
1625: PetscCheck(type == NORM_FROBENIUS || type == NORM_1 || type == NORM_INFINITY || type == NORM_2, PETSC_COMM_SELF, PETSC_ERR_SUP, "Unsupported norm type %s", NormTypes[type]);
1626: // NORM_2 is the largest singular value, which requires an SVD. There is no device gesvd()
1627: // wrapper available, so defer to MatNorm_SeqDense(), which copies the matrix to the host
1628: if (type == NORM_2) {
1629: PetscCall(MatNorm_SeqDense(A, type, nrm));
1630: PetscFunctionReturn(PETSC_SUCCESS);
1631: }
1632: if (!N) {
1633: *nrm = 0.0;
1634: PetscFunctionReturn(PETSC_SUCCESS);
1635: }
1636: PetscCall(PetscInfo(A, "Performing Norm %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m, n));
1637: PetscCall(GetHandles_(&dctx, &stream));
1638: {
1639: const auto da = DeviceArrayRead(dctx, A);
1640: const auto lda = static_cast<PetscInt>(MatIMPLCast(A)->lda);
1642: if (type == NORM_FROBENIUS) {
1643: if (lda > m) {
1644: PetscReal sum = 0.0;
1646: // clang-format off
1647: PetscCallThrust(
1648: const auto sub_mat = detail::make_submat_iterator(0, m, 0, n, lda, da.data());
1650: sum = THRUST_CALL(
1651: thrust::transform_reduce,
1652: stream,
1653: sub_mat.begin(), sub_mat.end(),
1654: detail::real_abs_squared{},
1655: PetscReal{0.0},
1656: thrust::plus<PetscReal>{}
1657: );
1658: );
1659: // clang-format on
1660: *nrm = PetscSqrtReal(sum);
1661: } else {
1662: cupmBlasHandle_t handle;
1664: PetscCall(GetHandlesFrom_(dctx, &handle));
1665: PetscCall(PetscLogGpuTimeBegin());
1666: PetscCallCUPMBLAS(cupmBlasXnrm2(handle, N, da.cupmdata(), 1, cupmRealPtrCast(nrm)));
1667: PetscCall(PetscLogGpuTimeEnd());
1668: }
1669: PetscCall(PetscLogGpuFlops(2.0 * N));
1670: } else {
1671: const auto ngroups = (type == NORM_1) ? n : m;
1672: const auto group_size = (type == NORM_1) ? m : n;
1673: const auto keys = thrust::make_transform_iterator(thrust::make_counting_iterator(PetscInt{0}), detail::GroupIndexFunctor<PetscInt>{group_size});
1674: PetscReal *sums = nullptr;
1676: PetscCall(PetscDeviceMalloc(dctx, PETSC_MEMTYPE_CUPM(), ngroups, &sums));
1677: if (type == NORM_1) {
1678: // the sub-matrix iterator enumerates the entries column-major, so each run of m
1679: // consecutive keys is exactly one column, for any lda
1680: // clang-format off
1681: PetscCallThrust(
1682: const auto sub_mat = detail::make_submat_iterator(0, m, 0, n, lda, da.data());
1684: THRUST_CALL(
1685: thrust::reduce_by_key,
1686: stream,
1687: keys, keys + N,
1688: thrust::make_transform_iterator(sub_mat.begin(), detail::real_abs{}),
1689: thrust::make_discard_iterator(),
1690: thrust::device_pointer_cast(sums)
1691: );
1692: );
1693: // clang-format on
1694: } else {
1695: // enumerate the entries row-major, so each run of n consecutive keys is exactly one row
1696: // clang-format off
1697: PetscCallThrust(
1698: const auto row_major = thrust::make_permutation_iterator(
1699: thrust::device_pointer_cast(da.data()),
1700: thrust::make_transform_iterator(thrust::make_counting_iterator(PetscInt{0}), detail::RowMajorIndexFunctor<PetscInt>{n, lda})
1701: );
1703: THRUST_CALL(
1704: thrust::reduce_by_key,
1705: stream,
1706: keys, keys + N,
1707: thrust::make_transform_iterator(row_major, detail::real_abs{}),
1708: thrust::make_discard_iterator(),
1709: thrust::device_pointer_cast(sums)
1710: );
1711: );
1712: // clang-format on
1713: }
1714: // clang-format off
1715: PetscCallThrust(
1716: const auto dsums = thrust::device_pointer_cast(sums);
1718: *nrm = THRUST_CALL(
1719: thrust::reduce,
1720: stream,
1721: dsums, dsums + ngroups,
1722: PetscReal{0.0},
1723: max_functor{}
1724: );
1725: );
1726: // clang-format on
1727: PetscCall(PetscDeviceFree(dctx, sums));
1728: PetscCall(PetscLogGpuFlops(1.0 * N));
1729: }
1730: }
1731: PetscFunctionReturn(PETSC_SUCCESS);
1732: }
1734: template <device::cupm::DeviceType T>
1735: inline PetscErrorCode MatDense_Seq_CUPM<T>::DiagonalScale(Mat A, Vec l, Vec r) noexcept
1736: {
1737: PetscDeviceContext dctx;
1738: cupmBlasHandle_t handle;
1739: const PetscScalar *dlr;
1740: PetscMemType mtype;
1741: auto m = A->rmap->n, n = A->cmap->n;
1743: PetscFunctionBegin;
1744: if (!m || !n) PetscFunctionReturn(PETSC_SUCCESS);
1745: PetscCall(PetscInfo(A, "Performing DiagonalScale %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m, n));
1746: PetscCall(GetHandles_(&dctx, &handle));
1747: {
1748: Vec lr;
1749: const auto da = DeviceArrayReadWrite(dctx, A);
1750: const auto lda = MatIMPLCast(A)->lda;
1752: if (l) {
1753: PetscCall(VecGetLocalSize(l, &m));
1754: PetscCheck(m == A->rmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Left scaling Vec of wrong size");
1755: PetscCall(VecGetArrayReadAndMemType(l, &dlr, &mtype));
1756: /* the array must live in memory this backend's BLAS can dereference, so test against this backend's
1757: memory type rather than with PetscMemTypeDevice(), which would also accept the other backend's
1758: device memory when PETSc is configured with both CUDA and HIP */
1759: if (!(T == device::cupm::DeviceType::CUDA ? PetscMemTypeCUDA(mtype) : PetscMemTypeHIP(mtype))) {
1760: PetscCall(VecRestoreArrayReadAndMemType(l, &dlr));
1761: /* l may be the parallel scaling Vec of a MATMPIDENSE, so the temporary is local and holds the local part */
1762: PetscCall(VecCreate(PETSC_COMM_SELF, &lr));
1763: PetscCall(VecSetSizes(lr, m, m));
1764: PetscCall(VecSetType(lr, VecSeq_CUPM::VECSEQCUPM()));
1765: PetscCall(VecCopy(l, lr));
1766: PetscCall(VecGetArrayReadAndMemType(lr, &dlr, nullptr));
1767: } else lr = l;
1768: {
1769: constexpr auto side = CUPMBLAS_SIDE_LEFT;
1771: PetscCall(PetscLogGpuTimeBegin());
1772: PetscCallCUPMBLAS(cupmBlasXdgmm(handle, side, m, n, da.cupmdata(), lda, cupmScalarPtrCast(dlr), 1, da.cupmdata(), lda));
1773: PetscCall(PetscLogGpuTimeEnd());
1774: PetscCall(PetscLogGpuFlops(1.0 * n * m));
1775: }
1776: PetscCall(VecRestoreArrayReadAndMemType(lr, &dlr));
1777: if (lr != l) PetscCall(VecDestroy(&lr));
1778: }
1779: if (r) {
1780: PetscCall(VecGetLocalSize(r, &n));
1781: PetscCheck(n == A->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Right scaling Vec of wrong size");
1782: PetscCall(VecGetArrayReadAndMemType(r, &dlr, &mtype));
1783: if (!(T == device::cupm::DeviceType::CUDA ? PetscMemTypeCUDA(mtype) : PetscMemTypeHIP(mtype))) {
1784: PetscCall(VecRestoreArrayReadAndMemType(r, &dlr));
1785: PetscCall(VecCreate(PETSC_COMM_SELF, &lr));
1786: PetscCall(VecSetSizes(lr, n, n));
1787: PetscCall(VecSetType(lr, VecSeq_CUPM::VECSEQCUPM()));
1788: PetscCall(VecCopy(r, lr));
1789: PetscCall(VecGetArrayReadAndMemType(lr, &dlr, nullptr));
1790: } else lr = r;
1791: {
1792: constexpr auto side = CUPMBLAS_SIDE_RIGHT;
1794: PetscCall(PetscLogGpuTimeBegin());
1795: PetscCallCUPMBLAS(cupmBlasXdgmm(handle, side, m, n, da.cupmdata(), lda, cupmScalarPtrCast(dlr), 1, da.cupmdata(), lda));
1796: PetscCall(PetscLogGpuTimeEnd());
1797: PetscCall(PetscLogGpuFlops(1.0 * n * m));
1798: }
1799: PetscCall(VecRestoreArrayReadAndMemType(lr, &dlr));
1800: if (lr != r) PetscCall(VecDestroy(&lr));
1801: }
1802: }
1803: PetscFunctionReturn(PETSC_SUCCESS);
1804: }
1806: template <device::cupm::DeviceType T>
1807: inline PetscErrorCode MatDense_Seq_CUPM<T>::AXPY(Mat Y, PetscScalar alpha, Mat X, MatStructure) noexcept
1808: {
1809: const auto m_x = X->rmap->n, m_y = Y->rmap->n;
1810: const auto n_x = X->cmap->n, n_y = Y->cmap->n;
1811: const auto N = m_x * n_x;
1812: PetscDeviceContext dctx;
1814: PetscFunctionBegin;
1815: if (!m_x || !n_x || alpha == (PetscScalar)0.0) PetscFunctionReturn(PETSC_SUCCESS);
1816: PetscCall(PetscInfo(Y, "Performing AXPY %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m_y, n_y));
1817: PetscCall(GetHandles_(&dctx));
1818: {
1819: const auto dx = DeviceArrayRead(dctx, X);
1820: const auto dy = DeviceArrayReadWrite(dctx, Y);
1821: const auto lda_x = MatIMPLCast(X)->lda;
1822: const auto lda_y = MatIMPLCast(Y)->lda;
1824: if (lda_x > m_x || lda_y > m_x) {
1825: cupmStream_t stream;
1827: PetscCall(GetHandlesFrom_(dctx, &stream));
1828: // clang-format off
1829: PetscCallThrust(
1830: const auto sub_mat_y = detail::make_submat_iterator(0, m_y, 0, n_y, lda_y, dy.data());
1831: const auto sub_mat_x = detail::make_submat_iterator(0, m_x, 0, n_x, lda_x, dx.data());
1833: THRUST_CALL(
1834: thrust::transform,
1835: stream,
1836: sub_mat_x.begin(), sub_mat_x.end(), sub_mat_y.begin(), sub_mat_y.begin(),
1837: device::cupm::functors::make_axpy(alpha)
1838: );
1839: );
1840: // clang-format on
1841: } else {
1842: const auto cu_alpha = cupmScalarCast(alpha);
1843: cupmBlasHandle_t handle;
1845: PetscCall(GetHandlesFrom_(dctx, &handle));
1846: PetscCall(PetscLogGpuTimeBegin());
1847: PetscCallCUPMBLAS(cupmBlasXaxpy(handle, N, &cu_alpha, dx.cupmdata(), 1, dy.cupmdata(), 1));
1848: PetscCall(PetscLogGpuTimeEnd());
1849: }
1850: }
1851: PetscCall(PetscLogGpuFlops(PetscMax(2 * N - 1, 0)));
1852: PetscFunctionReturn(PETSC_SUCCESS);
1853: }
1855: template <device::cupm::DeviceType T>
1856: inline PetscErrorCode MatDense_Seq_CUPM<T>::Duplicate(Mat A, MatDuplicateOption opt, Mat *B) noexcept
1857: {
1858: const auto hopt = (opt == MAT_COPY_VALUES && A->offloadmask != PETSC_OFFLOAD_CPU) ? MAT_DO_NOT_COPY_VALUES : opt;
1859: PetscDeviceContext dctx;
1861: PetscFunctionBegin;
1862: PetscCall(GetHandles_(&dctx));
1863: // do not call SetPreallocation() yet, we call it afterwards??
1864: PetscCall(MatCreateSeqDenseCUPM<T>(PetscObjectComm(PetscObjectCast(A)), A->rmap->n, A->cmap->n, nullptr, B, dctx, /* preallocate */ false));
1865: PetscCall(MatDuplicateNoCreate_SeqDense(*B, A, hopt));
1866: if (opt == MAT_COPY_VALUES && hopt != MAT_COPY_VALUES) PetscCall(Copy(A, *B, SAME_NONZERO_PATTERN));
1867: // allocate memory if needed
1868: if (opt != MAT_COPY_VALUES && !MatCUPMCast(*B)->d_v) PetscCall(SetPreallocation(*B, dctx, nullptr));
1869: PetscFunctionReturn(PETSC_SUCCESS);
1870: }
1872: template <device::cupm::DeviceType T>
1873: inline PetscErrorCode MatDense_Seq_CUPM<T>::SetRandom(Mat A, PetscRandom rng) noexcept
1874: {
1875: PetscBool device_rand_is_rander48;
1876: PetscBool device = PETSC_FALSE;
1878: PetscFunctionBegin;
1879: // CUPMObject::PETSCDEVICERAD() is PETSCRANDER48 until PetscRandom is implemented for hiprand
1880: PetscCall(PetscStrncmp(PETSCDEVICERAND(), PETSCRANDER48, sizeof(PETSCRANDER48), &device_rand_is_rander48));
1881: if (!device_rand_is_rander48) PetscCall(PetscObjectTypeCompare(PetscObjectCast(rng), PETSCDEVICERAND(), &device));
1882: if (device) {
1883: const auto m = A->rmap->n;
1884: const auto n = A->cmap->n;
1885: PetscDeviceContext dctx;
1887: PetscCall(GetHandles_(&dctx));
1888: {
1889: const auto a = DeviceArrayWrite(dctx, A);
1890: PetscInt lda;
1892: PetscCall(MatDenseGetLDA(A, &lda));
1893: if (lda > m) {
1894: for (PetscInt i = 0; i < n; i++) PetscCall(PetscRandomGetValues(rng, m, a.data() + i * lda));
1895: } else {
1896: PetscInt mn;
1898: PetscCall(PetscIntMultError(m, n, &mn));
1899: PetscCall(PetscRandomGetValues(rng, mn, a));
1900: }
1901: }
1902: } else {
1903: PetscCall(MatSetRandom_SeqDense(A, rng));
1904: }
1905: PetscFunctionReturn(PETSC_SUCCESS);
1906: }
1908: // ==========================================================================================
1910: template <device::cupm::DeviceType T>
1911: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetColumnVector(Mat A, Vec v, PetscInt col) noexcept
1912: {
1913: const auto offloadmask = A->offloadmask;
1914: const auto n = A->rmap->n;
1915: const auto col_offset = [&](const PetscScalar *ptr) { return ptr + col * MatIMPLCast(A)->lda; };
1916: PetscBool viscupm;
1917: PetscDeviceContext dctx;
1918: cupmStream_t stream;
1920: PetscFunctionBegin;
1921: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(v), &viscupm, VecSeq_CUPM::VECSEQCUPM(), VecSeq_CUPM::VECMPICUPM(), VecSeq_CUPM::VECCUPM(), ""));
1922: PetscCall(GetHandles_(&dctx, &stream));
1923: if (viscupm && !v->boundtocpu) {
1924: const auto x = VecSeq_CUPM::DeviceArrayWrite(dctx, v);
1926: // update device data
1927: if (PetscOffloadDevice(offloadmask)) {
1928: PetscCall(PetscCUPMMemcpyAsync(x.data(), col_offset(DeviceArrayRead(dctx, A)), n, cupmMemcpyDeviceToDevice, stream));
1929: } else {
1930: PetscCall(PetscCUPMMemcpyAsync(x.data(), col_offset(HostArrayRead(dctx, A)), n, cupmMemcpyHostToDevice, stream));
1931: }
1932: } else {
1933: PetscScalar *x;
1935: // update host data
1936: PetscCall(VecGetArrayWrite(v, &x));
1937: if (PetscOffloadUnallocated(offloadmask) || PetscOffloadHost(offloadmask)) {
1938: PetscCall(PetscArraycpy(x, col_offset(HostArrayRead(dctx, A)), n));
1939: } else if (PetscOffloadDevice(offloadmask)) {
1940: PetscCall(PetscCUPMMemcpyAsync(x, col_offset(DeviceArrayRead(dctx, A)), n, cupmMemcpyDeviceToHost, stream));
1941: }
1942: PetscCall(VecRestoreArrayWrite(v, &x));
1943: }
1944: PetscFunctionReturn(PETSC_SUCCESS);
1945: }
1947: // Place the device array a[] of a column of the matrix in its column Vec cvec. MatDenseCreateColumnVec_Private()
1948: // creates cvec from the VecType of the matrix, so it is a CUPM Vec, or a VECKOKKOS when the Kokkos backend runs on
1949: // this device (see MatCreateDenseFromVecType()); each needs its own place/reset routines
1950: template <device::cupm::DeviceType T>
1951: inline PetscErrorCode MatDense_Seq_CUPM<T>::PlaceColumnVecArray(Vec cvec, PetscScalar *a) noexcept
1952: {
1953: using namespace vec::cupm;
1954: PetscBool iscupm;
1956: PetscFunctionBegin;
1957: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(cvec), &iscupm, VecSeq_CUPM::VECSEQCUPM(), VecSeq_CUPM::VECMPICUPM(), ""));
1958: if (iscupm) PetscCall(VecCUPMPlaceArrayAsync<T>(cvec, a));
1959: else {
1960: #if PetscDefined(HAVE_KOKKOS_KERNELS)
1961: constexpr bool kokkos_on_this_device = (T == device::cupm::DeviceType::CUDA && PetscDefined(HAVE_MACRO_KOKKOS_ENABLE_CUDA)) || (T == device::cupm::DeviceType::HIP && PetscDefined(HAVE_MACRO_KOKKOS_ENABLE_HIP));
1962: PetscBool iskokkos;
1964: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(cvec), &iskokkos, VECSEQKOKKOS, VECMPIKOKKOS, ""));
1965: PetscCheck(iskokkos, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Column Vec of type %s cannot be used with a %s matrix", PetscObjectCast(cvec)->type_name, MATDENSECUPM());
1966: PetscCheck(kokkos_on_this_device, PETSC_COMM_SELF, PETSC_ERR_SUP, "The Kokkos backend does not run on the device of a %s matrix, so a %s cannot be its column Vec", MATDENSECUPM(), PetscObjectCast(cvec)->type_name);
1967: PetscCall(VecKokkosPlaceArray(cvec, a));
1968: #else
1969: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Column Vec of type %s cannot be used with a %s matrix", PetscObjectCast(cvec)->type_name, MATDENSECUPM());
1970: #endif
1971: }
1972: PetscFunctionReturn(PETSC_SUCCESS);
1973: }
1975: template <device::cupm::DeviceType T>
1976: inline PetscErrorCode MatDense_Seq_CUPM<T>::ResetColumnVecArray(Vec cvec) noexcept
1977: {
1978: using namespace vec::cupm;
1979: PetscBool iscupm;
1981: PetscFunctionBegin;
1982: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(cvec), &iscupm, VecSeq_CUPM::VECSEQCUPM(), VecSeq_CUPM::VECMPICUPM(), ""));
1983: if (iscupm) PetscCall(VecCUPMResetArrayAsync<T>(cvec));
1984: else {
1985: #if PetscDefined(HAVE_KOKKOS_KERNELS)
1986: PetscCall(VecKokkosResetArray(cvec)); // PlaceColumnVecArray() already checked cvec is a VECKOKKOS
1987: #else
1988: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Column Vec of type %s cannot be used with a %s matrix", PetscObjectCast(cvec)->type_name, MATDENSECUPM());
1989: #endif
1990: }
1991: PetscFunctionReturn(PETSC_SUCCESS);
1992: }
1994: template <device::cupm::DeviceType T>
1995: template <PetscMemoryAccessMode access>
1996: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetColumnVec(Mat A, PetscInt col, Vec *v) noexcept
1997: {
1998: using namespace vec::cupm;
1999: const auto mimpl = MatIMPLCast(A);
2000: PetscDeviceContext dctx;
2002: PetscFunctionBegin;
2003: PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
2004: PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2005: mimpl->vecinuse = col + 1;
2006: if (!mimpl->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &mimpl->cvec));
2007: PetscCall(GetHandles_(&dctx));
2008: PetscCall(GetArray<PETSC_MEMTYPE_DEVICE, access>(A, const_cast<PetscScalar **>(&mimpl->ptrinuse), dctx));
2009: PetscCall(PlaceColumnVecArray(mimpl->cvec, const_cast<PetscScalar *>(mimpl->ptrinuse) + static_cast<std::size_t>(col) * static_cast<std::size_t>(mimpl->lda)));
2010: if (access == PETSC_MEMORY_ACCESS_READ) PetscCall(VecLockReadPush(mimpl->cvec));
2011: *v = mimpl->cvec;
2012: PetscFunctionReturn(PETSC_SUCCESS);
2013: }
2015: template <device::cupm::DeviceType T>
2016: template <PetscMemoryAccessMode access>
2017: inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreColumnVec(Mat A, PetscInt, Vec *v) noexcept
2018: {
2019: using namespace vec::cupm;
2020: const auto mimpl = MatIMPLCast(A);
2021: const auto cvec = mimpl->cvec;
2022: PetscDeviceContext dctx;
2024: PetscFunctionBegin;
2025: PetscCheck(mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
2026: PetscCheck(cvec, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column vector");
2027: mimpl->vecinuse = 0;
2028: if (access == PETSC_MEMORY_ACCESS_READ) PetscCall(VecLockReadPop(cvec));
2029: PetscCall(ResetColumnVecArray(cvec));
2030: PetscCall(GetHandles_(&dctx));
2031: PetscCall(RestoreArray<PETSC_MEMTYPE_DEVICE, access>(A, const_cast<PetscScalar **>(&mimpl->ptrinuse), dctx));
2032: if (v) *v = nullptr;
2033: PetscFunctionReturn(PETSC_SUCCESS);
2034: }
2036: // ==========================================================================================
2038: template <device::cupm::DeviceType T>
2039: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetFactor(Mat A, MatFactorType ftype, Mat *fact_out) noexcept
2040: {
2041: Mat fact = nullptr;
2042: PetscDeviceContext dctx;
2044: PetscFunctionBegin;
2045: PetscCall(GetHandles_(&dctx));
2046: PetscCall(MatCreateSeqDenseCUPM<T>(PetscObjectComm(PetscObjectCast(A)), A->rmap->n, A->cmap->n, nullptr, &fact, dctx, /* preallocate */ false));
2047: fact->factortype = ftype;
2048: switch (ftype) {
2049: case MAT_FACTOR_LU:
2050: case MAT_FACTOR_ILU: // fall-through
2051: fact->ops->lufactorsymbolic = MatLUFactorSymbolic_SeqDense;
2052: fact->ops->ilufactorsymbolic = MatLUFactorSymbolic_SeqDense;
2053: break;
2054: case MAT_FACTOR_CHOLESKY:
2055: case MAT_FACTOR_ICC: // fall-through
2056: fact->ops->choleskyfactorsymbolic = MatCholeskyFactorSymbolic_SeqDense;
2057: break;
2058: case MAT_FACTOR_QR: {
2059: const auto pobj = PetscObjectCast(fact);
2061: PetscCall(PetscObjectComposeFunction(pobj, "MatQRFactor_C", MatQRFactor_SeqDense));
2062: PetscCall(PetscObjectComposeFunction(pobj, "MatQRFactorSymbolic_C", MatQRFactorSymbolic_SeqDense));
2063: } break;
2064: case MAT_FACTOR_NONE:
2065: case MAT_FACTOR_ILUDT: // fall-through
2066: case MAT_FACTOR_NUM_TYPES: // fall-through
2067: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "MatFactorType %s not supported", MatFactorTypes[ftype]);
2068: }
2069: PetscCall(PetscStrFreeAllocpy(MATSOLVERCUPM(), &fact->solvertype));
2070: PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_LU));
2071: PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_ILU));
2072: PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_CHOLESKY));
2073: PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_ICC));
2074: *fact_out = fact;
2075: PetscFunctionReturn(PETSC_SUCCESS);
2076: }
2078: template <device::cupm::DeviceType T>
2079: inline PetscErrorCode MatDense_Seq_CUPM<T>::InvertFactors(Mat A) noexcept
2080: {
2081: const auto mimpl = MatIMPLCast(A);
2082: const auto mcu = MatCUPMCast(A);
2083: const auto n = static_cast<cupmBlasInt_t>(A->cmap->n);
2084: cupmSolverHandle_t handle;
2085: PetscDeviceContext dctx;
2086: cupmStream_t stream;
2088: PetscFunctionBegin;
2089: #if PetscDefined(HAVE_CUDA) && PetscDefined(USING_NVCC)
2090: // HIP appears to have this by default??
2091: PetscCheck(PETSC_PKG_CUDA_VERSION_GE(10, 1, 0), PETSC_COMM_SELF, PETSC_ERR_SUP, "Upgrade to CUDA version 10.1.0 or higher");
2092: #endif
2093: if (!n || !A->rmap->n) PetscFunctionReturn(PETSC_SUCCESS);
2094: PetscCheck(A->factortype == MAT_FACTOR_CHOLESKY, PETSC_COMM_SELF, PETSC_ERR_LIB, "Factor type %s not implemented", MatFactorTypes[A->factortype]);
2095: // spd
2096: PetscCheck(!mcu->d_fact_ipiv, PETSC_COMM_SELF, PETSC_ERR_LIB, "%sDnsytri not implemented", cupmSolverName());
2098: PetscCall(GetHandles_(&dctx, &handle, &stream));
2099: {
2100: const auto da = DeviceArrayReadWrite(dctx, A);
2101: const auto lda = static_cast<cupmBlasInt_t>(mimpl->lda);
2102: cupmBlasInt_t il;
2104: PetscCallCUPMSOLVER(cupmSolverXpotri_bufferSize(handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, &il));
2105: if (il > mcu->d_fact_lwork) {
2106: mcu->d_fact_lwork = il;
2107: PetscCallCUPM(cupmFreeAsync(mcu->d_fact_work, stream));
2108: PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_work, il, stream));
2109: }
2110: PetscCall(PetscLogGpuTimeBegin());
2111: PetscCallCUPMSOLVER(cupmSolverXpotri(handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_info));
2112: PetscCall(PetscLogGpuTimeEnd());
2113: }
2114: PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream));
2115: // TODO (write cuda kernel)
2116: PetscCall(MatSeqDenseSymmetrize_Private(A, PETSC_TRUE));
2117: PetscCall(PetscLogGpuFlops(1.0 * n * n * n / 3.0));
2119: A->ops->solve = nullptr;
2120: A->ops->solvetranspose = nullptr;
2121: A->ops->matsolve = nullptr;
2122: A->factortype = MAT_FACTOR_NONE;
2124: PetscCall(PetscFree(A->solvertype));
2125: PetscFunctionReturn(PETSC_SUCCESS);
2126: }
2128: // ==========================================================================================
2130: template <device::cupm::DeviceType T>
2131: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetSubMatrix(Mat A, PetscInt rbegin, PetscInt rend, PetscInt cbegin, PetscInt cend, Mat *mat) noexcept
2132: {
2133: const auto mimpl = MatIMPLCast(A);
2134: const auto array_offset = [&](PetscScalar *ptr) { return ptr + rbegin + static_cast<std::size_t>(cbegin) * mimpl->lda; };
2135: const auto n = rend - rbegin;
2136: const auto m = cend - cbegin;
2137: auto &cmat = mimpl->cmat;
2138: PetscDeviceContext dctx;
2140: PetscFunctionBegin;
2141: PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
2142: PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2143: mimpl->matinuse = cbegin + 1;
2145: PetscCall(GetHandles_(&dctx));
2146: PetscCall(HostToDevice_(A, dctx));
2148: if (cmat && ((m != cmat->cmap->N) || (n != cmat->rmap->N))) PetscCall(MatDestroy(&cmat));
2149: {
2150: const auto device_array = array_offset(MatCUPMCast(A)->d_v);
2152: if (cmat) {
2153: PetscCall(PlaceArray(cmat, device_array));
2154: } else {
2155: PetscCall(MatCreateSeqDenseCUPM<T>(PetscObjectComm(PetscObjectCast(A)), n, m, device_array, &cmat, dctx));
2156: PetscCall(MatSetVecType(cmat, A->defaultvectype));
2157: }
2158: }
2159: PetscCall(MatDenseSetLDA(cmat, mimpl->lda));
2160: // place CPU array if present but do not copy any data
2161: if (const auto host_array = mimpl->v) {
2162: cmat->offloadmask = PETSC_OFFLOAD_GPU;
2163: PetscCall(MatDensePlaceArray(cmat, array_offset(host_array)));
2164: }
2166: cmat->offloadmask = A->offloadmask;
2167: *mat = cmat;
2168: PetscFunctionReturn(PETSC_SUCCESS);
2169: }
2171: template <device::cupm::DeviceType T>
2172: inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreSubMatrix(Mat A, Mat *m) noexcept
2173: {
2174: const auto mimpl = MatIMPLCast(A);
2175: const auto cmat = mimpl->cmat;
2176: const auto reset = static_cast<bool>(mimpl->v);
2177: bool copy, was_offload_host;
2179: PetscFunctionBegin;
2180: PetscCheck(mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetSubMatrix() first");
2181: PetscCheck(cmat, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column matrix");
2182: PetscCheck(*m == cmat, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Not the matrix obtained from MatDenseGetSubMatrix()");
2183: mimpl->matinuse = 0;
2185: // calls to ResetArray may change it, so save it here
2186: was_offload_host = cmat->offloadmask == PETSC_OFFLOAD_CPU;
2187: if (was_offload_host && !reset) {
2188: copy = true;
2189: PetscCall(MatSeqDenseSetPreallocation(A, nullptr));
2190: } else {
2191: copy = false;
2192: }
2194: PetscCall(ResetArray(cmat));
2195: if (reset) PetscCall(MatDenseResetArray(cmat));
2196: if (copy) {
2197: PetscDeviceContext dctx;
2199: PetscCall(GetHandles_(&dctx));
2200: PetscCall(DeviceToHost_(A, dctx));
2201: } else {
2202: A->offloadmask = was_offload_host ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
2203: }
2205: cmat->offloadmask = PETSC_OFFLOAD_UNALLOCATED;
2206: *m = nullptr;
2207: PetscFunctionReturn(PETSC_SUCCESS);
2208: }
2210: // ==========================================================================================
2212: namespace
2213: {
2215: template <device::cupm::DeviceType T>
2216: inline PetscErrorCode MatMatMultNumeric_SeqDenseCUPM_SeqDenseCUPM(Mat A, Mat B, Mat C, PetscBool TA, PetscBool TB) noexcept
2217: {
2218: PetscFunctionBegin;
2219: if (TA) {
2220: if (TB) {
2221: PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<true, true>(A, B, C));
2222: } else {
2223: PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<true, false>(A, B, C));
2224: }
2225: } else {
2226: if (TB) {
2227: PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<false, true>(A, B, C));
2228: } else {
2229: PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<false, false>(A, B, C));
2230: }
2231: }
2232: PetscFunctionReturn(PETSC_SUCCESS);
2233: }
2235: template <device::cupm::DeviceType T>
2236: inline PetscErrorCode MatSolverTypeRegister_DENSECUPM() noexcept
2237: {
2238: PetscFunctionBegin;
2239: for (auto ftype : util::make_array(MAT_FACTOR_LU, MAT_FACTOR_CHOLESKY, MAT_FACTOR_QR)) {
2240: PetscCall(MatSolverTypeRegister(MatDense_Seq_CUPM<T>::MATSOLVERCUPM(), MATSEQDENSE, ftype, MatDense_Seq_CUPM<T>::GetFactor));
2241: PetscCall(MatSolverTypeRegister(MatDense_Seq_CUPM<T>::MATSOLVERCUPM(), MatDense_Seq_CUPM<T>::MATSEQDENSECUPM(), ftype, MatDense_Seq_CUPM<T>::GetFactor));
2242: }
2243: PetscFunctionReturn(PETSC_SUCCESS);
2244: }
2246: } // anonymous namespace
2248: } // namespace impl
2250: } // namespace cupm
2252: } // namespace mat
2254: } // namespace Petsc