xref: /petsc/src/mat/impls/dense/seq/cupm/matseqdensecupm.hpp (revision a496304597bacff3545e802853d69e8765312868)
1*a4963045SJacob Faibussowitsch #pragma once
24742e46bSJacob Faibussowitsch 
34742e46bSJacob Faibussowitsch #include <petsc/private/matdensecupmimpl.h> /*I <petscmat.h> I*/
44742e46bSJacob Faibussowitsch #include <../src/mat/impls/dense/seq/dense.h>
54742e46bSJacob Faibussowitsch 
64742e46bSJacob Faibussowitsch #include <petsc/private/deviceimpl.h> // PetscDeviceContextGetOptionalNullContext_Internal()
74742e46bSJacob Faibussowitsch #include <petsc/private/randomimpl.h> // _p_PetscRandom
84742e46bSJacob Faibussowitsch #include <petsc/private/vecimpl.h>    // _p_Vec
94742e46bSJacob Faibussowitsch #include <petsc/private/cupmobject.hpp>
104742e46bSJacob Faibussowitsch #include <petsc/private/cupmsolverinterface.hpp>
114742e46bSJacob Faibussowitsch 
124742e46bSJacob Faibussowitsch #include <petsc/private/cpp/type_traits.hpp> // PetscObjectCast()
134742e46bSJacob Faibussowitsch #include <petsc/private/cpp/utility.hpp>     // util::exchange()
144742e46bSJacob Faibussowitsch 
154742e46bSJacob Faibussowitsch #include <../src/vec/vec/impls/seq/cupm/vecseqcupm.hpp> // for VecSeq_CUPM
164742e46bSJacob Faibussowitsch 
174742e46bSJacob Faibussowitsch namespace Petsc
184742e46bSJacob Faibussowitsch {
194742e46bSJacob Faibussowitsch 
204742e46bSJacob Faibussowitsch namespace mat
214742e46bSJacob Faibussowitsch {
224742e46bSJacob Faibussowitsch 
234742e46bSJacob Faibussowitsch namespace cupm
244742e46bSJacob Faibussowitsch {
254742e46bSJacob Faibussowitsch 
264742e46bSJacob Faibussowitsch namespace impl
274742e46bSJacob Faibussowitsch {
284742e46bSJacob Faibussowitsch 
294742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
304742e46bSJacob Faibussowitsch class MatDense_Seq_CUPM : MatDense_CUPM<T, MatDense_Seq_CUPM<T>> {
314742e46bSJacob Faibussowitsch public:
324742e46bSJacob Faibussowitsch   MATDENSECUPM_HEADER(T, MatDense_Seq_CUPM<T>);
334742e46bSJacob Faibussowitsch 
344742e46bSJacob Faibussowitsch private:
354742e46bSJacob Faibussowitsch   struct Mat_SeqDenseCUPM {
364742e46bSJacob Faibussowitsch     PetscScalar *d_v;           // pointer to the matrix on the GPU
374742e46bSJacob Faibussowitsch     PetscScalar *unplacedarray; // if one called MatCUPMDensePlaceArray(), this is where it stashed the original
384742e46bSJacob Faibussowitsch     bool         d_user_alloc;
394742e46bSJacob Faibussowitsch     bool         d_unplaced_user_alloc;
404742e46bSJacob Faibussowitsch     // factorization support
414742e46bSJacob Faibussowitsch     cupmBlasInt_t *d_fact_ipiv;  // device pivots
424742e46bSJacob Faibussowitsch     cupmScalar_t  *d_fact_tau;   // device QR tau vector
434742e46bSJacob Faibussowitsch     cupmBlasInt_t *d_fact_info;  // device info
444742e46bSJacob Faibussowitsch     cupmScalar_t  *d_fact_work;  // device workspace
454742e46bSJacob Faibussowitsch     cupmBlasInt_t  d_fact_lwork; // size of device workspace
464742e46bSJacob Faibussowitsch     // workspace
474742e46bSJacob Faibussowitsch     Vec workvec;
484742e46bSJacob Faibussowitsch   };
494742e46bSJacob Faibussowitsch 
504742e46bSJacob Faibussowitsch   static PetscErrorCode SetPreallocation_(Mat, PetscDeviceContext, PetscScalar *) noexcept;
514742e46bSJacob Faibussowitsch 
524742e46bSJacob Faibussowitsch   static PetscErrorCode HostToDevice_(Mat, PetscDeviceContext) noexcept;
534742e46bSJacob Faibussowitsch   static PetscErrorCode DeviceToHost_(Mat, PetscDeviceContext) noexcept;
544742e46bSJacob Faibussowitsch 
554742e46bSJacob Faibussowitsch   static PetscErrorCode CheckCUPMSolverInfo_(const cupmBlasInt_t *, cupmStream_t) noexcept;
564742e46bSJacob Faibussowitsch 
574742e46bSJacob Faibussowitsch   template <typename Derived>
584742e46bSJacob Faibussowitsch   struct SolveCommon;
594742e46bSJacob Faibussowitsch   struct SolveQR;
604742e46bSJacob Faibussowitsch   struct SolveCholesky;
614742e46bSJacob Faibussowitsch   struct SolveLU;
624742e46bSJacob Faibussowitsch 
634742e46bSJacob Faibussowitsch   template <typename Solver, bool transpose>
644742e46bSJacob Faibussowitsch   static PetscErrorCode MatSolve_Factored_Dispatch_(Mat, Vec, Vec) noexcept;
654742e46bSJacob Faibussowitsch   template <typename Solver, bool transpose>
664742e46bSJacob Faibussowitsch   static PetscErrorCode MatMatSolve_Factored_Dispatch_(Mat, Mat, Mat) noexcept;
674742e46bSJacob Faibussowitsch   template <bool transpose>
684742e46bSJacob Faibussowitsch   static PetscErrorCode MatMultAdd_Dispatch_(Mat, Vec, Vec, Vec) noexcept;
694742e46bSJacob Faibussowitsch 
704742e46bSJacob Faibussowitsch   template <bool to_host>
714742e46bSJacob Faibussowitsch   static PetscErrorCode Convert_Dispatch_(Mat, MatType, MatReuse, Mat *) noexcept;
724742e46bSJacob Faibussowitsch 
734742e46bSJacob Faibussowitsch   PETSC_NODISCARD static constexpr MatType       MATIMPLCUPM_() noexcept;
744742e46bSJacob Faibussowitsch   PETSC_NODISCARD static constexpr Mat_SeqDense *MatIMPLCast_(Mat) noexcept;
754742e46bSJacob Faibussowitsch 
764742e46bSJacob Faibussowitsch public:
774742e46bSJacob Faibussowitsch   PETSC_NODISCARD static constexpr Mat_SeqDenseCUPM *MatCUPMCast(Mat) noexcept;
784742e46bSJacob Faibussowitsch 
794742e46bSJacob Faibussowitsch   // define these by hand since they don't fit the above mold
804742e46bSJacob Faibussowitsch   PETSC_NODISCARD static constexpr const char *MatConvert_seqdensecupm_seqdense_C() noexcept;
814742e46bSJacob Faibussowitsch   PETSC_NODISCARD static constexpr const char *MatProductSetFromOptions_seqaij_seqdensecupm_C() noexcept;
824742e46bSJacob Faibussowitsch 
834742e46bSJacob Faibussowitsch   static PetscErrorCode Create(Mat) noexcept;
844742e46bSJacob Faibussowitsch   static PetscErrorCode Destroy(Mat) noexcept;
854742e46bSJacob Faibussowitsch   static PetscErrorCode SetUp(Mat) noexcept;
864742e46bSJacob Faibussowitsch   static PetscErrorCode Reset(Mat) noexcept;
874742e46bSJacob Faibussowitsch 
884742e46bSJacob Faibussowitsch   static PetscErrorCode BindToCPU(Mat, PetscBool) noexcept;
894742e46bSJacob Faibussowitsch   static PetscErrorCode Convert_SeqDense_SeqDenseCUPM(Mat, MatType, MatReuse, Mat *) noexcept;
904742e46bSJacob Faibussowitsch   static PetscErrorCode Convert_SeqDenseCUPM_SeqDense(Mat, MatType, MatReuse, Mat *) noexcept;
914742e46bSJacob Faibussowitsch 
924742e46bSJacob Faibussowitsch   template <PetscMemType, PetscMemoryAccessMode>
934742e46bSJacob Faibussowitsch   static PetscErrorCode GetArray(Mat, PetscScalar **, PetscDeviceContext) noexcept;
944742e46bSJacob Faibussowitsch   template <PetscMemType, PetscMemoryAccessMode>
954742e46bSJacob Faibussowitsch   static PetscErrorCode RestoreArray(Mat, PetscScalar **, PetscDeviceContext) noexcept;
964742e46bSJacob Faibussowitsch   template <PetscMemoryAccessMode>
974742e46bSJacob Faibussowitsch   static PetscErrorCode GetArrayAndMemType(Mat, PetscScalar **, PetscMemType *, PetscDeviceContext) noexcept;
984742e46bSJacob Faibussowitsch   template <PetscMemoryAccessMode>
994742e46bSJacob Faibussowitsch   static PetscErrorCode RestoreArrayAndMemType(Mat, PetscScalar **, PetscDeviceContext) noexcept;
1004742e46bSJacob Faibussowitsch 
1014742e46bSJacob Faibussowitsch private:
1024742e46bSJacob Faibussowitsch   template <PetscMemType mtype, PetscMemoryAccessMode mode>
1034742e46bSJacob Faibussowitsch   static PetscErrorCode GetArrayC_(Mat m, PetscScalar **p) noexcept
1044742e46bSJacob Faibussowitsch   {
1054742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
1064742e46bSJacob Faibussowitsch 
1074742e46bSJacob Faibussowitsch     PetscFunctionBegin;
1084742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
1094742e46bSJacob Faibussowitsch     PetscCall(GetArray<mtype, mode>(m, p, dctx));
1104742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
1114742e46bSJacob Faibussowitsch   }
1124742e46bSJacob Faibussowitsch 
1134742e46bSJacob Faibussowitsch   template <PetscMemType mtype, PetscMemoryAccessMode mode>
1144742e46bSJacob Faibussowitsch   static PetscErrorCode RestoreArrayC_(Mat m, PetscScalar **p) noexcept
1154742e46bSJacob Faibussowitsch   {
1164742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
1174742e46bSJacob Faibussowitsch 
1184742e46bSJacob Faibussowitsch     PetscFunctionBegin;
1194742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
1204742e46bSJacob Faibussowitsch     PetscCall(RestoreArray<mtype, mode>(m, p, dctx));
1214742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
1224742e46bSJacob Faibussowitsch   }
1234742e46bSJacob Faibussowitsch 
1244742e46bSJacob Faibussowitsch   template <PetscMemoryAccessMode mode>
1254742e46bSJacob Faibussowitsch   static PetscErrorCode GetArrayAndMemTypeC_(Mat m, PetscScalar **p, PetscMemType *tp) noexcept
1264742e46bSJacob Faibussowitsch   {
1274742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
1284742e46bSJacob Faibussowitsch 
1294742e46bSJacob Faibussowitsch     PetscFunctionBegin;
1304742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
1314742e46bSJacob Faibussowitsch     PetscCall(GetArrayAndMemType<mode>(m, p, tp, dctx));
1324742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
1334742e46bSJacob Faibussowitsch   }
1344742e46bSJacob Faibussowitsch 
1354742e46bSJacob Faibussowitsch   template <PetscMemoryAccessMode mode>
1364742e46bSJacob Faibussowitsch   static PetscErrorCode RestoreArrayAndMemTypeC_(Mat m, PetscScalar **p) noexcept
1374742e46bSJacob Faibussowitsch   {
1384742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
1394742e46bSJacob Faibussowitsch 
1404742e46bSJacob Faibussowitsch     PetscFunctionBegin;
1414742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
1424742e46bSJacob Faibussowitsch     PetscCall(RestoreArrayAndMemType<mode>(m, p, dctx));
1434742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
1444742e46bSJacob Faibussowitsch   }
1454742e46bSJacob Faibussowitsch 
1464742e46bSJacob Faibussowitsch public:
1474742e46bSJacob Faibussowitsch   static PetscErrorCode PlaceArray(Mat, const PetscScalar *) noexcept;
1484742e46bSJacob Faibussowitsch   static PetscErrorCode ReplaceArray(Mat, const PetscScalar *) noexcept;
1494742e46bSJacob Faibussowitsch   static PetscErrorCode ResetArray(Mat) noexcept;
1504742e46bSJacob Faibussowitsch 
1514742e46bSJacob Faibussowitsch   template <bool transpose_A, bool transpose_B>
1524742e46bSJacob Faibussowitsch   static PetscErrorCode MatMatMult_Numeric_Dispatch(Mat, Mat, Mat) noexcept;
1534742e46bSJacob Faibussowitsch   static PetscErrorCode Copy(Mat, Mat, MatStructure) noexcept;
1544742e46bSJacob Faibussowitsch   static PetscErrorCode ZeroEntries(Mat) noexcept;
1554742e46bSJacob Faibussowitsch   static PetscErrorCode Scale(Mat, PetscScalar) noexcept;
1564742e46bSJacob Faibussowitsch   static PetscErrorCode AXPY(Mat, PetscScalar, Mat, MatStructure) noexcept;
1574742e46bSJacob Faibussowitsch   static PetscErrorCode Duplicate(Mat, MatDuplicateOption, Mat *) noexcept;
1584742e46bSJacob Faibussowitsch   static PetscErrorCode SetRandom(Mat, PetscRandom) noexcept;
1594742e46bSJacob Faibussowitsch 
1604742e46bSJacob Faibussowitsch   static PetscErrorCode GetColumnVector(Mat, Vec, PetscInt) noexcept;
1614742e46bSJacob Faibussowitsch   template <PetscMemoryAccessMode>
1624742e46bSJacob Faibussowitsch   static PetscErrorCode GetColumnVec(Mat, PetscInt, Vec *) noexcept;
1634742e46bSJacob Faibussowitsch   template <PetscMemoryAccessMode>
1644742e46bSJacob Faibussowitsch   static PetscErrorCode RestoreColumnVec(Mat, PetscInt, Vec *) noexcept;
1654742e46bSJacob Faibussowitsch 
1664742e46bSJacob Faibussowitsch   static PetscErrorCode GetFactor(Mat, MatFactorType, Mat *) noexcept;
1674742e46bSJacob Faibussowitsch   static PetscErrorCode InvertFactors(Mat) noexcept;
1684742e46bSJacob Faibussowitsch 
1694742e46bSJacob Faibussowitsch   static PetscErrorCode GetSubMatrix(Mat, PetscInt, PetscInt, PetscInt, PetscInt, Mat *) noexcept;
1704742e46bSJacob Faibussowitsch   static PetscErrorCode RestoreSubMatrix(Mat, Mat *) noexcept;
1714742e46bSJacob Faibussowitsch };
1724742e46bSJacob Faibussowitsch 
1734742e46bSJacob Faibussowitsch } // namespace impl
1744742e46bSJacob Faibussowitsch 
1754742e46bSJacob Faibussowitsch namespace
1764742e46bSJacob Faibussowitsch {
1774742e46bSJacob Faibussowitsch 
1784742e46bSJacob Faibussowitsch // Declare this here so that the functions below can make use of it
1794742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
1804742e46bSJacob Faibussowitsch inline PetscErrorCode MatCreateSeqDenseCUPM(MPI_Comm comm, PetscInt m, PetscInt n, PetscScalar *data, Mat *A, PetscDeviceContext dctx = nullptr, bool preallocate = true) noexcept
1814742e46bSJacob Faibussowitsch {
1824742e46bSJacob Faibussowitsch   PetscFunctionBegin;
1834742e46bSJacob Faibussowitsch   PetscCall(impl::MatDense_Seq_CUPM<T>::CreateIMPLDenseCUPM(comm, m, n, m, n, data, A, dctx, preallocate));
1844742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1854742e46bSJacob Faibussowitsch }
1864742e46bSJacob Faibussowitsch 
1874742e46bSJacob Faibussowitsch } // anonymous namespace
1884742e46bSJacob Faibussowitsch 
1894742e46bSJacob Faibussowitsch namespace impl
1904742e46bSJacob Faibussowitsch {
1914742e46bSJacob Faibussowitsch 
1924742e46bSJacob Faibussowitsch // ==========================================================================================
1934742e46bSJacob Faibussowitsch // MatDense_Seq_CUPM - Private API - Utility
1944742e46bSJacob Faibussowitsch // ==========================================================================================
1954742e46bSJacob Faibussowitsch 
1964742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
1974742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::SetPreallocation_(Mat m, PetscDeviceContext dctx, PetscScalar *user_device_array) noexcept
1984742e46bSJacob Faibussowitsch {
1994742e46bSJacob Faibussowitsch   const auto   mcu   = MatCUPMCast(m);
2004742e46bSJacob Faibussowitsch   const auto   nrows = m->rmap->n;
2014742e46bSJacob Faibussowitsch   const auto   ncols = m->cmap->n;
2024742e46bSJacob Faibussowitsch   auto        &lda   = MatIMPLCast(m)->lda;
2034742e46bSJacob Faibussowitsch   cupmStream_t stream;
2044742e46bSJacob Faibussowitsch 
2054742e46bSJacob Faibussowitsch   PetscFunctionBegin;
2064742e46bSJacob Faibussowitsch   PetscCheckTypeName(m, MATSEQDENSECUPM());
2074742e46bSJacob Faibussowitsch   PetscValidDeviceContext(dctx, 2);
2084742e46bSJacob Faibussowitsch   PetscCall(checkCupmBlasIntCast(nrows));
2094742e46bSJacob Faibussowitsch   PetscCall(checkCupmBlasIntCast(ncols));
2104742e46bSJacob Faibussowitsch   PetscCall(GetHandlesFrom_(dctx, &stream));
2114742e46bSJacob Faibussowitsch   if (lda <= 0) lda = nrows;
2124742e46bSJacob Faibussowitsch   if (!mcu->d_user_alloc) PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream));
2134742e46bSJacob Faibussowitsch   if (user_device_array) {
2144742e46bSJacob Faibussowitsch     mcu->d_user_alloc = PETSC_TRUE;
2154742e46bSJacob Faibussowitsch     mcu->d_v          = user_device_array;
2164742e46bSJacob Faibussowitsch   } else {
2174742e46bSJacob Faibussowitsch     PetscInt size;
2184742e46bSJacob Faibussowitsch 
2194742e46bSJacob Faibussowitsch     mcu->d_user_alloc = PETSC_FALSE;
2204742e46bSJacob Faibussowitsch     PetscCall(PetscIntMultError(lda, ncols, &size));
2214742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMMallocAsync(&mcu->d_v, size, stream));
2224742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMMemsetAsync(mcu->d_v, 0, size, stream));
2234742e46bSJacob Faibussowitsch   }
2244742e46bSJacob Faibussowitsch   m->offloadmask = PETSC_OFFLOAD_GPU;
2254742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2264742e46bSJacob Faibussowitsch }
2274742e46bSJacob Faibussowitsch 
2284742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
2294742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::HostToDevice_(Mat m, PetscDeviceContext dctx) noexcept
2304742e46bSJacob Faibussowitsch {
2314742e46bSJacob Faibussowitsch   const auto nrows = m->rmap->n;
2324742e46bSJacob Faibussowitsch   const auto ncols = m->cmap->n;
2334742e46bSJacob Faibussowitsch   const auto copy  = m->offloadmask == PETSC_OFFLOAD_CPU || m->offloadmask == PETSC_OFFLOAD_UNALLOCATED;
2344742e46bSJacob Faibussowitsch 
2354742e46bSJacob Faibussowitsch   PetscFunctionBegin;
2364742e46bSJacob Faibussowitsch   PetscCheckTypeName(m, MATSEQDENSECUPM());
2374742e46bSJacob Faibussowitsch   if (m->boundtocpu) PetscFunctionReturn(PETSC_SUCCESS);
2384742e46bSJacob Faibussowitsch   PetscCall(PetscInfo(m, "%s matrix %" PetscInt_FMT " x %" PetscInt_FMT "\n", copy ? "Copy" : "Reusing", nrows, ncols));
2394742e46bSJacob Faibussowitsch   if (copy) {
2404742e46bSJacob Faibussowitsch     const auto   mcu = MatCUPMCast(m);
2414742e46bSJacob Faibussowitsch     cupmStream_t stream;
2424742e46bSJacob Faibussowitsch 
2434742e46bSJacob Faibussowitsch     // Allocate GPU memory if not present
2443d9668e3SJacob Faibussowitsch     if (!mcu->d_v) PetscCall(SetPreallocation(m, dctx, nullptr));
2454742e46bSJacob Faibussowitsch     PetscCall(GetHandlesFrom_(dctx, &stream));
2464742e46bSJacob Faibussowitsch     PetscCall(PetscLogEventBegin(MAT_DenseCopyToGPU, m, 0, 0, 0));
2474742e46bSJacob Faibussowitsch     {
2484742e46bSJacob Faibussowitsch       const auto mimpl = MatIMPLCast(m);
2494742e46bSJacob Faibussowitsch       const auto lda   = mimpl->lda;
2504742e46bSJacob Faibussowitsch       const auto src   = mimpl->v;
2514742e46bSJacob Faibussowitsch       const auto dest  = mcu->d_v;
2524742e46bSJacob Faibussowitsch 
2534742e46bSJacob Faibussowitsch       if (lda > nrows) {
2544742e46bSJacob Faibussowitsch         PetscCall(PetscCUPMMemcpy2DAsync(dest, lda, src, lda, nrows, ncols, cupmMemcpyHostToDevice, stream));
2554742e46bSJacob Faibussowitsch       } else {
2564742e46bSJacob Faibussowitsch         PetscCall(PetscCUPMMemcpyAsync(dest, src, lda * ncols, cupmMemcpyHostToDevice, stream));
2574742e46bSJacob Faibussowitsch       }
2584742e46bSJacob Faibussowitsch     }
2594742e46bSJacob Faibussowitsch     PetscCall(PetscLogEventEnd(MAT_DenseCopyToGPU, m, 0, 0, 0));
2604742e46bSJacob Faibussowitsch     // order important, ensure that offloadmask is PETSC_OFFLOAD_BOTH
2614742e46bSJacob Faibussowitsch     m->offloadmask = PETSC_OFFLOAD_BOTH;
2624742e46bSJacob Faibussowitsch   }
2634742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2644742e46bSJacob Faibussowitsch }
2654742e46bSJacob Faibussowitsch 
2664742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
2674742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::DeviceToHost_(Mat m, PetscDeviceContext dctx) noexcept
2684742e46bSJacob Faibussowitsch {
2694742e46bSJacob Faibussowitsch   const auto nrows = m->rmap->n;
2704742e46bSJacob Faibussowitsch   const auto ncols = m->cmap->n;
2714742e46bSJacob Faibussowitsch   const auto copy  = m->offloadmask == PETSC_OFFLOAD_GPU;
2724742e46bSJacob Faibussowitsch 
2734742e46bSJacob Faibussowitsch   PetscFunctionBegin;
2744742e46bSJacob Faibussowitsch   PetscCheckTypeName(m, MATSEQDENSECUPM());
2754742e46bSJacob Faibussowitsch   PetscCall(PetscInfo(m, "%s matrix %" PetscInt_FMT " x %" PetscInt_FMT "\n", copy ? "Copy" : "Reusing", nrows, ncols));
2764742e46bSJacob Faibussowitsch   if (copy) {
2774742e46bSJacob Faibussowitsch     const auto   mimpl = MatIMPLCast(m);
2784742e46bSJacob Faibussowitsch     cupmStream_t stream;
2794742e46bSJacob Faibussowitsch 
2804742e46bSJacob Faibussowitsch     // MatCreateSeqDenseCUPM may not allocate CPU memory. Allocate if needed
2814742e46bSJacob Faibussowitsch     if (!mimpl->v) PetscCall(MatSeqDenseSetPreallocation(m, nullptr));
2824742e46bSJacob Faibussowitsch     PetscCall(GetHandlesFrom_(dctx, &stream));
2834742e46bSJacob Faibussowitsch     PetscCall(PetscLogEventBegin(MAT_DenseCopyFromGPU, m, 0, 0, 0));
2844742e46bSJacob Faibussowitsch     {
2854742e46bSJacob Faibussowitsch       const auto lda  = mimpl->lda;
2864742e46bSJacob Faibussowitsch       const auto dest = mimpl->v;
2874742e46bSJacob Faibussowitsch       const auto src  = MatCUPMCast(m)->d_v;
2884742e46bSJacob Faibussowitsch 
2894742e46bSJacob Faibussowitsch       if (lda > nrows) {
2904742e46bSJacob Faibussowitsch         PetscCall(PetscCUPMMemcpy2DAsync(dest, lda, src, lda, nrows, ncols, cupmMemcpyDeviceToHost, stream));
2914742e46bSJacob Faibussowitsch       } else {
2924742e46bSJacob Faibussowitsch         PetscCall(PetscCUPMMemcpyAsync(dest, src, lda * ncols, cupmMemcpyDeviceToHost, stream));
2934742e46bSJacob Faibussowitsch       }
2944742e46bSJacob Faibussowitsch     }
2954742e46bSJacob Faibussowitsch     PetscCall(PetscLogEventEnd(MAT_DenseCopyFromGPU, m, 0, 0, 0));
2964742e46bSJacob Faibussowitsch     // order is important, MatSeqDenseSetPreallocation() might set offloadmask
2974742e46bSJacob Faibussowitsch     m->offloadmask = PETSC_OFFLOAD_BOTH;
2984742e46bSJacob Faibussowitsch   }
2994742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3004742e46bSJacob Faibussowitsch }
3014742e46bSJacob Faibussowitsch 
3024742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
3034742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::CheckCUPMSolverInfo_(const cupmBlasInt_t *fact_info, cupmStream_t stream) noexcept
3044742e46bSJacob Faibussowitsch {
3054742e46bSJacob Faibussowitsch   PetscFunctionBegin;
3064742e46bSJacob Faibussowitsch   if (PetscDefined(USE_DEBUG)) {
3074742e46bSJacob Faibussowitsch     cupmBlasInt_t info = 0;
3084742e46bSJacob Faibussowitsch 
3094742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMMemcpyAsync(&info, fact_info, 1, cupmMemcpyDeviceToHost, stream));
3104742e46bSJacob Faibussowitsch     if (stream) PetscCallCUPM(cupmStreamSynchronize(stream));
3114742e46bSJacob Faibussowitsch     static_assert(std::is_same<decltype(info), int>::value, "");
3124742e46bSJacob Faibussowitsch     PetscCheck(info <= 0, PETSC_COMM_SELF, PETSC_ERR_MAT_CH_ZRPVT, "Bad factorization: zero pivot in row %d", info - 1);
3134742e46bSJacob Faibussowitsch     PetscCheck(info >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Wrong argument to cupmSolver %d", -info);
3144742e46bSJacob Faibussowitsch   }
3154742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3164742e46bSJacob Faibussowitsch }
3174742e46bSJacob Faibussowitsch 
3184742e46bSJacob Faibussowitsch // ==========================================================================================
3194742e46bSJacob Faibussowitsch // MatDense_Seq_CUPM - Private API - Solver Dispatch
3204742e46bSJacob Faibussowitsch // ==========================================================================================
3214742e46bSJacob Faibussowitsch 
3224742e46bSJacob Faibussowitsch // specific solvers called through the dispatch_() family of functions
3234742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
3244742e46bSJacob Faibussowitsch template <typename Derived>
3254742e46bSJacob Faibussowitsch struct MatDense_Seq_CUPM<T>::SolveCommon {
3264742e46bSJacob Faibussowitsch   using derived_type = Derived;
3274742e46bSJacob Faibussowitsch 
3284742e46bSJacob Faibussowitsch   template <typename F>
3294742e46bSJacob Faibussowitsch   static PetscErrorCode ResizeFactLwork(Mat_SeqDenseCUPM *mcu, cupmStream_t stream, F &&cupmSolverComputeFactLwork) noexcept
3304742e46bSJacob Faibussowitsch   {
3314742e46bSJacob Faibussowitsch     cupmBlasInt_t lwork;
3324742e46bSJacob Faibussowitsch 
3334742e46bSJacob Faibussowitsch     PetscFunctionBegin;
3344742e46bSJacob Faibussowitsch     PetscCallCUPMSOLVER(cupmSolverComputeFactLwork(&lwork));
3354742e46bSJacob Faibussowitsch     if (lwork > mcu->d_fact_lwork) {
3364742e46bSJacob Faibussowitsch       mcu->d_fact_lwork = lwork;
3374742e46bSJacob Faibussowitsch       PetscCallCUPM(cupmFreeAsync(mcu->d_fact_work, stream));
3384742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_work, lwork, stream));
3394742e46bSJacob Faibussowitsch     }
3404742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
3414742e46bSJacob Faibussowitsch   }
3424742e46bSJacob Faibussowitsch 
3434742e46bSJacob Faibussowitsch   static PetscErrorCode FactorPrepare(Mat A, cupmStream_t stream) noexcept
3444742e46bSJacob Faibussowitsch   {
3454742e46bSJacob Faibussowitsch     const auto mcu = MatCUPMCast(A);
3464742e46bSJacob Faibussowitsch 
3474742e46bSJacob Faibussowitsch     PetscFunctionBegin;
3484742e46bSJacob Faibussowitsch     PetscCall(PetscInfo(A, "%s factor %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", derived_type::NAME(), A->rmap->n, A->cmap->n));
3494742e46bSJacob Faibussowitsch     A->factortype             = derived_type::MATFACTORTYPE();
3504742e46bSJacob Faibussowitsch     A->ops->solve             = MatSolve_Factored_Dispatch_<derived_type, false>;
3514742e46bSJacob Faibussowitsch     A->ops->solvetranspose    = MatSolve_Factored_Dispatch_<derived_type, true>;
3524742e46bSJacob Faibussowitsch     A->ops->matsolve          = MatMatSolve_Factored_Dispatch_<derived_type, false>;
3534742e46bSJacob Faibussowitsch     A->ops->matsolvetranspose = MatMatSolve_Factored_Dispatch_<derived_type, true>;
3544742e46bSJacob Faibussowitsch 
3554742e46bSJacob Faibussowitsch     PetscCall(PetscStrFreeAllocpy(MATSOLVERCUPM(), &A->solvertype));
3564742e46bSJacob Faibussowitsch     if (!mcu->d_fact_info) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_info, 1, stream));
3574742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
3584742e46bSJacob Faibussowitsch   }
3594742e46bSJacob Faibussowitsch };
3604742e46bSJacob Faibussowitsch 
3614742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
3624742e46bSJacob Faibussowitsch struct MatDense_Seq_CUPM<T>::SolveLU : SolveCommon<SolveLU> {
3634742e46bSJacob Faibussowitsch   using base_type = SolveCommon<SolveLU>;
3644742e46bSJacob Faibussowitsch 
3654742e46bSJacob Faibussowitsch   static constexpr const char   *NAME() noexcept { return "LU"; }
3664742e46bSJacob Faibussowitsch   static constexpr MatFactorType MATFACTORTYPE() noexcept { return MAT_FACTOR_LU; }
3674742e46bSJacob Faibussowitsch 
3684742e46bSJacob Faibussowitsch   static PetscErrorCode Factor(Mat A, IS, IS, const MatFactorInfo *) noexcept
3694742e46bSJacob Faibussowitsch   {
3704742e46bSJacob Faibussowitsch     const auto         m = static_cast<cupmBlasInt_t>(A->rmap->n);
3714742e46bSJacob Faibussowitsch     const auto         n = static_cast<cupmBlasInt_t>(A->cmap->n);
3724742e46bSJacob Faibussowitsch     cupmStream_t       stream;
3734742e46bSJacob Faibussowitsch     cupmSolverHandle_t handle;
3744742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
3754742e46bSJacob Faibussowitsch 
3764742e46bSJacob Faibussowitsch     PetscFunctionBegin;
3774742e46bSJacob Faibussowitsch     if (!m || !n) PetscFunctionReturn(PETSC_SUCCESS);
3784742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx, &handle, &stream));
3794742e46bSJacob Faibussowitsch     PetscCall(base_type::FactorPrepare(A, stream));
3804742e46bSJacob Faibussowitsch     {
3814742e46bSJacob Faibussowitsch       const auto mcu = MatCUPMCast(A);
3824742e46bSJacob Faibussowitsch       const auto da  = DeviceArrayReadWrite(dctx, A);
3834742e46bSJacob Faibussowitsch       const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);
3844742e46bSJacob Faibussowitsch 
3854742e46bSJacob Faibussowitsch       // clang-format off
3864742e46bSJacob Faibussowitsch       PetscCall(
3874742e46bSJacob Faibussowitsch         base_type::ResizeFactLwork(
3884742e46bSJacob Faibussowitsch           mcu, stream,
3894742e46bSJacob Faibussowitsch           [&](cupmBlasInt_t *fact_lwork)
3904742e46bSJacob Faibussowitsch           {
3914742e46bSJacob Faibussowitsch             return cupmSolverXgetrf_bufferSize(handle, m, n, da.cupmdata(), lda, fact_lwork);
3924742e46bSJacob Faibussowitsch           }
3934742e46bSJacob Faibussowitsch         )
3944742e46bSJacob Faibussowitsch       );
3954742e46bSJacob Faibussowitsch       // clang-format on
3964742e46bSJacob Faibussowitsch       if (!mcu->d_fact_ipiv) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_ipiv, n, stream));
3974742e46bSJacob Faibussowitsch 
3984742e46bSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeBegin());
3994742e46bSJacob Faibussowitsch       PetscCallCUPMSOLVER(cupmSolverXgetrf(handle, m, n, da.cupmdata(), lda, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_ipiv, mcu->d_fact_info));
4004742e46bSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeEnd());
4014742e46bSJacob Faibussowitsch       PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream));
4024742e46bSJacob Faibussowitsch     }
4034742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuFlops(2.0 * n * n * m / 3.0));
4044742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
4054742e46bSJacob Faibussowitsch   }
4064742e46bSJacob Faibussowitsch 
4074742e46bSJacob Faibussowitsch   template <bool transpose>
4084742e46bSJacob Faibussowitsch   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
4094742e46bSJacob Faibussowitsch   {
4104742e46bSJacob Faibussowitsch     const auto         mcu       = MatCUPMCast(A);
4114742e46bSJacob Faibussowitsch     const auto         fact_info = mcu->d_fact_info;
4124742e46bSJacob Faibussowitsch     const auto         fact_ipiv = mcu->d_fact_ipiv;
4134742e46bSJacob Faibussowitsch     cupmSolverHandle_t handle;
4144742e46bSJacob Faibussowitsch 
4154742e46bSJacob Faibussowitsch     PetscFunctionBegin;
4164742e46bSJacob Faibussowitsch     PetscCall(GetHandlesFrom_(dctx, &handle));
4174742e46bSJacob Faibussowitsch     PetscCall(PetscInfo(A, "%s solve %d x %d on backend\n", NAME(), m, k));
4184742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeBegin());
4194742e46bSJacob Faibussowitsch     {
4204742e46bSJacob Faibussowitsch       constexpr auto op  = transpose ? CUPMSOLVER_OP_T : CUPMSOLVER_OP_N;
4214742e46bSJacob Faibussowitsch       const auto     da  = DeviceArrayRead(dctx, A);
4224742e46bSJacob Faibussowitsch       const auto     lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);
4234742e46bSJacob Faibussowitsch 
4244742e46bSJacob Faibussowitsch       // clang-format off
4254742e46bSJacob Faibussowitsch       PetscCall(
4264742e46bSJacob Faibussowitsch         base_type::ResizeFactLwork(
4274742e46bSJacob Faibussowitsch           mcu, stream,
4284742e46bSJacob Faibussowitsch           [&](cupmBlasInt_t *lwork)
4294742e46bSJacob Faibussowitsch           {
4304742e46bSJacob Faibussowitsch             return cupmSolverXgetrs_bufferSize(
4314742e46bSJacob Faibussowitsch               handle, op, m, nrhs, da.cupmdata(), lda, fact_ipiv, x, ldx, lwork
4324742e46bSJacob Faibussowitsch             );
4334742e46bSJacob Faibussowitsch           }
4344742e46bSJacob Faibussowitsch         )
4354742e46bSJacob Faibussowitsch       );
4364742e46bSJacob Faibussowitsch       // clang-format on
4374742e46bSJacob Faibussowitsch       PetscCallCUPMSOLVER(cupmSolverXgetrs(handle, op, m, nrhs, da.cupmdata(), lda, fact_ipiv, x, ldx, mcu->d_fact_work, mcu->d_fact_lwork, fact_info));
4384742e46bSJacob Faibussowitsch       PetscCall(CheckCUPMSolverInfo_(fact_info, stream));
4394742e46bSJacob Faibussowitsch     }
4404742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeEnd());
4414742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuFlops(nrhs * (2.0 * m * m - m)));
4424742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
4434742e46bSJacob Faibussowitsch   }
4444742e46bSJacob Faibussowitsch };
4454742e46bSJacob Faibussowitsch 
4464742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
4474742e46bSJacob Faibussowitsch struct MatDense_Seq_CUPM<T>::SolveCholesky : SolveCommon<SolveCholesky> {
4484742e46bSJacob Faibussowitsch   using base_type = SolveCommon<SolveCholesky>;
4494742e46bSJacob Faibussowitsch 
4504742e46bSJacob Faibussowitsch   static constexpr const char   *NAME() noexcept { return "Cholesky"; }
4514742e46bSJacob Faibussowitsch   static constexpr MatFactorType MATFACTORTYPE() noexcept { return MAT_FACTOR_CHOLESKY; }
4524742e46bSJacob Faibussowitsch 
4534742e46bSJacob Faibussowitsch   static PetscErrorCode Factor(Mat A, IS, const MatFactorInfo *) noexcept
4544742e46bSJacob Faibussowitsch   {
4554742e46bSJacob Faibussowitsch     const auto         n = static_cast<cupmBlasInt_t>(A->rmap->n);
4564742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
4574742e46bSJacob Faibussowitsch     cupmSolverHandle_t handle;
4584742e46bSJacob Faibussowitsch     cupmStream_t       stream;
4594742e46bSJacob Faibussowitsch 
4604742e46bSJacob Faibussowitsch     PetscFunctionBegin;
4614742e46bSJacob Faibussowitsch     if (!n || !A->cmap->n) PetscFunctionReturn(PETSC_SUCCESS);
4624742e46bSJacob Faibussowitsch     PetscCheck(A->spd == PETSC_BOOL3_TRUE, PETSC_COMM_SELF, PETSC_ERR_SUP, "%ssytrs unavailable. Use MAT_FACTOR_LU", cupmSolverName());
4634742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx, &handle, &stream));
4644742e46bSJacob Faibussowitsch     PetscCall(base_type::FactorPrepare(A, stream));
4654742e46bSJacob Faibussowitsch     {
4664742e46bSJacob Faibussowitsch       const auto mcu = MatCUPMCast(A);
4674742e46bSJacob Faibussowitsch       const auto da  = DeviceArrayReadWrite(dctx, A);
4684742e46bSJacob Faibussowitsch       const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);
4694742e46bSJacob Faibussowitsch 
4704742e46bSJacob Faibussowitsch       // clang-format off
4714742e46bSJacob Faibussowitsch       PetscCall(
4724742e46bSJacob Faibussowitsch         base_type::ResizeFactLwork(
4734742e46bSJacob Faibussowitsch           mcu, stream,
4744742e46bSJacob Faibussowitsch           [&](cupmBlasInt_t *fact_lwork)
4754742e46bSJacob Faibussowitsch           {
4764742e46bSJacob Faibussowitsch             return cupmSolverXpotrf_bufferSize(
4774742e46bSJacob Faibussowitsch               handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, fact_lwork
4784742e46bSJacob Faibussowitsch             );
4794742e46bSJacob Faibussowitsch           }
4804742e46bSJacob Faibussowitsch         )
4814742e46bSJacob Faibussowitsch       );
4824742e46bSJacob Faibussowitsch       // clang-format on
4834742e46bSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeBegin());
4844742e46bSJacob Faibussowitsch       PetscCallCUPMSOLVER(cupmSolverXpotrf(handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_info));
4854742e46bSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeEnd());
4864742e46bSJacob Faibussowitsch       PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream));
4874742e46bSJacob Faibussowitsch     }
4884742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuFlops(1.0 * n * n * n / 3.0));
4894742e46bSJacob Faibussowitsch 
4904742e46bSJacob Faibussowitsch #if 0
4914742e46bSJacob Faibussowitsch     // At the time of writing this interface (cuda 10.0), cusolverDn does not implement *sytrs
4924742e46bSJacob Faibussowitsch     // and *hetr* routines. The code below should work, and it can be activated when *sytrs
4934742e46bSJacob Faibussowitsch     // routines will be available
4944742e46bSJacob Faibussowitsch     if (!mcu->d_fact_ipiv) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_ipiv, n, stream));
4954742e46bSJacob Faibussowitsch     if (!mcu->d_fact_lwork) {
4964742e46bSJacob Faibussowitsch       PetscCallCUPMSOLVER(cupmSolverDnXsytrf_bufferSize(handle, n, da.cupmdata(), lda, &mcu->d_fact_lwork));
4974742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_work, mcu->d_fact_lwork, stream));
4984742e46bSJacob Faibussowitsch     }
4994742e46bSJacob Faibussowitsch     if (mcu->d_fact_info) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_info, 1, stream));
5004742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeBegin());
5014742e46bSJacob Faibussowitsch     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));
5024742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeEnd());
5034742e46bSJacob Faibussowitsch #endif
5044742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
5054742e46bSJacob Faibussowitsch   }
5064742e46bSJacob Faibussowitsch 
5074742e46bSJacob Faibussowitsch   template <bool transpose>
5084742e46bSJacob Faibussowitsch   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
5094742e46bSJacob Faibussowitsch   {
5104742e46bSJacob Faibussowitsch     const auto         mcu       = MatCUPMCast(A);
5114742e46bSJacob Faibussowitsch     const auto         fact_info = mcu->d_fact_info;
5124742e46bSJacob Faibussowitsch     cupmSolverHandle_t handle;
5134742e46bSJacob Faibussowitsch 
5144742e46bSJacob Faibussowitsch     PetscFunctionBegin;
5154742e46bSJacob Faibussowitsch     PetscAssert(!mcu->d_fact_ipiv, PETSC_COMM_SELF, PETSC_ERR_LIB, "%ssytrs not implemented", cupmSolverName());
5164742e46bSJacob Faibussowitsch     PetscCall(GetHandlesFrom_(dctx, &handle));
5174742e46bSJacob Faibussowitsch     PetscCall(PetscInfo(A, "%s solve %d x %d on backend\n", NAME(), m, k));
5184742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeBegin());
5194742e46bSJacob Faibussowitsch     {
5204742e46bSJacob Faibussowitsch       const auto da  = DeviceArrayRead(dctx, A);
5214742e46bSJacob Faibussowitsch       const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);
5224742e46bSJacob Faibussowitsch 
5234742e46bSJacob Faibussowitsch       // clang-format off
5244742e46bSJacob Faibussowitsch       PetscCall(
5254742e46bSJacob Faibussowitsch         base_type::ResizeFactLwork(
5264742e46bSJacob Faibussowitsch           mcu, stream,
5274742e46bSJacob Faibussowitsch           [&](cupmBlasInt_t *lwork)
5284742e46bSJacob Faibussowitsch           {
5294742e46bSJacob Faibussowitsch             return cupmSolverXpotrs_bufferSize(
5304742e46bSJacob Faibussowitsch               handle, CUPMSOLVER_FILL_MODE_LOWER, m, nrhs, da.cupmdata(), lda, x, ldx, lwork
5314742e46bSJacob Faibussowitsch             );
5324742e46bSJacob Faibussowitsch           }
5334742e46bSJacob Faibussowitsch         )
5344742e46bSJacob Faibussowitsch       );
5354742e46bSJacob Faibussowitsch       // clang-format on
5364742e46bSJacob Faibussowitsch       PetscCallCUPMSOLVER(cupmSolverXpotrs(handle, CUPMSOLVER_FILL_MODE_LOWER, m, nrhs, da.cupmdata(), lda, x, ldx, mcu->d_fact_work, mcu->d_fact_lwork, fact_info));
5374742e46bSJacob Faibussowitsch       PetscCall(CheckCUPMSolverInfo_(fact_info, stream));
5384742e46bSJacob Faibussowitsch     }
5394742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeEnd());
5404742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuFlops(nrhs * (2.0 * m * m - m)));
5414742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
5424742e46bSJacob Faibussowitsch   }
5434742e46bSJacob Faibussowitsch };
5444742e46bSJacob Faibussowitsch 
5454742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
5464742e46bSJacob Faibussowitsch struct MatDense_Seq_CUPM<T>::SolveQR : SolveCommon<SolveQR> {
5474742e46bSJacob Faibussowitsch   using base_type = SolveCommon<SolveQR>;
5484742e46bSJacob Faibussowitsch 
5494742e46bSJacob Faibussowitsch   static constexpr const char   *NAME() noexcept { return "QR"; }
5504742e46bSJacob Faibussowitsch   static constexpr MatFactorType MATFACTORTYPE() noexcept { return MAT_FACTOR_QR; }
5514742e46bSJacob Faibussowitsch 
5524742e46bSJacob Faibussowitsch   static PetscErrorCode Factor(Mat A, IS, const MatFactorInfo *) noexcept
5534742e46bSJacob Faibussowitsch   {
5544742e46bSJacob Faibussowitsch     const auto         m     = static_cast<cupmBlasInt_t>(A->rmap->n);
5554742e46bSJacob Faibussowitsch     const auto         n     = static_cast<cupmBlasInt_t>(A->cmap->n);
5564742e46bSJacob Faibussowitsch     const auto         min   = std::min(m, n);
5574742e46bSJacob Faibussowitsch     const auto         mimpl = MatIMPLCast(A);
5584742e46bSJacob Faibussowitsch     cupmStream_t       stream;
5594742e46bSJacob Faibussowitsch     cupmSolverHandle_t handle;
5604742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
5614742e46bSJacob Faibussowitsch 
5624742e46bSJacob Faibussowitsch     PetscFunctionBegin;
5634742e46bSJacob Faibussowitsch     if (!m || !n) PetscFunctionReturn(PETSC_SUCCESS);
5644742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx, &handle, &stream));
5654742e46bSJacob Faibussowitsch     PetscCall(base_type::FactorPrepare(A, stream));
5664742e46bSJacob Faibussowitsch     mimpl->rank = min;
5674742e46bSJacob Faibussowitsch     {
5684742e46bSJacob Faibussowitsch       const auto mcu = MatCUPMCast(A);
5694742e46bSJacob Faibussowitsch       const auto da  = DeviceArrayReadWrite(dctx, A);
5704742e46bSJacob Faibussowitsch       const auto lda = static_cast<cupmBlasInt_t>(mimpl->lda);
5714742e46bSJacob Faibussowitsch 
5724742e46bSJacob Faibussowitsch       if (!mcu->workvec) PetscCall(vec::cupm::VecCreateSeqCUPMAsync<T>(PetscObjectComm(PetscObjectCast(A)), m, &mcu->workvec));
5734742e46bSJacob Faibussowitsch       if (!mcu->d_fact_tau) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_tau, min, stream));
5744742e46bSJacob Faibussowitsch       // clang-format off
5754742e46bSJacob Faibussowitsch       PetscCall(
5764742e46bSJacob Faibussowitsch         base_type::ResizeFactLwork(
5774742e46bSJacob Faibussowitsch           mcu, stream,
5784742e46bSJacob Faibussowitsch           [&](cupmBlasInt_t *fact_lwork)
5794742e46bSJacob Faibussowitsch           {
5804742e46bSJacob Faibussowitsch             return cupmSolverXgeqrf_bufferSize(handle, m, n, da.cupmdata(), lda, fact_lwork);
5814742e46bSJacob Faibussowitsch           }
5824742e46bSJacob Faibussowitsch         )
5834742e46bSJacob Faibussowitsch       );
5844742e46bSJacob Faibussowitsch       // clang-format on
5854742e46bSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeBegin());
5864742e46bSJacob Faibussowitsch       PetscCallCUPMSOLVER(cupmSolverXgeqrf(handle, m, n, da.cupmdata(), lda, mcu->d_fact_tau, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_info));
5874742e46bSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeEnd());
5884742e46bSJacob Faibussowitsch       PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream));
5894742e46bSJacob Faibussowitsch     }
5904742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuFlops(2.0 * min * min * (std::max(m, n) - min / 3.0)));
5914742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
5924742e46bSJacob Faibussowitsch   }
5934742e46bSJacob Faibussowitsch 
5944742e46bSJacob Faibussowitsch   template <bool transpose>
5954742e46bSJacob Faibussowitsch   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
5964742e46bSJacob Faibussowitsch   {
5974742e46bSJacob Faibussowitsch     const auto         mimpl      = MatIMPLCast(A);
5984742e46bSJacob Faibussowitsch     const auto         rank       = static_cast<cupmBlasInt_t>(mimpl->rank);
5994742e46bSJacob Faibussowitsch     const auto         mcu        = MatCUPMCast(A);
6004742e46bSJacob Faibussowitsch     const auto         fact_info  = mcu->d_fact_info;
6014742e46bSJacob Faibussowitsch     const auto         fact_tau   = mcu->d_fact_tau;
6024742e46bSJacob Faibussowitsch     const auto         fact_work  = mcu->d_fact_work;
6034742e46bSJacob Faibussowitsch     const auto         fact_lwork = mcu->d_fact_lwork;
6044742e46bSJacob Faibussowitsch     cupmSolverHandle_t solver_handle;
6054742e46bSJacob Faibussowitsch     cupmBlasHandle_t   blas_handle;
6064742e46bSJacob Faibussowitsch 
6074742e46bSJacob Faibussowitsch     PetscFunctionBegin;
6084742e46bSJacob Faibussowitsch     PetscCall(GetHandlesFrom_(dctx, &blas_handle, &solver_handle));
6094742e46bSJacob Faibussowitsch     PetscCall(PetscInfo(A, "%s solve %d x %d on backend\n", NAME(), m, k));
6104742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeBegin());
6114742e46bSJacob Faibussowitsch     {
6124742e46bSJacob Faibussowitsch       const auto da  = DeviceArrayRead(dctx, A);
6134742e46bSJacob Faibussowitsch       const auto one = cupmScalarCast(1.0);
6144742e46bSJacob Faibussowitsch       const auto lda = static_cast<cupmBlasInt_t>(mimpl->lda);
6154742e46bSJacob Faibussowitsch 
6164742e46bSJacob Faibussowitsch       if (transpose) {
6174742e46bSJacob Faibussowitsch         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));
6184742e46bSJacob Faibussowitsch         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));
6194742e46bSJacob Faibussowitsch         PetscCall(CheckCUPMSolverInfo_(fact_info, stream));
6204742e46bSJacob Faibussowitsch       } else {
6214742e46bSJacob Faibussowitsch         constexpr auto op = PetscDefined(USE_COMPLEX) ? CUPMSOLVER_OP_C : CUPMSOLVER_OP_T;
6224742e46bSJacob Faibussowitsch 
6234742e46bSJacob Faibussowitsch         PetscCallCUPMSOLVER(cupmSolverXormqr(solver_handle, CUPMSOLVER_SIDE_LEFT, op, m, nrhs, rank, da.cupmdata(), lda, fact_tau, x, ldx, fact_work, fact_lwork, fact_info));
6244742e46bSJacob Faibussowitsch         PetscCall(CheckCUPMSolverInfo_(fact_info, stream));
6254742e46bSJacob Faibussowitsch         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));
6264742e46bSJacob Faibussowitsch       }
6274742e46bSJacob Faibussowitsch     }
6284742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeEnd());
6294742e46bSJacob Faibussowitsch     PetscCall(PetscLogFlops(nrhs * (4.0 * m * rank - (rank * rank))));
6304742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
6314742e46bSJacob Faibussowitsch   }
6324742e46bSJacob Faibussowitsch };
6334742e46bSJacob Faibussowitsch 
6344742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
6354742e46bSJacob Faibussowitsch template <typename Solver, bool transpose>
6364742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::MatSolve_Factored_Dispatch_(Mat A, Vec x, Vec y) noexcept
6374742e46bSJacob Faibussowitsch {
6384742e46bSJacob Faibussowitsch   using namespace vec::cupm;
6394742e46bSJacob Faibussowitsch   const auto         pobj_A  = PetscObjectCast(A);
6404742e46bSJacob Faibussowitsch   const auto         m       = static_cast<cupmBlasInt_t>(A->rmap->n);
6414742e46bSJacob Faibussowitsch   const auto         k       = static_cast<cupmBlasInt_t>(A->cmap->n);
6424742e46bSJacob Faibussowitsch   auto              &workvec = MatCUPMCast(A)->workvec;
6434742e46bSJacob Faibussowitsch   PetscScalar       *y_array = nullptr;
6444742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
6454742e46bSJacob Faibussowitsch   PetscBool          xiscupm, yiscupm, aiscupm;
6464742e46bSJacob Faibussowitsch   bool               use_y_array_directly;
6474742e46bSJacob Faibussowitsch   cupmStream_t       stream;
6484742e46bSJacob Faibussowitsch 
6494742e46bSJacob Faibussowitsch   PetscFunctionBegin;
6504742e46bSJacob Faibussowitsch   PetscCheck(A->factortype != MAT_FACTOR_NONE, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix must be factored to solve");
6514742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(PetscObjectCast(x), VecSeq_CUPM::VECSEQCUPM(), &xiscupm));
6524742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(PetscObjectCast(y), VecSeq_CUPM::VECSEQCUPM(), &yiscupm));
6534742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(pobj_A, MATSEQDENSECUPM(), &aiscupm));
6544742e46bSJacob Faibussowitsch   PetscAssert(aiscupm, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Matrix A is somehow not CUPM?????????????????????????????");
6554742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx, &stream));
6564742e46bSJacob Faibussowitsch   use_y_array_directly = yiscupm && (k >= m);
6574742e46bSJacob Faibussowitsch   {
6584742e46bSJacob Faibussowitsch     const PetscScalar *x_array;
6594742e46bSJacob Faibussowitsch     const auto         xisdevice = xiscupm && PetscOffloadDevice(x->offloadmask);
6604742e46bSJacob Faibussowitsch     const auto         copy_mode = xisdevice ? cupmMemcpyDeviceToDevice : cupmMemcpyHostToDevice;
6614742e46bSJacob Faibussowitsch 
6624742e46bSJacob Faibussowitsch     if (!use_y_array_directly && !workvec) PetscCall(VecCreateSeqCUPMAsync<T>(PetscObjectComm(pobj_A), m, &workvec));
6634742e46bSJacob Faibussowitsch     // The logic here is to try to minimize the amount of memory copying:
6644742e46bSJacob Faibussowitsch     //
6654742e46bSJacob Faibussowitsch     // If we call VecCUPMGetArrayRead(X, &x) every time xiscupm and the data is not offloaded
6664742e46bSJacob Faibussowitsch     // to the GPU yet, then the data is copied to the GPU. But we are only trying to get the
6674742e46bSJacob Faibussowitsch     // data in order to copy it into the y array. So the array x will be wherever the data
6684742e46bSJacob Faibussowitsch     // already is so that only one memcpy is performed
6694742e46bSJacob Faibussowitsch     if (xisdevice) {
6704742e46bSJacob Faibussowitsch       PetscCall(VecCUPMGetArrayReadAsync<T>(x, &x_array, dctx));
6714742e46bSJacob Faibussowitsch     } else {
6724742e46bSJacob Faibussowitsch       PetscCall(VecGetArrayRead(x, &x_array));
6734742e46bSJacob Faibussowitsch     }
6744742e46bSJacob Faibussowitsch     PetscCall(VecCUPMGetArrayWriteAsync<T>(use_y_array_directly ? y : workvec, &y_array, dctx));
6754742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMMemcpyAsync(y_array, x_array, m, copy_mode, stream));
6764742e46bSJacob Faibussowitsch     if (xisdevice) {
6774742e46bSJacob Faibussowitsch       PetscCall(VecCUPMRestoreArrayReadAsync<T>(x, &x_array, dctx));
6784742e46bSJacob Faibussowitsch     } else {
6794742e46bSJacob Faibussowitsch       PetscCall(VecRestoreArrayRead(x, &x_array));
6804742e46bSJacob Faibussowitsch     }
6814742e46bSJacob Faibussowitsch   }
6824742e46bSJacob Faibussowitsch 
6834742e46bSJacob Faibussowitsch   if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
6844742e46bSJacob Faibussowitsch   PetscCall(Solver{}.template Solve<transpose>(A, cupmScalarPtrCast(y_array), m, m, 1, k, dctx, stream));
6854742e46bSJacob Faibussowitsch   if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A));
6864742e46bSJacob Faibussowitsch 
6874742e46bSJacob Faibussowitsch   if (use_y_array_directly) {
6884742e46bSJacob Faibussowitsch     PetscCall(VecCUPMRestoreArrayWriteAsync<T>(y, &y_array, dctx));
6894742e46bSJacob Faibussowitsch   } else {
6904742e46bSJacob Faibussowitsch     const auto   copy_mode = yiscupm ? cupmMemcpyDeviceToDevice : cupmMemcpyDeviceToHost;
6914742e46bSJacob Faibussowitsch     PetscScalar *yv;
6924742e46bSJacob Faibussowitsch 
6934742e46bSJacob Faibussowitsch     // The logic here is that the data is not yet in either y's GPU array or its CPU array.
6944742e46bSJacob Faibussowitsch     // There is nothing in the interface to say where the user would like it to end up. So we
6954742e46bSJacob Faibussowitsch     // choose the GPU, because it is the faster option
6964742e46bSJacob Faibussowitsch     if (yiscupm) {
6974742e46bSJacob Faibussowitsch       PetscCall(VecCUPMGetArrayWriteAsync<T>(y, &yv, dctx));
6984742e46bSJacob Faibussowitsch     } else {
6994742e46bSJacob Faibussowitsch       PetscCall(VecGetArray(y, &yv));
7004742e46bSJacob Faibussowitsch     }
7014742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMMemcpyAsync(yv, y_array, k, copy_mode, stream));
7024742e46bSJacob Faibussowitsch     if (yiscupm) {
7034742e46bSJacob Faibussowitsch       PetscCall(VecCUPMRestoreArrayWriteAsync<T>(y, &yv, dctx));
7044742e46bSJacob Faibussowitsch     } else {
7054742e46bSJacob Faibussowitsch       PetscCall(VecRestoreArray(y, &yv));
7064742e46bSJacob Faibussowitsch     }
7074742e46bSJacob Faibussowitsch     PetscCall(VecCUPMRestoreArrayWriteAsync<T>(workvec, &y_array));
7084742e46bSJacob Faibussowitsch   }
7094742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
7104742e46bSJacob Faibussowitsch }
7114742e46bSJacob Faibussowitsch 
7124742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
7134742e46bSJacob Faibussowitsch template <typename Solver, bool transpose>
7144742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMatSolve_Factored_Dispatch_(Mat A, Mat B, Mat X) noexcept
7154742e46bSJacob Faibussowitsch {
7164742e46bSJacob Faibussowitsch   const auto         m = static_cast<cupmBlasInt_t>(A->rmap->n);
7174742e46bSJacob Faibussowitsch   const auto         k = static_cast<cupmBlasInt_t>(A->cmap->n);
7184742e46bSJacob Faibussowitsch   cupmBlasInt_t      nrhs, ldb, ldx, ldy;
7194742e46bSJacob Faibussowitsch   PetscScalar       *y;
7204742e46bSJacob Faibussowitsch   PetscBool          biscupm, xiscupm, aiscupm;
7214742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
7224742e46bSJacob Faibussowitsch   cupmStream_t       stream;
7234742e46bSJacob Faibussowitsch 
7244742e46bSJacob Faibussowitsch   PetscFunctionBegin;
7254742e46bSJacob Faibussowitsch   PetscCheck(A->factortype != MAT_FACTOR_NONE, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix must be factored to solve");
7264742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(PetscObjectCast(B), MATSEQDENSECUPM(), &biscupm));
7274742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(PetscObjectCast(X), MATSEQDENSECUPM(), &xiscupm));
7284742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(PetscObjectCast(A), MATSEQDENSECUPM(), &aiscupm));
7294742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx, &stream));
7304742e46bSJacob Faibussowitsch   {
7314742e46bSJacob Faibussowitsch     PetscInt n;
7324742e46bSJacob Faibussowitsch 
7334742e46bSJacob Faibussowitsch     PetscCall(MatGetSize(B, nullptr, &n));
7344742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMBlasIntCast(n, &nrhs));
7354742e46bSJacob Faibussowitsch     PetscCall(MatDenseGetLDA(B, &n));
7364742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMBlasIntCast(n, &ldb));
7374742e46bSJacob Faibussowitsch     PetscCall(MatDenseGetLDA(X, &n));
7384742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMBlasIntCast(n, &ldx));
7394742e46bSJacob Faibussowitsch   }
7404742e46bSJacob Faibussowitsch   {
7414742e46bSJacob Faibussowitsch     // The logic here is to try to minimize the amount of memory copying:
7424742e46bSJacob Faibussowitsch     //
7434742e46bSJacob Faibussowitsch     // If we call MatDenseCUPMGetArrayRead(B, &b) every time biscupm and the data is not
7444742e46bSJacob Faibussowitsch     // offloaded to the GPU yet, then the data is copied to the GPU. But we are only trying to
7454742e46bSJacob Faibussowitsch     // get the data in order to copy it into the y array. So the array b will be wherever the
7464742e46bSJacob Faibussowitsch     // data already is so that only one memcpy is performed
7474742e46bSJacob Faibussowitsch     const auto         bisdevice = biscupm && PetscOffloadDevice(B->offloadmask);
7484742e46bSJacob Faibussowitsch     const auto         copy_mode = bisdevice ? cupmMemcpyDeviceToDevice : cupmMemcpyHostToDevice;
7494742e46bSJacob Faibussowitsch     const PetscScalar *b;
7504742e46bSJacob Faibussowitsch 
7514742e46bSJacob Faibussowitsch     if (bisdevice) {
7524742e46bSJacob Faibussowitsch       b = DeviceArrayRead(dctx, B);
7534742e46bSJacob Faibussowitsch     } else if (biscupm) {
7544742e46bSJacob Faibussowitsch       b = HostArrayRead(dctx, B);
7554742e46bSJacob Faibussowitsch     } else {
7564742e46bSJacob Faibussowitsch       PetscCall(MatDenseGetArrayRead(B, &b));
7574742e46bSJacob Faibussowitsch     }
7584742e46bSJacob Faibussowitsch 
7594742e46bSJacob Faibussowitsch     if (ldx < m || !xiscupm) {
7604742e46bSJacob Faibussowitsch       // X's array cannot serve as the array (too small or not on device), B's array cannot
7614742e46bSJacob Faibussowitsch       // serve as the array (const), so allocate a new array
7624742e46bSJacob Faibussowitsch       ldy = m;
7634742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMallocAsync(&y, nrhs * m));
7644742e46bSJacob Faibussowitsch     } else {
7654742e46bSJacob Faibussowitsch       // X's array should serve as the array
7664742e46bSJacob Faibussowitsch       ldy = ldx;
7674742e46bSJacob Faibussowitsch       y   = DeviceArrayWrite(dctx, X);
7684742e46bSJacob Faibussowitsch     }
7694742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMMemcpy2DAsync(y, ldy, b, ldb, m, nrhs, copy_mode, stream));
7704742e46bSJacob Faibussowitsch     if (!bisdevice && !biscupm) PetscCall(MatDenseRestoreArrayRead(B, &b));
7714742e46bSJacob Faibussowitsch   }
7724742e46bSJacob Faibussowitsch 
7734742e46bSJacob Faibussowitsch   // convert to CUPM twice??????????????????????????????????
7744742e46bSJacob Faibussowitsch   // but A should already be CUPM??????????????????????????????????????
7754742e46bSJacob Faibussowitsch   if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
7764742e46bSJacob Faibussowitsch   PetscCall(Solver{}.template Solve<transpose>(A, cupmScalarPtrCast(y), ldy, m, nrhs, k, dctx, stream));
7774742e46bSJacob Faibussowitsch   if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
7784742e46bSJacob Faibussowitsch 
7794742e46bSJacob Faibussowitsch   if (ldx < m || !xiscupm) {
7804742e46bSJacob Faibussowitsch     const auto   copy_mode = xiscupm ? cupmMemcpyDeviceToDevice : cupmMemcpyDeviceToHost;
7814742e46bSJacob Faibussowitsch     PetscScalar *x;
7824742e46bSJacob Faibussowitsch 
7834742e46bSJacob Faibussowitsch     // The logic here is that the data is not yet in either X's GPU array or its CPU
7844742e46bSJacob Faibussowitsch     // array. There is nothing in the interface to say where the user would like it to end up.
7854742e46bSJacob Faibussowitsch     // So we choose the GPU, because it is the faster option
7864742e46bSJacob Faibussowitsch     if (xiscupm) {
7874742e46bSJacob Faibussowitsch       x = DeviceArrayWrite(dctx, X);
7884742e46bSJacob Faibussowitsch     } else {
7894742e46bSJacob Faibussowitsch       PetscCall(MatDenseGetArray(X, &x));
7904742e46bSJacob Faibussowitsch     }
7914742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMMemcpy2DAsync(x, ldx, y, ldy, k, nrhs, copy_mode, stream));
7924742e46bSJacob Faibussowitsch     if (!xiscupm) PetscCall(MatDenseRestoreArray(X, &x));
7934742e46bSJacob Faibussowitsch     PetscCallCUPM(cupmFreeAsync(y, stream));
7944742e46bSJacob Faibussowitsch   }
7954742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
7964742e46bSJacob Faibussowitsch }
7974742e46bSJacob Faibussowitsch 
7984742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
7994742e46bSJacob Faibussowitsch template <bool transpose>
8004742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMultAdd_Dispatch_(Mat A, Vec xx, Vec yy, Vec zz) noexcept
8014742e46bSJacob Faibussowitsch {
8024742e46bSJacob Faibussowitsch   const auto         m = static_cast<cupmBlasInt_t>(A->rmap->n);
8034742e46bSJacob Faibussowitsch   const auto         n = static_cast<cupmBlasInt_t>(A->cmap->n);
8044742e46bSJacob Faibussowitsch   cupmBlasHandle_t   handle;
8054742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
8064742e46bSJacob Faibussowitsch 
8074742e46bSJacob Faibussowitsch   PetscFunctionBegin;
8084742e46bSJacob Faibussowitsch   if (yy && yy != zz) PetscCall(VecSeq_CUPM::Copy(yy, zz)); // mult add
8094742e46bSJacob Faibussowitsch   if (!m || !n) {
8104742e46bSJacob Faibussowitsch     // mult only
8114742e46bSJacob Faibussowitsch     if (!yy) PetscCall(VecSeq_CUPM::Set(zz, 0.0));
8124742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
8134742e46bSJacob Faibussowitsch   }
8144742e46bSJacob Faibussowitsch   PetscCall(PetscInfo(A, "Matrix-vector product %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " on backend\n", m, n));
8154742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx, &handle));
8164742e46bSJacob Faibussowitsch   {
8174742e46bSJacob Faibussowitsch     constexpr auto op   = transpose ? CUPMBLAS_OP_T : CUPMBLAS_OP_N;
8184742e46bSJacob Faibussowitsch     const auto     one  = cupmScalarCast(1.0);
8194742e46bSJacob Faibussowitsch     const auto     zero = cupmScalarCast(0.0);
8204742e46bSJacob Faibussowitsch     const auto     da   = DeviceArrayRead(dctx, A);
8214742e46bSJacob Faibussowitsch     const auto     dxx  = VecSeq_CUPM::DeviceArrayRead(dctx, xx);
8224742e46bSJacob Faibussowitsch     const auto     dzz  = VecSeq_CUPM::DeviceArrayReadWrite(dctx, zz);
8234742e46bSJacob Faibussowitsch 
8244742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeBegin());
8254742e46bSJacob Faibussowitsch     PetscCallCUPMBLAS(cupmBlasXgemv(handle, op, m, n, &one, da.cupmdata(), static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda), dxx.cupmdata(), 1, (yy ? &one : &zero), dzz.cupmdata(), 1));
8264742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeEnd());
8274742e46bSJacob Faibussowitsch   }
8284742e46bSJacob Faibussowitsch   PetscCall(PetscLogGpuFlops(2.0 * m * n - (yy ? 0 : m)));
8294742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
8304742e46bSJacob Faibussowitsch }
8314742e46bSJacob Faibussowitsch 
8324742e46bSJacob Faibussowitsch // ==========================================================================================
8334742e46bSJacob Faibussowitsch // MatDense_Seq_CUPM - Private API - Conversion Dispatch
8344742e46bSJacob Faibussowitsch // ==========================================================================================
8354742e46bSJacob Faibussowitsch 
8364742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
8374742e46bSJacob Faibussowitsch template <bool to_host>
8384742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_Dispatch_(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept
8394742e46bSJacob Faibussowitsch {
8404742e46bSJacob Faibussowitsch   PetscFunctionBegin;
8414742e46bSJacob Faibussowitsch   if (reuse == MAT_REUSE_MATRIX || reuse == MAT_INITIAL_MATRIX) {
8424742e46bSJacob Faibussowitsch     // TODO these cases should be optimized
8434742e46bSJacob Faibussowitsch     PetscCall(MatConvert_Basic(M, type, reuse, newmat));
8444742e46bSJacob Faibussowitsch   } else {
8454742e46bSJacob Faibussowitsch     const auto B    = *newmat;
8464742e46bSJacob Faibussowitsch     const auto pobj = PetscObjectCast(B);
8474742e46bSJacob Faibussowitsch 
8484742e46bSJacob Faibussowitsch     if (to_host) {
8494742e46bSJacob Faibussowitsch       PetscCall(BindToCPU(B, PETSC_TRUE));
8504742e46bSJacob Faibussowitsch       PetscCall(Reset(B));
8514742e46bSJacob Faibussowitsch     } else {
8524742e46bSJacob Faibussowitsch       PetscCall(PetscDeviceInitialize(PETSC_DEVICE_CUPM()));
8534742e46bSJacob Faibussowitsch     }
8544742e46bSJacob Faibussowitsch 
8554742e46bSJacob Faibussowitsch     PetscCall(PetscStrFreeAllocpy(to_host ? VECSTANDARD : VecSeq_CUPM::VECCUPM(), &B->defaultvectype));
8564742e46bSJacob Faibussowitsch     PetscCall(PetscObjectChangeTypeName(pobj, to_host ? MATSEQDENSE : MATSEQDENSECUPM()));
8574742e46bSJacob Faibussowitsch     // cvec might be the wrong VecType, destroy and rebuild it if necessary
8584742e46bSJacob Faibussowitsch     // REVIEW ME: this is possibly very inefficient
8594742e46bSJacob Faibussowitsch     PetscCall(VecDestroy(&MatIMPLCast(B)->cvec));
8604742e46bSJacob Faibussowitsch 
8614742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatConvert_seqdensecupm_seqdense_C(), nullptr, Convert_SeqDenseCUPM_SeqDense);
8624742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArray_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ_WRITE>);
8634742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArrayRead_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ>);
8644742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArrayWrite_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_WRITE>);
8654742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArray_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ_WRITE>);
8664742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArrayRead_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ>);
8674742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArrayWrite_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_WRITE>);
8684742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMPlaceArray_C(), nullptr, PlaceArray);
8694742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMResetArray_C(), nullptr, ResetArray);
8704742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMReplaceArray_C(), nullptr, ReplaceArray);
8714742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatProductSetFromOptions_seqaij_seqdensecupm_C(), nullptr, MatProductSetFromOptions_SeqAIJ_SeqDense);
8723d9668e3SJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMSetPreallocation_C(), nullptr, SetPreallocation);
8734742e46bSJacob Faibussowitsch 
8744742e46bSJacob Faibussowitsch     if (to_host) {
8754742e46bSJacob Faibussowitsch       B->offloadmask = PETSC_OFFLOAD_CPU;
8764742e46bSJacob Faibussowitsch     } else {
8774742e46bSJacob Faibussowitsch       Mat_SeqDenseCUPM *mcu;
8784742e46bSJacob Faibussowitsch 
8794742e46bSJacob Faibussowitsch       PetscCall(PetscNew(&mcu));
8804742e46bSJacob Faibussowitsch       B->spptr       = mcu;
8814742e46bSJacob Faibussowitsch       B->offloadmask = PETSC_OFFLOAD_UNALLOCATED; // REVIEW ME: why not offload host??
8824742e46bSJacob Faibussowitsch       PetscCall(BindToCPU(B, PETSC_FALSE));
8834742e46bSJacob Faibussowitsch     }
8844742e46bSJacob Faibussowitsch 
8854742e46bSJacob Faibussowitsch     MatSetOp_CUPM(to_host, B, bindtocpu, nullptr, BindToCPU);
8864742e46bSJacob Faibussowitsch     MatSetOp_CUPM(to_host, B, destroy, MatDestroy_SeqDense, Destroy);
8874742e46bSJacob Faibussowitsch   }
8884742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
8894742e46bSJacob Faibussowitsch }
8904742e46bSJacob Faibussowitsch 
8914742e46bSJacob Faibussowitsch // ==========================================================================================
8924742e46bSJacob Faibussowitsch // MatDense_Seq_CUPM - Public API
8934742e46bSJacob Faibussowitsch // ==========================================================================================
8944742e46bSJacob Faibussowitsch 
8954742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
8964742e46bSJacob Faibussowitsch inline constexpr MatType MatDense_Seq_CUPM<T>::MATIMPLCUPM_() noexcept
8974742e46bSJacob Faibussowitsch {
8984742e46bSJacob Faibussowitsch   return MATSEQDENSECUPM();
8994742e46bSJacob Faibussowitsch }
9004742e46bSJacob Faibussowitsch 
9014742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9024742e46bSJacob Faibussowitsch inline constexpr typename MatDense_Seq_CUPM<T>::Mat_SeqDenseCUPM *MatDense_Seq_CUPM<T>::MatCUPMCast(Mat m) noexcept
9034742e46bSJacob Faibussowitsch {
9044742e46bSJacob Faibussowitsch   return static_cast<Mat_SeqDenseCUPM *>(m->spptr);
9054742e46bSJacob Faibussowitsch }
9064742e46bSJacob Faibussowitsch 
9074742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9084742e46bSJacob Faibussowitsch inline constexpr Mat_SeqDense *MatDense_Seq_CUPM<T>::MatIMPLCast_(Mat m) noexcept
9094742e46bSJacob Faibussowitsch {
9104742e46bSJacob Faibussowitsch   return static_cast<Mat_SeqDense *>(m->data);
9114742e46bSJacob Faibussowitsch }
9124742e46bSJacob Faibussowitsch 
9134742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9144742e46bSJacob Faibussowitsch inline constexpr const char *MatDense_Seq_CUPM<T>::MatConvert_seqdensecupm_seqdense_C() noexcept
9154742e46bSJacob Faibussowitsch {
9164742e46bSJacob Faibussowitsch   return T == device::cupm::DeviceType::CUDA ? "MatConvert_seqdensecuda_seqdense_C" : "MatConvert_seqdensehip_seqdense_C";
9174742e46bSJacob Faibussowitsch }
9184742e46bSJacob Faibussowitsch 
9194742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9204742e46bSJacob Faibussowitsch inline constexpr const char *MatDense_Seq_CUPM<T>::MatProductSetFromOptions_seqaij_seqdensecupm_C() noexcept
9214742e46bSJacob Faibussowitsch {
9224742e46bSJacob Faibussowitsch   return T == device::cupm::DeviceType::CUDA ? "MatProductSetFromOptions_seqaij_seqdensecuda_C" : "MatProductSetFromOptions_seqaij_seqdensehip_C";
9234742e46bSJacob Faibussowitsch }
9244742e46bSJacob Faibussowitsch 
9254742e46bSJacob Faibussowitsch // ==========================================================================================
9264742e46bSJacob Faibussowitsch 
9274742e46bSJacob Faibussowitsch // MatCreate_SeqDenseCUPM()
9284742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9294742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Create(Mat A) noexcept
9304742e46bSJacob Faibussowitsch {
9314742e46bSJacob Faibussowitsch   PetscFunctionBegin;
9324742e46bSJacob Faibussowitsch   PetscCall(PetscDeviceInitialize(PETSC_DEVICE_CUPM()));
9334742e46bSJacob Faibussowitsch   PetscCall(MatCreate_SeqDense(A));
9344742e46bSJacob Faibussowitsch   PetscCall(Convert_SeqDense_SeqDenseCUPM(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
9354742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
9364742e46bSJacob Faibussowitsch }
9374742e46bSJacob Faibussowitsch 
9384742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9394742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Destroy(Mat A) noexcept
9404742e46bSJacob Faibussowitsch {
9414742e46bSJacob Faibussowitsch   PetscFunctionBegin;
9424742e46bSJacob Faibussowitsch   // prevent copying back data if we own the data pointer
9434742e46bSJacob Faibussowitsch   if (!MatIMPLCast(A)->user_alloc) A->offloadmask = PETSC_OFFLOAD_CPU;
9444742e46bSJacob Faibussowitsch   PetscCall(Convert_SeqDenseCUPM_SeqDense(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A));
9454742e46bSJacob Faibussowitsch   PetscCall(MatDestroy_SeqDense(A));
9464742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
9474742e46bSJacob Faibussowitsch }
9484742e46bSJacob Faibussowitsch 
9494742e46bSJacob Faibussowitsch // obj->ops->setup()
9504742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9514742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::SetUp(Mat A) noexcept
9524742e46bSJacob Faibussowitsch {
9534742e46bSJacob Faibussowitsch   PetscFunctionBegin;
9544742e46bSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(A->rmap));
9554742e46bSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(A->cmap));
9564742e46bSJacob Faibussowitsch   if (!A->preallocated) {
9574742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
9584742e46bSJacob Faibussowitsch 
9594742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
9603d9668e3SJacob Faibussowitsch     PetscCall(SetPreallocation(A, dctx, nullptr));
9614742e46bSJacob Faibussowitsch   }
9624742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
9634742e46bSJacob Faibussowitsch }
9644742e46bSJacob Faibussowitsch 
9654742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9664742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Reset(Mat A) noexcept
9674742e46bSJacob Faibussowitsch {
9684742e46bSJacob Faibussowitsch   PetscFunctionBegin;
9694742e46bSJacob Faibussowitsch   if (const auto mcu = MatCUPMCast(A)) {
9704742e46bSJacob Faibussowitsch     cupmStream_t stream;
9714742e46bSJacob Faibussowitsch 
9724742e46bSJacob Faibussowitsch     PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME());
9734742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&stream));
9744742e46bSJacob Faibussowitsch     if (!mcu->d_user_alloc) PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream));
9754742e46bSJacob Faibussowitsch     PetscCallCUPM(cupmFreeAsync(mcu->d_fact_tau, stream));
9764742e46bSJacob Faibussowitsch     PetscCallCUPM(cupmFreeAsync(mcu->d_fact_ipiv, stream));
9774742e46bSJacob Faibussowitsch     PetscCallCUPM(cupmFreeAsync(mcu->d_fact_info, stream));
9784742e46bSJacob Faibussowitsch     PetscCallCUPM(cupmFreeAsync(mcu->d_fact_work, stream));
9794742e46bSJacob Faibussowitsch     PetscCall(VecDestroy(&mcu->workvec));
9804742e46bSJacob Faibussowitsch     PetscCall(PetscFree(A->spptr /* mcu */));
9814742e46bSJacob Faibussowitsch   }
9824742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
9834742e46bSJacob Faibussowitsch }
9844742e46bSJacob Faibussowitsch 
9854742e46bSJacob Faibussowitsch // ==========================================================================================
9864742e46bSJacob Faibussowitsch 
9874742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9884742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::BindToCPU(Mat A, PetscBool to_host) noexcept
9894742e46bSJacob Faibussowitsch {
9904742e46bSJacob Faibussowitsch   const auto mimpl = MatIMPLCast(A);
9914742e46bSJacob Faibussowitsch   const auto pobj  = PetscObjectCast(A);
9924742e46bSJacob Faibussowitsch 
9934742e46bSJacob Faibussowitsch   PetscFunctionBegin;
9944742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
9954742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
9964742e46bSJacob Faibussowitsch   A->boundtocpu = to_host;
9974742e46bSJacob Faibussowitsch   PetscCall(PetscStrFreeAllocpy(to_host ? PETSCRANDER48 : PETSCDEVICERAND(), &A->defaultrandtype));
9984742e46bSJacob Faibussowitsch   if (to_host) {
9994742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
10004742e46bSJacob Faibussowitsch 
10014742e46bSJacob Faibussowitsch     // make sure we have an up-to-date copy on the CPU
10024742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
10034742e46bSJacob Faibussowitsch     PetscCall(DeviceToHost_(A, dctx));
10044742e46bSJacob Faibussowitsch   } else {
10054742e46bSJacob Faibussowitsch     PetscBool iscupm;
10064742e46bSJacob Faibussowitsch 
10074742e46bSJacob Faibussowitsch     if (auto &cvec = mimpl->cvec) {
10084742e46bSJacob Faibussowitsch       PetscCall(PetscObjectTypeCompare(PetscObjectCast(cvec), VecSeq_CUPM::VECSEQCUPM(), &iscupm));
10094742e46bSJacob Faibussowitsch       if (!iscupm) PetscCall(VecDestroy(&cvec));
10104742e46bSJacob Faibussowitsch     }
10114742e46bSJacob Faibussowitsch     if (auto &cmat = mimpl->cmat) {
10124742e46bSJacob Faibussowitsch       PetscCall(PetscObjectTypeCompare(PetscObjectCast(cmat), MATSEQDENSECUPM(), &iscupm));
10134742e46bSJacob Faibussowitsch       if (!iscupm) PetscCall(MatDestroy(&cmat));
10144742e46bSJacob Faibussowitsch     }
10154742e46bSJacob Faibussowitsch   }
10164742e46bSJacob Faibussowitsch 
10174742e46bSJacob Faibussowitsch   // ============================================================
10184742e46bSJacob Faibussowitsch   // Composed ops
10194742e46bSJacob Faibussowitsch   // ============================================================
10204742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArray_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ_WRITE>);
10214742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayRead_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ>);
10224742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayWrite_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_WRITE>);
10234742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ_WRITE>);
10244742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ_WRITE>);
10254742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayReadAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ>);
10264742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayReadAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ>);
10274742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayWriteAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_WRITE>);
10284742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayWriteAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_WRITE>);
10294742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVec_C", MatDenseGetColumnVec_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_READ_WRITE>);
10304742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVec_C", MatDenseRestoreColumnVec_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_READ_WRITE>);
10314742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVecRead_C", MatDenseGetColumnVecRead_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_READ>);
10324742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVecRead_C", MatDenseRestoreColumnVecRead_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_READ>);
10334742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVecWrite_C", MatDenseGetColumnVecWrite_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_WRITE>);
10344742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVecWrite_C", MatDenseRestoreColumnVecWrite_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_WRITE>);
10354742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetSubMatrix_C", MatDenseGetSubMatrix_SeqDense, GetSubMatrix);
10364742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreSubMatrix_C", MatDenseRestoreSubMatrix_SeqDense, RestoreSubMatrix);
10374742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatQRFactor_C", MatQRFactor_SeqDense, SolveQR::Factor);
10384742e46bSJacob Faibussowitsch   // always the same
10394742e46bSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction(pobj, "MatDenseSetLDA_C", MatDenseSetLDA_SeqDense));
10404742e46bSJacob Faibussowitsch 
10414742e46bSJacob Faibussowitsch   // ============================================================
10424742e46bSJacob Faibussowitsch   // Function pointer ops
10434742e46bSJacob Faibussowitsch   // ============================================================
10444742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, duplicate, MatDuplicate_SeqDense, Duplicate);
10454742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, mult, MatMult_SeqDense, [](Mat A, Vec xx, Vec yy) { return MatMultAdd_Dispatch_</* transpose */ false>(A, xx, nullptr, yy); });
10464742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, multtranspose, MatMultTranspose_SeqDense, [](Mat A, Vec xx, Vec yy) { return MatMultAdd_Dispatch_</* transpose */ true>(A, xx, nullptr, yy); });
10474742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, multadd, MatMultAdd_SeqDense, MatMultAdd_Dispatch_</* transpose */ false>);
10484742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, multtransposeadd, MatMultTransposeAdd_SeqDense, MatMultAdd_Dispatch_</* transpose */ true>);
10494742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, matmultnumeric, MatMatMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ false, /* transpose_B */ false>);
10504742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, mattransposemultnumeric, MatMatTransposeMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ false, /* transpose_B */ true>);
10514742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, transposematmultnumeric, MatTransposeMatMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ true, /* transpose_B */ false>);
10524742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, axpy, MatAXPY_SeqDense, AXPY);
10534742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, choleskyfactor, MatCholeskyFactor_SeqDense, SolveCholesky::Factor);
10544742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, lufactor, MatLUFactor_SeqDense, SolveLU::Factor);
10554742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, getcolumnvector, MatGetColumnVector_SeqDense, GetColumnVector);
10564742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, scale, MatScale_SeqDense, Scale);
10574742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, shift, MatShift_SeqDense, Shift);
10584742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, copy, MatCopy_SeqDense, Copy);
10594742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, zeroentries, MatZeroEntries_SeqDense, ZeroEntries);
10604742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, setup, MatSetUp_SeqDense, SetUp);
10614742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, setrandom, MatSetRandom_SeqDense, SetRandom);
106214277c92SJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, getdiagonal, MatGetDiagonal_SeqDense, GetDiagonal);
10634742e46bSJacob Faibussowitsch   // seemingly always the same
10644742e46bSJacob Faibussowitsch   A->ops->productsetfromoptions = MatProductSetFromOptions_SeqDense;
10654742e46bSJacob Faibussowitsch 
10664742e46bSJacob Faibussowitsch   if (const auto cmat = mimpl->cmat) PetscCall(MatBindToCPU(cmat, to_host));
10674742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
10684742e46bSJacob Faibussowitsch }
10694742e46bSJacob Faibussowitsch 
10704742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
10714742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_SeqDenseCUPM_SeqDense(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept
10724742e46bSJacob Faibussowitsch {
10734742e46bSJacob Faibussowitsch   PetscFunctionBegin;
10744742e46bSJacob Faibussowitsch   PetscCall(Convert_Dispatch_</* to host */ true>(M, type, reuse, newmat));
10754742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
10764742e46bSJacob Faibussowitsch }
10774742e46bSJacob Faibussowitsch 
10784742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
10794742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_SeqDense_SeqDenseCUPM(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept
10804742e46bSJacob Faibussowitsch {
10814742e46bSJacob Faibussowitsch   PetscFunctionBegin;
10824742e46bSJacob Faibussowitsch   PetscCall(Convert_Dispatch_</* to host */ false>(M, type, reuse, newmat));
10834742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
10844742e46bSJacob Faibussowitsch }
10854742e46bSJacob Faibussowitsch 
10864742e46bSJacob Faibussowitsch // ==========================================================================================
10874742e46bSJacob Faibussowitsch 
10884742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
10894742e46bSJacob Faibussowitsch template <PetscMemType mtype, PetscMemoryAccessMode access>
10904742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::GetArray(Mat m, PetscScalar **array, PetscDeviceContext dctx) noexcept
10914742e46bSJacob Faibussowitsch {
10924742e46bSJacob Faibussowitsch   constexpr auto hostmem     = PetscMemTypeHost(mtype);
10934742e46bSJacob Faibussowitsch   constexpr auto read_access = PetscMemoryAccessRead(access);
10944742e46bSJacob Faibussowitsch 
10954742e46bSJacob Faibussowitsch   PetscFunctionBegin;
10964742e46bSJacob Faibussowitsch   static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), "");
10974742e46bSJacob Faibussowitsch   PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
10984742e46bSJacob Faibussowitsch   if (hostmem) {
10994742e46bSJacob Faibussowitsch     if (read_access) {
11004742e46bSJacob Faibussowitsch       PetscCall(DeviceToHost_(m, dctx));
11014742e46bSJacob Faibussowitsch     } else if (!MatIMPLCast(m)->v) {
11024742e46bSJacob Faibussowitsch       // MatCreateSeqDenseCUPM may not allocate CPU memory. Allocate if needed
11034742e46bSJacob Faibussowitsch       PetscCall(MatSeqDenseSetPreallocation(m, nullptr));
11044742e46bSJacob Faibussowitsch     }
11054742e46bSJacob Faibussowitsch     *array = MatIMPLCast(m)->v;
11064742e46bSJacob Faibussowitsch   } else {
11074742e46bSJacob Faibussowitsch     if (read_access) {
11084742e46bSJacob Faibussowitsch       PetscCall(HostToDevice_(m, dctx));
11094742e46bSJacob Faibussowitsch     } else if (!MatCUPMCast(m)->d_v) {
11104742e46bSJacob Faibussowitsch       // write-only
11114742e46bSJacob Faibussowitsch       PetscCall(SetPreallocation(m, dctx, nullptr));
11124742e46bSJacob Faibussowitsch     }
11134742e46bSJacob Faibussowitsch     *array = MatCUPMCast(m)->d_v;
11144742e46bSJacob Faibussowitsch   }
11154742e46bSJacob Faibussowitsch   if (PetscMemoryAccessWrite(access)) {
11164742e46bSJacob Faibussowitsch     m->offloadmask = hostmem ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
11174742e46bSJacob Faibussowitsch     PetscCall(PetscObjectStateIncrease(PetscObjectCast(m)));
11184742e46bSJacob Faibussowitsch   }
11194742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
11204742e46bSJacob Faibussowitsch }
11214742e46bSJacob Faibussowitsch 
11224742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
11234742e46bSJacob Faibussowitsch template <PetscMemType mtype, PetscMemoryAccessMode access>
11244742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreArray(Mat m, PetscScalar **array, PetscDeviceContext) noexcept
11254742e46bSJacob Faibussowitsch {
11264742e46bSJacob Faibussowitsch   PetscFunctionBegin;
11274742e46bSJacob Faibussowitsch   static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), "");
11284742e46bSJacob Faibussowitsch   if (PetscMemoryAccessWrite(access)) {
11294742e46bSJacob Faibussowitsch     // WRITE or READ_WRITE
11304742e46bSJacob Faibussowitsch     m->offloadmask = PetscMemTypeHost(mtype) ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
11314742e46bSJacob Faibussowitsch     PetscCall(PetscObjectStateIncrease(PetscObjectCast(m)));
11324742e46bSJacob Faibussowitsch   }
113301329d7dSJacob Faibussowitsch   if (array) {
113401329d7dSJacob Faibussowitsch     PetscCall(CheckPointerMatchesMemType_(*array, mtype));
11354742e46bSJacob Faibussowitsch     *array = nullptr;
113601329d7dSJacob Faibussowitsch   }
11374742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
11384742e46bSJacob Faibussowitsch }
11394742e46bSJacob Faibussowitsch 
11404742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
11414742e46bSJacob Faibussowitsch template <PetscMemoryAccessMode access>
11424742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::GetArrayAndMemType(Mat m, PetscScalar **array, PetscMemType *mtype, PetscDeviceContext dctx) noexcept
11434742e46bSJacob Faibussowitsch {
11444742e46bSJacob Faibussowitsch   PetscFunctionBegin;
11454742e46bSJacob Faibussowitsch   PetscCall(GetArray<PETSC_MEMTYPE_DEVICE, access>(m, array, dctx));
11464742e46bSJacob Faibussowitsch   if (mtype) *mtype = PETSC_MEMTYPE_CUPM();
11474742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
11484742e46bSJacob Faibussowitsch }
11494742e46bSJacob Faibussowitsch 
11504742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
11514742e46bSJacob Faibussowitsch template <PetscMemoryAccessMode access>
11524742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreArrayAndMemType(Mat m, PetscScalar **array, PetscDeviceContext dctx) noexcept
11534742e46bSJacob Faibussowitsch {
11544742e46bSJacob Faibussowitsch   PetscFunctionBegin;
11554742e46bSJacob Faibussowitsch   PetscCall(RestoreArray<PETSC_MEMTYPE_DEVICE, access>(m, array, dctx));
11564742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
11574742e46bSJacob Faibussowitsch }
11584742e46bSJacob Faibussowitsch 
11594742e46bSJacob Faibussowitsch // ==========================================================================================
11604742e46bSJacob Faibussowitsch 
11614742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
11624742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::PlaceArray(Mat A, const PetscScalar *array) noexcept
11634742e46bSJacob Faibussowitsch {
11644742e46bSJacob Faibussowitsch   const auto mimpl = MatIMPLCast(A);
11654742e46bSJacob Faibussowitsch   const auto mcu   = MatCUPMCast(A);
11664742e46bSJacob Faibussowitsch 
11674742e46bSJacob Faibussowitsch   PetscFunctionBegin;
11684742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
11694742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
11704742e46bSJacob Faibussowitsch   PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME());
11714742e46bSJacob Faibussowitsch   if (mimpl->v) {
11724742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
11734742e46bSJacob Faibussowitsch 
11744742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
11754742e46bSJacob Faibussowitsch     PetscCall(HostToDevice_(A, dctx));
11764742e46bSJacob Faibussowitsch   }
11774742e46bSJacob Faibussowitsch   mcu->unplacedarray         = util::exchange(mcu->d_v, const_cast<PetscScalar *>(array));
11784742e46bSJacob Faibussowitsch   mcu->d_unplaced_user_alloc = util::exchange(mcu->d_user_alloc, PETSC_TRUE);
11794742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
11804742e46bSJacob Faibussowitsch }
11814742e46bSJacob Faibussowitsch 
11824742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
11834742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::ReplaceArray(Mat A, const PetscScalar *array) noexcept
11844742e46bSJacob Faibussowitsch {
11854742e46bSJacob Faibussowitsch   const auto mimpl = MatIMPLCast(A);
11864742e46bSJacob Faibussowitsch   const auto mcu   = MatCUPMCast(A);
11874742e46bSJacob Faibussowitsch 
11884742e46bSJacob Faibussowitsch   PetscFunctionBegin;
11894742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
11904742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
11914742e46bSJacob Faibussowitsch   PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME());
11924742e46bSJacob Faibussowitsch   if (!mcu->d_user_alloc) {
11934742e46bSJacob Faibussowitsch     cupmStream_t stream;
11944742e46bSJacob Faibussowitsch 
11954742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&stream));
11964742e46bSJacob Faibussowitsch     PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream));
11974742e46bSJacob Faibussowitsch   }
11984742e46bSJacob Faibussowitsch   mcu->d_v          = const_cast<PetscScalar *>(array);
11994742e46bSJacob Faibussowitsch   mcu->d_user_alloc = PETSC_FALSE;
12004742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
12014742e46bSJacob Faibussowitsch }
12024742e46bSJacob Faibussowitsch 
12034742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
12044742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::ResetArray(Mat A) noexcept
12054742e46bSJacob Faibussowitsch {
12064742e46bSJacob Faibussowitsch   const auto mimpl = MatIMPLCast(A);
12074742e46bSJacob Faibussowitsch   const auto mcu   = MatCUPMCast(A);
12084742e46bSJacob Faibussowitsch 
12094742e46bSJacob Faibussowitsch   PetscFunctionBegin;
12104742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
12114742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
12124742e46bSJacob Faibussowitsch   if (mimpl->v) {
12134742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
12144742e46bSJacob Faibussowitsch 
12154742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
12164742e46bSJacob Faibussowitsch     PetscCall(HostToDevice_(A, dctx));
12174742e46bSJacob Faibussowitsch   }
12184742e46bSJacob Faibussowitsch   mcu->d_v          = util::exchange(mcu->unplacedarray, nullptr);
12194742e46bSJacob Faibussowitsch   mcu->d_user_alloc = mcu->d_unplaced_user_alloc;
12204742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
12214742e46bSJacob Faibussowitsch }
12224742e46bSJacob Faibussowitsch 
12234742e46bSJacob Faibussowitsch // ==========================================================================================
12244742e46bSJacob Faibussowitsch 
12254742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
12264742e46bSJacob Faibussowitsch template <bool transpose_A, bool transpose_B>
12274742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMatMult_Numeric_Dispatch(Mat A, Mat B, Mat C) noexcept
12284742e46bSJacob Faibussowitsch {
12294742e46bSJacob Faibussowitsch   cupmBlasInt_t      m, n, k;
12304742e46bSJacob Faibussowitsch   PetscBool          Aiscupm, Biscupm;
12314742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
12324742e46bSJacob Faibussowitsch   cupmBlasHandle_t   handle;
12334742e46bSJacob Faibussowitsch 
12344742e46bSJacob Faibussowitsch   PetscFunctionBegin;
12354742e46bSJacob Faibussowitsch   PetscCall(PetscCUPMBlasIntCast(C->rmap->n, &m));
12364742e46bSJacob Faibussowitsch   PetscCall(PetscCUPMBlasIntCast(C->cmap->n, &n));
12374742e46bSJacob Faibussowitsch   PetscCall(PetscCUPMBlasIntCast(transpose_A ? A->rmap->n : A->cmap->n, &k));
12384742e46bSJacob Faibussowitsch   if (!m || !n || !k) PetscFunctionReturn(PETSC_SUCCESS);
12394742e46bSJacob Faibussowitsch 
12404742e46bSJacob Faibussowitsch   // we may end up with SEQDENSE as one of the arguments
12414742e46bSJacob Faibussowitsch   // REVIEW ME: how? and why is it not B and C????????
12424742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(PetscObjectCast(A), MATSEQDENSECUPM(), &Aiscupm));
12434742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(PetscObjectCast(B), MATSEQDENSECUPM(), &Biscupm));
12444742e46bSJacob Faibussowitsch   if (!Aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
12454742e46bSJacob Faibussowitsch   if (!Biscupm) PetscCall(MatConvert(B, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &B));
12464742e46bSJacob Faibussowitsch   PetscCall(PetscInfo(C, "Matrix-Matrix product %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " on backend\n", m, k, n));
12474742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx, &handle));
12484742e46bSJacob Faibussowitsch 
12494742e46bSJacob Faibussowitsch   PetscCall(PetscLogGpuTimeBegin());
12504742e46bSJacob Faibussowitsch   {
12514742e46bSJacob Faibussowitsch     const auto one  = cupmScalarCast(1.0);
12524742e46bSJacob Faibussowitsch     const auto zero = cupmScalarCast(0.0);
12534742e46bSJacob Faibussowitsch     const auto da   = DeviceArrayRead(dctx, A);
12544742e46bSJacob Faibussowitsch     const auto db   = DeviceArrayRead(dctx, B);
12554742e46bSJacob Faibussowitsch     const auto dc   = DeviceArrayWrite(dctx, C);
12564742e46bSJacob Faibussowitsch     PetscInt   alda, blda, clda;
12574742e46bSJacob Faibussowitsch 
12584742e46bSJacob Faibussowitsch     PetscCall(MatDenseGetLDA(A, &alda));
12594742e46bSJacob Faibussowitsch     PetscCall(MatDenseGetLDA(B, &blda));
12604742e46bSJacob Faibussowitsch     PetscCall(MatDenseGetLDA(C, &clda));
12614742e46bSJacob Faibussowitsch     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));
12624742e46bSJacob Faibussowitsch   }
12634742e46bSJacob Faibussowitsch   PetscCall(PetscLogGpuTimeEnd());
12644742e46bSJacob Faibussowitsch 
12654742e46bSJacob Faibussowitsch   PetscCall(PetscLogGpuFlops(1.0 * m * n * k + 1.0 * m * n * (k - 1)));
12664742e46bSJacob Faibussowitsch   if (!Aiscupm) PetscCall(MatConvert(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A));
12674742e46bSJacob Faibussowitsch   if (!Biscupm) PetscCall(MatConvert(B, MATSEQDENSE, MAT_INPLACE_MATRIX, &B));
12684742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
12694742e46bSJacob Faibussowitsch }
12704742e46bSJacob Faibussowitsch 
12714742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
12724742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Copy(Mat A, Mat B, MatStructure str) noexcept
12734742e46bSJacob Faibussowitsch {
12744742e46bSJacob Faibussowitsch   const auto m = A->rmap->n;
12754742e46bSJacob Faibussowitsch   const auto n = A->cmap->n;
12764742e46bSJacob Faibussowitsch 
12774742e46bSJacob Faibussowitsch   PetscFunctionBegin;
12784742e46bSJacob Faibussowitsch   PetscAssert(m == B->rmap->n && n == B->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "size(B) != size(A)");
12794742e46bSJacob Faibussowitsch   // The two matrices must have the same copy implementation to be eligible for fast copy
12804742e46bSJacob Faibussowitsch   if (A->ops->copy == B->ops->copy) {
12814742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
12824742e46bSJacob Faibussowitsch     cupmStream_t       stream;
12834742e46bSJacob Faibussowitsch 
12844742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx, &stream));
12854742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeBegin());
12864742e46bSJacob Faibussowitsch     {
12874742e46bSJacob Faibussowitsch       const auto va = DeviceArrayRead(dctx, A);
12884742e46bSJacob Faibussowitsch       const auto vb = DeviceArrayWrite(dctx, B);
12894742e46bSJacob Faibussowitsch       // order is important, DeviceArrayRead/Write() might call SetPreallocation() which sets
12904742e46bSJacob Faibussowitsch       // lda!
12914742e46bSJacob Faibussowitsch       const auto lda_a = MatIMPLCast(A)->lda;
12924742e46bSJacob Faibussowitsch       const auto lda_b = MatIMPLCast(B)->lda;
12934742e46bSJacob Faibussowitsch 
12944742e46bSJacob Faibussowitsch       if (lda_a > m || lda_b > m) {
12954742e46bSJacob Faibussowitsch         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());
12964742e46bSJacob Faibussowitsch         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());
12974742e46bSJacob Faibussowitsch         PetscCall(PetscCUPMMemcpy2DAsync(vb.data(), lda_b, va.data(), lda_a, m, n, cupmMemcpyDeviceToDevice, stream));
12984742e46bSJacob Faibussowitsch       } else {
12994742e46bSJacob Faibussowitsch         PetscCall(PetscCUPMMemcpyAsync(vb.data(), va.data(), m * n, cupmMemcpyDeviceToDevice, stream));
13004742e46bSJacob Faibussowitsch       }
13014742e46bSJacob Faibussowitsch     }
13024742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeEnd());
13034742e46bSJacob Faibussowitsch   } else {
13044742e46bSJacob Faibussowitsch     PetscCall(MatCopy_Basic(A, B, str));
13054742e46bSJacob Faibussowitsch   }
13064742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
13074742e46bSJacob Faibussowitsch }
13084742e46bSJacob Faibussowitsch 
13094742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
13104742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::ZeroEntries(Mat m) noexcept
13114742e46bSJacob Faibussowitsch {
13124742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
13134742e46bSJacob Faibussowitsch   cupmStream_t       stream;
13144742e46bSJacob Faibussowitsch 
13154742e46bSJacob Faibussowitsch   PetscFunctionBegin;
13164742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx, &stream));
13174742e46bSJacob Faibussowitsch   PetscCall(PetscLogGpuTimeBegin());
13184742e46bSJacob Faibussowitsch   {
13194742e46bSJacob Faibussowitsch     const auto va  = DeviceArrayWrite(dctx, m);
13204742e46bSJacob Faibussowitsch     const auto lda = MatIMPLCast(m)->lda;
13214742e46bSJacob Faibussowitsch     const auto ma  = m->rmap->n;
13224742e46bSJacob Faibussowitsch     const auto na  = m->cmap->n;
13234742e46bSJacob Faibussowitsch 
13244742e46bSJacob Faibussowitsch     if (lda > ma) {
13254742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMemset2DAsync(va.data(), lda, 0, ma, na, stream));
13264742e46bSJacob Faibussowitsch     } else {
13274742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMemsetAsync(va.data(), 0, ma * na, stream));
13284742e46bSJacob Faibussowitsch     }
13294742e46bSJacob Faibussowitsch   }
13304742e46bSJacob Faibussowitsch   PetscCall(PetscLogGpuTimeEnd());
13314742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
13324742e46bSJacob Faibussowitsch }
13334742e46bSJacob Faibussowitsch 
13342ea277ceSJacob Faibussowitsch namespace detail
13352ea277ceSJacob Faibussowitsch {
13362ea277ceSJacob Faibussowitsch 
13372ea277ceSJacob Faibussowitsch // ==========================================================================================
13382ea277ceSJacob Faibussowitsch // SubMatIndexFunctor
13392ea277ceSJacob Faibussowitsch //
13402ea277ceSJacob Faibussowitsch // Iterator which permutes a linear index range into matrix indices for am nrows x ncols
13412ea277ceSJacob Faibussowitsch // submat with leading dimension lda. Essentially SubMatIndexFunctor(i) returns the index for
13422ea277ceSJacob Faibussowitsch // the i'th sequential entry in the matrix.
13432ea277ceSJacob Faibussowitsch // ==========================================================================================
13442ea277ceSJacob Faibussowitsch template <typename T>
13452ea277ceSJacob Faibussowitsch struct SubMatIndexFunctor {
13462ea277ceSJacob Faibussowitsch   PETSC_HOSTDEVICE_INLINE_DECL T operator()(T x) const noexcept { return ((x / nrows) * lda) + (x % nrows); }
13472ea277ceSJacob Faibussowitsch 
13482ea277ceSJacob Faibussowitsch   PetscInt nrows;
13492ea277ceSJacob Faibussowitsch   PetscInt ncols;
13502ea277ceSJacob Faibussowitsch   PetscInt lda;
13512ea277ceSJacob Faibussowitsch };
13522ea277ceSJacob Faibussowitsch 
13532ea277ceSJacob Faibussowitsch template <typename Iterator>
13542ea277ceSJacob Faibussowitsch struct SubMatrixIterator : MatrixIteratorBase<Iterator, SubMatIndexFunctor<typename thrust::iterator_difference<Iterator>::type>> {
13552ea277ceSJacob Faibussowitsch   using base_type = MatrixIteratorBase<Iterator, SubMatIndexFunctor<typename thrust::iterator_difference<Iterator>::type>>;
13562ea277ceSJacob Faibussowitsch 
13572ea277ceSJacob Faibussowitsch   using iterator = typename base_type::iterator;
13582ea277ceSJacob Faibussowitsch 
13592ea277ceSJacob Faibussowitsch   constexpr SubMatrixIterator(Iterator first, Iterator last, PetscInt nrows, PetscInt ncols, PetscInt lda) noexcept :
13602ea277ceSJacob Faibussowitsch     base_type{
13612ea277ceSJacob Faibussowitsch       std::move(first), std::move(last), {nrows, ncols, lda}
13622ea277ceSJacob Faibussowitsch   }
13632ea277ceSJacob Faibussowitsch   {
13642ea277ceSJacob Faibussowitsch   }
13652ea277ceSJacob Faibussowitsch 
13662ea277ceSJacob Faibussowitsch   PETSC_NODISCARD iterator end() const noexcept { return this->begin() + (this->func.nrows * this->func.ncols); }
13672ea277ceSJacob Faibussowitsch };
13682ea277ceSJacob Faibussowitsch 
13692ea277ceSJacob Faibussowitsch namespace
13702ea277ceSJacob Faibussowitsch {
13712ea277ceSJacob Faibussowitsch 
13722ea277ceSJacob Faibussowitsch template <typename T>
13732ea277ceSJacob Faibussowitsch 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
13742ea277ceSJacob Faibussowitsch {
13752ea277ceSJacob Faibussowitsch   const auto nrows = rend - rstart;
13762ea277ceSJacob Faibussowitsch   const auto ncols = cend - cstart;
13772ea277ceSJacob Faibussowitsch   const auto dptr  = thrust::device_pointer_cast(ptr);
13782ea277ceSJacob Faibussowitsch 
13792ea277ceSJacob Faibussowitsch   return {dptr + (rstart * lda) + cstart, dptr + ((rstart + nrows) * lda) + cstart, nrows, ncols, lda};
13802ea277ceSJacob Faibussowitsch }
13812ea277ceSJacob Faibussowitsch 
13822ea277ceSJacob Faibussowitsch } // namespace
13832ea277ceSJacob Faibussowitsch 
13842ea277ceSJacob Faibussowitsch } // namespace detail
13852ea277ceSJacob Faibussowitsch 
13864742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
13874742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Scale(Mat A, PetscScalar alpha) noexcept
13884742e46bSJacob Faibussowitsch {
13892ea277ceSJacob Faibussowitsch   const auto         m = A->rmap->n;
13902ea277ceSJacob Faibussowitsch   const auto         n = A->cmap->n;
13914742e46bSJacob Faibussowitsch   const auto         N = m * n;
13924742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
13934742e46bSJacob Faibussowitsch 
13944742e46bSJacob Faibussowitsch   PetscFunctionBegin;
13952ea277ceSJacob Faibussowitsch   PetscCall(PetscInfo(A, "Performing Scale %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m, n));
13962ea277ceSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx));
13974742e46bSJacob Faibussowitsch   {
13984742e46bSJacob Faibussowitsch     const auto da  = DeviceArrayReadWrite(dctx, A);
13992ea277ceSJacob Faibussowitsch     const auto lda = MatIMPLCast(A)->lda;
14004742e46bSJacob Faibussowitsch 
14014742e46bSJacob Faibussowitsch     if (lda > m) {
14022ea277ceSJacob Faibussowitsch       cupmStream_t stream;
14032ea277ceSJacob Faibussowitsch 
14042ea277ceSJacob Faibussowitsch       PetscCall(GetHandlesFrom_(dctx, &stream));
14052ea277ceSJacob Faibussowitsch       // clang-format off
14062ea277ceSJacob Faibussowitsch       PetscCallThrust(
14072ea277ceSJacob Faibussowitsch         const auto sub_mat = detail::make_submat_iterator(0, m, 0, n, lda, da.data());
14082ea277ceSJacob Faibussowitsch 
14092ea277ceSJacob Faibussowitsch         THRUST_CALL(
14102ea277ceSJacob Faibussowitsch           thrust::transform,
14112ea277ceSJacob Faibussowitsch           stream,
14122ea277ceSJacob Faibussowitsch           sub_mat.begin(), sub_mat.end(), sub_mat.begin(),
14132ea277ceSJacob Faibussowitsch           device::cupm::functors::make_times_equals(alpha)
14142ea277ceSJacob Faibussowitsch         )
14152ea277ceSJacob Faibussowitsch       );
14162ea277ceSJacob Faibussowitsch       // clang-format on
14174742e46bSJacob Faibussowitsch     } else {
14182ea277ceSJacob Faibussowitsch       const auto       cu_alpha = cupmScalarCast(alpha);
14192ea277ceSJacob Faibussowitsch       cupmBlasHandle_t handle;
14202ea277ceSJacob Faibussowitsch 
14212ea277ceSJacob Faibussowitsch       PetscCall(GetHandlesFrom_(dctx, &handle));
14222ea277ceSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeBegin());
14234742e46bSJacob Faibussowitsch       PetscCallCUPMBLAS(cupmBlasXscal(handle, N, &cu_alpha, da.cupmdata(), 1));
14244742e46bSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeEnd());
14252ea277ceSJacob Faibussowitsch     }
14262ea277ceSJacob Faibussowitsch   }
14274742e46bSJacob Faibussowitsch   PetscCall(PetscLogGpuFlops(N));
14284742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
14294742e46bSJacob Faibussowitsch }
14304742e46bSJacob Faibussowitsch 
14314742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
14324742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::AXPY(Mat Y, PetscScalar alpha, Mat X, MatStructure) noexcept
14334742e46bSJacob Faibussowitsch {
14344742e46bSJacob Faibussowitsch   const auto         m_x = X->rmap->n, m_y = Y->rmap->n;
14354742e46bSJacob Faibussowitsch   const auto         n_x = X->cmap->n, n_y = Y->cmap->n;
1436025e0618SJacob Faibussowitsch   const auto         N = m_x * n_x;
14374742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
14384742e46bSJacob Faibussowitsch 
14394742e46bSJacob Faibussowitsch   PetscFunctionBegin;
1440025e0618SJacob Faibussowitsch   if (!m_x || !n_x || alpha == (PetscScalar)0.0) PetscFunctionReturn(PETSC_SUCCESS);
14414742e46bSJacob Faibussowitsch   PetscCall(PetscInfo(Y, "Performing AXPY %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m_y, n_y));
1442025e0618SJacob Faibussowitsch   PetscCall(GetHandles_(&dctx));
14434742e46bSJacob Faibussowitsch   {
14444742e46bSJacob Faibussowitsch     const auto dx    = DeviceArrayRead(dctx, X);
1445025e0618SJacob Faibussowitsch     const auto dy    = DeviceArrayReadWrite(dctx, Y);
1446025e0618SJacob Faibussowitsch     const auto lda_x = MatIMPLCast(X)->lda;
1447025e0618SJacob Faibussowitsch     const auto lda_y = MatIMPLCast(Y)->lda;
14484742e46bSJacob Faibussowitsch 
1449025e0618SJacob Faibussowitsch     if (lda_x > m_x || lda_y > m_x) {
1450025e0618SJacob Faibussowitsch       cupmStream_t stream;
1451025e0618SJacob Faibussowitsch 
1452025e0618SJacob Faibussowitsch       PetscCall(GetHandlesFrom_(dctx, &stream));
1453025e0618SJacob Faibussowitsch       // clang-format off
1454025e0618SJacob Faibussowitsch       PetscCallThrust(
1455025e0618SJacob Faibussowitsch         const auto sub_mat_y = detail::make_submat_iterator(0, m_y, 0, n_y, lda_y, dy.data());
1456025e0618SJacob Faibussowitsch         const auto sub_mat_x = detail::make_submat_iterator(0, m_x, 0, n_x, lda_x, dx.data());
1457025e0618SJacob Faibussowitsch 
1458025e0618SJacob Faibussowitsch         THRUST_CALL(
1459025e0618SJacob Faibussowitsch           thrust::transform,
1460025e0618SJacob Faibussowitsch           stream,
1461025e0618SJacob Faibussowitsch           sub_mat_x.begin(), sub_mat_x.end(), sub_mat_y.begin(), sub_mat_y.begin(),
1462025e0618SJacob Faibussowitsch           device::cupm::functors::make_axpy(alpha)
1463025e0618SJacob Faibussowitsch         );
1464025e0618SJacob Faibussowitsch       );
1465025e0618SJacob Faibussowitsch       // clang-format on
14664742e46bSJacob Faibussowitsch     } else {
1467025e0618SJacob Faibussowitsch       const auto       cu_alpha = cupmScalarCast(alpha);
1468025e0618SJacob Faibussowitsch       cupmBlasHandle_t handle;
1469025e0618SJacob Faibussowitsch 
1470025e0618SJacob Faibussowitsch       PetscCall(GetHandlesFrom_(dctx, &handle));
1471025e0618SJacob Faibussowitsch       PetscCall(PetscLogGpuTimeBegin());
1472025e0618SJacob Faibussowitsch       PetscCallCUPMBLAS(cupmBlasXaxpy(handle, N, &cu_alpha, dx.cupmdata(), 1, dy.cupmdata(), 1));
14734742e46bSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeEnd());
14744742e46bSJacob Faibussowitsch     }
1475025e0618SJacob Faibussowitsch   }
1476025e0618SJacob Faibussowitsch   PetscCall(PetscLogGpuFlops(PetscMax(2 * N - 1, 0)));
14774742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
14784742e46bSJacob Faibussowitsch }
14794742e46bSJacob Faibussowitsch 
14804742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
14814742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Duplicate(Mat A, MatDuplicateOption opt, Mat *B) noexcept
14824742e46bSJacob Faibussowitsch {
14834742e46bSJacob Faibussowitsch   const auto         hopt = (opt == MAT_COPY_VALUES && A->offloadmask != PETSC_OFFLOAD_CPU) ? MAT_DO_NOT_COPY_VALUES : opt;
14844742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
14854742e46bSJacob Faibussowitsch 
14864742e46bSJacob Faibussowitsch   PetscFunctionBegin;
14874742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx));
14884742e46bSJacob Faibussowitsch   // do not call SetPreallocation() yet, we call it afterwards??
14894742e46bSJacob Faibussowitsch   PetscCall(MatCreateSeqDenseCUPM<T>(PetscObjectComm(PetscObjectCast(A)), A->rmap->n, A->cmap->n, nullptr, B, dctx, /* preallocate */ false));
14904742e46bSJacob Faibussowitsch   PetscCall(MatDuplicateNoCreate_SeqDense(*B, A, hopt));
14914742e46bSJacob Faibussowitsch   if (opt == MAT_COPY_VALUES && hopt != MAT_COPY_VALUES) PetscCall(Copy(A, *B, SAME_NONZERO_PATTERN));
14924742e46bSJacob Faibussowitsch   // allocate memory if needed
14933d9668e3SJacob Faibussowitsch   if (opt != MAT_COPY_VALUES && !MatCUPMCast(*B)->d_v) PetscCall(SetPreallocation(*B, dctx, nullptr));
14944742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
14954742e46bSJacob Faibussowitsch }
14964742e46bSJacob Faibussowitsch 
14974742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
14984742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::SetRandom(Mat A, PetscRandom rng) noexcept
14994742e46bSJacob Faibussowitsch {
15004742e46bSJacob Faibussowitsch   PetscBool device;
15014742e46bSJacob Faibussowitsch 
15024742e46bSJacob Faibussowitsch   PetscFunctionBegin;
15034742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(PetscObjectCast(rng), PETSCDEVICERAND(), &device));
15044742e46bSJacob Faibussowitsch   if (device) {
15054742e46bSJacob Faibussowitsch     const auto         m = A->rmap->n;
15064742e46bSJacob Faibussowitsch     const auto         n = A->cmap->n;
15074742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
15084742e46bSJacob Faibussowitsch 
15094742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
15104742e46bSJacob Faibussowitsch     {
15114742e46bSJacob Faibussowitsch       const auto a = DeviceArrayWrite(dctx, A);
15124742e46bSJacob Faibussowitsch       PetscInt   lda;
15134742e46bSJacob Faibussowitsch 
15144742e46bSJacob Faibussowitsch       PetscCall(MatDenseGetLDA(A, &lda));
15154742e46bSJacob Faibussowitsch       if (lda > m) {
15164742e46bSJacob Faibussowitsch         for (PetscInt i = 0; i < n; i++) PetscCall(PetscRandomGetValues(rng, m, a.data() + i * lda));
15174742e46bSJacob Faibussowitsch       } else {
15184742e46bSJacob Faibussowitsch         PetscInt mn;
15194742e46bSJacob Faibussowitsch 
15204742e46bSJacob Faibussowitsch         PetscCall(PetscIntMultError(m, n, &mn));
15214742e46bSJacob Faibussowitsch         PetscCall(PetscRandomGetValues(rng, mn, a));
15224742e46bSJacob Faibussowitsch       }
15234742e46bSJacob Faibussowitsch     }
15244742e46bSJacob Faibussowitsch   } else {
15254742e46bSJacob Faibussowitsch     PetscCall(MatSetRandom_SeqDense(A, rng));
15264742e46bSJacob Faibussowitsch   }
15274742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15284742e46bSJacob Faibussowitsch }
15294742e46bSJacob Faibussowitsch 
15304742e46bSJacob Faibussowitsch // ==========================================================================================
15314742e46bSJacob Faibussowitsch 
15324742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
15334742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::GetColumnVector(Mat A, Vec v, PetscInt col) noexcept
15344742e46bSJacob Faibussowitsch {
15354742e46bSJacob Faibussowitsch   const auto         offloadmask = A->offloadmask;
15364742e46bSJacob Faibussowitsch   const auto         n           = A->rmap->n;
15374742e46bSJacob Faibussowitsch   const auto         col_offset  = [&](const PetscScalar *ptr) { return ptr + col * MatIMPLCast(A)->lda; };
15384742e46bSJacob Faibussowitsch   PetscBool          viscupm;
15394742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
15404742e46bSJacob Faibussowitsch   cupmStream_t       stream;
15414742e46bSJacob Faibussowitsch 
15424742e46bSJacob Faibussowitsch   PetscFunctionBegin;
15434742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(v), &viscupm, VecSeq_CUPM::VECSEQCUPM(), VecSeq_CUPM::VECMPICUPM(), VecSeq_CUPM::VECCUPM(), ""));
15444742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx, &stream));
15454742e46bSJacob Faibussowitsch   if (viscupm && !v->boundtocpu) {
15464742e46bSJacob Faibussowitsch     const auto x = VecSeq_CUPM::DeviceArrayWrite(dctx, v);
15474742e46bSJacob Faibussowitsch 
15484742e46bSJacob Faibussowitsch     // update device data
15494742e46bSJacob Faibussowitsch     if (PetscOffloadDevice(offloadmask)) {
15504742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMemcpyAsync(x.data(), col_offset(DeviceArrayRead(dctx, A)), n, cupmMemcpyDeviceToDevice, stream));
15514742e46bSJacob Faibussowitsch     } else {
15524742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMemcpyAsync(x.data(), col_offset(HostArrayRead(dctx, A)), n, cupmMemcpyHostToDevice, stream));
15534742e46bSJacob Faibussowitsch     }
15544742e46bSJacob Faibussowitsch   } else {
15554742e46bSJacob Faibussowitsch     PetscScalar *x;
15564742e46bSJacob Faibussowitsch 
15574742e46bSJacob Faibussowitsch     // update host data
15584742e46bSJacob Faibussowitsch     PetscCall(VecGetArrayWrite(v, &x));
15594742e46bSJacob Faibussowitsch     if (PetscOffloadUnallocated(offloadmask) || PetscOffloadHost(offloadmask)) {
15604742e46bSJacob Faibussowitsch       PetscCall(PetscArraycpy(x, col_offset(HostArrayRead(dctx, A)), n));
15614742e46bSJacob Faibussowitsch     } else if (PetscOffloadDevice(offloadmask)) {
15624742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMemcpyAsync(x, col_offset(DeviceArrayRead(dctx, A)), n, cupmMemcpyDeviceToHost, stream));
15634742e46bSJacob Faibussowitsch     }
15644742e46bSJacob Faibussowitsch     PetscCall(VecRestoreArrayWrite(v, &x));
15654742e46bSJacob Faibussowitsch   }
15664742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15674742e46bSJacob Faibussowitsch }
15684742e46bSJacob Faibussowitsch 
15694742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
15704742e46bSJacob Faibussowitsch template <PetscMemoryAccessMode access>
15714742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::GetColumnVec(Mat A, PetscInt col, Vec *v) noexcept
15724742e46bSJacob Faibussowitsch {
15734742e46bSJacob Faibussowitsch   using namespace vec::cupm;
15744742e46bSJacob Faibussowitsch   const auto         mimpl = MatIMPLCast(A);
15754742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
15764742e46bSJacob Faibussowitsch 
15774742e46bSJacob Faibussowitsch   PetscFunctionBegin;
15784742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
15794742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
15804742e46bSJacob Faibussowitsch   mimpl->vecinuse = col + 1;
1581d16ceb75SStefano Zampini   if (!mimpl->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &mimpl->cvec));
15824742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx));
15834742e46bSJacob Faibussowitsch   PetscCall(GetArray<PETSC_MEMTYPE_DEVICE, access>(A, const_cast<PetscScalar **>(&mimpl->ptrinuse), dctx));
15844742e46bSJacob Faibussowitsch   PetscCall(VecCUPMPlaceArrayAsync<T>(mimpl->cvec, mimpl->ptrinuse + static_cast<std::size_t>(col) * static_cast<std::size_t>(mimpl->lda)));
15854742e46bSJacob Faibussowitsch   if (access == PETSC_MEMORY_ACCESS_READ) PetscCall(VecLockReadPush(mimpl->cvec));
15864742e46bSJacob Faibussowitsch   *v = mimpl->cvec;
15874742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15884742e46bSJacob Faibussowitsch }
15894742e46bSJacob Faibussowitsch 
15904742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
15914742e46bSJacob Faibussowitsch template <PetscMemoryAccessMode access>
15924742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreColumnVec(Mat A, PetscInt, Vec *v) noexcept
15934742e46bSJacob Faibussowitsch {
15944742e46bSJacob Faibussowitsch   using namespace vec::cupm;
15954742e46bSJacob Faibussowitsch   const auto         mimpl = MatIMPLCast(A);
15964742e46bSJacob Faibussowitsch   const auto         cvec  = mimpl->cvec;
15974742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
15984742e46bSJacob Faibussowitsch 
15994742e46bSJacob Faibussowitsch   PetscFunctionBegin;
16004742e46bSJacob Faibussowitsch   PetscCheck(mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
16014742e46bSJacob Faibussowitsch   PetscCheck(cvec, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column vector");
16024742e46bSJacob Faibussowitsch   mimpl->vecinuse = 0;
16034742e46bSJacob Faibussowitsch   if (access == PETSC_MEMORY_ACCESS_READ) PetscCall(VecLockReadPop(cvec));
16044742e46bSJacob Faibussowitsch   PetscCall(VecCUPMResetArrayAsync<T>(cvec));
16054742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx));
16064742e46bSJacob Faibussowitsch   PetscCall(RestoreArray<PETSC_MEMTYPE_DEVICE, access>(A, const_cast<PetscScalar **>(&mimpl->ptrinuse), dctx));
16074742e46bSJacob Faibussowitsch   if (v) *v = nullptr;
16084742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
16094742e46bSJacob Faibussowitsch }
16104742e46bSJacob Faibussowitsch 
16114742e46bSJacob Faibussowitsch // ==========================================================================================
16124742e46bSJacob Faibussowitsch 
16134742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
16144742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::GetFactor(Mat A, MatFactorType ftype, Mat *fact_out) noexcept
16154742e46bSJacob Faibussowitsch {
1616c72d0b4cSJacob Faibussowitsch   Mat                fact = nullptr;
16174742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
16184742e46bSJacob Faibussowitsch 
16194742e46bSJacob Faibussowitsch   PetscFunctionBegin;
16204742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx));
16214742e46bSJacob Faibussowitsch   PetscCall(MatCreateSeqDenseCUPM<T>(PetscObjectComm(PetscObjectCast(A)), A->rmap->n, A->cmap->n, nullptr, &fact, dctx, /* preallocate */ false));
16224742e46bSJacob Faibussowitsch   fact->factortype = ftype;
16234742e46bSJacob Faibussowitsch   switch (ftype) {
16244742e46bSJacob Faibussowitsch   case MAT_FACTOR_LU:
16254742e46bSJacob Faibussowitsch   case MAT_FACTOR_ILU: // fall-through
16264742e46bSJacob Faibussowitsch     fact->ops->lufactorsymbolic  = MatLUFactorSymbolic_SeqDense;
16274742e46bSJacob Faibussowitsch     fact->ops->ilufactorsymbolic = MatLUFactorSymbolic_SeqDense;
16284742e46bSJacob Faibussowitsch     break;
16294742e46bSJacob Faibussowitsch   case MAT_FACTOR_CHOLESKY:
16304742e46bSJacob Faibussowitsch   case MAT_FACTOR_ICC: // fall-through
16314742e46bSJacob Faibussowitsch     fact->ops->choleskyfactorsymbolic = MatCholeskyFactorSymbolic_SeqDense;
16324742e46bSJacob Faibussowitsch     break;
16334742e46bSJacob Faibussowitsch   case MAT_FACTOR_QR: {
16344742e46bSJacob Faibussowitsch     const auto pobj = PetscObjectCast(fact);
16354742e46bSJacob Faibussowitsch 
16364742e46bSJacob Faibussowitsch     PetscCall(PetscObjectComposeFunction(pobj, "MatQRFactor_C", MatQRFactor_SeqDense));
16374742e46bSJacob Faibussowitsch     PetscCall(PetscObjectComposeFunction(pobj, "MatQRFactorSymbolic_C", MatQRFactorSymbolic_SeqDense));
16384742e46bSJacob Faibussowitsch   } break;
16394742e46bSJacob Faibussowitsch   case MAT_FACTOR_NONE:
16404742e46bSJacob Faibussowitsch   case MAT_FACTOR_ILUDT:     // fall-through
16414742e46bSJacob Faibussowitsch   case MAT_FACTOR_NUM_TYPES: // fall-through
16424742e46bSJacob Faibussowitsch     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "MatFactorType %s not supported", MatFactorTypes[ftype]);
16434742e46bSJacob Faibussowitsch   }
16444742e46bSJacob Faibussowitsch   PetscCall(PetscStrFreeAllocpy(MATSOLVERCUPM(), &fact->solvertype));
16454742e46bSJacob Faibussowitsch   PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_LU));
16464742e46bSJacob Faibussowitsch   PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_ILU));
16474742e46bSJacob Faibussowitsch   PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_CHOLESKY));
16484742e46bSJacob Faibussowitsch   PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_ICC));
16494742e46bSJacob Faibussowitsch   *fact_out = fact;
16504742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
16514742e46bSJacob Faibussowitsch }
16524742e46bSJacob Faibussowitsch 
16534742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
16544742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::InvertFactors(Mat A) noexcept
16554742e46bSJacob Faibussowitsch {
16564742e46bSJacob Faibussowitsch   const auto         mimpl = MatIMPLCast(A);
16574742e46bSJacob Faibussowitsch   const auto         mcu   = MatCUPMCast(A);
16584742e46bSJacob Faibussowitsch   const auto         n     = static_cast<cupmBlasInt_t>(A->cmap->n);
16594742e46bSJacob Faibussowitsch   cupmSolverHandle_t handle;
16604742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
16614742e46bSJacob Faibussowitsch   cupmStream_t       stream;
16624742e46bSJacob Faibussowitsch 
16634742e46bSJacob Faibussowitsch   PetscFunctionBegin;
16644742e46bSJacob Faibussowitsch #if PetscDefined(HAVE_CUDA) && PetscDefined(USING_NVCC)
16654742e46bSJacob Faibussowitsch   // HIP appears to have this by default??
16664742e46bSJacob Faibussowitsch   PetscCheck(PETSC_PKG_CUDA_VERSION_GE(10, 1, 0), PETSC_COMM_SELF, PETSC_ERR_SUP, "Upgrade to CUDA version 10.1.0 or higher");
16674742e46bSJacob Faibussowitsch #endif
16684742e46bSJacob Faibussowitsch   if (!n || !A->rmap->n) PetscFunctionReturn(PETSC_SUCCESS);
16694742e46bSJacob Faibussowitsch   PetscCheck(A->factortype == MAT_FACTOR_CHOLESKY, PETSC_COMM_SELF, PETSC_ERR_LIB, "Factor type %s not implemented", MatFactorTypes[A->factortype]);
16704742e46bSJacob Faibussowitsch   // spd
16714742e46bSJacob Faibussowitsch   PetscCheck(!mcu->d_fact_ipiv, PETSC_COMM_SELF, PETSC_ERR_LIB, "%sDnsytri not implemented", cupmSolverName());
16724742e46bSJacob Faibussowitsch 
16734742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx, &handle, &stream));
16744742e46bSJacob Faibussowitsch   {
16754742e46bSJacob Faibussowitsch     const auto    da  = DeviceArrayReadWrite(dctx, A);
16764742e46bSJacob Faibussowitsch     const auto    lda = static_cast<cupmBlasInt_t>(mimpl->lda);
16774742e46bSJacob Faibussowitsch     cupmBlasInt_t il;
16784742e46bSJacob Faibussowitsch 
16794742e46bSJacob Faibussowitsch     PetscCallCUPMSOLVER(cupmSolverXpotri_bufferSize(handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, &il));
16804742e46bSJacob Faibussowitsch     if (il > mcu->d_fact_lwork) {
16814742e46bSJacob Faibussowitsch       mcu->d_fact_lwork = il;
16824742e46bSJacob Faibussowitsch       PetscCallCUPM(cupmFreeAsync(mcu->d_fact_work, stream));
16834742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_work, il, stream));
16844742e46bSJacob Faibussowitsch     }
16854742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeBegin());
16864742e46bSJacob Faibussowitsch     PetscCallCUPMSOLVER(cupmSolverXpotri(handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_info));
16874742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeEnd());
16884742e46bSJacob Faibussowitsch   }
16894742e46bSJacob Faibussowitsch   PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream));
16904742e46bSJacob Faibussowitsch   // TODO (write cuda kernel)
16914742e46bSJacob Faibussowitsch   PetscCall(MatSeqDenseSymmetrize_Private(A, PETSC_TRUE));
16924742e46bSJacob Faibussowitsch   PetscCall(PetscLogGpuFlops(1.0 * n * n * n / 3.0));
16934742e46bSJacob Faibussowitsch 
16944742e46bSJacob Faibussowitsch   A->ops->solve          = nullptr;
16954742e46bSJacob Faibussowitsch   A->ops->solvetranspose = nullptr;
16964742e46bSJacob Faibussowitsch   A->ops->matsolve       = nullptr;
16974742e46bSJacob Faibussowitsch   A->factortype          = MAT_FACTOR_NONE;
16984742e46bSJacob Faibussowitsch 
16994742e46bSJacob Faibussowitsch   PetscCall(PetscFree(A->solvertype));
17004742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
17014742e46bSJacob Faibussowitsch }
17024742e46bSJacob Faibussowitsch 
17034742e46bSJacob Faibussowitsch // ==========================================================================================
17044742e46bSJacob Faibussowitsch 
17054742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
17064742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::GetSubMatrix(Mat A, PetscInt rbegin, PetscInt rend, PetscInt cbegin, PetscInt cend, Mat *mat) noexcept
17074742e46bSJacob Faibussowitsch {
17084742e46bSJacob Faibussowitsch   const auto         mimpl        = MatIMPLCast(A);
17094742e46bSJacob Faibussowitsch   const auto         array_offset = [&](PetscScalar *ptr) { return ptr + rbegin + static_cast<std::size_t>(cbegin) * mimpl->lda; };
17104742e46bSJacob Faibussowitsch   const auto         n            = rend - rbegin;
17114742e46bSJacob Faibussowitsch   const auto         m            = cend - cbegin;
17124742e46bSJacob Faibussowitsch   auto              &cmat         = mimpl->cmat;
17134742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
17144742e46bSJacob Faibussowitsch 
17154742e46bSJacob Faibussowitsch   PetscFunctionBegin;
17164742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
17174742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
17184742e46bSJacob Faibussowitsch   mimpl->matinuse = cbegin + 1;
17194742e46bSJacob Faibussowitsch 
17204742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx));
17214742e46bSJacob Faibussowitsch   PetscCall(HostToDevice_(A, dctx));
17224742e46bSJacob Faibussowitsch 
17234742e46bSJacob Faibussowitsch   if (cmat && ((m != cmat->cmap->N) || (n != cmat->rmap->N))) PetscCall(MatDestroy(&cmat));
17244742e46bSJacob Faibussowitsch   {
17254742e46bSJacob Faibussowitsch     const auto device_array = array_offset(MatCUPMCast(A)->d_v);
17264742e46bSJacob Faibussowitsch 
17274742e46bSJacob Faibussowitsch     if (cmat) {
17284742e46bSJacob Faibussowitsch       PetscCall(PlaceArray(cmat, device_array));
17294742e46bSJacob Faibussowitsch     } else {
17304742e46bSJacob Faibussowitsch       PetscCall(MatCreateSeqDenseCUPM<T>(PetscObjectComm(PetscObjectCast(A)), n, m, device_array, &cmat, dctx));
17314742e46bSJacob Faibussowitsch     }
17324742e46bSJacob Faibussowitsch   }
17334742e46bSJacob Faibussowitsch   PetscCall(MatDenseSetLDA(cmat, mimpl->lda));
17344742e46bSJacob Faibussowitsch   // place CPU array if present but do not copy any data
17354742e46bSJacob Faibussowitsch   if (const auto host_array = mimpl->v) {
17364742e46bSJacob Faibussowitsch     cmat->offloadmask = PETSC_OFFLOAD_GPU;
17374742e46bSJacob Faibussowitsch     PetscCall(MatDensePlaceArray(cmat, array_offset(host_array)));
17384742e46bSJacob Faibussowitsch   }
17394742e46bSJacob Faibussowitsch 
17404742e46bSJacob Faibussowitsch   cmat->offloadmask = A->offloadmask;
17414742e46bSJacob Faibussowitsch   *mat              = cmat;
17424742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
17434742e46bSJacob Faibussowitsch }
17444742e46bSJacob Faibussowitsch 
17454742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
17464742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreSubMatrix(Mat A, Mat *m) noexcept
17474742e46bSJacob Faibussowitsch {
17484742e46bSJacob Faibussowitsch   const auto mimpl = MatIMPLCast(A);
17494742e46bSJacob Faibussowitsch   const auto cmat  = mimpl->cmat;
17504742e46bSJacob Faibussowitsch   const auto reset = static_cast<bool>(mimpl->v);
17514742e46bSJacob Faibussowitsch   bool       copy, was_offload_host;
17524742e46bSJacob Faibussowitsch 
17534742e46bSJacob Faibussowitsch   PetscFunctionBegin;
17544742e46bSJacob Faibussowitsch   PetscCheck(mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetSubMatrix() first");
17554742e46bSJacob Faibussowitsch   PetscCheck(cmat, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column matrix");
17564742e46bSJacob Faibussowitsch   PetscCheck(*m == cmat, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Not the matrix obtained from MatDenseGetSubMatrix()");
17574742e46bSJacob Faibussowitsch   mimpl->matinuse = 0;
17584742e46bSJacob Faibussowitsch 
17594742e46bSJacob Faibussowitsch   // calls to ResetArray may change it, so save it here
17604742e46bSJacob Faibussowitsch   was_offload_host = cmat->offloadmask == PETSC_OFFLOAD_CPU;
17614742e46bSJacob Faibussowitsch   if (was_offload_host && !reset) {
17624742e46bSJacob Faibussowitsch     copy = true;
17634742e46bSJacob Faibussowitsch     PetscCall(MatSeqDenseSetPreallocation(A, nullptr));
17644742e46bSJacob Faibussowitsch   } else {
17654742e46bSJacob Faibussowitsch     copy = false;
17664742e46bSJacob Faibussowitsch   }
17674742e46bSJacob Faibussowitsch 
17684742e46bSJacob Faibussowitsch   PetscCall(ResetArray(cmat));
17694742e46bSJacob Faibussowitsch   if (reset) PetscCall(MatDenseResetArray(cmat));
17704742e46bSJacob Faibussowitsch   if (copy) {
17714742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
17724742e46bSJacob Faibussowitsch 
17734742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
17744742e46bSJacob Faibussowitsch     PetscCall(DeviceToHost_(A, dctx));
17754742e46bSJacob Faibussowitsch   } else {
17764742e46bSJacob Faibussowitsch     A->offloadmask = was_offload_host ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
17774742e46bSJacob Faibussowitsch   }
17784742e46bSJacob Faibussowitsch 
17794742e46bSJacob Faibussowitsch   cmat->offloadmask = PETSC_OFFLOAD_UNALLOCATED;
17804742e46bSJacob Faibussowitsch   *m                = nullptr;
17814742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
17824742e46bSJacob Faibussowitsch }
17834742e46bSJacob Faibussowitsch 
17844742e46bSJacob Faibussowitsch // ==========================================================================================
17854742e46bSJacob Faibussowitsch 
17864742e46bSJacob Faibussowitsch namespace
17874742e46bSJacob Faibussowitsch {
17884742e46bSJacob Faibussowitsch 
17894742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
17904742e46bSJacob Faibussowitsch inline PetscErrorCode MatMatMultNumeric_SeqDenseCUPM_SeqDenseCUPM(Mat A, Mat B, Mat C, PetscBool TA, PetscBool TB) noexcept
17914742e46bSJacob Faibussowitsch {
17924742e46bSJacob Faibussowitsch   PetscFunctionBegin;
17934742e46bSJacob Faibussowitsch   if (TA) {
17944742e46bSJacob Faibussowitsch     if (TB) {
17954742e46bSJacob Faibussowitsch       PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<true, true>(A, B, C));
17964742e46bSJacob Faibussowitsch     } else {
17974742e46bSJacob Faibussowitsch       PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<true, false>(A, B, C));
17984742e46bSJacob Faibussowitsch     }
17994742e46bSJacob Faibussowitsch   } else {
18004742e46bSJacob Faibussowitsch     if (TB) {
18014742e46bSJacob Faibussowitsch       PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<false, true>(A, B, C));
18024742e46bSJacob Faibussowitsch     } else {
18034742e46bSJacob Faibussowitsch       PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<false, false>(A, B, C));
18044742e46bSJacob Faibussowitsch     }
18054742e46bSJacob Faibussowitsch   }
18064742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
18074742e46bSJacob Faibussowitsch }
18084742e46bSJacob Faibussowitsch 
18094742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
18104742e46bSJacob Faibussowitsch inline PetscErrorCode MatSolverTypeRegister_DENSECUPM() noexcept
18114742e46bSJacob Faibussowitsch {
18124742e46bSJacob Faibussowitsch   PetscFunctionBegin;
18134742e46bSJacob Faibussowitsch   for (auto ftype : util::make_array(MAT_FACTOR_LU, MAT_FACTOR_CHOLESKY, MAT_FACTOR_QR)) {
18144742e46bSJacob Faibussowitsch     PetscCall(MatSolverTypeRegister(MatDense_Seq_CUPM<T>::MATSOLVERCUPM(), MATSEQDENSE, ftype, MatDense_Seq_CUPM<T>::GetFactor));
18154742e46bSJacob Faibussowitsch     PetscCall(MatSolverTypeRegister(MatDense_Seq_CUPM<T>::MATSOLVERCUPM(), MatDense_Seq_CUPM<T>::MATSEQDENSECUPM(), ftype, MatDense_Seq_CUPM<T>::GetFactor));
18164742e46bSJacob Faibussowitsch   }
18174742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
18184742e46bSJacob Faibussowitsch }
18194742e46bSJacob Faibussowitsch 
18204742e46bSJacob Faibussowitsch } // anonymous namespace
18214742e46bSJacob Faibussowitsch 
18224742e46bSJacob Faibussowitsch } // namespace impl
18234742e46bSJacob Faibussowitsch 
18244742e46bSJacob Faibussowitsch } // namespace cupm
18254742e46bSJacob Faibussowitsch 
18264742e46bSJacob Faibussowitsch } // namespace mat
18274742e46bSJacob Faibussowitsch 
18284742e46bSJacob Faibussowitsch } // namespace Petsc
1829