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: PetscFunctionReturn(PETSC_SUCCESS);
883: }
885: template <device::cupm::DeviceType T>
886: template <bool transpose, bool hermitian>
887: inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMultColumnRange_Dispatch_(Mat A, Vec xx, Vec yy, PetscInt c_start, PetscInt c_end) noexcept
888: {
889: PetscFunctionBegin;
890: PetscCall(MatMultAddColumnRange_Dispatch_<transpose, hermitian>(A, xx, nullptr, yy, c_start, c_end));
891: PetscFunctionReturn(PETSC_SUCCESS);
892: }
894: template <device::cupm::DeviceType T>
895: template <bool transpose, bool hermitian>
896: inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMultAdd_Dispatch_(Mat A, Vec xx, Vec yy, Vec zz) noexcept
897: {
898: PetscFunctionBegin;
899: PetscCall(MatMultAddColumnRange_Dispatch_<transpose, hermitian>(A, xx, yy, zz, 0, A->cmap->n));
900: PetscFunctionReturn(PETSC_SUCCESS);
901: }
903: // ==========================================================================================
904: // MatDense_Seq_CUPM - Private API - Conversion Dispatch
905: // ==========================================================================================
907: template <device::cupm::DeviceType T>
908: template <bool to_host>
909: inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_Dispatch_(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept
910: {
911: PetscFunctionBegin;
912: if (reuse == MAT_REUSE_MATRIX || reuse == MAT_INITIAL_MATRIX) {
913: // TODO these cases should be optimized
914: PetscCall(MatConvert_Basic(M, type, reuse, newmat));
915: } else {
916: const auto B = *newmat;
917: const auto pobj = PetscObjectCast(B);
919: if (to_host) {
920: PetscCall(BindToCPU(B, PETSC_TRUE));
921: PetscCall(Reset(B));
922: } else {
923: PetscCall(PetscDeviceInitialize(PETSC_DEVICE_CUPM()));
924: }
926: PetscCall(PetscStrFreeAllocpy(to_host ? VECSTANDARD : VecSeq_CUPM::VECCUPM(), &B->defaultvectype));
927: PetscCall(PetscObjectChangeTypeName(pobj, to_host ? MATSEQDENSE : MATSEQDENSECUPM()));
928: // cvec might be the wrong VecType, destroy and rebuild it if necessary
929: // REVIEW ME: this is possibly very inefficient
930: PetscCall(VecDestroy(&MatIMPLCast(B)->cvec));
932: MatComposeOp_CUPM(to_host, pobj, MatConvert_seqdensecupm_seqdense_C(), nullptr, Convert_SeqDenseCUPM_SeqDense);
933: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArray_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ_WRITE>);
934: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArrayRead_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ>);
935: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArrayWrite_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_WRITE>);
936: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArray_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ_WRITE>);
937: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArrayRead_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ>);
938: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArrayWrite_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_WRITE>);
939: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMPlaceArray_C(), nullptr, PlaceArray);
940: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMResetArray_C(), nullptr, ResetArray);
941: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMReplaceArray_C(), nullptr, ReplaceArray);
942: MatComposeOp_CUPM(to_host, pobj, MatProductSetFromOptions_seqaij_seqdensecupm_C(), nullptr, MatProductSetFromOptions_SeqAIJ_SeqDense);
943: MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMSetPreallocation_C(), nullptr, SetPreallocation);
945: if (to_host) {
946: B->offloadmask = PETSC_OFFLOAD_CPU;
947: } else {
948: Mat_SeqDenseCUPM *mcu;
950: PetscCall(PetscNew(&mcu));
951: B->spptr = mcu;
952: B->offloadmask = PETSC_OFFLOAD_UNALLOCATED; // REVIEW ME: why not offload host??
953: PetscCall(BindToCPU(B, PETSC_FALSE));
954: }
956: MatSetOp_CUPM(to_host, B, bindtocpu, nullptr, BindToCPU);
957: MatSetOp_CUPM(to_host, B, destroy, MatDestroy_SeqDense, Destroy);
958: }
959: PetscFunctionReturn(PETSC_SUCCESS);
960: }
962: // ==========================================================================================
963: // MatDense_Seq_CUPM - Public API
964: // ==========================================================================================
966: template <device::cupm::DeviceType T>
967: inline constexpr MatType MatDense_Seq_CUPM<T>::MATIMPLCUPM_() noexcept
968: {
969: return MATSEQDENSECUPM();
970: }
972: template <device::cupm::DeviceType T>
973: inline constexpr typename MatDense_Seq_CUPM<T>::Mat_SeqDenseCUPM *MatDense_Seq_CUPM<T>::MatCUPMCast(Mat m) noexcept
974: {
975: return static_cast<Mat_SeqDenseCUPM *>(m->spptr);
976: }
978: template <device::cupm::DeviceType T>
979: inline constexpr Mat_SeqDense *MatDense_Seq_CUPM<T>::MatIMPLCast_(Mat m) noexcept
980: {
981: return static_cast<Mat_SeqDense *>(m->data);
982: }
984: template <device::cupm::DeviceType T>
985: inline constexpr const char *MatDense_Seq_CUPM<T>::MatConvert_seqdensecupm_seqdense_C() noexcept
986: {
987: return T == device::cupm::DeviceType::CUDA ? "MatConvert_seqdensecuda_seqdense_C" : "MatConvert_seqdensehip_seqdense_C";
988: }
990: template <device::cupm::DeviceType T>
991: inline constexpr const char *MatDense_Seq_CUPM<T>::MatProductSetFromOptions_seqaij_seqdensecupm_C() noexcept
992: {
993: return T == device::cupm::DeviceType::CUDA ? "MatProductSetFromOptions_seqaij_seqdensecuda_C" : "MatProductSetFromOptions_seqaij_seqdensehip_C";
994: }
996: // ==========================================================================================
998: // MatCreate_SeqDenseCUPM()
999: template <device::cupm::DeviceType T>
1000: inline PetscErrorCode MatDense_Seq_CUPM<T>::Create(Mat A) noexcept
1001: {
1002: PetscFunctionBegin;
1003: PetscCall(PetscDeviceInitialize(PETSC_DEVICE_CUPM()));
1004: PetscCall(MatCreate_SeqDense(A));
1005: PetscCall(Convert_SeqDense_SeqDenseCUPM(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
1006: PetscFunctionReturn(PETSC_SUCCESS);
1007: }
1009: template <device::cupm::DeviceType T>
1010: inline PetscErrorCode MatDense_Seq_CUPM<T>::Destroy(Mat A) noexcept
1011: {
1012: PetscFunctionBegin;
1013: // prevent copying back data if we own the data pointer
1014: if (!MatIMPLCast(A)->user_alloc) A->offloadmask = PETSC_OFFLOAD_CPU;
1015: PetscCall(Convert_SeqDenseCUPM_SeqDense(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A));
1016: PetscCall(MatDestroy_SeqDense(A));
1017: PetscFunctionReturn(PETSC_SUCCESS);
1018: }
1020: // obj->ops->setup()
1021: template <device::cupm::DeviceType T>
1022: inline PetscErrorCode MatDense_Seq_CUPM<T>::SetUp(Mat A) noexcept
1023: {
1024: PetscFunctionBegin;
1025: PetscCall(PetscLayoutSetUp(A->rmap));
1026: PetscCall(PetscLayoutSetUp(A->cmap));
1027: if (!A->preallocated) {
1028: PetscDeviceContext dctx;
1030: PetscCall(GetHandles_(&dctx));
1031: PetscCall(SetPreallocation(A, dctx, nullptr));
1032: }
1033: PetscFunctionReturn(PETSC_SUCCESS);
1034: }
1036: template <device::cupm::DeviceType T>
1037: inline PetscErrorCode MatDense_Seq_CUPM<T>::Reset(Mat A) noexcept
1038: {
1039: PetscFunctionBegin;
1040: if (const auto mcu = MatCUPMCast(A)) {
1041: cupmStream_t stream;
1043: PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME());
1044: PetscCall(GetHandles_(&stream));
1045: if (!mcu->d_user_alloc) PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream));
1046: PetscCallCUPM(cupmFreeAsync(mcu->d_fact_tau, stream));
1047: PetscCallCUPM(cupmFreeAsync(mcu->d_fact_ipiv, stream));
1048: PetscCallCUPM(cupmFreeAsync(mcu->d_fact_info, stream));
1049: PetscCallCUPM(cupmFreeAsync(mcu->d_fact_work, stream));
1050: PetscCall(VecDestroy(&mcu->workvec));
1051: PetscCall(PetscFree(A->spptr /* mcu */));
1052: }
1053: PetscFunctionReturn(PETSC_SUCCESS);
1054: }
1056: // ==========================================================================================
1058: template <device::cupm::DeviceType T>
1059: inline PetscErrorCode MatDense_Seq_CUPM<T>::BindToCPU(Mat A, PetscBool to_host) noexcept
1060: {
1061: const auto mimpl = MatIMPLCast(A);
1062: const auto pobj = PetscObjectCast(A);
1064: PetscFunctionBegin;
1065: PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
1066: PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1067: A->boundtocpu = to_host;
1068: PetscCall(PetscStrFreeAllocpy(to_host ? PETSCRANDER48 : PETSCDEVICERAND(), &A->defaultrandtype));
1069: if (to_host) {
1070: PetscDeviceContext dctx;
1072: // make sure we have an up-to-date copy on the CPU
1073: PetscCall(GetHandles_(&dctx));
1074: PetscCall(DeviceToHost_(A, dctx));
1075: } else {
1076: PetscBool iscupm;
1078: if (auto &cvec = mimpl->cvec) {
1079: PetscCall(PetscObjectTypeCompare(PetscObjectCast(cvec), VecSeq_CUPM::VECSEQCUPM(), &iscupm));
1080: if (!iscupm) PetscCall(VecDestroy(&cvec));
1081: }
1082: if (auto &cmat = mimpl->cmat) {
1083: PetscCall(PetscObjectTypeCompare(PetscObjectCast(cmat), MATSEQDENSECUPM(), &iscupm));
1084: if (!iscupm) PetscCall(MatDestroy(&cmat));
1085: }
1086: }
1088: // ============================================================
1089: // Composed ops
1090: // ============================================================
1091: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArray_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ_WRITE>);
1092: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayRead_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ>);
1093: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayWrite_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_WRITE>);
1094: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ_WRITE>);
1095: MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ_WRITE>);
1096: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayReadAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ>);
1097: MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayReadAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ>);
1098: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayWriteAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_WRITE>);
1099: MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayWriteAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_WRITE>);
1100: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVec_C", MatDenseGetColumnVec_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_READ_WRITE>);
1101: MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVec_C", MatDenseRestoreColumnVec_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_READ_WRITE>);
1102: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVecRead_C", MatDenseGetColumnVecRead_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_READ>);
1103: MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVecRead_C", MatDenseRestoreColumnVecRead_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_READ>);
1104: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVecWrite_C", MatDenseGetColumnVecWrite_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_WRITE>);
1105: MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVecWrite_C", MatDenseRestoreColumnVecWrite_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_WRITE>);
1106: MatComposeOp_CUPM(to_host, pobj, "MatDenseGetSubMatrix_C", MatDenseGetSubMatrix_SeqDense, GetSubMatrix);
1107: MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreSubMatrix_C", MatDenseRestoreSubMatrix_SeqDense, RestoreSubMatrix);
1108: MatComposeOp_CUPM(to_host, pobj, "MatQRFactor_C", MatQRFactor_SeqDense, SolveQR::Factor);
1109: MatComposeOp_CUPM(to_host, pobj, "MatMultColumnRange_C", MatMultColumnRange_SeqDense, MatMultColumnRange_Dispatch_</* transpose */ false, /* hermitian */ false>);
1110: MatComposeOp_CUPM(to_host, pobj, "MatMultAddColumnRange_C", MatMultAddColumnRange_SeqDense, MatMultAddColumnRange_Dispatch_</* transpose */ false, /* hermitian */ false>);
1111: MatComposeOp_CUPM(to_host, pobj, "MatMultHermitianTransposeColumnRange_C", MatMultHermitianTransposeColumnRange_SeqDense, MatMultColumnRange_Dispatch_</* transpose */ true, /* hermitian */ true>);
1112: MatComposeOp_CUPM(to_host, pobj, "MatMultHermitianTransposeAddColumnRange_C", MatMultHermitianTransposeAddColumnRange_SeqDense, MatMultAddColumnRange_Dispatch_</* transpose */ true, /* hermitian */ true>);
1113: // always the same
1114: PetscCall(PetscObjectComposeFunction(pobj, "MatDenseSetLDA_C", MatDenseSetLDA_SeqDense));
1116: // ============================================================
1117: // Function pointer ops
1118: // ============================================================
1119: MatSetOp_CUPM(to_host, A, duplicate, MatDuplicate_SeqDense, Duplicate);
1120: 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); });
1121: 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); });
1122: 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); });
1123: MatSetOp_CUPM(to_host, A, multadd, MatMultAdd_SeqDense, MatMultAdd_Dispatch_</* transpose */ false, /* hermitian */ false>);
1124: MatSetOp_CUPM(to_host, A, multtransposeadd, MatMultTransposeAdd_SeqDense, MatMultAdd_Dispatch_</* transpose */ true, /* hermitian */ false>);
1125: MatSetOp_CUPM(to_host, A, multhermitiantransposeadd, MatMultHermitianTransposeAdd_SeqDense, MatMultAdd_Dispatch_</* transpose */ true, /* hermitian */ true>);
1126: MatSetOp_CUPM(to_host, A, matmultnumeric, MatMatMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ false, /* transpose_B */ false>);
1127: MatSetOp_CUPM(to_host, A, mattransposemultnumeric, MatMatTransposeMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ false, /* transpose_B */ true>);
1128: MatSetOp_CUPM(to_host, A, transposematmultnumeric, MatTransposeMatMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ true, /* transpose_B */ false>);
1129: MatSetOp_CUPM(to_host, A, axpy, MatAXPY_SeqDense, AXPY);
1130: MatSetOp_CUPM(to_host, A, choleskyfactor, MatCholeskyFactor_SeqDense, SolveCholesky::Factor);
1131: MatSetOp_CUPM(to_host, A, lufactor, MatLUFactor_SeqDense, SolveLU::Factor);
1132: MatSetOp_CUPM(to_host, A, getcolumnvector, MatGetColumnVector_SeqDense, GetColumnVector);
1133: MatSetOp_CUPM(to_host, A, conjugate, MatConjugate_SeqDense, Conjugate);
1134: MatSetOp_CUPM(to_host, A, scale, MatScale_SeqDense, Scale);
1135: MatSetOp_CUPM(to_host, A, norm, MatNorm_SeqDense, Norm);
1136: MatSetOp_CUPM(to_host, A, diagonalscale, MatDiagonalScale_SeqDense, DiagonalScale);
1137: MatSetOp_CUPM(to_host, A, shift, MatShift_SeqDense, Shift);
1138: MatSetOp_CUPM(to_host, A, copy, MatCopy_SeqDense, Copy);
1139: MatSetOp_CUPM(to_host, A, zeroentries, MatZeroEntries_SeqDense, ZeroEntries);
1140: MatSetOp_CUPM(to_host, A, setup, MatSetUp_SeqDense, SetUp);
1141: MatSetOp_CUPM(to_host, A, setrandom, MatSetRandom_SeqDense, SetRandom);
1142: MatSetOp_CUPM(to_host, A, getdiagonal, MatGetDiagonal_SeqDense, GetDiagonal);
1143: // seemingly always the same
1144: A->ops->productsetfromoptions = MatProductSetFromOptions_SeqDense;
1146: if (const auto cmat = mimpl->cmat) PetscCall(MatBindToCPU(cmat, to_host));
1147: PetscFunctionReturn(PETSC_SUCCESS);
1148: }
1150: template <device::cupm::DeviceType T>
1151: inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_SeqDenseCUPM_SeqDense(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept
1152: {
1153: PetscFunctionBegin;
1154: PetscCall(Convert_Dispatch_</* to host */ true>(M, type, reuse, newmat));
1155: PetscFunctionReturn(PETSC_SUCCESS);
1156: }
1158: template <device::cupm::DeviceType T>
1159: inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_SeqDense_SeqDenseCUPM(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept
1160: {
1161: PetscFunctionBegin;
1162: PetscCall(Convert_Dispatch_</* to host */ false>(M, type, reuse, newmat));
1163: PetscFunctionReturn(PETSC_SUCCESS);
1164: }
1166: // ==========================================================================================
1168: template <device::cupm::DeviceType T>
1169: template <PetscMemType mtype, PetscMemoryAccessMode access>
1170: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetArray(Mat m, PetscScalar **array, PetscDeviceContext dctx) noexcept
1171: {
1172: constexpr auto hostmem = PetscMemTypeHost(mtype);
1173: constexpr auto read_access = PetscMemoryAccessRead(access);
1175: PetscFunctionBegin;
1176: static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), "");
1177: PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
1178: if (hostmem) {
1179: if (read_access) {
1180: PetscCall(DeviceToHost_(m, dctx));
1181: } else if (!MatIMPLCast(m)->v) {
1182: // MatCreateSeqDenseCUPM may not allocate CPU memory. Allocate if needed
1183: PetscCall(MatSeqDenseSetPreallocation(m, nullptr));
1184: }
1185: *array = MatIMPLCast(m)->v;
1186: } else {
1187: if (read_access) {
1188: PetscCall(HostToDevice_(m, dctx));
1189: } else if (!MatCUPMCast(m)->d_v) {
1190: // write-only
1191: PetscCall(SetPreallocation(m, dctx, nullptr));
1192: }
1193: *array = MatCUPMCast(m)->d_v;
1194: }
1195: if (PetscMemoryAccessWrite(access)) {
1196: m->offloadmask = hostmem ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
1197: PetscCall(PetscObjectStateIncrease(PetscObjectCast(m)));
1198: }
1199: PetscFunctionReturn(PETSC_SUCCESS);
1200: }
1202: template <device::cupm::DeviceType T>
1203: template <PetscMemType mtype, PetscMemoryAccessMode access>
1204: inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreArray(Mat m, PetscScalar **array, PetscDeviceContext) noexcept
1205: {
1206: PetscFunctionBegin;
1207: static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), "");
1208: if (PetscMemoryAccessWrite(access)) {
1209: // WRITE or READ_WRITE
1210: m->offloadmask = PetscMemTypeHost(mtype) ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
1211: PetscCall(PetscObjectStateIncrease(PetscObjectCast(m)));
1212: }
1213: if (array) {
1214: PetscCall(CheckPointerMatchesMemType_(*array, mtype));
1215: *array = nullptr;
1216: }
1217: PetscFunctionReturn(PETSC_SUCCESS);
1218: }
1220: template <device::cupm::DeviceType T>
1221: template <PetscMemoryAccessMode access>
1222: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetArrayAndMemType(Mat m, PetscScalar **array, PetscMemType *mtype, PetscDeviceContext dctx) noexcept
1223: {
1224: PetscFunctionBegin;
1225: PetscCall(GetArray<PETSC_MEMTYPE_DEVICE, access>(m, array, dctx));
1226: if (mtype) *mtype = PETSC_MEMTYPE_CUPM();
1227: PetscFunctionReturn(PETSC_SUCCESS);
1228: }
1230: template <device::cupm::DeviceType T>
1231: template <PetscMemoryAccessMode access>
1232: inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreArrayAndMemType(Mat m, PetscScalar **array, PetscDeviceContext dctx) noexcept
1233: {
1234: PetscFunctionBegin;
1235: PetscCall(RestoreArray<PETSC_MEMTYPE_DEVICE, access>(m, array, dctx));
1236: PetscFunctionReturn(PETSC_SUCCESS);
1237: }
1239: // ==========================================================================================
1241: template <device::cupm::DeviceType T>
1242: inline PetscErrorCode MatDense_Seq_CUPM<T>::PlaceArray(Mat A, const PetscScalar *array) noexcept
1243: {
1244: const auto mimpl = MatIMPLCast(A);
1245: const auto mcu = MatCUPMCast(A);
1247: PetscFunctionBegin;
1248: PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
1249: PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1250: PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME());
1251: if (mimpl->v) {
1252: PetscDeviceContext dctx;
1254: PetscCall(GetHandles_(&dctx));
1255: PetscCall(HostToDevice_(A, dctx));
1256: }
1257: mcu->unplacedarray = util::exchange(mcu->d_v, const_cast<PetscScalar *>(array));
1258: mcu->d_unplaced_user_alloc = util::exchange(mcu->d_user_alloc, PETSC_TRUE);
1259: PetscFunctionReturn(PETSC_SUCCESS);
1260: }
1262: template <device::cupm::DeviceType T>
1263: inline PetscErrorCode MatDense_Seq_CUPM<T>::ReplaceArray(Mat A, const PetscScalar *array) noexcept
1264: {
1265: const auto mimpl = MatIMPLCast(A);
1266: const auto mcu = MatCUPMCast(A);
1268: PetscFunctionBegin;
1269: PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
1270: PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1271: PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME());
1272: if (!mcu->d_user_alloc) {
1273: cupmStream_t stream;
1275: PetscCall(GetHandles_(&stream));
1276: PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream));
1277: }
1278: mcu->d_v = const_cast<PetscScalar *>(array);
1279: mcu->d_user_alloc = PETSC_FALSE;
1280: PetscFunctionReturn(PETSC_SUCCESS);
1281: }
1283: template <device::cupm::DeviceType T>
1284: inline PetscErrorCode MatDense_Seq_CUPM<T>::ResetArray(Mat A) noexcept
1285: {
1286: const auto mimpl = MatIMPLCast(A);
1287: const auto mcu = MatCUPMCast(A);
1289: PetscFunctionBegin;
1290: PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
1291: PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1292: if (mimpl->v) {
1293: PetscDeviceContext dctx;
1295: PetscCall(GetHandles_(&dctx));
1296: PetscCall(HostToDevice_(A, dctx));
1297: }
1298: mcu->d_v = util::exchange(mcu->unplacedarray, nullptr);
1299: mcu->d_user_alloc = mcu->d_unplaced_user_alloc;
1300: PetscFunctionReturn(PETSC_SUCCESS);
1301: }
1303: // ==========================================================================================
1305: template <device::cupm::DeviceType T>
1306: template <bool transpose_A, bool transpose_B>
1307: inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMatMult_Numeric_Dispatch(Mat A, Mat B, Mat C) noexcept
1308: {
1309: cupmBlasInt_t m, n, k;
1310: PetscBool Aiscupm, Biscupm;
1311: PetscDeviceContext dctx;
1312: cupmBlasHandle_t handle;
1314: PetscFunctionBegin;
1315: PetscCall(PetscCUPMBlasIntCast(C->rmap->n, &m));
1316: PetscCall(PetscCUPMBlasIntCast(C->cmap->n, &n));
1317: PetscCall(PetscCUPMBlasIntCast(transpose_A ? A->rmap->n : A->cmap->n, &k));
1318: if (!m || !n || !k) {
1319: PetscCall(ZeroEntries(C));
1320: PetscFunctionReturn(PETSC_SUCCESS);
1321: }
1323: // we may end up with SEQDENSE as one of the arguments
1324: // REVIEW ME: how? and why is it not B and C????????
1325: PetscCall(PetscObjectTypeCompare(PetscObjectCast(A), MATSEQDENSECUPM(), &Aiscupm));
1326: PetscCall(PetscObjectTypeCompare(PetscObjectCast(B), MATSEQDENSECUPM(), &Biscupm));
1327: if (!Aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
1328: if (!Biscupm) PetscCall(MatConvert(B, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &B));
1329: PetscCall(PetscInfo(C, "Matrix-Matrix product %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " on backend\n", m, k, n));
1330: PetscCall(GetHandles_(&dctx, &handle));
1332: PetscCall(PetscLogGpuTimeBegin());
1333: {
1334: const auto one = cupmScalarCast(1.0);
1335: const auto zero = cupmScalarCast(0.0);
1336: const auto da = DeviceArrayRead(dctx, A);
1337: const auto db = DeviceArrayRead(dctx, B);
1338: const auto dc = DeviceArrayWrite(dctx, C);
1339: PetscInt alda, blda, clda;
1341: PetscCall(MatDenseGetLDA(A, &alda));
1342: PetscCall(MatDenseGetLDA(B, &blda));
1343: PetscCall(MatDenseGetLDA(C, &clda));
1344: 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));
1345: }
1346: PetscCall(PetscLogGpuTimeEnd());
1348: PetscCall(PetscLogGpuFlops(1.0 * m * n * k + 1.0 * m * n * (k - 1)));
1349: if (!Aiscupm) PetscCall(MatConvert(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A));
1350: if (!Biscupm) PetscCall(MatConvert(B, MATSEQDENSE, MAT_INPLACE_MATRIX, &B));
1351: PetscFunctionReturn(PETSC_SUCCESS);
1352: }
1354: template <device::cupm::DeviceType T>
1355: inline PetscErrorCode MatDense_Seq_CUPM<T>::Copy(Mat A, Mat B, MatStructure str) noexcept
1356: {
1357: const auto m = A->rmap->n;
1358: const auto n = A->cmap->n;
1360: PetscFunctionBegin;
1361: PetscAssert(m == B->rmap->n && n == B->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "size(B) != size(A)");
1362: // The two matrices must have the same copy implementation to be eligible for fast copy
1363: if (A->ops->copy == B->ops->copy) {
1364: PetscDeviceContext dctx;
1365: cupmStream_t stream;
1367: PetscCall(GetHandles_(&dctx, &stream));
1368: PetscCall(PetscLogGpuTimeBegin());
1369: {
1370: const auto va = DeviceArrayRead(dctx, A);
1371: const auto vb = DeviceArrayWrite(dctx, B);
1372: // order is important, DeviceArrayRead/Write() might call SetPreallocation() which sets
1373: // lda!
1374: const auto lda_a = MatIMPLCast(A)->lda;
1375: const auto lda_b = MatIMPLCast(B)->lda;
1377: if (lda_a > m || lda_b > m) {
1378: 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());
1379: 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());
1380: PetscCall(PetscCUPMMemcpy2DAsync(vb.data(), lda_b, va.data(), lda_a, m, n, cupmMemcpyDeviceToDevice, stream));
1381: } else {
1382: PetscCall(PetscCUPMMemcpyAsync(vb.data(), va.data(), m * n, cupmMemcpyDeviceToDevice, stream));
1383: }
1384: }
1385: PetscCall(PetscLogGpuTimeEnd());
1386: } else {
1387: PetscCall(MatCopy_Basic(A, B, str));
1388: }
1389: PetscFunctionReturn(PETSC_SUCCESS);
1390: }
1392: template <device::cupm::DeviceType T>
1393: inline PetscErrorCode MatDense_Seq_CUPM<T>::ZeroEntries(Mat m) noexcept
1394: {
1395: PetscDeviceContext dctx;
1396: cupmStream_t stream;
1398: PetscFunctionBegin;
1399: PetscCall(GetHandles_(&dctx, &stream));
1400: PetscCall(PetscLogGpuTimeBegin());
1401: {
1402: const auto va = DeviceArrayWrite(dctx, m);
1403: const auto lda = MatIMPLCast(m)->lda;
1404: const auto ma = m->rmap->n;
1405: const auto na = m->cmap->n;
1407: if (lda > ma) {
1408: PetscCall(PetscCUPMMemset2DAsync(va.data(), lda, 0, ma, na, stream));
1409: } else {
1410: PetscCall(PetscCUPMMemsetAsync(va.data(), 0, ma * na, stream));
1411: }
1412: }
1413: PetscCall(PetscLogGpuTimeEnd());
1414: PetscFunctionReturn(PETSC_SUCCESS);
1415: }
1417: namespace detail
1418: {
1420: // ==========================================================================================
1421: // SubMatIndexFunctor
1422: //
1423: // Iterator which permutes a linear index range into matrix indices for an nrows x ncols
1424: // submat with leading dimension lda. Essentially SubMatIndexFunctor(i) returns the index for
1425: // the i'th sequential entry in the matrix.
1426: // ==========================================================================================
1427: template <typename T>
1428: struct SubMatIndexFunctor {
1429: PETSC_HOSTDEVICE_INLINE_DECL T operator()(T x) const noexcept { return ((x / nrows) * lda) + (x % nrows); }
1431: PetscInt nrows;
1432: PetscInt ncols;
1433: PetscInt lda;
1434: };
1436: template <typename Iterator>
1437: struct SubMatrixIterator : MatrixIteratorBase<Iterator, SubMatIndexFunctor<iter_difference_t<Iterator>>> {
1438: using base_type = MatrixIteratorBase<Iterator, SubMatIndexFunctor<iter_difference_t<Iterator>>>;
1440: using iterator = typename base_type::iterator;
1442: constexpr SubMatrixIterator(Iterator first, Iterator last, PetscInt nrows, PetscInt ncols, PetscInt lda) noexcept :
1443: base_type{
1444: std::move(first), std::move(last), {nrows, ncols, lda}
1445: }
1446: {
1447: }
1449: PETSC_NODISCARD iterator end() const noexcept { return this->begin() + (this->func.nrows * this->func.ncols); }
1450: };
1452: namespace
1453: {
1455: template <typename T>
1456: 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
1457: {
1458: const auto nrows = rend - rstart;
1459: const auto ncols = cend - cstart;
1460: const auto dptr = thrust::device_pointer_cast(ptr);
1462: return {dptr + (rstart * lda) + cstart, dptr + ((rstart + nrows) * lda) + cstart, nrows, ncols, lda};
1463: }
1465: } // namespace
1467: struct conjugate {
1468: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &x) const noexcept { return PetscConj(x); }
1469: };
1471: struct real_abs {
1472: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL PetscReal operator()(const PetscScalar &x) const noexcept { return PetscAbsScalar(x); }
1473: };
1475: struct real_abs_squared {
1476: PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL PetscReal operator()(const PetscScalar &x) const noexcept { return PetscRealPart(PetscConj(x) * x); }
1477: };
1479: // ==========================================================================================
1480: // RowMajorIndexFunctor
1481: //
1482: // Iterator which permutes a linear row-major index range into the memory offsets of a matrix
1483: // with ncols columns which is stored column-major with leading dimension lda. Essentially
1484: // RowMajorIndexFunctor(k) returns the index of the k'th entry of the matrix when it is
1485: // traversed one row at a time.
1486: // ==========================================================================================
1487: template <typename T>
1488: struct RowMajorIndexFunctor {
1489: PETSC_HOSTDEVICE_INLINE_DECL T operator()(T x) const noexcept { return ((x % ncols) * lda) + (x / ncols); }
1491: PetscInt ncols;
1492: PetscInt lda;
1493: };
1495: // ==========================================================================================
1496: // GroupIndexFunctor
1497: //
1498: // Maps a linear index range onto the index of the group of group_size consecutive entries
1499: // that each index belongs to. Used to generate the keys of a segmented reduction, where each
1500: // group is one column (or one row) of the matrix.
1501: // ==========================================================================================
1502: template <typename T>
1503: struct GroupIndexFunctor {
1504: PETSC_HOSTDEVICE_INLINE_DECL T operator()(T x) const noexcept { return x / group_size; }
1506: PetscInt group_size;
1507: };
1509: } // namespace detail
1511: template <device::cupm::DeviceType T>
1512: inline PetscErrorCode MatDense_Seq_CUPM<T>::Conjugate(Mat A) noexcept
1513: {
1514: const auto m = A->rmap->n;
1515: const auto n = A->cmap->n;
1516: const auto N = m * n;
1517: PetscDeviceContext dctx;
1518: cupmStream_t stream;
1520: PetscFunctionBegin;
1521: if (PetscDefined(USE_COMPLEX)) {
1522: PetscCall(GetHandles_(&dctx, &stream));
1523: PetscCall(PetscLogGpuTimeBegin());
1524: {
1525: const auto da = DeviceArrayReadWrite(dctx, A);
1526: const auto lda = MatIMPLCast(A)->lda;
1527: cupmStream_t stream;
1528: PetscCall(GetHandlesFrom_(dctx, &stream));
1530: if (lda > m) {
1531: // clang-format off
1532: PetscCallThrust(
1533: const auto sub_mat = detail::make_submat_iterator(0, m, 0, n, lda, da.data());
1535: THRUST_CALL(
1536: thrust::transform,
1537: stream,
1538: sub_mat.begin(), sub_mat.end(), sub_mat.begin(),
1539: detail::conjugate{}
1540: )
1541: );
1542: // clang-format on
1543: } else {
1544: // clang-format off
1545: PetscCallThrust(
1546: const auto aptr = thrust::device_pointer_cast(da.data());
1548: THRUST_CALL(
1549: thrust::transform,
1550: stream,
1551: aptr, aptr + N, aptr,
1552: detail::conjugate{}
1553: )
1554: );
1555: // clang-format on
1556: }
1557: }
1558: PetscCall(PetscLogGpuTimeEnd());
1559: }
1560: PetscFunctionReturn(PETSC_SUCCESS);
1561: }
1563: template <device::cupm::DeviceType T>
1564: inline PetscErrorCode MatDense_Seq_CUPM<T>::Scale(Mat A, PetscScalar alpha) noexcept
1565: {
1566: const auto m = A->rmap->n;
1567: const auto n = A->cmap->n;
1568: const auto N = m * n;
1569: PetscDeviceContext dctx;
1571: PetscFunctionBegin;
1572: PetscCall(PetscInfo(A, "Performing Scale %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m, n));
1573: PetscCall(GetHandles_(&dctx));
1574: {
1575: const auto da = DeviceArrayReadWrite(dctx, A);
1576: const auto lda = MatIMPLCast(A)->lda;
1578: if (lda > m) {
1579: cupmStream_t stream;
1581: PetscCall(GetHandlesFrom_(dctx, &stream));
1582: // clang-format off
1583: PetscCallThrust(
1584: const auto sub_mat = detail::make_submat_iterator(0, m, 0, n, lda, da.data());
1586: THRUST_CALL(
1587: thrust::transform,
1588: stream,
1589: sub_mat.begin(), sub_mat.end(), sub_mat.begin(),
1590: device::cupm::functors::make_times_equals(alpha)
1591: )
1592: );
1593: // clang-format on
1594: } else {
1595: const auto cu_alpha = cupmScalarCast(alpha);
1596: cupmBlasHandle_t handle;
1598: PetscCall(GetHandlesFrom_(dctx, &handle));
1599: PetscCall(PetscLogGpuTimeBegin());
1600: PetscCallCUPMBLAS(cupmBlasXscal(handle, N, &cu_alpha, da.cupmdata(), 1));
1601: PetscCall(PetscLogGpuTimeEnd());
1602: }
1603: }
1604: PetscCall(PetscLogGpuFlops(N));
1605: PetscFunctionReturn(PETSC_SUCCESS);
1606: }
1608: template <device::cupm::DeviceType T>
1609: inline PetscErrorCode MatDense_Seq_CUPM<T>::Norm(Mat A, NormType type, PetscReal *nrm) noexcept
1610: {
1611: #if PetscDefined(USING_NVCC) && CCCL_VERSION >= 3001000
1612: using max_functor = cuda::maximum<PetscReal>;
1613: #else
1614: using max_functor = thrust::maximum<PetscReal>;
1615: #endif
1616: const auto m = A->rmap->n;
1617: const auto n = A->cmap->n;
1618: const auto N = m * n;
1619: PetscDeviceContext dctx;
1620: cupmStream_t stream;
1622: PetscFunctionBegin;
1623: 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]);
1624: // NORM_2 is the largest singular value, which requires an SVD. There is no device gesvd()
1625: // wrapper available, so defer to MatNorm_SeqDense(), which copies the matrix to the host
1626: if (type == NORM_2) {
1627: PetscCall(MatNorm_SeqDense(A, type, nrm));
1628: PetscFunctionReturn(PETSC_SUCCESS);
1629: }
1630: if (!N) {
1631: *nrm = 0.0;
1632: PetscFunctionReturn(PETSC_SUCCESS);
1633: }
1634: PetscCall(PetscInfo(A, "Performing Norm %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m, n));
1635: PetscCall(GetHandles_(&dctx, &stream));
1636: {
1637: const auto da = DeviceArrayRead(dctx, A);
1638: const auto lda = static_cast<PetscInt>(MatIMPLCast(A)->lda);
1640: if (type == NORM_FROBENIUS) {
1641: if (lda > m) {
1642: PetscReal sum = 0.0;
1644: // clang-format off
1645: PetscCallThrust(
1646: const auto sub_mat = detail::make_submat_iterator(0, m, 0, n, lda, da.data());
1648: sum = THRUST_CALL(
1649: thrust::transform_reduce,
1650: stream,
1651: sub_mat.begin(), sub_mat.end(),
1652: detail::real_abs_squared{},
1653: PetscReal{0.0},
1654: thrust::plus<PetscReal>{}
1655: );
1656: );
1657: // clang-format on
1658: *nrm = PetscSqrtReal(sum);
1659: } else {
1660: cupmBlasHandle_t handle;
1662: PetscCall(GetHandlesFrom_(dctx, &handle));
1663: PetscCall(PetscLogGpuTimeBegin());
1664: PetscCallCUPMBLAS(cupmBlasXnrm2(handle, N, da.cupmdata(), 1, cupmRealPtrCast(nrm)));
1665: PetscCall(PetscLogGpuTimeEnd());
1666: }
1667: PetscCall(PetscLogGpuFlops(2.0 * N));
1668: } else {
1669: const auto ngroups = (type == NORM_1) ? n : m;
1670: const auto group_size = (type == NORM_1) ? m : n;
1671: const auto keys = thrust::make_transform_iterator(thrust::make_counting_iterator(PetscInt{0}), detail::GroupIndexFunctor<PetscInt>{group_size});
1672: PetscReal *sums = nullptr;
1674: PetscCall(PetscDeviceMalloc(dctx, PETSC_MEMTYPE_CUPM(), ngroups, &sums));
1675: if (type == NORM_1) {
1676: // the sub-matrix iterator enumerates the entries column-major, so each run of m
1677: // consecutive keys is exactly one column, for any lda
1678: // clang-format off
1679: PetscCallThrust(
1680: const auto sub_mat = detail::make_submat_iterator(0, m, 0, n, lda, da.data());
1682: THRUST_CALL(
1683: thrust::reduce_by_key,
1684: stream,
1685: keys, keys + N,
1686: thrust::make_transform_iterator(sub_mat.begin(), detail::real_abs{}),
1687: thrust::make_discard_iterator(),
1688: thrust::device_pointer_cast(sums)
1689: );
1690: );
1691: // clang-format on
1692: } else {
1693: // enumerate the entries row-major, so each run of n consecutive keys is exactly one row
1694: // clang-format off
1695: PetscCallThrust(
1696: const auto row_major = thrust::make_permutation_iterator(
1697: thrust::device_pointer_cast(da.data()),
1698: thrust::make_transform_iterator(thrust::make_counting_iterator(PetscInt{0}), detail::RowMajorIndexFunctor<PetscInt>{n, lda})
1699: );
1701: THRUST_CALL(
1702: thrust::reduce_by_key,
1703: stream,
1704: keys, keys + N,
1705: thrust::make_transform_iterator(row_major, detail::real_abs{}),
1706: thrust::make_discard_iterator(),
1707: thrust::device_pointer_cast(sums)
1708: );
1709: );
1710: // clang-format on
1711: }
1712: // clang-format off
1713: PetscCallThrust(
1714: const auto dsums = thrust::device_pointer_cast(sums);
1716: *nrm = THRUST_CALL(
1717: thrust::reduce,
1718: stream,
1719: dsums, dsums + ngroups,
1720: PetscReal{0.0},
1721: max_functor{}
1722: );
1723: );
1724: // clang-format on
1725: PetscCall(PetscDeviceFree(dctx, sums));
1726: PetscCall(PetscLogGpuFlops(1.0 * N));
1727: }
1728: }
1729: PetscFunctionReturn(PETSC_SUCCESS);
1730: }
1732: template <device::cupm::DeviceType T>
1733: inline PetscErrorCode MatDense_Seq_CUPM<T>::DiagonalScale(Mat A, Vec l, Vec r) noexcept
1734: {
1735: PetscDeviceContext dctx;
1736: cupmBlasHandle_t handle;
1737: const PetscScalar *dlr;
1738: PetscMemType mtype;
1739: auto m = A->rmap->n, n = A->cmap->n;
1741: PetscFunctionBegin;
1742: if (!m || !n) PetscFunctionReturn(PETSC_SUCCESS);
1743: PetscCall(PetscInfo(A, "Performing DiagonalScale %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m, n));
1744: PetscCall(GetHandles_(&dctx, &handle));
1745: {
1746: Vec lr;
1747: const auto da = DeviceArrayReadWrite(dctx, A);
1748: const auto lda = MatIMPLCast(A)->lda;
1750: if (l) {
1751: PetscCall(VecGetLocalSize(l, &m));
1752: PetscCheck(m == A->rmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Left scaling Vec of wrong size");
1753: PetscCall(VecGetArrayReadAndMemType(l, &dlr, &mtype));
1754: /* the array must live in memory this backend's BLAS can dereference, so test against this backend's
1755: memory type rather than with PetscMemTypeDevice(), which would also accept the other backend's
1756: device memory when PETSc is configured with both CUDA and HIP */
1757: if (!(T == device::cupm::DeviceType::CUDA ? PetscMemTypeCUDA(mtype) : PetscMemTypeHIP(mtype))) {
1758: PetscCall(VecRestoreArrayReadAndMemType(l, &dlr));
1759: /* l may be the parallel scaling Vec of a MATMPIDENSE, so the temporary is local and holds the local part */
1760: PetscCall(VecCreate(PETSC_COMM_SELF, &lr));
1761: PetscCall(VecSetSizes(lr, m, m));
1762: PetscCall(VecSetType(lr, VecSeq_CUPM::VECSEQCUPM()));
1763: PetscCall(VecCopy(l, lr));
1764: PetscCall(VecGetArrayReadAndMemType(lr, &dlr, nullptr));
1765: } else lr = l;
1766: {
1767: constexpr auto side = CUPMBLAS_SIDE_LEFT;
1769: PetscCall(PetscLogGpuTimeBegin());
1770: PetscCallCUPMBLAS(cupmBlasXdgmm(handle, side, m, n, da.cupmdata(), lda, cupmScalarPtrCast(dlr), 1, da.cupmdata(), lda));
1771: PetscCall(PetscLogGpuTimeEnd());
1772: PetscCall(PetscLogGpuFlops(1.0 * n * m));
1773: }
1774: PetscCall(VecRestoreArrayReadAndMemType(lr, &dlr));
1775: if (lr != l) PetscCall(VecDestroy(&lr));
1776: }
1777: if (r) {
1778: PetscCall(VecGetLocalSize(r, &n));
1779: PetscCheck(n == A->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Right scaling Vec of wrong size");
1780: PetscCall(VecGetArrayReadAndMemType(r, &dlr, &mtype));
1781: if (!(T == device::cupm::DeviceType::CUDA ? PetscMemTypeCUDA(mtype) : PetscMemTypeHIP(mtype))) {
1782: PetscCall(VecRestoreArrayReadAndMemType(r, &dlr));
1783: PetscCall(VecCreate(PETSC_COMM_SELF, &lr));
1784: PetscCall(VecSetSizes(lr, n, n));
1785: PetscCall(VecSetType(lr, VecSeq_CUPM::VECSEQCUPM()));
1786: PetscCall(VecCopy(r, lr));
1787: PetscCall(VecGetArrayReadAndMemType(lr, &dlr, nullptr));
1788: } else lr = r;
1789: {
1790: constexpr auto side = CUPMBLAS_SIDE_RIGHT;
1792: PetscCall(PetscLogGpuTimeBegin());
1793: PetscCallCUPMBLAS(cupmBlasXdgmm(handle, side, m, n, da.cupmdata(), lda, cupmScalarPtrCast(dlr), 1, da.cupmdata(), lda));
1794: PetscCall(PetscLogGpuTimeEnd());
1795: PetscCall(PetscLogGpuFlops(1.0 * n * m));
1796: }
1797: PetscCall(VecRestoreArrayReadAndMemType(lr, &dlr));
1798: if (lr != r) PetscCall(VecDestroy(&lr));
1799: }
1800: }
1801: PetscFunctionReturn(PETSC_SUCCESS);
1802: }
1804: template <device::cupm::DeviceType T>
1805: inline PetscErrorCode MatDense_Seq_CUPM<T>::AXPY(Mat Y, PetscScalar alpha, Mat X, MatStructure) noexcept
1806: {
1807: const auto m_x = X->rmap->n, m_y = Y->rmap->n;
1808: const auto n_x = X->cmap->n, n_y = Y->cmap->n;
1809: const auto N = m_x * n_x;
1810: PetscDeviceContext dctx;
1812: PetscFunctionBegin;
1813: if (!m_x || !n_x || alpha == (PetscScalar)0.0) PetscFunctionReturn(PETSC_SUCCESS);
1814: PetscCall(PetscInfo(Y, "Performing AXPY %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m_y, n_y));
1815: PetscCall(GetHandles_(&dctx));
1816: {
1817: const auto dx = DeviceArrayRead(dctx, X);
1818: const auto dy = DeviceArrayReadWrite(dctx, Y);
1819: const auto lda_x = MatIMPLCast(X)->lda;
1820: const auto lda_y = MatIMPLCast(Y)->lda;
1822: if (lda_x > m_x || lda_y > m_x) {
1823: cupmStream_t stream;
1825: PetscCall(GetHandlesFrom_(dctx, &stream));
1826: // clang-format off
1827: PetscCallThrust(
1828: const auto sub_mat_y = detail::make_submat_iterator(0, m_y, 0, n_y, lda_y, dy.data());
1829: const auto sub_mat_x = detail::make_submat_iterator(0, m_x, 0, n_x, lda_x, dx.data());
1831: THRUST_CALL(
1832: thrust::transform,
1833: stream,
1834: sub_mat_x.begin(), sub_mat_x.end(), sub_mat_y.begin(), sub_mat_y.begin(),
1835: device::cupm::functors::make_axpy(alpha)
1836: );
1837: );
1838: // clang-format on
1839: } else {
1840: const auto cu_alpha = cupmScalarCast(alpha);
1841: cupmBlasHandle_t handle;
1843: PetscCall(GetHandlesFrom_(dctx, &handle));
1844: PetscCall(PetscLogGpuTimeBegin());
1845: PetscCallCUPMBLAS(cupmBlasXaxpy(handle, N, &cu_alpha, dx.cupmdata(), 1, dy.cupmdata(), 1));
1846: PetscCall(PetscLogGpuTimeEnd());
1847: }
1848: }
1849: PetscCall(PetscLogGpuFlops(PetscMax(2 * N - 1, 0)));
1850: PetscFunctionReturn(PETSC_SUCCESS);
1851: }
1853: template <device::cupm::DeviceType T>
1854: inline PetscErrorCode MatDense_Seq_CUPM<T>::Duplicate(Mat A, MatDuplicateOption opt, Mat *B) noexcept
1855: {
1856: const auto hopt = (opt == MAT_COPY_VALUES && A->offloadmask != PETSC_OFFLOAD_CPU) ? MAT_DO_NOT_COPY_VALUES : opt;
1857: PetscDeviceContext dctx;
1859: PetscFunctionBegin;
1860: PetscCall(GetHandles_(&dctx));
1861: // do not call SetPreallocation() yet, we call it afterwards??
1862: PetscCall(MatCreateSeqDenseCUPM<T>(PetscObjectComm(PetscObjectCast(A)), A->rmap->n, A->cmap->n, nullptr, B, dctx, /* preallocate */ false));
1863: PetscCall(MatDuplicateNoCreate_SeqDense(*B, A, hopt));
1864: if (opt == MAT_COPY_VALUES && hopt != MAT_COPY_VALUES) PetscCall(Copy(A, *B, SAME_NONZERO_PATTERN));
1865: // allocate memory if needed
1866: if (opt != MAT_COPY_VALUES && !MatCUPMCast(*B)->d_v) PetscCall(SetPreallocation(*B, dctx, nullptr));
1867: PetscFunctionReturn(PETSC_SUCCESS);
1868: }
1870: template <device::cupm::DeviceType T>
1871: inline PetscErrorCode MatDense_Seq_CUPM<T>::SetRandom(Mat A, PetscRandom rng) noexcept
1872: {
1873: PetscBool device_rand_is_rander48;
1874: PetscBool device = PETSC_FALSE;
1876: PetscFunctionBegin;
1877: // CUPMObject::PETSCDEVICERAD() is PETSCRANDER48 until PetscRandom is implemented for hiprand
1878: PetscCall(PetscStrncmp(PETSCDEVICERAND(), PETSCRANDER48, sizeof(PETSCRANDER48), &device_rand_is_rander48));
1879: if (!device_rand_is_rander48) PetscCall(PetscObjectTypeCompare(PetscObjectCast(rng), PETSCDEVICERAND(), &device));
1880: if (device) {
1881: const auto m = A->rmap->n;
1882: const auto n = A->cmap->n;
1883: PetscDeviceContext dctx;
1885: PetscCall(GetHandles_(&dctx));
1886: {
1887: const auto a = DeviceArrayWrite(dctx, A);
1888: PetscInt lda;
1890: PetscCall(MatDenseGetLDA(A, &lda));
1891: if (lda > m) {
1892: for (PetscInt i = 0; i < n; i++) PetscCall(PetscRandomGetValues(rng, m, a.data() + i * lda));
1893: } else {
1894: PetscInt mn;
1896: PetscCall(PetscIntMultError(m, n, &mn));
1897: PetscCall(PetscRandomGetValues(rng, mn, a));
1898: }
1899: }
1900: } else {
1901: PetscCall(MatSetRandom_SeqDense(A, rng));
1902: }
1903: PetscFunctionReturn(PETSC_SUCCESS);
1904: }
1906: // ==========================================================================================
1908: template <device::cupm::DeviceType T>
1909: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetColumnVector(Mat A, Vec v, PetscInt col) noexcept
1910: {
1911: const auto offloadmask = A->offloadmask;
1912: const auto n = A->rmap->n;
1913: const auto col_offset = [&](const PetscScalar *ptr) { return ptr + col * MatIMPLCast(A)->lda; };
1914: PetscBool viscupm;
1915: PetscDeviceContext dctx;
1916: cupmStream_t stream;
1918: PetscFunctionBegin;
1919: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(v), &viscupm, VecSeq_CUPM::VECSEQCUPM(), VecSeq_CUPM::VECMPICUPM(), VecSeq_CUPM::VECCUPM(), ""));
1920: PetscCall(GetHandles_(&dctx, &stream));
1921: if (viscupm && !v->boundtocpu) {
1922: const auto x = VecSeq_CUPM::DeviceArrayWrite(dctx, v);
1924: // update device data
1925: if (PetscOffloadDevice(offloadmask)) {
1926: PetscCall(PetscCUPMMemcpyAsync(x.data(), col_offset(DeviceArrayRead(dctx, A)), n, cupmMemcpyDeviceToDevice, stream));
1927: } else {
1928: PetscCall(PetscCUPMMemcpyAsync(x.data(), col_offset(HostArrayRead(dctx, A)), n, cupmMemcpyHostToDevice, stream));
1929: }
1930: } else {
1931: PetscScalar *x;
1933: // update host data
1934: PetscCall(VecGetArrayWrite(v, &x));
1935: if (PetscOffloadUnallocated(offloadmask) || PetscOffloadHost(offloadmask)) {
1936: PetscCall(PetscArraycpy(x, col_offset(HostArrayRead(dctx, A)), n));
1937: } else if (PetscOffloadDevice(offloadmask)) {
1938: PetscCall(PetscCUPMMemcpyAsync(x, col_offset(DeviceArrayRead(dctx, A)), n, cupmMemcpyDeviceToHost, stream));
1939: }
1940: PetscCall(VecRestoreArrayWrite(v, &x));
1941: }
1942: PetscFunctionReturn(PETSC_SUCCESS);
1943: }
1945: // Place the device array a[] of a column of the matrix in its column Vec cvec. MatDenseCreateColumnVec_Private()
1946: // creates cvec from the VecType of the matrix, so it is a CUPM Vec, or a VECKOKKOS when the Kokkos backend runs on
1947: // this device (see MatCreateDenseFromVecType()); each needs its own place/reset routines
1948: template <device::cupm::DeviceType T>
1949: inline PetscErrorCode MatDense_Seq_CUPM<T>::PlaceColumnVecArray(Vec cvec, PetscScalar *a) noexcept
1950: {
1951: using namespace vec::cupm;
1952: PetscBool iscupm;
1954: PetscFunctionBegin;
1955: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(cvec), &iscupm, VecSeq_CUPM::VECSEQCUPM(), VecSeq_CUPM::VECMPICUPM(), ""));
1956: if (iscupm) PetscCall(VecCUPMPlaceArrayAsync<T>(cvec, a));
1957: else {
1958: #if PetscDefined(HAVE_KOKKOS_KERNELS)
1959: 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));
1960: PetscBool iskokkos;
1962: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(cvec), &iskokkos, VECSEQKOKKOS, VECMPIKOKKOS, ""));
1963: 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());
1964: 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);
1965: PetscCall(VecKokkosPlaceArray(cvec, a));
1966: #else
1967: 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());
1968: #endif
1969: }
1970: PetscFunctionReturn(PETSC_SUCCESS);
1971: }
1973: template <device::cupm::DeviceType T>
1974: inline PetscErrorCode MatDense_Seq_CUPM<T>::ResetColumnVecArray(Vec cvec) noexcept
1975: {
1976: using namespace vec::cupm;
1977: PetscBool iscupm;
1979: PetscFunctionBegin;
1980: PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(cvec), &iscupm, VecSeq_CUPM::VECSEQCUPM(), VecSeq_CUPM::VECMPICUPM(), ""));
1981: if (iscupm) PetscCall(VecCUPMResetArrayAsync<T>(cvec));
1982: else {
1983: #if PetscDefined(HAVE_KOKKOS_KERNELS)
1984: PetscCall(VecKokkosResetArray(cvec)); // PlaceColumnVecArray() already checked cvec is a VECKOKKOS
1985: #else
1986: 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());
1987: #endif
1988: }
1989: PetscFunctionReturn(PETSC_SUCCESS);
1990: }
1992: template <device::cupm::DeviceType T>
1993: template <PetscMemoryAccessMode access>
1994: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetColumnVec(Mat A, PetscInt col, Vec *v) noexcept
1995: {
1996: using namespace vec::cupm;
1997: const auto mimpl = MatIMPLCast(A);
1998: PetscDeviceContext dctx;
2000: PetscFunctionBegin;
2001: PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
2002: PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2003: mimpl->vecinuse = col + 1;
2004: if (!mimpl->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &mimpl->cvec));
2005: PetscCall(GetHandles_(&dctx));
2006: PetscCall(GetArray<PETSC_MEMTYPE_DEVICE, access>(A, const_cast<PetscScalar **>(&mimpl->ptrinuse), dctx));
2007: PetscCall(PlaceColumnVecArray(mimpl->cvec, const_cast<PetscScalar *>(mimpl->ptrinuse) + static_cast<std::size_t>(col) * static_cast<std::size_t>(mimpl->lda)));
2008: if (access == PETSC_MEMORY_ACCESS_READ) PetscCall(VecLockReadPush(mimpl->cvec));
2009: *v = mimpl->cvec;
2010: PetscFunctionReturn(PETSC_SUCCESS);
2011: }
2013: template <device::cupm::DeviceType T>
2014: template <PetscMemoryAccessMode access>
2015: inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreColumnVec(Mat A, PetscInt, Vec *v) noexcept
2016: {
2017: using namespace vec::cupm;
2018: const auto mimpl = MatIMPLCast(A);
2019: const auto cvec = mimpl->cvec;
2020: PetscDeviceContext dctx;
2022: PetscFunctionBegin;
2023: PetscCheck(mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
2024: PetscCheck(cvec, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column vector");
2025: mimpl->vecinuse = 0;
2026: if (access == PETSC_MEMORY_ACCESS_READ) PetscCall(VecLockReadPop(cvec));
2027: PetscCall(ResetColumnVecArray(cvec));
2028: PetscCall(GetHandles_(&dctx));
2029: PetscCall(RestoreArray<PETSC_MEMTYPE_DEVICE, access>(A, const_cast<PetscScalar **>(&mimpl->ptrinuse), dctx));
2030: if (v) *v = nullptr;
2031: PetscFunctionReturn(PETSC_SUCCESS);
2032: }
2034: // ==========================================================================================
2036: template <device::cupm::DeviceType T>
2037: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetFactor(Mat A, MatFactorType ftype, Mat *fact_out) noexcept
2038: {
2039: Mat fact = nullptr;
2040: PetscDeviceContext dctx;
2042: PetscFunctionBegin;
2043: PetscCall(GetHandles_(&dctx));
2044: PetscCall(MatCreateSeqDenseCUPM<T>(PetscObjectComm(PetscObjectCast(A)), A->rmap->n, A->cmap->n, nullptr, &fact, dctx, /* preallocate */ false));
2045: fact->factortype = ftype;
2046: switch (ftype) {
2047: case MAT_FACTOR_LU:
2048: case MAT_FACTOR_ILU: // fall-through
2049: fact->ops->lufactorsymbolic = MatLUFactorSymbolic_SeqDense;
2050: fact->ops->ilufactorsymbolic = MatLUFactorSymbolic_SeqDense;
2051: break;
2052: case MAT_FACTOR_CHOLESKY:
2053: case MAT_FACTOR_ICC: // fall-through
2054: fact->ops->choleskyfactorsymbolic = MatCholeskyFactorSymbolic_SeqDense;
2055: break;
2056: case MAT_FACTOR_QR: {
2057: const auto pobj = PetscObjectCast(fact);
2059: PetscCall(PetscObjectComposeFunction(pobj, "MatQRFactor_C", MatQRFactor_SeqDense));
2060: PetscCall(PetscObjectComposeFunction(pobj, "MatQRFactorSymbolic_C", MatQRFactorSymbolic_SeqDense));
2061: } break;
2062: case MAT_FACTOR_NONE:
2063: case MAT_FACTOR_ILUDT: // fall-through
2064: case MAT_FACTOR_NUM_TYPES: // fall-through
2065: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "MatFactorType %s not supported", MatFactorTypes[ftype]);
2066: }
2067: PetscCall(PetscStrFreeAllocpy(MATSOLVERCUPM(), &fact->solvertype));
2068: PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_LU));
2069: PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_ILU));
2070: PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_CHOLESKY));
2071: PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_ICC));
2072: *fact_out = fact;
2073: PetscFunctionReturn(PETSC_SUCCESS);
2074: }
2076: template <device::cupm::DeviceType T>
2077: inline PetscErrorCode MatDense_Seq_CUPM<T>::InvertFactors(Mat A) noexcept
2078: {
2079: const auto mimpl = MatIMPLCast(A);
2080: const auto mcu = MatCUPMCast(A);
2081: const auto n = static_cast<cupmBlasInt_t>(A->cmap->n);
2082: cupmSolverHandle_t handle;
2083: PetscDeviceContext dctx;
2084: cupmStream_t stream;
2086: PetscFunctionBegin;
2087: #if PetscDefined(HAVE_CUDA) && PetscDefined(USING_NVCC)
2088: // HIP appears to have this by default??
2089: PetscCheck(PETSC_PKG_CUDA_VERSION_GE(10, 1, 0), PETSC_COMM_SELF, PETSC_ERR_SUP, "Upgrade to CUDA version 10.1.0 or higher");
2090: #endif
2091: if (!n || !A->rmap->n) PetscFunctionReturn(PETSC_SUCCESS);
2092: PetscCheck(A->factortype == MAT_FACTOR_CHOLESKY, PETSC_COMM_SELF, PETSC_ERR_LIB, "Factor type %s not implemented", MatFactorTypes[A->factortype]);
2093: // spd
2094: PetscCheck(!mcu->d_fact_ipiv, PETSC_COMM_SELF, PETSC_ERR_LIB, "%sDnsytri not implemented", cupmSolverName());
2096: PetscCall(GetHandles_(&dctx, &handle, &stream));
2097: {
2098: const auto da = DeviceArrayReadWrite(dctx, A);
2099: const auto lda = static_cast<cupmBlasInt_t>(mimpl->lda);
2100: cupmBlasInt_t il;
2102: PetscCallCUPMSOLVER(cupmSolverXpotri_bufferSize(handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, &il));
2103: if (il > mcu->d_fact_lwork) {
2104: mcu->d_fact_lwork = il;
2105: PetscCallCUPM(cupmFreeAsync(mcu->d_fact_work, stream));
2106: PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_work, il, stream));
2107: }
2108: PetscCall(PetscLogGpuTimeBegin());
2109: PetscCallCUPMSOLVER(cupmSolverXpotri(handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_info));
2110: PetscCall(PetscLogGpuTimeEnd());
2111: }
2112: PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream));
2113: // TODO (write cuda kernel)
2114: PetscCall(MatSeqDenseSymmetrize_Private(A, PETSC_TRUE));
2115: PetscCall(PetscLogGpuFlops(1.0 * n * n * n / 3.0));
2117: A->ops->solve = nullptr;
2118: A->ops->solvetranspose = nullptr;
2119: A->ops->matsolve = nullptr;
2120: A->factortype = MAT_FACTOR_NONE;
2122: PetscCall(PetscFree(A->solvertype));
2123: PetscFunctionReturn(PETSC_SUCCESS);
2124: }
2126: // ==========================================================================================
2128: template <device::cupm::DeviceType T>
2129: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetSubMatrix(Mat A, PetscInt rbegin, PetscInt rend, PetscInt cbegin, PetscInt cend, Mat *mat) noexcept
2130: {
2131: const auto mimpl = MatIMPLCast(A);
2132: const auto array_offset = [&](PetscScalar *ptr) { return ptr + rbegin + static_cast<std::size_t>(cbegin) * mimpl->lda; };
2133: const auto n = rend - rbegin;
2134: const auto m = cend - cbegin;
2135: auto &cmat = mimpl->cmat;
2136: PetscDeviceContext dctx;
2138: PetscFunctionBegin;
2139: PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
2140: PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2141: mimpl->matinuse = cbegin + 1;
2143: PetscCall(GetHandles_(&dctx));
2144: PetscCall(HostToDevice_(A, dctx));
2146: if (cmat && ((m != cmat->cmap->N) || (n != cmat->rmap->N))) PetscCall(MatDestroy(&cmat));
2147: {
2148: const auto device_array = array_offset(MatCUPMCast(A)->d_v);
2150: if (cmat) {
2151: PetscCall(PlaceArray(cmat, device_array));
2152: } else {
2153: PetscCall(MatCreateSeqDenseCUPM<T>(PetscObjectComm(PetscObjectCast(A)), n, m, device_array, &cmat, dctx));
2154: PetscCall(MatSetVecType(cmat, A->defaultvectype));
2155: }
2156: }
2157: PetscCall(MatDenseSetLDA(cmat, mimpl->lda));
2158: // place CPU array if present but do not copy any data
2159: if (const auto host_array = mimpl->v) {
2160: cmat->offloadmask = PETSC_OFFLOAD_GPU;
2161: PetscCall(MatDensePlaceArray(cmat, array_offset(host_array)));
2162: }
2164: cmat->offloadmask = A->offloadmask;
2165: *mat = cmat;
2166: PetscFunctionReturn(PETSC_SUCCESS);
2167: }
2169: template <device::cupm::DeviceType T>
2170: inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreSubMatrix(Mat A, Mat *m) noexcept
2171: {
2172: const auto mimpl = MatIMPLCast(A);
2173: const auto cmat = mimpl->cmat;
2174: const auto reset = static_cast<bool>(mimpl->v);
2175: bool copy, was_offload_host;
2177: PetscFunctionBegin;
2178: PetscCheck(mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetSubMatrix() first");
2179: PetscCheck(cmat, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column matrix");
2180: PetscCheck(*m == cmat, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Not the matrix obtained from MatDenseGetSubMatrix()");
2181: mimpl->matinuse = 0;
2183: // calls to ResetArray may change it, so save it here
2184: was_offload_host = cmat->offloadmask == PETSC_OFFLOAD_CPU;
2185: if (was_offload_host && !reset) {
2186: copy = true;
2187: PetscCall(MatSeqDenseSetPreallocation(A, nullptr));
2188: } else {
2189: copy = false;
2190: }
2192: PetscCall(ResetArray(cmat));
2193: if (reset) PetscCall(MatDenseResetArray(cmat));
2194: if (copy) {
2195: PetscDeviceContext dctx;
2197: PetscCall(GetHandles_(&dctx));
2198: PetscCall(DeviceToHost_(A, dctx));
2199: } else {
2200: A->offloadmask = was_offload_host ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
2201: }
2203: cmat->offloadmask = PETSC_OFFLOAD_UNALLOCATED;
2204: *m = nullptr;
2205: PetscFunctionReturn(PETSC_SUCCESS);
2206: }
2208: // ==========================================================================================
2210: namespace
2211: {
2213: template <device::cupm::DeviceType T>
2214: inline PetscErrorCode MatMatMultNumeric_SeqDenseCUPM_SeqDenseCUPM(Mat A, Mat B, Mat C, PetscBool TA, PetscBool TB) noexcept
2215: {
2216: PetscFunctionBegin;
2217: if (TA) {
2218: if (TB) {
2219: PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<true, true>(A, B, C));
2220: } else {
2221: PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<true, false>(A, B, C));
2222: }
2223: } else {
2224: if (TB) {
2225: PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<false, true>(A, B, C));
2226: } else {
2227: PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<false, false>(A, B, C));
2228: }
2229: }
2230: PetscFunctionReturn(PETSC_SUCCESS);
2231: }
2233: template <device::cupm::DeviceType T>
2234: inline PetscErrorCode MatSolverTypeRegister_DENSECUPM() noexcept
2235: {
2236: PetscFunctionBegin;
2237: for (auto ftype : util::make_array(MAT_FACTOR_LU, MAT_FACTOR_CHOLESKY, MAT_FACTOR_QR)) {
2238: PetscCall(MatSolverTypeRegister(MatDense_Seq_CUPM<T>::MATSOLVERCUPM(), MATSEQDENSE, ftype, MatDense_Seq_CUPM<T>::GetFactor));
2239: PetscCall(MatSolverTypeRegister(MatDense_Seq_CUPM<T>::MATSOLVERCUPM(), MatDense_Seq_CUPM<T>::MATSEQDENSECUPM(), ftype, MatDense_Seq_CUPM<T>::GetFactor));
2240: }
2241: PetscFunctionReturn(PETSC_SUCCESS);
2242: }
2244: } // anonymous namespace
2246: } // namespace impl
2248: } // namespace cupm
2250: } // namespace mat
2252: } // namespace Petsc