xref: /petsc/src/mat/impls/dense/seq/cupm/matseqdensecupm.hpp (revision 2ea277ce7ccf004135b03702db4a237d7886e1ba)
14742e46bSJacob Faibussowitsch #ifndef PETSCMATSEQDENSECUPM_HPP
24742e46bSJacob Faibussowitsch #define PETSCMATSEQDENSECUPM_HPP
34742e46bSJacob Faibussowitsch 
44742e46bSJacob Faibussowitsch #include <petsc/private/matdensecupmimpl.h> /*I <petscmat.h> I*/
54742e46bSJacob Faibussowitsch #include <../src/mat/impls/dense/seq/dense.h>
64742e46bSJacob Faibussowitsch 
74742e46bSJacob Faibussowitsch #if defined(__cplusplus)
84742e46bSJacob Faibussowitsch   #include <petsc/private/deviceimpl.h> // PetscDeviceContextGetOptionalNullContext_Internal()
94742e46bSJacob Faibussowitsch   #include <petsc/private/randomimpl.h> // _p_PetscRandom
104742e46bSJacob Faibussowitsch   #include <petsc/private/vecimpl.h>    // _p_Vec
114742e46bSJacob Faibussowitsch   #include <petsc/private/cupmobject.hpp>
124742e46bSJacob Faibussowitsch   #include <petsc/private/cupmsolverinterface.hpp>
134742e46bSJacob Faibussowitsch 
144742e46bSJacob Faibussowitsch   #include <petsc/private/cpp/type_traits.hpp> // PetscObjectCast()
154742e46bSJacob Faibussowitsch   #include <petsc/private/cpp/utility.hpp>     // util::exchange()
164742e46bSJacob Faibussowitsch 
174742e46bSJacob Faibussowitsch   #include <../src/vec/vec/impls/seq/cupm/vecseqcupm.hpp> // for VecSeq_CUPM
184742e46bSJacob Faibussowitsch 
194742e46bSJacob Faibussowitsch namespace Petsc
204742e46bSJacob Faibussowitsch {
214742e46bSJacob Faibussowitsch 
224742e46bSJacob Faibussowitsch namespace mat
234742e46bSJacob Faibussowitsch {
244742e46bSJacob Faibussowitsch 
254742e46bSJacob Faibussowitsch namespace cupm
264742e46bSJacob Faibussowitsch {
274742e46bSJacob Faibussowitsch 
284742e46bSJacob Faibussowitsch namespace impl
294742e46bSJacob Faibussowitsch {
304742e46bSJacob Faibussowitsch 
314742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
324742e46bSJacob Faibussowitsch class MatDense_Seq_CUPM : MatDense_CUPM<T, MatDense_Seq_CUPM<T>> {
334742e46bSJacob Faibussowitsch public:
344742e46bSJacob Faibussowitsch   MATDENSECUPM_HEADER(T, MatDense_Seq_CUPM<T>);
354742e46bSJacob Faibussowitsch 
364742e46bSJacob Faibussowitsch private:
374742e46bSJacob Faibussowitsch   struct Mat_SeqDenseCUPM {
384742e46bSJacob Faibussowitsch     PetscScalar *d_v;           // pointer to the matrix on the GPU
394742e46bSJacob Faibussowitsch     PetscScalar *unplacedarray; // if one called MatCUPMDensePlaceArray(), this is where it stashed the original
404742e46bSJacob Faibussowitsch     bool         d_user_alloc;
414742e46bSJacob Faibussowitsch     bool         d_unplaced_user_alloc;
424742e46bSJacob Faibussowitsch     // factorization support
434742e46bSJacob Faibussowitsch     cupmBlasInt_t *d_fact_ipiv;  // device pivots
444742e46bSJacob Faibussowitsch     cupmScalar_t  *d_fact_tau;   // device QR tau vector
454742e46bSJacob Faibussowitsch     cupmBlasInt_t *d_fact_info;  // device info
464742e46bSJacob Faibussowitsch     cupmScalar_t  *d_fact_work;  // device workspace
474742e46bSJacob Faibussowitsch     cupmBlasInt_t  d_fact_lwork; // size of device workspace
484742e46bSJacob Faibussowitsch     // workspace
494742e46bSJacob Faibussowitsch     Vec workvec;
504742e46bSJacob Faibussowitsch   };
514742e46bSJacob Faibussowitsch 
524742e46bSJacob Faibussowitsch   static PetscErrorCode SetPreallocation_(Mat, PetscDeviceContext, PetscScalar *) noexcept;
534742e46bSJacob Faibussowitsch 
544742e46bSJacob Faibussowitsch   static PetscErrorCode HostToDevice_(Mat, PetscDeviceContext) noexcept;
554742e46bSJacob Faibussowitsch   static PetscErrorCode DeviceToHost_(Mat, PetscDeviceContext) noexcept;
564742e46bSJacob Faibussowitsch 
574742e46bSJacob Faibussowitsch   static PetscErrorCode CheckCUPMSolverInfo_(const cupmBlasInt_t *, cupmStream_t) noexcept;
584742e46bSJacob Faibussowitsch 
594742e46bSJacob Faibussowitsch   template <typename Derived>
604742e46bSJacob Faibussowitsch   struct SolveCommon;
614742e46bSJacob Faibussowitsch   struct SolveQR;
624742e46bSJacob Faibussowitsch   struct SolveCholesky;
634742e46bSJacob Faibussowitsch   struct SolveLU;
644742e46bSJacob Faibussowitsch 
654742e46bSJacob Faibussowitsch   template <typename Solver, bool transpose>
664742e46bSJacob Faibussowitsch   static PetscErrorCode MatSolve_Factored_Dispatch_(Mat, Vec, Vec) noexcept;
674742e46bSJacob Faibussowitsch   template <typename Solver, bool transpose>
684742e46bSJacob Faibussowitsch   static PetscErrorCode MatMatSolve_Factored_Dispatch_(Mat, Mat, Mat) noexcept;
694742e46bSJacob Faibussowitsch   template <bool transpose>
704742e46bSJacob Faibussowitsch   static PetscErrorCode MatMultAdd_Dispatch_(Mat, Vec, Vec, Vec) noexcept;
714742e46bSJacob Faibussowitsch 
724742e46bSJacob Faibussowitsch   template <bool to_host>
734742e46bSJacob Faibussowitsch   static PetscErrorCode Convert_Dispatch_(Mat, MatType, MatReuse, Mat *) noexcept;
744742e46bSJacob Faibussowitsch 
754742e46bSJacob Faibussowitsch   PETSC_NODISCARD static constexpr MatType       MATIMPLCUPM_() noexcept;
764742e46bSJacob Faibussowitsch   PETSC_NODISCARD static constexpr Mat_SeqDense *MatIMPLCast_(Mat) noexcept;
774742e46bSJacob Faibussowitsch 
784742e46bSJacob Faibussowitsch public:
794742e46bSJacob Faibussowitsch   PETSC_NODISCARD static constexpr Mat_SeqDenseCUPM *MatCUPMCast(Mat) noexcept;
804742e46bSJacob Faibussowitsch 
814742e46bSJacob Faibussowitsch   // define these by hand since they don't fit the above mold
824742e46bSJacob Faibussowitsch   PETSC_NODISCARD static constexpr const char *MatConvert_seqdensecupm_seqdense_C() noexcept;
834742e46bSJacob Faibussowitsch   PETSC_NODISCARD static constexpr const char *MatProductSetFromOptions_seqaij_seqdensecupm_C() noexcept;
844742e46bSJacob Faibussowitsch 
854742e46bSJacob Faibussowitsch   static PetscErrorCode Create(Mat) noexcept;
864742e46bSJacob Faibussowitsch   static PetscErrorCode Destroy(Mat) noexcept;
874742e46bSJacob Faibussowitsch   static PetscErrorCode SetUp(Mat) noexcept;
884742e46bSJacob Faibussowitsch   static PetscErrorCode Reset(Mat) noexcept;
894742e46bSJacob Faibussowitsch 
904742e46bSJacob Faibussowitsch   static PetscErrorCode BindToCPU(Mat, PetscBool) noexcept;
914742e46bSJacob Faibussowitsch   static PetscErrorCode Convert_SeqDense_SeqDenseCUPM(Mat, MatType, MatReuse, Mat *) noexcept;
924742e46bSJacob Faibussowitsch   static PetscErrorCode Convert_SeqDenseCUPM_SeqDense(Mat, MatType, MatReuse, Mat *) noexcept;
934742e46bSJacob Faibussowitsch 
944742e46bSJacob Faibussowitsch   template <PetscMemType, PetscMemoryAccessMode>
954742e46bSJacob Faibussowitsch   static PetscErrorCode GetArray(Mat, PetscScalar **, PetscDeviceContext) noexcept;
964742e46bSJacob Faibussowitsch   template <PetscMemType, PetscMemoryAccessMode>
974742e46bSJacob Faibussowitsch   static PetscErrorCode RestoreArray(Mat, PetscScalar **, PetscDeviceContext) noexcept;
984742e46bSJacob Faibussowitsch   template <PetscMemoryAccessMode>
994742e46bSJacob Faibussowitsch   static PetscErrorCode GetArrayAndMemType(Mat, PetscScalar **, PetscMemType *, PetscDeviceContext) noexcept;
1004742e46bSJacob Faibussowitsch   template <PetscMemoryAccessMode>
1014742e46bSJacob Faibussowitsch   static PetscErrorCode RestoreArrayAndMemType(Mat, PetscScalar **, PetscDeviceContext) noexcept;
1024742e46bSJacob Faibussowitsch 
1034742e46bSJacob Faibussowitsch private:
1044742e46bSJacob Faibussowitsch   template <PetscMemType mtype, PetscMemoryAccessMode mode>
1054742e46bSJacob Faibussowitsch   static PetscErrorCode GetArrayC_(Mat m, PetscScalar **p) noexcept
1064742e46bSJacob Faibussowitsch   {
1074742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
1084742e46bSJacob Faibussowitsch 
1094742e46bSJacob Faibussowitsch     PetscFunctionBegin;
1104742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
1114742e46bSJacob Faibussowitsch     PetscCall(GetArray<mtype, mode>(m, p, dctx));
1124742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
1134742e46bSJacob Faibussowitsch   }
1144742e46bSJacob Faibussowitsch 
1154742e46bSJacob Faibussowitsch   template <PetscMemType mtype, PetscMemoryAccessMode mode>
1164742e46bSJacob Faibussowitsch   static PetscErrorCode RestoreArrayC_(Mat m, PetscScalar **p) noexcept
1174742e46bSJacob Faibussowitsch   {
1184742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
1194742e46bSJacob Faibussowitsch 
1204742e46bSJacob Faibussowitsch     PetscFunctionBegin;
1214742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
1224742e46bSJacob Faibussowitsch     PetscCall(RestoreArray<mtype, mode>(m, p, dctx));
1234742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
1244742e46bSJacob Faibussowitsch   }
1254742e46bSJacob Faibussowitsch 
1264742e46bSJacob Faibussowitsch   template <PetscMemoryAccessMode mode>
1274742e46bSJacob Faibussowitsch   static PetscErrorCode GetArrayAndMemTypeC_(Mat m, PetscScalar **p, PetscMemType *tp) noexcept
1284742e46bSJacob Faibussowitsch   {
1294742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
1304742e46bSJacob Faibussowitsch 
1314742e46bSJacob Faibussowitsch     PetscFunctionBegin;
1324742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
1334742e46bSJacob Faibussowitsch     PetscCall(GetArrayAndMemType<mode>(m, p, tp, dctx));
1344742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
1354742e46bSJacob Faibussowitsch   }
1364742e46bSJacob Faibussowitsch 
1374742e46bSJacob Faibussowitsch   template <PetscMemoryAccessMode mode>
1384742e46bSJacob Faibussowitsch   static PetscErrorCode RestoreArrayAndMemTypeC_(Mat m, PetscScalar **p) noexcept
1394742e46bSJacob Faibussowitsch   {
1404742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
1414742e46bSJacob Faibussowitsch 
1424742e46bSJacob Faibussowitsch     PetscFunctionBegin;
1434742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
1444742e46bSJacob Faibussowitsch     PetscCall(RestoreArrayAndMemType<mode>(m, p, dctx));
1454742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
1464742e46bSJacob Faibussowitsch   }
1474742e46bSJacob Faibussowitsch 
1484742e46bSJacob Faibussowitsch public:
1494742e46bSJacob Faibussowitsch   static PetscErrorCode PlaceArray(Mat, const PetscScalar *) noexcept;
1504742e46bSJacob Faibussowitsch   static PetscErrorCode ReplaceArray(Mat, const PetscScalar *) noexcept;
1514742e46bSJacob Faibussowitsch   static PetscErrorCode ResetArray(Mat) noexcept;
1524742e46bSJacob Faibussowitsch 
1534742e46bSJacob Faibussowitsch   template <bool transpose_A, bool transpose_B>
1544742e46bSJacob Faibussowitsch   static PetscErrorCode MatMatMult_Numeric_Dispatch(Mat, Mat, Mat) noexcept;
1554742e46bSJacob Faibussowitsch   static PetscErrorCode Copy(Mat, Mat, MatStructure) noexcept;
1564742e46bSJacob Faibussowitsch   static PetscErrorCode ZeroEntries(Mat) noexcept;
1574742e46bSJacob Faibussowitsch   static PetscErrorCode Scale(Mat, PetscScalar) noexcept;
1584742e46bSJacob Faibussowitsch   static PetscErrorCode Shift(Mat, PetscScalar) noexcept;
1594742e46bSJacob Faibussowitsch   static PetscErrorCode AXPY(Mat, PetscScalar, Mat, MatStructure) noexcept;
1604742e46bSJacob Faibussowitsch   static PetscErrorCode Duplicate(Mat, MatDuplicateOption, Mat *) noexcept;
1614742e46bSJacob Faibussowitsch   static PetscErrorCode SetRandom(Mat, PetscRandom) noexcept;
1624742e46bSJacob Faibussowitsch 
1634742e46bSJacob Faibussowitsch   static PetscErrorCode GetColumnVector(Mat, Vec, PetscInt) noexcept;
1644742e46bSJacob Faibussowitsch   template <PetscMemoryAccessMode>
1654742e46bSJacob Faibussowitsch   static PetscErrorCode GetColumnVec(Mat, PetscInt, Vec *) noexcept;
1664742e46bSJacob Faibussowitsch   template <PetscMemoryAccessMode>
1674742e46bSJacob Faibussowitsch   static PetscErrorCode RestoreColumnVec(Mat, PetscInt, Vec *) noexcept;
1684742e46bSJacob Faibussowitsch 
1694742e46bSJacob Faibussowitsch   static PetscErrorCode GetFactor(Mat, MatFactorType, Mat *) noexcept;
1704742e46bSJacob Faibussowitsch   static PetscErrorCode InvertFactors(Mat) noexcept;
1714742e46bSJacob Faibussowitsch 
1724742e46bSJacob Faibussowitsch   static PetscErrorCode GetSubMatrix(Mat, PetscInt, PetscInt, PetscInt, PetscInt, Mat *) noexcept;
1734742e46bSJacob Faibussowitsch   static PetscErrorCode RestoreSubMatrix(Mat, Mat *) noexcept;
1744742e46bSJacob Faibussowitsch };
1754742e46bSJacob Faibussowitsch 
1764742e46bSJacob Faibussowitsch } // namespace impl
1774742e46bSJacob Faibussowitsch 
1784742e46bSJacob Faibussowitsch namespace
1794742e46bSJacob Faibussowitsch {
1804742e46bSJacob Faibussowitsch 
1814742e46bSJacob Faibussowitsch // Declare this here so that the functions below can make use of it
1824742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
1834742e46bSJacob Faibussowitsch inline PetscErrorCode MatCreateSeqDenseCUPM(MPI_Comm comm, PetscInt m, PetscInt n, PetscScalar *data, Mat *A, PetscDeviceContext dctx = nullptr, bool preallocate = true) noexcept
1844742e46bSJacob Faibussowitsch {
1854742e46bSJacob Faibussowitsch   PetscFunctionBegin;
1864742e46bSJacob Faibussowitsch   PetscCall(impl::MatDense_Seq_CUPM<T>::CreateIMPLDenseCUPM(comm, m, n, m, n, data, A, dctx, preallocate));
1874742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1884742e46bSJacob Faibussowitsch }
1894742e46bSJacob Faibussowitsch 
1904742e46bSJacob Faibussowitsch } // anonymous namespace
1914742e46bSJacob Faibussowitsch 
1924742e46bSJacob Faibussowitsch namespace impl
1934742e46bSJacob Faibussowitsch {
1944742e46bSJacob Faibussowitsch 
1954742e46bSJacob Faibussowitsch // ==========================================================================================
1964742e46bSJacob Faibussowitsch // MatDense_Seq_CUPM - Private API - Utility
1974742e46bSJacob Faibussowitsch // ==========================================================================================
1984742e46bSJacob Faibussowitsch 
1994742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
2004742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::SetPreallocation_(Mat m, PetscDeviceContext dctx, PetscScalar *user_device_array) noexcept
2014742e46bSJacob Faibussowitsch {
2024742e46bSJacob Faibussowitsch   const auto   mcu   = MatCUPMCast(m);
2034742e46bSJacob Faibussowitsch   const auto   nrows = m->rmap->n;
2044742e46bSJacob Faibussowitsch   const auto   ncols = m->cmap->n;
2054742e46bSJacob Faibussowitsch   auto        &lda   = MatIMPLCast(m)->lda;
2064742e46bSJacob Faibussowitsch   cupmStream_t stream;
2074742e46bSJacob Faibussowitsch 
2084742e46bSJacob Faibussowitsch   PetscFunctionBegin;
2094742e46bSJacob Faibussowitsch   PetscCheckTypeName(m, MATSEQDENSECUPM());
2104742e46bSJacob Faibussowitsch   PetscValidDeviceContext(dctx, 2);
2114742e46bSJacob Faibussowitsch   PetscCall(checkCupmBlasIntCast(nrows));
2124742e46bSJacob Faibussowitsch   PetscCall(checkCupmBlasIntCast(ncols));
2134742e46bSJacob Faibussowitsch   PetscCall(GetHandlesFrom_(dctx, &stream));
2144742e46bSJacob Faibussowitsch   if (lda <= 0) lda = nrows;
2154742e46bSJacob Faibussowitsch   if (!mcu->d_user_alloc) PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream));
2164742e46bSJacob Faibussowitsch   if (user_device_array) {
2174742e46bSJacob Faibussowitsch     mcu->d_user_alloc = PETSC_TRUE;
2184742e46bSJacob Faibussowitsch     mcu->d_v          = user_device_array;
2194742e46bSJacob Faibussowitsch   } else {
2204742e46bSJacob Faibussowitsch     PetscInt size;
2214742e46bSJacob Faibussowitsch 
2224742e46bSJacob Faibussowitsch     mcu->d_user_alloc = PETSC_FALSE;
2234742e46bSJacob Faibussowitsch     PetscCall(PetscIntMultError(lda, ncols, &size));
2244742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMMallocAsync(&mcu->d_v, size, stream));
2254742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMMemsetAsync(mcu->d_v, 0, size, stream));
2264742e46bSJacob Faibussowitsch   }
2274742e46bSJacob Faibussowitsch   m->offloadmask = PETSC_OFFLOAD_GPU;
2284742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2294742e46bSJacob Faibussowitsch }
2304742e46bSJacob Faibussowitsch 
2314742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
2324742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::HostToDevice_(Mat m, PetscDeviceContext dctx) noexcept
2334742e46bSJacob Faibussowitsch {
2344742e46bSJacob Faibussowitsch   const auto nrows = m->rmap->n;
2354742e46bSJacob Faibussowitsch   const auto ncols = m->cmap->n;
2364742e46bSJacob Faibussowitsch   const auto copy  = m->offloadmask == PETSC_OFFLOAD_CPU || m->offloadmask == PETSC_OFFLOAD_UNALLOCATED;
2374742e46bSJacob Faibussowitsch 
2384742e46bSJacob Faibussowitsch   PetscFunctionBegin;
2394742e46bSJacob Faibussowitsch   PetscCheckTypeName(m, MATSEQDENSECUPM());
2404742e46bSJacob Faibussowitsch   if (m->boundtocpu) PetscFunctionReturn(PETSC_SUCCESS);
2414742e46bSJacob Faibussowitsch   PetscCall(PetscInfo(m, "%s matrix %" PetscInt_FMT " x %" PetscInt_FMT "\n", copy ? "Copy" : "Reusing", nrows, ncols));
2424742e46bSJacob Faibussowitsch   if (copy) {
2434742e46bSJacob Faibussowitsch     const auto   mcu = MatCUPMCast(m);
2444742e46bSJacob Faibussowitsch     cupmStream_t stream;
2454742e46bSJacob Faibussowitsch 
2464742e46bSJacob Faibussowitsch     // Allocate GPU memory if not present
2474742e46bSJacob Faibussowitsch     if (!mcu->d_v) PetscCall(SetPreallocation(m, dctx));
2484742e46bSJacob Faibussowitsch     PetscCall(GetHandlesFrom_(dctx, &stream));
2494742e46bSJacob Faibussowitsch     PetscCall(PetscLogEventBegin(MAT_DenseCopyToGPU, m, 0, 0, 0));
2504742e46bSJacob Faibussowitsch     {
2514742e46bSJacob Faibussowitsch       const auto mimpl = MatIMPLCast(m);
2524742e46bSJacob Faibussowitsch       const auto lda   = mimpl->lda;
2534742e46bSJacob Faibussowitsch       const auto src   = mimpl->v;
2544742e46bSJacob Faibussowitsch       const auto dest  = mcu->d_v;
2554742e46bSJacob Faibussowitsch 
2564742e46bSJacob Faibussowitsch       if (lda > nrows) {
2574742e46bSJacob Faibussowitsch         PetscCall(PetscCUPMMemcpy2DAsync(dest, lda, src, lda, nrows, ncols, cupmMemcpyHostToDevice, stream));
2584742e46bSJacob Faibussowitsch       } else {
2594742e46bSJacob Faibussowitsch         PetscCall(PetscCUPMMemcpyAsync(dest, src, lda * ncols, cupmMemcpyHostToDevice, stream));
2604742e46bSJacob Faibussowitsch       }
2614742e46bSJacob Faibussowitsch     }
2624742e46bSJacob Faibussowitsch     PetscCall(PetscLogEventEnd(MAT_DenseCopyToGPU, m, 0, 0, 0));
2634742e46bSJacob Faibussowitsch     // order important, ensure that offloadmask is PETSC_OFFLOAD_BOTH
2644742e46bSJacob Faibussowitsch     m->offloadmask = PETSC_OFFLOAD_BOTH;
2654742e46bSJacob Faibussowitsch   }
2664742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2674742e46bSJacob Faibussowitsch }
2684742e46bSJacob Faibussowitsch 
2694742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
2704742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::DeviceToHost_(Mat m, PetscDeviceContext dctx) noexcept
2714742e46bSJacob Faibussowitsch {
2724742e46bSJacob Faibussowitsch   const auto nrows = m->rmap->n;
2734742e46bSJacob Faibussowitsch   const auto ncols = m->cmap->n;
2744742e46bSJacob Faibussowitsch   const auto copy  = m->offloadmask == PETSC_OFFLOAD_GPU;
2754742e46bSJacob Faibussowitsch 
2764742e46bSJacob Faibussowitsch   PetscFunctionBegin;
2774742e46bSJacob Faibussowitsch   PetscCheckTypeName(m, MATSEQDENSECUPM());
2784742e46bSJacob Faibussowitsch   PetscCall(PetscInfo(m, "%s matrix %" PetscInt_FMT " x %" PetscInt_FMT "\n", copy ? "Copy" : "Reusing", nrows, ncols));
2794742e46bSJacob Faibussowitsch   if (copy) {
2804742e46bSJacob Faibussowitsch     const auto   mimpl = MatIMPLCast(m);
2814742e46bSJacob Faibussowitsch     cupmStream_t stream;
2824742e46bSJacob Faibussowitsch 
2834742e46bSJacob Faibussowitsch     // MatCreateSeqDenseCUPM may not allocate CPU memory. Allocate if needed
2844742e46bSJacob Faibussowitsch     if (!mimpl->v) PetscCall(MatSeqDenseSetPreallocation(m, nullptr));
2854742e46bSJacob Faibussowitsch     PetscCall(GetHandlesFrom_(dctx, &stream));
2864742e46bSJacob Faibussowitsch     PetscCall(PetscLogEventBegin(MAT_DenseCopyFromGPU, m, 0, 0, 0));
2874742e46bSJacob Faibussowitsch     {
2884742e46bSJacob Faibussowitsch       const auto lda  = mimpl->lda;
2894742e46bSJacob Faibussowitsch       const auto dest = mimpl->v;
2904742e46bSJacob Faibussowitsch       const auto src  = MatCUPMCast(m)->d_v;
2914742e46bSJacob Faibussowitsch 
2924742e46bSJacob Faibussowitsch       if (lda > nrows) {
2934742e46bSJacob Faibussowitsch         PetscCall(PetscCUPMMemcpy2DAsync(dest, lda, src, lda, nrows, ncols, cupmMemcpyDeviceToHost, stream));
2944742e46bSJacob Faibussowitsch       } else {
2954742e46bSJacob Faibussowitsch         PetscCall(PetscCUPMMemcpyAsync(dest, src, lda * ncols, cupmMemcpyDeviceToHost, stream));
2964742e46bSJacob Faibussowitsch       }
2974742e46bSJacob Faibussowitsch     }
2984742e46bSJacob Faibussowitsch     PetscCall(PetscLogEventEnd(MAT_DenseCopyFromGPU, m, 0, 0, 0));
2994742e46bSJacob Faibussowitsch     // order is important, MatSeqDenseSetPreallocation() might set offloadmask
3004742e46bSJacob Faibussowitsch     m->offloadmask = PETSC_OFFLOAD_BOTH;
3014742e46bSJacob Faibussowitsch   }
3024742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3034742e46bSJacob Faibussowitsch }
3044742e46bSJacob Faibussowitsch 
3054742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
3064742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::CheckCUPMSolverInfo_(const cupmBlasInt_t *fact_info, cupmStream_t stream) noexcept
3074742e46bSJacob Faibussowitsch {
3084742e46bSJacob Faibussowitsch   PetscFunctionBegin;
3094742e46bSJacob Faibussowitsch   if (PetscDefined(USE_DEBUG)) {
3104742e46bSJacob Faibussowitsch     cupmBlasInt_t info = 0;
3114742e46bSJacob Faibussowitsch 
3124742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMMemcpyAsync(&info, fact_info, 1, cupmMemcpyDeviceToHost, stream));
3134742e46bSJacob Faibussowitsch     if (stream) PetscCallCUPM(cupmStreamSynchronize(stream));
3144742e46bSJacob Faibussowitsch     static_assert(std::is_same<decltype(info), int>::value, "");
3154742e46bSJacob Faibussowitsch     PetscCheck(info <= 0, PETSC_COMM_SELF, PETSC_ERR_MAT_CH_ZRPVT, "Bad factorization: zero pivot in row %d", info - 1);
3164742e46bSJacob Faibussowitsch     PetscCheck(info >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Wrong argument to cupmSolver %d", -info);
3174742e46bSJacob Faibussowitsch   }
3184742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3194742e46bSJacob Faibussowitsch }
3204742e46bSJacob Faibussowitsch 
3214742e46bSJacob Faibussowitsch // ==========================================================================================
3224742e46bSJacob Faibussowitsch // MatDense_Seq_CUPM - Private API - Solver Dispatch
3234742e46bSJacob Faibussowitsch // ==========================================================================================
3244742e46bSJacob Faibussowitsch 
3254742e46bSJacob Faibussowitsch // specific solvers called through the dispatch_() family of functions
3264742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
3274742e46bSJacob Faibussowitsch template <typename Derived>
3284742e46bSJacob Faibussowitsch struct MatDense_Seq_CUPM<T>::SolveCommon {
3294742e46bSJacob Faibussowitsch   using derived_type = Derived;
3304742e46bSJacob Faibussowitsch 
3314742e46bSJacob Faibussowitsch   template <typename F>
3324742e46bSJacob Faibussowitsch   static PetscErrorCode ResizeFactLwork(Mat_SeqDenseCUPM *mcu, cupmStream_t stream, F &&cupmSolverComputeFactLwork) noexcept
3334742e46bSJacob Faibussowitsch   {
3344742e46bSJacob Faibussowitsch     cupmBlasInt_t lwork;
3354742e46bSJacob Faibussowitsch 
3364742e46bSJacob Faibussowitsch     PetscFunctionBegin;
3374742e46bSJacob Faibussowitsch     PetscCallCUPMSOLVER(cupmSolverComputeFactLwork(&lwork));
3384742e46bSJacob Faibussowitsch     if (lwork > mcu->d_fact_lwork) {
3394742e46bSJacob Faibussowitsch       mcu->d_fact_lwork = lwork;
3404742e46bSJacob Faibussowitsch       PetscCallCUPM(cupmFreeAsync(mcu->d_fact_work, stream));
3414742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_work, lwork, stream));
3424742e46bSJacob Faibussowitsch     }
3434742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
3444742e46bSJacob Faibussowitsch   }
3454742e46bSJacob Faibussowitsch 
3464742e46bSJacob Faibussowitsch   static PetscErrorCode FactorPrepare(Mat A, cupmStream_t stream) noexcept
3474742e46bSJacob Faibussowitsch   {
3484742e46bSJacob Faibussowitsch     const auto mcu = MatCUPMCast(A);
3494742e46bSJacob Faibussowitsch 
3504742e46bSJacob Faibussowitsch     PetscFunctionBegin;
3514742e46bSJacob Faibussowitsch     PetscCall(PetscInfo(A, "%s factor %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", derived_type::NAME(), A->rmap->n, A->cmap->n));
3524742e46bSJacob Faibussowitsch     A->factortype             = derived_type::MATFACTORTYPE();
3534742e46bSJacob Faibussowitsch     A->ops->solve             = MatSolve_Factored_Dispatch_<derived_type, false>;
3544742e46bSJacob Faibussowitsch     A->ops->solvetranspose    = MatSolve_Factored_Dispatch_<derived_type, true>;
3554742e46bSJacob Faibussowitsch     A->ops->matsolve          = MatMatSolve_Factored_Dispatch_<derived_type, false>;
3564742e46bSJacob Faibussowitsch     A->ops->matsolvetranspose = MatMatSolve_Factored_Dispatch_<derived_type, true>;
3574742e46bSJacob Faibussowitsch 
3584742e46bSJacob Faibussowitsch     PetscCall(PetscStrFreeAllocpy(MATSOLVERCUPM(), &A->solvertype));
3594742e46bSJacob Faibussowitsch     if (!mcu->d_fact_info) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_info, 1, stream));
3604742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
3614742e46bSJacob Faibussowitsch   }
3624742e46bSJacob Faibussowitsch };
3634742e46bSJacob Faibussowitsch 
3644742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
3654742e46bSJacob Faibussowitsch struct MatDense_Seq_CUPM<T>::SolveLU : SolveCommon<SolveLU> {
3664742e46bSJacob Faibussowitsch   using base_type = SolveCommon<SolveLU>;
3674742e46bSJacob Faibussowitsch 
3684742e46bSJacob Faibussowitsch   static constexpr const char   *NAME() noexcept { return "LU"; }
3694742e46bSJacob Faibussowitsch   static constexpr MatFactorType MATFACTORTYPE() noexcept { return MAT_FACTOR_LU; }
3704742e46bSJacob Faibussowitsch 
3714742e46bSJacob Faibussowitsch   static PetscErrorCode Factor(Mat A, IS, IS, const MatFactorInfo *) noexcept
3724742e46bSJacob Faibussowitsch   {
3734742e46bSJacob Faibussowitsch     const auto         m = static_cast<cupmBlasInt_t>(A->rmap->n);
3744742e46bSJacob Faibussowitsch     const auto         n = static_cast<cupmBlasInt_t>(A->cmap->n);
3754742e46bSJacob Faibussowitsch     cupmStream_t       stream;
3764742e46bSJacob Faibussowitsch     cupmSolverHandle_t handle;
3774742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
3784742e46bSJacob Faibussowitsch 
3794742e46bSJacob Faibussowitsch     PetscFunctionBegin;
3804742e46bSJacob Faibussowitsch     if (!m || !n) PetscFunctionReturn(PETSC_SUCCESS);
3814742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx, &handle, &stream));
3824742e46bSJacob Faibussowitsch     PetscCall(base_type::FactorPrepare(A, stream));
3834742e46bSJacob Faibussowitsch     {
3844742e46bSJacob Faibussowitsch       const auto mcu = MatCUPMCast(A);
3854742e46bSJacob Faibussowitsch       const auto da  = DeviceArrayReadWrite(dctx, A);
3864742e46bSJacob Faibussowitsch       const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);
3874742e46bSJacob Faibussowitsch 
3884742e46bSJacob Faibussowitsch       // clang-format off
3894742e46bSJacob Faibussowitsch       PetscCall(
3904742e46bSJacob Faibussowitsch         base_type::ResizeFactLwork(
3914742e46bSJacob Faibussowitsch           mcu, stream,
3924742e46bSJacob Faibussowitsch           [&](cupmBlasInt_t *fact_lwork)
3934742e46bSJacob Faibussowitsch           {
3944742e46bSJacob Faibussowitsch             return cupmSolverXgetrf_bufferSize(handle, m, n, da.cupmdata(), lda, fact_lwork);
3954742e46bSJacob Faibussowitsch           }
3964742e46bSJacob Faibussowitsch         )
3974742e46bSJacob Faibussowitsch       );
3984742e46bSJacob Faibussowitsch       // clang-format on
3994742e46bSJacob Faibussowitsch       if (!mcu->d_fact_ipiv) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_ipiv, n, stream));
4004742e46bSJacob Faibussowitsch 
4014742e46bSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeBegin());
4024742e46bSJacob 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));
4034742e46bSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeEnd());
4044742e46bSJacob Faibussowitsch       PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream));
4054742e46bSJacob Faibussowitsch     }
4064742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuFlops(2.0 * n * n * m / 3.0));
4074742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
4084742e46bSJacob Faibussowitsch   }
4094742e46bSJacob Faibussowitsch 
4104742e46bSJacob Faibussowitsch   template <bool transpose>
4114742e46bSJacob 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
4124742e46bSJacob Faibussowitsch   {
4134742e46bSJacob Faibussowitsch     const auto         mcu       = MatCUPMCast(A);
4144742e46bSJacob Faibussowitsch     const auto         fact_info = mcu->d_fact_info;
4154742e46bSJacob Faibussowitsch     const auto         fact_ipiv = mcu->d_fact_ipiv;
4164742e46bSJacob Faibussowitsch     const auto         lda       = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);
4174742e46bSJacob Faibussowitsch     cupmSolverHandle_t handle;
4184742e46bSJacob Faibussowitsch 
4194742e46bSJacob Faibussowitsch     PetscFunctionBegin;
4204742e46bSJacob Faibussowitsch     PetscCall(GetHandlesFrom_(dctx, &handle));
4214742e46bSJacob Faibussowitsch     PetscCall(PetscInfo(A, "%s solve %d x %d on backend\n", NAME(), m, k));
4224742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeBegin());
4234742e46bSJacob Faibussowitsch     {
4244742e46bSJacob Faibussowitsch       constexpr auto op  = transpose ? CUPMSOLVER_OP_T : CUPMSOLVER_OP_N;
4254742e46bSJacob Faibussowitsch       const auto     da  = DeviceArrayRead(dctx, A);
4264742e46bSJacob Faibussowitsch       const auto     lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);
4274742e46bSJacob Faibussowitsch 
4284742e46bSJacob Faibussowitsch       // clang-format off
4294742e46bSJacob Faibussowitsch       PetscCall(
4304742e46bSJacob Faibussowitsch         base_type::ResizeFactLwork(
4314742e46bSJacob Faibussowitsch           mcu, stream,
4324742e46bSJacob Faibussowitsch           [&](cupmBlasInt_t *lwork)
4334742e46bSJacob Faibussowitsch           {
4344742e46bSJacob Faibussowitsch             return cupmSolverXgetrs_bufferSize(
4354742e46bSJacob Faibussowitsch               handle, op, m, nrhs, da.cupmdata(), lda, fact_ipiv, x, ldx, lwork
4364742e46bSJacob Faibussowitsch             );
4374742e46bSJacob Faibussowitsch           }
4384742e46bSJacob Faibussowitsch         )
4394742e46bSJacob Faibussowitsch       );
4404742e46bSJacob Faibussowitsch       // clang-format on
4414742e46bSJacob Faibussowitsch       PetscCallCUPMSOLVER(cupmSolverXgetrs(handle, op, m, nrhs, da.cupmdata(), lda, fact_ipiv, x, ldx, mcu->d_fact_work, mcu->d_fact_lwork, fact_info));
4424742e46bSJacob Faibussowitsch       PetscCall(CheckCUPMSolverInfo_(fact_info, stream));
4434742e46bSJacob Faibussowitsch     }
4444742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeEnd());
4454742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuFlops(nrhs * (2.0 * m * m - m)));
4464742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
4474742e46bSJacob Faibussowitsch   }
4484742e46bSJacob Faibussowitsch };
4494742e46bSJacob Faibussowitsch 
4504742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
4514742e46bSJacob Faibussowitsch struct MatDense_Seq_CUPM<T>::SolveCholesky : SolveCommon<SolveCholesky> {
4524742e46bSJacob Faibussowitsch   using base_type = SolveCommon<SolveCholesky>;
4534742e46bSJacob Faibussowitsch 
4544742e46bSJacob Faibussowitsch   static constexpr const char   *NAME() noexcept { return "Cholesky"; }
4554742e46bSJacob Faibussowitsch   static constexpr MatFactorType MATFACTORTYPE() noexcept { return MAT_FACTOR_CHOLESKY; }
4564742e46bSJacob Faibussowitsch 
4574742e46bSJacob Faibussowitsch   static PetscErrorCode Factor(Mat A, IS, const MatFactorInfo *) noexcept
4584742e46bSJacob Faibussowitsch   {
4594742e46bSJacob Faibussowitsch     const auto         n = static_cast<cupmBlasInt_t>(A->rmap->n);
4604742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
4614742e46bSJacob Faibussowitsch     cupmSolverHandle_t handle;
4624742e46bSJacob Faibussowitsch     cupmStream_t       stream;
4634742e46bSJacob Faibussowitsch 
4644742e46bSJacob Faibussowitsch     PetscFunctionBegin;
4654742e46bSJacob Faibussowitsch     if (!n || !A->cmap->n) PetscFunctionReturn(PETSC_SUCCESS);
4664742e46bSJacob Faibussowitsch     PetscCheck(A->spd == PETSC_BOOL3_TRUE, PETSC_COMM_SELF, PETSC_ERR_SUP, "%ssytrs unavailable. Use MAT_FACTOR_LU", cupmSolverName());
4674742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx, &handle, &stream));
4684742e46bSJacob Faibussowitsch     PetscCall(base_type::FactorPrepare(A, stream));
4694742e46bSJacob Faibussowitsch     {
4704742e46bSJacob Faibussowitsch       const auto mcu = MatCUPMCast(A);
4714742e46bSJacob Faibussowitsch       const auto da  = DeviceArrayReadWrite(dctx, A);
4724742e46bSJacob Faibussowitsch       const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);
4734742e46bSJacob Faibussowitsch 
4744742e46bSJacob Faibussowitsch       // clang-format off
4754742e46bSJacob Faibussowitsch       PetscCall(
4764742e46bSJacob Faibussowitsch         base_type::ResizeFactLwork(
4774742e46bSJacob Faibussowitsch           mcu, stream,
4784742e46bSJacob Faibussowitsch           [&](cupmBlasInt_t *fact_lwork)
4794742e46bSJacob Faibussowitsch           {
4804742e46bSJacob Faibussowitsch             return cupmSolverXpotrf_bufferSize(
4814742e46bSJacob Faibussowitsch               handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, fact_lwork
4824742e46bSJacob Faibussowitsch             );
4834742e46bSJacob Faibussowitsch           }
4844742e46bSJacob Faibussowitsch         )
4854742e46bSJacob Faibussowitsch       );
4864742e46bSJacob Faibussowitsch       // clang-format on
4874742e46bSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeBegin());
4884742e46bSJacob Faibussowitsch       PetscCallCUPMSOLVER(cupmSolverXpotrf(handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_info));
4894742e46bSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeEnd());
4904742e46bSJacob Faibussowitsch       PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream));
4914742e46bSJacob Faibussowitsch     }
4924742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuFlops(1.0 * n * n * n / 3.0));
4934742e46bSJacob Faibussowitsch 
4944742e46bSJacob Faibussowitsch   #if 0
4954742e46bSJacob Faibussowitsch     // At the time of writing this interface (cuda 10.0), cusolverDn does not implement *sytrs
4964742e46bSJacob Faibussowitsch     // and *hetr* routines. The code below should work, and it can be activated when *sytrs
4974742e46bSJacob Faibussowitsch     // routines will be available
4984742e46bSJacob Faibussowitsch     if (!mcu->d_fact_ipiv) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_ipiv, n, stream));
4994742e46bSJacob Faibussowitsch     if (!mcu->d_fact_lwork) {
5004742e46bSJacob Faibussowitsch       PetscCallCUPMSOLVER(cupmSolverDnXsytrf_bufferSize(handle, n, da.cupmdata(), lda, &mcu->d_fact_lwork));
5014742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_work, mcu->d_fact_lwork, stream));
5024742e46bSJacob Faibussowitsch     }
5034742e46bSJacob Faibussowitsch     if (mcu->d_fact_info) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_info, 1, stream));
5044742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeBegin());
5054742e46bSJacob 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));
5064742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeEnd());
5074742e46bSJacob Faibussowitsch   #endif
5084742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
5094742e46bSJacob Faibussowitsch   }
5104742e46bSJacob Faibussowitsch 
5114742e46bSJacob Faibussowitsch   template <bool transpose>
5124742e46bSJacob 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
5134742e46bSJacob Faibussowitsch   {
5144742e46bSJacob Faibussowitsch     const auto         mcu       = MatCUPMCast(A);
5154742e46bSJacob Faibussowitsch     const auto         fact_info = mcu->d_fact_info;
5164742e46bSJacob Faibussowitsch     cupmSolverHandle_t handle;
5174742e46bSJacob Faibussowitsch 
5184742e46bSJacob Faibussowitsch     PetscFunctionBegin;
5194742e46bSJacob Faibussowitsch     PetscAssert(!mcu->d_fact_ipiv, PETSC_COMM_SELF, PETSC_ERR_LIB, "%ssytrs not implemented", cupmSolverName());
5204742e46bSJacob Faibussowitsch     PetscCall(GetHandlesFrom_(dctx, &handle));
5214742e46bSJacob Faibussowitsch     PetscCall(PetscInfo(A, "%s solve %d x %d on backend\n", NAME(), m, k));
5224742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeBegin());
5234742e46bSJacob Faibussowitsch     {
5244742e46bSJacob Faibussowitsch       const auto da  = DeviceArrayRead(dctx, A);
5254742e46bSJacob Faibussowitsch       const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);
5264742e46bSJacob Faibussowitsch 
5274742e46bSJacob Faibussowitsch       // clang-format off
5284742e46bSJacob Faibussowitsch       PetscCall(
5294742e46bSJacob Faibussowitsch         base_type::ResizeFactLwork(
5304742e46bSJacob Faibussowitsch           mcu, stream,
5314742e46bSJacob Faibussowitsch           [&](cupmBlasInt_t *lwork)
5324742e46bSJacob Faibussowitsch           {
5334742e46bSJacob Faibussowitsch             return cupmSolverXpotrs_bufferSize(
5344742e46bSJacob Faibussowitsch               handle, CUPMSOLVER_FILL_MODE_LOWER, m, nrhs, da.cupmdata(), lda, x, ldx, lwork
5354742e46bSJacob Faibussowitsch             );
5364742e46bSJacob Faibussowitsch           }
5374742e46bSJacob Faibussowitsch         )
5384742e46bSJacob Faibussowitsch       );
5394742e46bSJacob Faibussowitsch       // clang-format on
5404742e46bSJacob 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));
5414742e46bSJacob Faibussowitsch       PetscCall(CheckCUPMSolverInfo_(fact_info, stream));
5424742e46bSJacob Faibussowitsch     }
5434742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeEnd());
5444742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuFlops(nrhs * (2.0 * m * m - m)));
5454742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
5464742e46bSJacob Faibussowitsch   }
5474742e46bSJacob Faibussowitsch };
5484742e46bSJacob Faibussowitsch 
5494742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
5504742e46bSJacob Faibussowitsch struct MatDense_Seq_CUPM<T>::SolveQR : SolveCommon<SolveQR> {
5514742e46bSJacob Faibussowitsch   using base_type = SolveCommon<SolveQR>;
5524742e46bSJacob Faibussowitsch 
5534742e46bSJacob Faibussowitsch   static constexpr const char   *NAME() noexcept { return "QR"; }
5544742e46bSJacob Faibussowitsch   static constexpr MatFactorType MATFACTORTYPE() noexcept { return MAT_FACTOR_QR; }
5554742e46bSJacob Faibussowitsch 
5564742e46bSJacob Faibussowitsch   static PetscErrorCode Factor(Mat A, IS, const MatFactorInfo *) noexcept
5574742e46bSJacob Faibussowitsch   {
5584742e46bSJacob Faibussowitsch     const auto         m     = static_cast<cupmBlasInt_t>(A->rmap->n);
5594742e46bSJacob Faibussowitsch     const auto         n     = static_cast<cupmBlasInt_t>(A->cmap->n);
5604742e46bSJacob Faibussowitsch     const auto         min   = std::min(m, n);
5614742e46bSJacob Faibussowitsch     const auto         mimpl = MatIMPLCast(A);
5624742e46bSJacob Faibussowitsch     cupmStream_t       stream;
5634742e46bSJacob Faibussowitsch     cupmSolverHandle_t handle;
5644742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
5654742e46bSJacob Faibussowitsch 
5664742e46bSJacob Faibussowitsch     PetscFunctionBegin;
5674742e46bSJacob Faibussowitsch     if (!m || !n) PetscFunctionReturn(PETSC_SUCCESS);
5684742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx, &handle, &stream));
5694742e46bSJacob Faibussowitsch     PetscCall(base_type::FactorPrepare(A, stream));
5704742e46bSJacob Faibussowitsch     mimpl->rank = min;
5714742e46bSJacob Faibussowitsch     {
5724742e46bSJacob Faibussowitsch       const auto mcu = MatCUPMCast(A);
5734742e46bSJacob Faibussowitsch       const auto da  = DeviceArrayReadWrite(dctx, A);
5744742e46bSJacob Faibussowitsch       const auto lda = static_cast<cupmBlasInt_t>(mimpl->lda);
5754742e46bSJacob Faibussowitsch 
5764742e46bSJacob Faibussowitsch       if (!mcu->workvec) PetscCall(vec::cupm::VecCreateSeqCUPMAsync<T>(PetscObjectComm(PetscObjectCast(A)), m, &mcu->workvec));
5774742e46bSJacob Faibussowitsch       if (!mcu->d_fact_tau) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_tau, min, stream));
5784742e46bSJacob Faibussowitsch       // clang-format off
5794742e46bSJacob Faibussowitsch       PetscCall(
5804742e46bSJacob Faibussowitsch         base_type::ResizeFactLwork(
5814742e46bSJacob Faibussowitsch           mcu, stream,
5824742e46bSJacob Faibussowitsch           [&](cupmBlasInt_t *fact_lwork)
5834742e46bSJacob Faibussowitsch           {
5844742e46bSJacob Faibussowitsch             return cupmSolverXgeqrf_bufferSize(handle, m, n, da.cupmdata(), lda, fact_lwork);
5854742e46bSJacob Faibussowitsch           }
5864742e46bSJacob Faibussowitsch         )
5874742e46bSJacob Faibussowitsch       );
5884742e46bSJacob Faibussowitsch       // clang-format on
5894742e46bSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeBegin());
5904742e46bSJacob 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));
5914742e46bSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeEnd());
5924742e46bSJacob Faibussowitsch       PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream));
5934742e46bSJacob Faibussowitsch     }
5944742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuFlops(2.0 * min * min * (std::max(m, n) - min / 3.0)));
5954742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
5964742e46bSJacob Faibussowitsch   }
5974742e46bSJacob Faibussowitsch 
5984742e46bSJacob Faibussowitsch   template <bool transpose>
5994742e46bSJacob 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
6004742e46bSJacob Faibussowitsch   {
6014742e46bSJacob Faibussowitsch     const auto         mimpl      = MatIMPLCast(A);
6024742e46bSJacob Faibussowitsch     const auto         rank       = static_cast<cupmBlasInt_t>(mimpl->rank);
6034742e46bSJacob Faibussowitsch     const auto         mcu        = MatCUPMCast(A);
6044742e46bSJacob Faibussowitsch     const auto         fact_info  = mcu->d_fact_info;
6054742e46bSJacob Faibussowitsch     const auto         fact_tau   = mcu->d_fact_tau;
6064742e46bSJacob Faibussowitsch     const auto         fact_work  = mcu->d_fact_work;
6074742e46bSJacob Faibussowitsch     const auto         fact_lwork = mcu->d_fact_lwork;
6084742e46bSJacob Faibussowitsch     cupmSolverHandle_t solver_handle;
6094742e46bSJacob Faibussowitsch     cupmBlasHandle_t   blas_handle;
6104742e46bSJacob Faibussowitsch 
6114742e46bSJacob Faibussowitsch     PetscFunctionBegin;
6124742e46bSJacob Faibussowitsch     PetscCall(GetHandlesFrom_(dctx, &blas_handle, &solver_handle));
6134742e46bSJacob Faibussowitsch     PetscCall(PetscInfo(A, "%s solve %d x %d on backend\n", NAME(), m, k));
6144742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeBegin());
6154742e46bSJacob Faibussowitsch     {
6164742e46bSJacob Faibussowitsch       const auto da  = DeviceArrayRead(dctx, A);
6174742e46bSJacob Faibussowitsch       const auto one = cupmScalarCast(1.0);
6184742e46bSJacob Faibussowitsch       const auto lda = static_cast<cupmBlasInt_t>(mimpl->lda);
6194742e46bSJacob Faibussowitsch 
6204742e46bSJacob Faibussowitsch       if (transpose) {
6214742e46bSJacob 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));
6224742e46bSJacob 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));
6234742e46bSJacob Faibussowitsch         PetscCall(CheckCUPMSolverInfo_(fact_info, stream));
6244742e46bSJacob Faibussowitsch       } else {
6254742e46bSJacob Faibussowitsch         constexpr auto op = PetscDefined(USE_COMPLEX) ? CUPMSOLVER_OP_C : CUPMSOLVER_OP_T;
6264742e46bSJacob Faibussowitsch 
6274742e46bSJacob 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));
6284742e46bSJacob Faibussowitsch         PetscCall(CheckCUPMSolverInfo_(fact_info, stream));
6294742e46bSJacob 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));
6304742e46bSJacob Faibussowitsch       }
6314742e46bSJacob Faibussowitsch     }
6324742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeEnd());
6334742e46bSJacob Faibussowitsch     PetscCall(PetscLogFlops(nrhs * (4.0 * m * rank - (rank * rank))));
6344742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
6354742e46bSJacob Faibussowitsch   }
6364742e46bSJacob Faibussowitsch };
6374742e46bSJacob Faibussowitsch 
6384742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
6394742e46bSJacob Faibussowitsch template <typename Solver, bool transpose>
6404742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::MatSolve_Factored_Dispatch_(Mat A, Vec x, Vec y) noexcept
6414742e46bSJacob Faibussowitsch {
6424742e46bSJacob Faibussowitsch   using namespace vec::cupm;
6434742e46bSJacob Faibussowitsch   const auto         pobj_A  = PetscObjectCast(A);
6444742e46bSJacob Faibussowitsch   const auto         m       = static_cast<cupmBlasInt_t>(A->rmap->n);
6454742e46bSJacob Faibussowitsch   const auto         k       = static_cast<cupmBlasInt_t>(A->cmap->n);
6464742e46bSJacob Faibussowitsch   auto              &workvec = MatCUPMCast(A)->workvec;
6474742e46bSJacob Faibussowitsch   PetscScalar       *y_array = nullptr;
6484742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
6494742e46bSJacob Faibussowitsch   PetscBool          xiscupm, yiscupm, aiscupm;
6504742e46bSJacob Faibussowitsch   bool               use_y_array_directly;
6514742e46bSJacob Faibussowitsch   cupmStream_t       stream;
6524742e46bSJacob Faibussowitsch 
6534742e46bSJacob Faibussowitsch   PetscFunctionBegin;
6544742e46bSJacob Faibussowitsch   PetscCheck(A->factortype != MAT_FACTOR_NONE, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix must be factored to solve");
6554742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(PetscObjectCast(x), VecSeq_CUPM::VECSEQCUPM(), &xiscupm));
6564742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(PetscObjectCast(y), VecSeq_CUPM::VECSEQCUPM(), &yiscupm));
6574742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(pobj_A, MATSEQDENSECUPM(), &aiscupm));
6584742e46bSJacob Faibussowitsch   PetscAssert(aiscupm, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Matrix A is somehow not CUPM?????????????????????????????");
6594742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx, &stream));
6604742e46bSJacob Faibussowitsch   use_y_array_directly = yiscupm && (k >= m);
6614742e46bSJacob Faibussowitsch   {
6624742e46bSJacob Faibussowitsch     const PetscScalar *x_array;
6634742e46bSJacob Faibussowitsch     const auto         xisdevice = xiscupm && PetscOffloadDevice(x->offloadmask);
6644742e46bSJacob Faibussowitsch     const auto         copy_mode = xisdevice ? cupmMemcpyDeviceToDevice : cupmMemcpyHostToDevice;
6654742e46bSJacob Faibussowitsch 
6664742e46bSJacob Faibussowitsch     if (!use_y_array_directly && !workvec) PetscCall(VecCreateSeqCUPMAsync<T>(PetscObjectComm(pobj_A), m, &workvec));
6674742e46bSJacob Faibussowitsch     // The logic here is to try to minimize the amount of memory copying:
6684742e46bSJacob Faibussowitsch     //
6694742e46bSJacob Faibussowitsch     // If we call VecCUPMGetArrayRead(X, &x) every time xiscupm and the data is not offloaded
6704742e46bSJacob Faibussowitsch     // to the GPU yet, then the data is copied to the GPU. But we are only trying to get the
6714742e46bSJacob Faibussowitsch     // data in order to copy it into the y array. So the array x will be wherever the data
6724742e46bSJacob Faibussowitsch     // already is so that only one memcpy is performed
6734742e46bSJacob Faibussowitsch     if (xisdevice) {
6744742e46bSJacob Faibussowitsch       PetscCall(VecCUPMGetArrayReadAsync<T>(x, &x_array, dctx));
6754742e46bSJacob Faibussowitsch     } else {
6764742e46bSJacob Faibussowitsch       PetscCall(VecGetArrayRead(x, &x_array));
6774742e46bSJacob Faibussowitsch     }
6784742e46bSJacob Faibussowitsch     PetscCall(VecCUPMGetArrayWriteAsync<T>(use_y_array_directly ? y : workvec, &y_array, dctx));
6794742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMMemcpyAsync(y_array, x_array, m, copy_mode, stream));
6804742e46bSJacob Faibussowitsch     if (xisdevice) {
6814742e46bSJacob Faibussowitsch       PetscCall(VecCUPMRestoreArrayReadAsync<T>(x, &x_array, dctx));
6824742e46bSJacob Faibussowitsch     } else {
6834742e46bSJacob Faibussowitsch       PetscCall(VecRestoreArrayRead(x, &x_array));
6844742e46bSJacob Faibussowitsch     }
6854742e46bSJacob Faibussowitsch   }
6864742e46bSJacob Faibussowitsch 
6874742e46bSJacob Faibussowitsch   if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
6884742e46bSJacob Faibussowitsch   PetscCall(Solver{}.template Solve<transpose>(A, cupmScalarPtrCast(y_array), m, m, 1, k, dctx, stream));
6894742e46bSJacob Faibussowitsch   if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A));
6904742e46bSJacob Faibussowitsch 
6914742e46bSJacob Faibussowitsch   if (use_y_array_directly) {
6924742e46bSJacob Faibussowitsch     PetscCall(VecCUPMRestoreArrayWriteAsync<T>(y, &y_array, dctx));
6934742e46bSJacob Faibussowitsch   } else {
6944742e46bSJacob Faibussowitsch     const auto   copy_mode = yiscupm ? cupmMemcpyDeviceToDevice : cupmMemcpyDeviceToHost;
6954742e46bSJacob Faibussowitsch     PetscScalar *yv;
6964742e46bSJacob Faibussowitsch 
6974742e46bSJacob Faibussowitsch     // The logic here is that the data is not yet in either y's GPU array or its CPU array.
6984742e46bSJacob Faibussowitsch     // There is nothing in the interface to say where the user would like it to end up. So we
6994742e46bSJacob Faibussowitsch     // choose the GPU, because it is the faster option
7004742e46bSJacob Faibussowitsch     if (yiscupm) {
7014742e46bSJacob Faibussowitsch       PetscCall(VecCUPMGetArrayWriteAsync<T>(y, &yv, dctx));
7024742e46bSJacob Faibussowitsch     } else {
7034742e46bSJacob Faibussowitsch       PetscCall(VecGetArray(y, &yv));
7044742e46bSJacob Faibussowitsch     }
7054742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMMemcpyAsync(yv, y_array, k, copy_mode, stream));
7064742e46bSJacob Faibussowitsch     if (yiscupm) {
7074742e46bSJacob Faibussowitsch       PetscCall(VecCUPMRestoreArrayWriteAsync<T>(y, &yv, dctx));
7084742e46bSJacob Faibussowitsch     } else {
7094742e46bSJacob Faibussowitsch       PetscCall(VecRestoreArray(y, &yv));
7104742e46bSJacob Faibussowitsch     }
7114742e46bSJacob Faibussowitsch     PetscCall(VecCUPMRestoreArrayWriteAsync<T>(workvec, &y_array));
7124742e46bSJacob Faibussowitsch   }
7134742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
7144742e46bSJacob Faibussowitsch }
7154742e46bSJacob Faibussowitsch 
7164742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
7174742e46bSJacob Faibussowitsch template <typename Solver, bool transpose>
7184742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMatSolve_Factored_Dispatch_(Mat A, Mat B, Mat X) noexcept
7194742e46bSJacob Faibussowitsch {
7204742e46bSJacob Faibussowitsch   const auto         m = static_cast<cupmBlasInt_t>(A->rmap->n);
7214742e46bSJacob Faibussowitsch   const auto         k = static_cast<cupmBlasInt_t>(A->cmap->n);
7224742e46bSJacob Faibussowitsch   cupmBlasInt_t      nrhs, ldb, ldx, ldy;
7234742e46bSJacob Faibussowitsch   PetscScalar       *y;
7244742e46bSJacob Faibussowitsch   PetscBool          biscupm, xiscupm, aiscupm;
7254742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
7264742e46bSJacob Faibussowitsch   cupmStream_t       stream;
7274742e46bSJacob Faibussowitsch 
7284742e46bSJacob Faibussowitsch   PetscFunctionBegin;
7294742e46bSJacob Faibussowitsch   PetscCheck(A->factortype != MAT_FACTOR_NONE, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix must be factored to solve");
7304742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(PetscObjectCast(B), MATSEQDENSECUPM(), &biscupm));
7314742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(PetscObjectCast(X), MATSEQDENSECUPM(), &xiscupm));
7324742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(PetscObjectCast(A), MATSEQDENSECUPM(), &aiscupm));
7334742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx, &stream));
7344742e46bSJacob Faibussowitsch   {
7354742e46bSJacob Faibussowitsch     PetscInt n;
7364742e46bSJacob Faibussowitsch 
7374742e46bSJacob Faibussowitsch     PetscCall(MatGetSize(B, nullptr, &n));
7384742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMBlasIntCast(n, &nrhs));
7394742e46bSJacob Faibussowitsch     PetscCall(MatDenseGetLDA(B, &n));
7404742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMBlasIntCast(n, &ldb));
7414742e46bSJacob Faibussowitsch     PetscCall(MatDenseGetLDA(X, &n));
7424742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMBlasIntCast(n, &ldx));
7434742e46bSJacob Faibussowitsch   }
7444742e46bSJacob Faibussowitsch   {
7454742e46bSJacob Faibussowitsch     // The logic here is to try to minimize the amount of memory copying:
7464742e46bSJacob Faibussowitsch     //
7474742e46bSJacob Faibussowitsch     // If we call MatDenseCUPMGetArrayRead(B, &b) every time biscupm and the data is not
7484742e46bSJacob Faibussowitsch     // offloaded to the GPU yet, then the data is copied to the GPU. But we are only trying to
7494742e46bSJacob Faibussowitsch     // get the data in order to copy it into the y array. So the array b will be wherever the
7504742e46bSJacob Faibussowitsch     // data already is so that only one memcpy is performed
7514742e46bSJacob Faibussowitsch     const auto         bisdevice = biscupm && PetscOffloadDevice(B->offloadmask);
7524742e46bSJacob Faibussowitsch     const auto         copy_mode = bisdevice ? cupmMemcpyDeviceToDevice : cupmMemcpyHostToDevice;
7534742e46bSJacob Faibussowitsch     const PetscScalar *b;
7544742e46bSJacob Faibussowitsch 
7554742e46bSJacob Faibussowitsch     if (bisdevice) {
7564742e46bSJacob Faibussowitsch       b = DeviceArrayRead(dctx, B);
7574742e46bSJacob Faibussowitsch     } else if (biscupm) {
7584742e46bSJacob Faibussowitsch       b = HostArrayRead(dctx, B);
7594742e46bSJacob Faibussowitsch     } else {
7604742e46bSJacob Faibussowitsch       PetscCall(MatDenseGetArrayRead(B, &b));
7614742e46bSJacob Faibussowitsch     }
7624742e46bSJacob Faibussowitsch 
7634742e46bSJacob Faibussowitsch     if (ldx < m || !xiscupm) {
7644742e46bSJacob Faibussowitsch       // X's array cannot serve as the array (too small or not on device), B's array cannot
7654742e46bSJacob Faibussowitsch       // serve as the array (const), so allocate a new array
7664742e46bSJacob Faibussowitsch       ldy = m;
7674742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMallocAsync(&y, nrhs * m));
7684742e46bSJacob Faibussowitsch     } else {
7694742e46bSJacob Faibussowitsch       // X's array should serve as the array
7704742e46bSJacob Faibussowitsch       ldy = ldx;
7714742e46bSJacob Faibussowitsch       y   = DeviceArrayWrite(dctx, X);
7724742e46bSJacob Faibussowitsch     }
7734742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMMemcpy2DAsync(y, ldy, b, ldb, m, nrhs, copy_mode, stream));
7744742e46bSJacob Faibussowitsch     if (!bisdevice && !biscupm) PetscCall(MatDenseRestoreArrayRead(B, &b));
7754742e46bSJacob Faibussowitsch   }
7764742e46bSJacob Faibussowitsch 
7774742e46bSJacob Faibussowitsch   // convert to CUPM twice??????????????????????????????????
7784742e46bSJacob Faibussowitsch   // but A should already be CUPM??????????????????????????????????????
7794742e46bSJacob Faibussowitsch   if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
7804742e46bSJacob Faibussowitsch   PetscCall(Solver{}.template Solve<transpose>(A, cupmScalarPtrCast(y), ldy, m, nrhs, k, dctx, stream));
7814742e46bSJacob Faibussowitsch   if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
7824742e46bSJacob Faibussowitsch 
7834742e46bSJacob Faibussowitsch   if (ldx < m || !xiscupm) {
7844742e46bSJacob Faibussowitsch     const auto   copy_mode = xiscupm ? cupmMemcpyDeviceToDevice : cupmMemcpyDeviceToHost;
7854742e46bSJacob Faibussowitsch     PetscScalar *x;
7864742e46bSJacob Faibussowitsch 
7874742e46bSJacob Faibussowitsch     // The logic here is that the data is not yet in either X's GPU array or its CPU
7884742e46bSJacob Faibussowitsch     // array. There is nothing in the interface to say where the user would like it to end up.
7894742e46bSJacob Faibussowitsch     // So we choose the GPU, because it is the faster option
7904742e46bSJacob Faibussowitsch     if (xiscupm) {
7914742e46bSJacob Faibussowitsch       x = DeviceArrayWrite(dctx, X);
7924742e46bSJacob Faibussowitsch     } else {
7934742e46bSJacob Faibussowitsch       PetscCall(MatDenseGetArray(X, &x));
7944742e46bSJacob Faibussowitsch     }
7954742e46bSJacob Faibussowitsch     PetscCall(PetscCUPMMemcpy2DAsync(x, ldx, y, ldy, k, nrhs, copy_mode, stream));
7964742e46bSJacob Faibussowitsch     if (!xiscupm) PetscCall(MatDenseRestoreArray(X, &x));
7974742e46bSJacob Faibussowitsch     PetscCallCUPM(cupmFreeAsync(y, stream));
7984742e46bSJacob Faibussowitsch   }
7994742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
8004742e46bSJacob Faibussowitsch }
8014742e46bSJacob Faibussowitsch 
8024742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
8034742e46bSJacob Faibussowitsch template <bool transpose>
8044742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMultAdd_Dispatch_(Mat A, Vec xx, Vec yy, Vec zz) noexcept
8054742e46bSJacob Faibussowitsch {
8064742e46bSJacob Faibussowitsch   const auto         m = static_cast<cupmBlasInt_t>(A->rmap->n);
8074742e46bSJacob Faibussowitsch   const auto         n = static_cast<cupmBlasInt_t>(A->cmap->n);
8084742e46bSJacob Faibussowitsch   cupmBlasHandle_t   handle;
8094742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
8104742e46bSJacob Faibussowitsch 
8114742e46bSJacob Faibussowitsch   PetscFunctionBegin;
8124742e46bSJacob Faibussowitsch   if (yy && yy != zz) PetscCall(VecSeq_CUPM::Copy(yy, zz)); // mult add
8134742e46bSJacob Faibussowitsch   if (!m || !n) {
8144742e46bSJacob Faibussowitsch     // mult only
8154742e46bSJacob Faibussowitsch     if (!yy) PetscCall(VecSeq_CUPM::Set(zz, 0.0));
8164742e46bSJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
8174742e46bSJacob Faibussowitsch   }
8184742e46bSJacob Faibussowitsch   PetscCall(PetscInfo(A, "Matrix-vector product %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " on backend\n", m, n));
8194742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx, &handle));
8204742e46bSJacob Faibussowitsch   {
8214742e46bSJacob Faibussowitsch     constexpr auto op   = transpose ? CUPMBLAS_OP_T : CUPMBLAS_OP_N;
8224742e46bSJacob Faibussowitsch     const auto     one  = cupmScalarCast(1.0);
8234742e46bSJacob Faibussowitsch     const auto     zero = cupmScalarCast(0.0);
8244742e46bSJacob Faibussowitsch     const auto     da   = DeviceArrayRead(dctx, A);
8254742e46bSJacob Faibussowitsch     const auto     dxx  = VecSeq_CUPM::DeviceArrayRead(dctx, xx);
8264742e46bSJacob Faibussowitsch     const auto     dzz  = VecSeq_CUPM::DeviceArrayReadWrite(dctx, zz);
8274742e46bSJacob Faibussowitsch 
8284742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeBegin());
8294742e46bSJacob 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));
8304742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeEnd());
8314742e46bSJacob Faibussowitsch   }
8324742e46bSJacob Faibussowitsch   PetscCall(PetscLogGpuFlops(2.0 * m * n - (yy ? 0 : m)));
8334742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
8344742e46bSJacob Faibussowitsch }
8354742e46bSJacob Faibussowitsch 
8364742e46bSJacob Faibussowitsch // ==========================================================================================
8374742e46bSJacob Faibussowitsch // MatDense_Seq_CUPM - Private API - Conversion Dispatch
8384742e46bSJacob Faibussowitsch // ==========================================================================================
8394742e46bSJacob Faibussowitsch 
8404742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
8414742e46bSJacob Faibussowitsch template <bool to_host>
8424742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_Dispatch_(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept
8434742e46bSJacob Faibussowitsch {
8444742e46bSJacob Faibussowitsch   PetscFunctionBegin;
8454742e46bSJacob Faibussowitsch   if (reuse == MAT_REUSE_MATRIX || reuse == MAT_INITIAL_MATRIX) {
8464742e46bSJacob Faibussowitsch     // TODO these cases should be optimized
8474742e46bSJacob Faibussowitsch     PetscCall(MatConvert_Basic(M, type, reuse, newmat));
8484742e46bSJacob Faibussowitsch   } else {
8494742e46bSJacob Faibussowitsch     const auto B    = *newmat;
8504742e46bSJacob Faibussowitsch     const auto pobj = PetscObjectCast(B);
8514742e46bSJacob Faibussowitsch 
8524742e46bSJacob Faibussowitsch     if (to_host) {
8534742e46bSJacob Faibussowitsch       PetscCall(BindToCPU(B, PETSC_TRUE));
8544742e46bSJacob Faibussowitsch       PetscCall(Reset(B));
8554742e46bSJacob Faibussowitsch     } else {
8564742e46bSJacob Faibussowitsch       PetscCall(PetscDeviceInitialize(PETSC_DEVICE_CUPM()));
8574742e46bSJacob Faibussowitsch     }
8584742e46bSJacob Faibussowitsch 
8594742e46bSJacob Faibussowitsch     PetscCall(PetscStrFreeAllocpy(to_host ? VECSTANDARD : VecSeq_CUPM::VECCUPM(), &B->defaultvectype));
8604742e46bSJacob Faibussowitsch     PetscCall(PetscObjectChangeTypeName(pobj, to_host ? MATSEQDENSE : MATSEQDENSECUPM()));
8614742e46bSJacob Faibussowitsch     // cvec might be the wrong VecType, destroy and rebuild it if necessary
8624742e46bSJacob Faibussowitsch     // REVIEW ME: this is possibly very inefficient
8634742e46bSJacob Faibussowitsch     PetscCall(VecDestroy(&MatIMPLCast(B)->cvec));
8644742e46bSJacob Faibussowitsch 
8654742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatConvert_seqdensecupm_seqdense_C(), nullptr, Convert_SeqDenseCUPM_SeqDense);
8664742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArray_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ_WRITE>);
8674742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArrayRead_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ>);
8684742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArrayWrite_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_WRITE>);
8694742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArray_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ_WRITE>);
8704742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArrayRead_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ>);
8714742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArrayWrite_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_WRITE>);
8724742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMPlaceArray_C(), nullptr, PlaceArray);
8734742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMResetArray_C(), nullptr, ResetArray);
8744742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMReplaceArray_C(), nullptr, ReplaceArray);
8754742e46bSJacob Faibussowitsch     MatComposeOp_CUPM(to_host, pobj, MatProductSetFromOptions_seqaij_seqdensecupm_C(), nullptr, MatProductSetFromOptions_SeqAIJ_SeqDense);
8764742e46bSJacob Faibussowitsch 
8774742e46bSJacob Faibussowitsch     if (to_host) {
8784742e46bSJacob Faibussowitsch       B->offloadmask = PETSC_OFFLOAD_CPU;
8794742e46bSJacob Faibussowitsch     } else {
8804742e46bSJacob Faibussowitsch       Mat_SeqDenseCUPM *mcu;
8814742e46bSJacob Faibussowitsch 
8824742e46bSJacob Faibussowitsch       PetscCall(PetscNew(&mcu));
8834742e46bSJacob Faibussowitsch       B->spptr       = mcu;
8844742e46bSJacob Faibussowitsch       B->offloadmask = PETSC_OFFLOAD_UNALLOCATED; // REVIEW ME: why not offload host??
8854742e46bSJacob Faibussowitsch       PetscCall(BindToCPU(B, PETSC_FALSE));
8864742e46bSJacob Faibussowitsch     }
8874742e46bSJacob Faibussowitsch 
8884742e46bSJacob Faibussowitsch     MatSetOp_CUPM(to_host, B, bindtocpu, nullptr, BindToCPU);
8894742e46bSJacob Faibussowitsch     MatSetOp_CUPM(to_host, B, destroy, MatDestroy_SeqDense, Destroy);
8904742e46bSJacob Faibussowitsch   }
8914742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
8924742e46bSJacob Faibussowitsch }
8934742e46bSJacob Faibussowitsch 
8944742e46bSJacob Faibussowitsch // ==========================================================================================
8954742e46bSJacob Faibussowitsch // MatDense_Seq_CUPM - Public API
8964742e46bSJacob Faibussowitsch // ==========================================================================================
8974742e46bSJacob Faibussowitsch 
8984742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
8994742e46bSJacob Faibussowitsch inline constexpr MatType MatDense_Seq_CUPM<T>::MATIMPLCUPM_() noexcept
9004742e46bSJacob Faibussowitsch {
9014742e46bSJacob Faibussowitsch   return MATSEQDENSECUPM();
9024742e46bSJacob Faibussowitsch }
9034742e46bSJacob Faibussowitsch 
9044742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9054742e46bSJacob Faibussowitsch inline constexpr typename MatDense_Seq_CUPM<T>::Mat_SeqDenseCUPM *MatDense_Seq_CUPM<T>::MatCUPMCast(Mat m) noexcept
9064742e46bSJacob Faibussowitsch {
9074742e46bSJacob Faibussowitsch   return static_cast<Mat_SeqDenseCUPM *>(m->spptr);
9084742e46bSJacob Faibussowitsch }
9094742e46bSJacob Faibussowitsch 
9104742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9114742e46bSJacob Faibussowitsch inline constexpr Mat_SeqDense *MatDense_Seq_CUPM<T>::MatIMPLCast_(Mat m) noexcept
9124742e46bSJacob Faibussowitsch {
9134742e46bSJacob Faibussowitsch   return static_cast<Mat_SeqDense *>(m->data);
9144742e46bSJacob Faibussowitsch }
9154742e46bSJacob Faibussowitsch 
9164742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9174742e46bSJacob Faibussowitsch inline constexpr const char *MatDense_Seq_CUPM<T>::MatConvert_seqdensecupm_seqdense_C() noexcept
9184742e46bSJacob Faibussowitsch {
9194742e46bSJacob Faibussowitsch   return T == device::cupm::DeviceType::CUDA ? "MatConvert_seqdensecuda_seqdense_C" : "MatConvert_seqdensehip_seqdense_C";
9204742e46bSJacob Faibussowitsch }
9214742e46bSJacob Faibussowitsch 
9224742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9234742e46bSJacob Faibussowitsch inline constexpr const char *MatDense_Seq_CUPM<T>::MatProductSetFromOptions_seqaij_seqdensecupm_C() noexcept
9244742e46bSJacob Faibussowitsch {
9254742e46bSJacob Faibussowitsch   return T == device::cupm::DeviceType::CUDA ? "MatProductSetFromOptions_seqaij_seqdensecuda_C" : "MatProductSetFromOptions_seqaij_seqdensehip_C";
9264742e46bSJacob Faibussowitsch }
9274742e46bSJacob Faibussowitsch 
9284742e46bSJacob Faibussowitsch // ==========================================================================================
9294742e46bSJacob Faibussowitsch 
9304742e46bSJacob Faibussowitsch // MatCreate_SeqDenseCUPM()
9314742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9324742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Create(Mat A) noexcept
9334742e46bSJacob Faibussowitsch {
9344742e46bSJacob Faibussowitsch   PetscFunctionBegin;
9354742e46bSJacob Faibussowitsch   PetscCall(PetscDeviceInitialize(PETSC_DEVICE_CUPM()));
9364742e46bSJacob Faibussowitsch   PetscCall(MatCreate_SeqDense(A));
9374742e46bSJacob Faibussowitsch   PetscCall(Convert_SeqDense_SeqDenseCUPM(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
9384742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
9394742e46bSJacob Faibussowitsch }
9404742e46bSJacob Faibussowitsch 
9414742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9424742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Destroy(Mat A) noexcept
9434742e46bSJacob Faibussowitsch {
9444742e46bSJacob Faibussowitsch   PetscFunctionBegin;
9454742e46bSJacob Faibussowitsch   // prevent copying back data if we own the data pointer
9464742e46bSJacob Faibussowitsch   if (!MatIMPLCast(A)->user_alloc) A->offloadmask = PETSC_OFFLOAD_CPU;
9474742e46bSJacob Faibussowitsch   PetscCall(Convert_SeqDenseCUPM_SeqDense(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A));
9484742e46bSJacob Faibussowitsch   PetscCall(MatDestroy_SeqDense(A));
9494742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
9504742e46bSJacob Faibussowitsch }
9514742e46bSJacob Faibussowitsch 
9524742e46bSJacob Faibussowitsch // obj->ops->setup()
9534742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9544742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::SetUp(Mat A) noexcept
9554742e46bSJacob Faibussowitsch {
9564742e46bSJacob Faibussowitsch   PetscFunctionBegin;
9574742e46bSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(A->rmap));
9584742e46bSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(A->cmap));
9594742e46bSJacob Faibussowitsch   if (!A->preallocated) {
9604742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
9614742e46bSJacob Faibussowitsch 
9624742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
9634742e46bSJacob Faibussowitsch     PetscCall(SetPreallocation(A, dctx));
9644742e46bSJacob Faibussowitsch   }
9654742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
9664742e46bSJacob Faibussowitsch }
9674742e46bSJacob Faibussowitsch 
9684742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9694742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Reset(Mat A) noexcept
9704742e46bSJacob Faibussowitsch {
9714742e46bSJacob Faibussowitsch   PetscFunctionBegin;
9724742e46bSJacob Faibussowitsch   if (const auto mcu = MatCUPMCast(A)) {
9734742e46bSJacob Faibussowitsch     cupmStream_t stream;
9744742e46bSJacob Faibussowitsch 
9754742e46bSJacob Faibussowitsch     PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME());
9764742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&stream));
9774742e46bSJacob Faibussowitsch     if (!mcu->d_user_alloc) PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream));
9784742e46bSJacob Faibussowitsch     PetscCallCUPM(cupmFreeAsync(mcu->d_fact_tau, stream));
9794742e46bSJacob Faibussowitsch     PetscCallCUPM(cupmFreeAsync(mcu->d_fact_ipiv, stream));
9804742e46bSJacob Faibussowitsch     PetscCallCUPM(cupmFreeAsync(mcu->d_fact_info, stream));
9814742e46bSJacob Faibussowitsch     PetscCallCUPM(cupmFreeAsync(mcu->d_fact_work, stream));
9824742e46bSJacob Faibussowitsch     PetscCall(VecDestroy(&mcu->workvec));
9834742e46bSJacob Faibussowitsch     PetscCall(PetscFree(A->spptr /* mcu */));
9844742e46bSJacob Faibussowitsch   }
9854742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
9864742e46bSJacob Faibussowitsch }
9874742e46bSJacob Faibussowitsch 
9884742e46bSJacob Faibussowitsch // ==========================================================================================
9894742e46bSJacob Faibussowitsch 
9904742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
9914742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::BindToCPU(Mat A, PetscBool to_host) noexcept
9924742e46bSJacob Faibussowitsch {
9934742e46bSJacob Faibussowitsch   const auto mimpl = MatIMPLCast(A);
9944742e46bSJacob Faibussowitsch   const auto pobj  = PetscObjectCast(A);
9954742e46bSJacob Faibussowitsch 
9964742e46bSJacob Faibussowitsch   PetscFunctionBegin;
9974742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
9984742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
9994742e46bSJacob Faibussowitsch   A->boundtocpu = to_host;
10004742e46bSJacob Faibussowitsch   PetscCall(PetscStrFreeAllocpy(to_host ? PETSCRANDER48 : PETSCDEVICERAND(), &A->defaultrandtype));
10014742e46bSJacob Faibussowitsch   if (to_host) {
10024742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
10034742e46bSJacob Faibussowitsch 
10044742e46bSJacob Faibussowitsch     // make sure we have an up-to-date copy on the CPU
10054742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
10064742e46bSJacob Faibussowitsch     PetscCall(DeviceToHost_(A, dctx));
10074742e46bSJacob Faibussowitsch   } else {
10084742e46bSJacob Faibussowitsch     PetscBool iscupm;
10094742e46bSJacob Faibussowitsch 
10104742e46bSJacob Faibussowitsch     if (auto &cvec = mimpl->cvec) {
10114742e46bSJacob Faibussowitsch       PetscCall(PetscObjectTypeCompare(PetscObjectCast(cvec), VecSeq_CUPM::VECSEQCUPM(), &iscupm));
10124742e46bSJacob Faibussowitsch       if (!iscupm) PetscCall(VecDestroy(&cvec));
10134742e46bSJacob Faibussowitsch     }
10144742e46bSJacob Faibussowitsch     if (auto &cmat = mimpl->cmat) {
10154742e46bSJacob Faibussowitsch       PetscCall(PetscObjectTypeCompare(PetscObjectCast(cmat), MATSEQDENSECUPM(), &iscupm));
10164742e46bSJacob Faibussowitsch       if (!iscupm) PetscCall(MatDestroy(&cmat));
10174742e46bSJacob Faibussowitsch     }
10184742e46bSJacob Faibussowitsch   }
10194742e46bSJacob Faibussowitsch 
10204742e46bSJacob Faibussowitsch   // ============================================================
10214742e46bSJacob Faibussowitsch   // Composed ops
10224742e46bSJacob Faibussowitsch   // ============================================================
10234742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArray_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ_WRITE>);
10244742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayRead_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ>);
10254742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayWrite_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_WRITE>);
10264742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ_WRITE>);
10274742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ_WRITE>);
10284742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayReadAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ>);
10294742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayReadAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ>);
10304742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayWriteAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_WRITE>);
10314742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayWriteAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_WRITE>);
10324742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVec_C", MatDenseGetColumnVec_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_READ_WRITE>);
10334742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVec_C", MatDenseRestoreColumnVec_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_READ_WRITE>);
10344742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVecRead_C", MatDenseGetColumnVecRead_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_READ>);
10354742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVecRead_C", MatDenseRestoreColumnVecRead_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_READ>);
10364742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVecWrite_C", MatDenseGetColumnVecWrite_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_WRITE>);
10374742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVecWrite_C", MatDenseRestoreColumnVecWrite_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_WRITE>);
10384742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetSubMatrix_C", MatDenseGetSubMatrix_SeqDense, GetSubMatrix);
10394742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreSubMatrix_C", MatDenseRestoreSubMatrix_SeqDense, RestoreSubMatrix);
10404742e46bSJacob Faibussowitsch   MatComposeOp_CUPM(to_host, pobj, "MatQRFactor_C", MatQRFactor_SeqDense, SolveQR::Factor);
10414742e46bSJacob Faibussowitsch   // always the same
10424742e46bSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction(pobj, "MatDenseSetLDA_C", MatDenseSetLDA_SeqDense));
10434742e46bSJacob Faibussowitsch 
10444742e46bSJacob Faibussowitsch   // ============================================================
10454742e46bSJacob Faibussowitsch   // Function pointer ops
10464742e46bSJacob Faibussowitsch   // ============================================================
10474742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, duplicate, MatDuplicate_SeqDense, Duplicate);
10484742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, mult, MatMult_SeqDense, [](Mat A, Vec xx, Vec yy) { return MatMultAdd_Dispatch_</* transpose */ false>(A, xx, nullptr, yy); });
10494742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, multtranspose, MatMultTranspose_SeqDense, [](Mat A, Vec xx, Vec yy) { return MatMultAdd_Dispatch_</* transpose */ true>(A, xx, nullptr, yy); });
10504742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, multadd, MatMultAdd_SeqDense, MatMultAdd_Dispatch_</* transpose */ false>);
10514742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, multtransposeadd, MatMultTransposeAdd_SeqDense, MatMultAdd_Dispatch_</* transpose */ true>);
10524742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, matmultnumeric, MatMatMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ false, /* transpose_B */ false>);
10534742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, mattransposemultnumeric, MatMatTransposeMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ false, /* transpose_B */ true>);
10544742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, transposematmultnumeric, MatTransposeMatMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ true, /* transpose_B */ false>);
10554742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, axpy, MatAXPY_SeqDense, AXPY);
10564742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, choleskyfactor, MatCholeskyFactor_SeqDense, SolveCholesky::Factor);
10574742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, lufactor, MatLUFactor_SeqDense, SolveLU::Factor);
10584742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, getcolumnvector, MatGetColumnVector_SeqDense, GetColumnVector);
10594742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, scale, MatScale_SeqDense, Scale);
10604742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, shift, MatShift_SeqDense, Shift);
10614742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, copy, MatCopy_SeqDense, Copy);
10624742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, zeroentries, MatZeroEntries_SeqDense, ZeroEntries);
10634742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, setup, MatSetUp_SeqDense, SetUp);
10644742e46bSJacob Faibussowitsch   MatSetOp_CUPM(to_host, A, setrandom, MatSetRandom_SeqDense, SetRandom);
10654742e46bSJacob Faibussowitsch   // seemingly always the same
10664742e46bSJacob Faibussowitsch   A->ops->productsetfromoptions = MatProductSetFromOptions_SeqDense;
10674742e46bSJacob Faibussowitsch 
10684742e46bSJacob Faibussowitsch   if (const auto cmat = mimpl->cmat) PetscCall(MatBindToCPU(cmat, to_host));
10694742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
10704742e46bSJacob Faibussowitsch }
10714742e46bSJacob Faibussowitsch 
10724742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
10734742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_SeqDenseCUPM_SeqDense(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept
10744742e46bSJacob Faibussowitsch {
10754742e46bSJacob Faibussowitsch   PetscFunctionBegin;
10764742e46bSJacob Faibussowitsch   PetscCall(Convert_Dispatch_</* to host */ true>(M, type, reuse, newmat));
10774742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
10784742e46bSJacob Faibussowitsch }
10794742e46bSJacob Faibussowitsch 
10804742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
10814742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_SeqDense_SeqDenseCUPM(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept
10824742e46bSJacob Faibussowitsch {
10834742e46bSJacob Faibussowitsch   PetscFunctionBegin;
10844742e46bSJacob Faibussowitsch   PetscCall(Convert_Dispatch_</* to host */ false>(M, type, reuse, newmat));
10854742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
10864742e46bSJacob Faibussowitsch }
10874742e46bSJacob Faibussowitsch 
10884742e46bSJacob Faibussowitsch // ==========================================================================================
10894742e46bSJacob Faibussowitsch 
10904742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
10914742e46bSJacob Faibussowitsch template <PetscMemType mtype, PetscMemoryAccessMode access>
10924742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::GetArray(Mat m, PetscScalar **array, PetscDeviceContext dctx) noexcept
10934742e46bSJacob Faibussowitsch {
10944742e46bSJacob Faibussowitsch   constexpr auto hostmem     = PetscMemTypeHost(mtype);
10954742e46bSJacob Faibussowitsch   constexpr auto read_access = PetscMemoryAccessRead(access);
10964742e46bSJacob Faibussowitsch 
10974742e46bSJacob Faibussowitsch   PetscFunctionBegin;
10984742e46bSJacob Faibussowitsch   static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), "");
10994742e46bSJacob Faibussowitsch   PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
11004742e46bSJacob Faibussowitsch   if (hostmem) {
11014742e46bSJacob Faibussowitsch     if (read_access) {
11024742e46bSJacob Faibussowitsch       PetscCall(DeviceToHost_(m, dctx));
11034742e46bSJacob Faibussowitsch     } else if (!MatIMPLCast(m)->v) {
11044742e46bSJacob Faibussowitsch       // MatCreateSeqDenseCUPM may not allocate CPU memory. Allocate if needed
11054742e46bSJacob Faibussowitsch       PetscCall(MatSeqDenseSetPreallocation(m, nullptr));
11064742e46bSJacob Faibussowitsch     }
11074742e46bSJacob Faibussowitsch     *array = MatIMPLCast(m)->v;
11084742e46bSJacob Faibussowitsch   } else {
11094742e46bSJacob Faibussowitsch     if (read_access) {
11104742e46bSJacob Faibussowitsch       PetscCall(HostToDevice_(m, dctx));
11114742e46bSJacob Faibussowitsch     } else if (!MatCUPMCast(m)->d_v) {
11124742e46bSJacob Faibussowitsch       // write-only
11134742e46bSJacob Faibussowitsch       PetscCall(SetPreallocation(m, dctx, nullptr));
11144742e46bSJacob Faibussowitsch     }
11154742e46bSJacob Faibussowitsch     *array = MatCUPMCast(m)->d_v;
11164742e46bSJacob Faibussowitsch   }
11174742e46bSJacob Faibussowitsch   if (PetscMemoryAccessWrite(access)) {
11184742e46bSJacob Faibussowitsch     m->offloadmask = hostmem ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
11194742e46bSJacob Faibussowitsch     PetscCall(PetscObjectStateIncrease(PetscObjectCast(m)));
11204742e46bSJacob Faibussowitsch   }
11214742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
11224742e46bSJacob Faibussowitsch }
11234742e46bSJacob Faibussowitsch 
11244742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
11254742e46bSJacob Faibussowitsch template <PetscMemType mtype, PetscMemoryAccessMode access>
11264742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreArray(Mat m, PetscScalar **array, PetscDeviceContext) noexcept
11274742e46bSJacob Faibussowitsch {
11284742e46bSJacob Faibussowitsch   PetscFunctionBegin;
11294742e46bSJacob Faibussowitsch   static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), "");
11304742e46bSJacob Faibussowitsch   if (PetscMemoryAccessWrite(access)) {
11314742e46bSJacob Faibussowitsch     // WRITE or READ_WRITE
11324742e46bSJacob Faibussowitsch     m->offloadmask = PetscMemTypeHost(mtype) ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
11334742e46bSJacob Faibussowitsch     PetscCall(PetscObjectStateIncrease(PetscObjectCast(m)));
11344742e46bSJacob Faibussowitsch   }
11354742e46bSJacob Faibussowitsch   *array = nullptr;
11364742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
11374742e46bSJacob Faibussowitsch }
11384742e46bSJacob Faibussowitsch 
11394742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
11404742e46bSJacob Faibussowitsch template <PetscMemoryAccessMode access>
11414742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::GetArrayAndMemType(Mat m, PetscScalar **array, PetscMemType *mtype, PetscDeviceContext dctx) noexcept
11424742e46bSJacob Faibussowitsch {
11434742e46bSJacob Faibussowitsch   PetscFunctionBegin;
11444742e46bSJacob Faibussowitsch   PetscCall(GetArray<PETSC_MEMTYPE_DEVICE, access>(m, array, dctx));
11454742e46bSJacob Faibussowitsch   if (mtype) *mtype = PETSC_MEMTYPE_CUPM();
11464742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
11474742e46bSJacob Faibussowitsch }
11484742e46bSJacob Faibussowitsch 
11494742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
11504742e46bSJacob Faibussowitsch template <PetscMemoryAccessMode access>
11514742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreArrayAndMemType(Mat m, PetscScalar **array, PetscDeviceContext dctx) noexcept
11524742e46bSJacob Faibussowitsch {
11534742e46bSJacob Faibussowitsch   PetscFunctionBegin;
11544742e46bSJacob Faibussowitsch   PetscCall(RestoreArray<PETSC_MEMTYPE_DEVICE, access>(m, array, dctx));
11554742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
11564742e46bSJacob Faibussowitsch }
11574742e46bSJacob Faibussowitsch 
11584742e46bSJacob Faibussowitsch // ==========================================================================================
11594742e46bSJacob Faibussowitsch 
11604742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
11614742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::PlaceArray(Mat A, const PetscScalar *array) noexcept
11624742e46bSJacob Faibussowitsch {
11634742e46bSJacob Faibussowitsch   const auto mimpl = MatIMPLCast(A);
11644742e46bSJacob Faibussowitsch   const auto mcu   = MatCUPMCast(A);
11654742e46bSJacob Faibussowitsch 
11664742e46bSJacob Faibussowitsch   PetscFunctionBegin;
11674742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
11684742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
11694742e46bSJacob Faibussowitsch   PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME());
11704742e46bSJacob Faibussowitsch   if (mimpl->v) {
11714742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
11724742e46bSJacob Faibussowitsch 
11734742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
11744742e46bSJacob Faibussowitsch     PetscCall(HostToDevice_(A, dctx));
11754742e46bSJacob Faibussowitsch   }
11764742e46bSJacob Faibussowitsch   mcu->unplacedarray         = util::exchange(mcu->d_v, const_cast<PetscScalar *>(array));
11774742e46bSJacob Faibussowitsch   mcu->d_unplaced_user_alloc = util::exchange(mcu->d_user_alloc, PETSC_TRUE);
11784742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
11794742e46bSJacob Faibussowitsch }
11804742e46bSJacob Faibussowitsch 
11814742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
11824742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::ReplaceArray(Mat A, const PetscScalar *array) noexcept
11834742e46bSJacob Faibussowitsch {
11844742e46bSJacob Faibussowitsch   const auto mimpl = MatIMPLCast(A);
11854742e46bSJacob Faibussowitsch   const auto mcu   = MatCUPMCast(A);
11864742e46bSJacob Faibussowitsch 
11874742e46bSJacob Faibussowitsch   PetscFunctionBegin;
11884742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
11894742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
11904742e46bSJacob Faibussowitsch   PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME());
11914742e46bSJacob Faibussowitsch   if (!mcu->d_user_alloc) {
11924742e46bSJacob Faibussowitsch     cupmStream_t stream;
11934742e46bSJacob Faibussowitsch 
11944742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&stream));
11954742e46bSJacob Faibussowitsch     PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream));
11964742e46bSJacob Faibussowitsch   }
11974742e46bSJacob Faibussowitsch   mcu->d_v          = const_cast<PetscScalar *>(array);
11984742e46bSJacob Faibussowitsch   mcu->d_user_alloc = PETSC_FALSE;
11994742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
12004742e46bSJacob Faibussowitsch }
12014742e46bSJacob Faibussowitsch 
12024742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
12034742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::ResetArray(Mat A) noexcept
12044742e46bSJacob Faibussowitsch {
12054742e46bSJacob Faibussowitsch   const auto mimpl = MatIMPLCast(A);
12064742e46bSJacob Faibussowitsch   const auto mcu   = MatCUPMCast(A);
12074742e46bSJacob Faibussowitsch 
12084742e46bSJacob Faibussowitsch   PetscFunctionBegin;
12094742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
12104742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
12114742e46bSJacob Faibussowitsch   if (mimpl->v) {
12124742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
12134742e46bSJacob Faibussowitsch 
12144742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
12154742e46bSJacob Faibussowitsch     PetscCall(HostToDevice_(A, dctx));
12164742e46bSJacob Faibussowitsch   }
12174742e46bSJacob Faibussowitsch   mcu->d_v          = util::exchange(mcu->unplacedarray, nullptr);
12184742e46bSJacob Faibussowitsch   mcu->d_user_alloc = mcu->d_unplaced_user_alloc;
12194742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
12204742e46bSJacob Faibussowitsch }
12214742e46bSJacob Faibussowitsch 
12224742e46bSJacob Faibussowitsch // ==========================================================================================
12234742e46bSJacob Faibussowitsch 
12244742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
12254742e46bSJacob Faibussowitsch template <bool transpose_A, bool transpose_B>
12264742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMatMult_Numeric_Dispatch(Mat A, Mat B, Mat C) noexcept
12274742e46bSJacob Faibussowitsch {
12284742e46bSJacob Faibussowitsch   cupmBlasInt_t      m, n, k;
12294742e46bSJacob Faibussowitsch   PetscBool          Aiscupm, Biscupm;
12304742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
12314742e46bSJacob Faibussowitsch   cupmBlasHandle_t   handle;
12324742e46bSJacob Faibussowitsch 
12334742e46bSJacob Faibussowitsch   PetscFunctionBegin;
12344742e46bSJacob Faibussowitsch   PetscCall(PetscCUPMBlasIntCast(C->rmap->n, &m));
12354742e46bSJacob Faibussowitsch   PetscCall(PetscCUPMBlasIntCast(C->cmap->n, &n));
12364742e46bSJacob Faibussowitsch   PetscCall(PetscCUPMBlasIntCast(transpose_A ? A->rmap->n : A->cmap->n, &k));
12374742e46bSJacob Faibussowitsch   if (!m || !n || !k) PetscFunctionReturn(PETSC_SUCCESS);
12384742e46bSJacob Faibussowitsch 
12394742e46bSJacob Faibussowitsch   // we may end up with SEQDENSE as one of the arguments
12404742e46bSJacob Faibussowitsch   // REVIEW ME: how? and why is it not B and C????????
12414742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(PetscObjectCast(A), MATSEQDENSECUPM(), &Aiscupm));
12424742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(PetscObjectCast(B), MATSEQDENSECUPM(), &Biscupm));
12434742e46bSJacob Faibussowitsch   if (!Aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
12444742e46bSJacob Faibussowitsch   if (!Biscupm) PetscCall(MatConvert(B, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &B));
12454742e46bSJacob Faibussowitsch   PetscCall(PetscInfo(C, "Matrix-Matrix product %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " on backend\n", m, k, n));
12464742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx, &handle));
12474742e46bSJacob Faibussowitsch 
12484742e46bSJacob Faibussowitsch   PetscCall(PetscLogGpuTimeBegin());
12494742e46bSJacob Faibussowitsch   {
12504742e46bSJacob Faibussowitsch     const auto one  = cupmScalarCast(1.0);
12514742e46bSJacob Faibussowitsch     const auto zero = cupmScalarCast(0.0);
12524742e46bSJacob Faibussowitsch     const auto da   = DeviceArrayRead(dctx, A);
12534742e46bSJacob Faibussowitsch     const auto db   = DeviceArrayRead(dctx, B);
12544742e46bSJacob Faibussowitsch     const auto dc   = DeviceArrayWrite(dctx, C);
12554742e46bSJacob Faibussowitsch     PetscInt   alda, blda, clda;
12564742e46bSJacob Faibussowitsch 
12574742e46bSJacob Faibussowitsch     PetscCall(MatDenseGetLDA(A, &alda));
12584742e46bSJacob Faibussowitsch     PetscCall(MatDenseGetLDA(B, &blda));
12594742e46bSJacob Faibussowitsch     PetscCall(MatDenseGetLDA(C, &clda));
12604742e46bSJacob 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));
12614742e46bSJacob Faibussowitsch   }
12624742e46bSJacob Faibussowitsch   PetscCall(PetscLogGpuTimeEnd());
12634742e46bSJacob Faibussowitsch 
12644742e46bSJacob Faibussowitsch   PetscCall(PetscLogGpuFlops(1.0 * m * n * k + 1.0 * m * n * (k - 1)));
12654742e46bSJacob Faibussowitsch   if (!Aiscupm) PetscCall(MatConvert(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A));
12664742e46bSJacob Faibussowitsch   if (!Biscupm) PetscCall(MatConvert(B, MATSEQDENSE, MAT_INPLACE_MATRIX, &B));
12674742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
12684742e46bSJacob Faibussowitsch }
12694742e46bSJacob Faibussowitsch 
12704742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
12714742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Copy(Mat A, Mat B, MatStructure str) noexcept
12724742e46bSJacob Faibussowitsch {
12734742e46bSJacob Faibussowitsch   const auto m = A->rmap->n;
12744742e46bSJacob Faibussowitsch   const auto n = A->cmap->n;
12754742e46bSJacob Faibussowitsch 
12764742e46bSJacob Faibussowitsch   PetscFunctionBegin;
12774742e46bSJacob Faibussowitsch   PetscAssert(m == B->rmap->n && n == B->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "size(B) != size(A)");
12784742e46bSJacob Faibussowitsch   // The two matrices must have the same copy implementation to be eligible for fast copy
12794742e46bSJacob Faibussowitsch   if (A->ops->copy == B->ops->copy) {
12804742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
12814742e46bSJacob Faibussowitsch     cupmStream_t       stream;
12824742e46bSJacob Faibussowitsch 
12834742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx, &stream));
12844742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeBegin());
12854742e46bSJacob Faibussowitsch     {
12864742e46bSJacob Faibussowitsch       const auto va = DeviceArrayRead(dctx, A);
12874742e46bSJacob Faibussowitsch       const auto vb = DeviceArrayWrite(dctx, B);
12884742e46bSJacob Faibussowitsch       // order is important, DeviceArrayRead/Write() might call SetPreallocation() which sets
12894742e46bSJacob Faibussowitsch       // lda!
12904742e46bSJacob Faibussowitsch       const auto lda_a = MatIMPLCast(A)->lda;
12914742e46bSJacob Faibussowitsch       const auto lda_b = MatIMPLCast(B)->lda;
12924742e46bSJacob Faibussowitsch 
12934742e46bSJacob Faibussowitsch       if (lda_a > m || lda_b > m) {
12944742e46bSJacob 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());
12954742e46bSJacob 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());
12964742e46bSJacob Faibussowitsch         PetscCall(PetscCUPMMemcpy2DAsync(vb.data(), lda_b, va.data(), lda_a, m, n, cupmMemcpyDeviceToDevice, stream));
12974742e46bSJacob Faibussowitsch       } else {
12984742e46bSJacob Faibussowitsch         PetscCall(PetscCUPMMemcpyAsync(vb.data(), va.data(), m * n, cupmMemcpyDeviceToDevice, stream));
12994742e46bSJacob Faibussowitsch       }
13004742e46bSJacob Faibussowitsch     }
13014742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeEnd());
13024742e46bSJacob Faibussowitsch   } else {
13034742e46bSJacob Faibussowitsch     PetscCall(MatCopy_Basic(A, B, str));
13044742e46bSJacob Faibussowitsch   }
13054742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
13064742e46bSJacob Faibussowitsch }
13074742e46bSJacob Faibussowitsch 
13084742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
13094742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::ZeroEntries(Mat m) noexcept
13104742e46bSJacob Faibussowitsch {
13114742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
13124742e46bSJacob Faibussowitsch   cupmStream_t       stream;
13134742e46bSJacob Faibussowitsch 
13144742e46bSJacob Faibussowitsch   PetscFunctionBegin;
13154742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx, &stream));
13164742e46bSJacob Faibussowitsch   PetscCall(PetscLogGpuTimeBegin());
13174742e46bSJacob Faibussowitsch   {
13184742e46bSJacob Faibussowitsch     const auto va  = DeviceArrayWrite(dctx, m);
13194742e46bSJacob Faibussowitsch     const auto lda = MatIMPLCast(m)->lda;
13204742e46bSJacob Faibussowitsch     const auto ma  = m->rmap->n;
13214742e46bSJacob Faibussowitsch     const auto na  = m->cmap->n;
13224742e46bSJacob Faibussowitsch 
13234742e46bSJacob Faibussowitsch     if (lda > ma) {
13244742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMemset2DAsync(va.data(), lda, 0, ma, na, stream));
13254742e46bSJacob Faibussowitsch     } else {
13264742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMemsetAsync(va.data(), 0, ma * na, stream));
13274742e46bSJacob Faibussowitsch     }
13284742e46bSJacob Faibussowitsch   }
13294742e46bSJacob Faibussowitsch   PetscCall(PetscLogGpuTimeEnd());
13304742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
13314742e46bSJacob Faibussowitsch }
13324742e46bSJacob Faibussowitsch 
1333*2ea277ceSJacob Faibussowitsch namespace detail
1334*2ea277ceSJacob Faibussowitsch {
1335*2ea277ceSJacob Faibussowitsch 
1336*2ea277ceSJacob Faibussowitsch // ==========================================================================================
1337*2ea277ceSJacob Faibussowitsch // SubMatIndexFunctor
1338*2ea277ceSJacob Faibussowitsch //
1339*2ea277ceSJacob Faibussowitsch // Iterator which permutes a linear index range into matrix indices for am nrows x ncols
1340*2ea277ceSJacob Faibussowitsch // submat with leading dimension lda. Essentially SubMatIndexFunctor(i) returns the index for
1341*2ea277ceSJacob Faibussowitsch // the i'th sequential entry in the matrix.
1342*2ea277ceSJacob Faibussowitsch // ==========================================================================================
1343*2ea277ceSJacob Faibussowitsch template <typename T>
1344*2ea277ceSJacob Faibussowitsch struct SubMatIndexFunctor {
1345*2ea277ceSJacob Faibussowitsch   PETSC_HOSTDEVICE_INLINE_DECL T operator()(T x) const noexcept { return ((x / nrows) * lda) + (x % nrows); }
1346*2ea277ceSJacob Faibussowitsch 
1347*2ea277ceSJacob Faibussowitsch   PetscInt nrows;
1348*2ea277ceSJacob Faibussowitsch   PetscInt ncols;
1349*2ea277ceSJacob Faibussowitsch   PetscInt lda;
1350*2ea277ceSJacob Faibussowitsch };
1351*2ea277ceSJacob Faibussowitsch 
1352*2ea277ceSJacob Faibussowitsch template <typename Iterator>
1353*2ea277ceSJacob Faibussowitsch struct SubMatrixIterator : MatrixIteratorBase<Iterator, SubMatIndexFunctor<typename thrust::iterator_difference<Iterator>::type>> {
1354*2ea277ceSJacob Faibussowitsch   using base_type = MatrixIteratorBase<Iterator, SubMatIndexFunctor<typename thrust::iterator_difference<Iterator>::type>>;
1355*2ea277ceSJacob Faibussowitsch 
1356*2ea277ceSJacob Faibussowitsch   using iterator = typename base_type::iterator;
1357*2ea277ceSJacob Faibussowitsch 
1358*2ea277ceSJacob Faibussowitsch   constexpr SubMatrixIterator(Iterator first, Iterator last, PetscInt nrows, PetscInt ncols, PetscInt lda) noexcept :
1359*2ea277ceSJacob Faibussowitsch     base_type{
1360*2ea277ceSJacob Faibussowitsch       std::move(first), std::move(last), {nrows, ncols, lda}
1361*2ea277ceSJacob Faibussowitsch   }
1362*2ea277ceSJacob Faibussowitsch   {
1363*2ea277ceSJacob Faibussowitsch   }
1364*2ea277ceSJacob Faibussowitsch 
1365*2ea277ceSJacob Faibussowitsch   PETSC_NODISCARD iterator end() const noexcept { return this->begin() + (this->func.nrows * this->func.ncols); }
1366*2ea277ceSJacob Faibussowitsch };
1367*2ea277ceSJacob Faibussowitsch 
1368*2ea277ceSJacob Faibussowitsch namespace
1369*2ea277ceSJacob Faibussowitsch {
1370*2ea277ceSJacob Faibussowitsch 
1371*2ea277ceSJacob Faibussowitsch template <typename T>
1372*2ea277ceSJacob 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
1373*2ea277ceSJacob Faibussowitsch {
1374*2ea277ceSJacob Faibussowitsch   const auto nrows = rend - rstart;
1375*2ea277ceSJacob Faibussowitsch   const auto ncols = cend - cstart;
1376*2ea277ceSJacob Faibussowitsch   const auto dptr  = thrust::device_pointer_cast(ptr);
1377*2ea277ceSJacob Faibussowitsch 
1378*2ea277ceSJacob Faibussowitsch   return {dptr + (rstart * lda) + cstart, dptr + ((rstart + nrows) * lda) + cstart, nrows, ncols, lda};
1379*2ea277ceSJacob Faibussowitsch }
1380*2ea277ceSJacob Faibussowitsch 
1381*2ea277ceSJacob Faibussowitsch } // namespace
1382*2ea277ceSJacob Faibussowitsch 
1383*2ea277ceSJacob Faibussowitsch } // namespace detail
1384*2ea277ceSJacob Faibussowitsch 
13854742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
13864742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Scale(Mat A, PetscScalar alpha) noexcept
13874742e46bSJacob Faibussowitsch {
1388*2ea277ceSJacob Faibussowitsch   const auto         m = A->rmap->n;
1389*2ea277ceSJacob Faibussowitsch   const auto         n = A->cmap->n;
13904742e46bSJacob Faibussowitsch   const auto         N = m * n;
13914742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
13924742e46bSJacob Faibussowitsch 
13934742e46bSJacob Faibussowitsch   PetscFunctionBegin;
1394*2ea277ceSJacob Faibussowitsch   PetscCall(PetscInfo(A, "Performing Scale %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m, n));
1395*2ea277ceSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx));
13964742e46bSJacob Faibussowitsch   {
13974742e46bSJacob Faibussowitsch     const auto da  = DeviceArrayReadWrite(dctx, A);
1398*2ea277ceSJacob Faibussowitsch     const auto lda = MatIMPLCast(A)->lda;
13994742e46bSJacob Faibussowitsch 
14004742e46bSJacob Faibussowitsch     if (lda > m) {
1401*2ea277ceSJacob Faibussowitsch       cupmStream_t stream;
1402*2ea277ceSJacob Faibussowitsch 
1403*2ea277ceSJacob Faibussowitsch       PetscCall(GetHandlesFrom_(dctx, &stream));
1404*2ea277ceSJacob Faibussowitsch       // clang-format off
1405*2ea277ceSJacob Faibussowitsch       PetscCallThrust(
1406*2ea277ceSJacob Faibussowitsch         const auto sub_mat = detail::make_submat_iterator(0, m, 0, n, lda, da.data());
1407*2ea277ceSJacob Faibussowitsch 
1408*2ea277ceSJacob Faibussowitsch         THRUST_CALL(
1409*2ea277ceSJacob Faibussowitsch           thrust::transform,
1410*2ea277ceSJacob Faibussowitsch           stream,
1411*2ea277ceSJacob Faibussowitsch           sub_mat.begin(), sub_mat.end(), sub_mat.begin(),
1412*2ea277ceSJacob Faibussowitsch           device::cupm::functors::make_times_equals(alpha)
1413*2ea277ceSJacob Faibussowitsch         )
1414*2ea277ceSJacob Faibussowitsch       );
1415*2ea277ceSJacob Faibussowitsch       // clang-format on
14164742e46bSJacob Faibussowitsch     } else {
1417*2ea277ceSJacob Faibussowitsch       const auto       cu_alpha = cupmScalarCast(alpha);
1418*2ea277ceSJacob Faibussowitsch       cupmBlasHandle_t handle;
1419*2ea277ceSJacob Faibussowitsch 
1420*2ea277ceSJacob Faibussowitsch       PetscCall(GetHandlesFrom_(dctx, &handle));
1421*2ea277ceSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeBegin());
14224742e46bSJacob Faibussowitsch       PetscCallCUPMBLAS(cupmBlasXscal(handle, N, &cu_alpha, da.cupmdata(), 1));
14234742e46bSJacob Faibussowitsch       PetscCall(PetscLogGpuTimeEnd());
1424*2ea277ceSJacob Faibussowitsch     }
1425*2ea277ceSJacob Faibussowitsch   }
14264742e46bSJacob Faibussowitsch   PetscCall(PetscLogGpuFlops(N));
14274742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
14284742e46bSJacob Faibussowitsch }
14294742e46bSJacob Faibussowitsch 
14304742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
14314742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Shift(Mat A, PetscScalar alpha) noexcept
14324742e46bSJacob Faibussowitsch {
14334742e46bSJacob Faibussowitsch   const auto         m = A->rmap->n;
14344742e46bSJacob Faibussowitsch   const auto         n = A->cmap->n;
14354742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
14364742e46bSJacob Faibussowitsch 
14374742e46bSJacob Faibussowitsch   PetscFunctionBegin;
14384742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx));
14394742e46bSJacob Faibussowitsch   PetscCall(PetscInfo(A, "Performing Shift %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m, n));
1440*2ea277ceSJacob Faibussowitsch   PetscCall(DiagonalUnaryTransform(A, 0, m, n, dctx, device::cupm::functors::make_plus_equals(alpha)));
14414742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
14424742e46bSJacob Faibussowitsch }
14434742e46bSJacob Faibussowitsch 
14444742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
14454742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::AXPY(Mat Y, PetscScalar alpha, Mat X, MatStructure) noexcept
14464742e46bSJacob Faibussowitsch {
14474742e46bSJacob Faibussowitsch   const auto         m_x = X->rmap->n, m_y = Y->rmap->n;
14484742e46bSJacob Faibussowitsch   const auto         n_x = X->cmap->n, n_y = Y->cmap->n;
14494742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
14504742e46bSJacob Faibussowitsch   cupmBlasHandle_t   handle;
14514742e46bSJacob Faibussowitsch 
14524742e46bSJacob Faibussowitsch   PetscFunctionBegin;
14534742e46bSJacob Faibussowitsch   if (!m_x || !n_x) PetscFunctionReturn(PETSC_SUCCESS);
14544742e46bSJacob Faibussowitsch   PetscCall(PetscInfo(Y, "Performing AXPY %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m_y, n_y));
14554742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx, &handle));
14564742e46bSJacob Faibussowitsch   {
14574742e46bSJacob Faibussowitsch     const auto N        = m_x * n_x;
14584742e46bSJacob Faibussowitsch     const auto dx       = DeviceArrayRead(dctx, X);
14594742e46bSJacob Faibussowitsch     const auto dy       = alpha == 0.0 ? DeviceArrayWrite(dctx, Y).cupmdata() : DeviceArrayReadWrite(dctx, Y).cupmdata();
14604742e46bSJacob Faibussowitsch     const auto ldax     = static_cast<cupmBlasInt_t>(MatIMPLCast(X)->lda);
14614742e46bSJacob Faibussowitsch     const auto lday     = static_cast<cupmBlasInt_t>(MatIMPLCast(Y)->lda);
14624742e46bSJacob Faibussowitsch     const auto cu_alpha = cupmScalarCast(alpha);
14634742e46bSJacob Faibussowitsch 
14644742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeBegin());
14654742e46bSJacob Faibussowitsch     if (ldax > m_x || lday > m_x) {
14664742e46bSJacob Faibussowitsch       for (cupmBlasInt_t j = 0; j < n_x; j++) PetscCallCUPMBLAS(cupmBlasXaxpy(handle, m_x, &cu_alpha, dx.cupmdata() + j * ldax, 1, dy + j * lday, 1));
14674742e46bSJacob Faibussowitsch     } else {
14684742e46bSJacob Faibussowitsch       PetscCallCUPMBLAS(cupmBlasXaxpy(handle, N, &cu_alpha, dx.cupmdata(), 1, dy, 1));
14694742e46bSJacob Faibussowitsch     }
14704742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuTimeEnd());
14714742e46bSJacob Faibussowitsch     PetscCall(PetscLogGpuFlops(PetscMax(2 * N - 1, 0)));
14724742e46bSJacob Faibussowitsch   }
14734742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
14744742e46bSJacob Faibussowitsch }
14754742e46bSJacob Faibussowitsch 
14764742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
14774742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::Duplicate(Mat A, MatDuplicateOption opt, Mat *B) noexcept
14784742e46bSJacob Faibussowitsch {
14794742e46bSJacob Faibussowitsch   const auto         hopt = (opt == MAT_COPY_VALUES && A->offloadmask != PETSC_OFFLOAD_CPU) ? MAT_DO_NOT_COPY_VALUES : opt;
14804742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
14814742e46bSJacob Faibussowitsch 
14824742e46bSJacob Faibussowitsch   PetscFunctionBegin;
14834742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx));
14844742e46bSJacob Faibussowitsch   // do not call SetPreallocation() yet, we call it afterwards??
14854742e46bSJacob Faibussowitsch   PetscCall(MatCreateSeqDenseCUPM<T>(PetscObjectComm(PetscObjectCast(A)), A->rmap->n, A->cmap->n, nullptr, B, dctx, /* preallocate */ false));
14864742e46bSJacob Faibussowitsch   PetscCall(MatDuplicateNoCreate_SeqDense(*B, A, hopt));
14874742e46bSJacob Faibussowitsch   if (opt == MAT_COPY_VALUES && hopt != MAT_COPY_VALUES) PetscCall(Copy(A, *B, SAME_NONZERO_PATTERN));
14884742e46bSJacob Faibussowitsch   // allocate memory if needed
14894742e46bSJacob Faibussowitsch   if (opt != MAT_COPY_VALUES && !MatCUPMCast(*B)->d_v) PetscCall(SetPreallocation(*B, dctx));
14904742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
14914742e46bSJacob Faibussowitsch }
14924742e46bSJacob Faibussowitsch 
14934742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
14944742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::SetRandom(Mat A, PetscRandom rng) noexcept
14954742e46bSJacob Faibussowitsch {
14964742e46bSJacob Faibussowitsch   PetscBool device;
14974742e46bSJacob Faibussowitsch 
14984742e46bSJacob Faibussowitsch   PetscFunctionBegin;
14994742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare(PetscObjectCast(rng), PETSCDEVICERAND(), &device));
15004742e46bSJacob Faibussowitsch   if (device) {
15014742e46bSJacob Faibussowitsch     const auto         m = A->rmap->n;
15024742e46bSJacob Faibussowitsch     const auto         n = A->cmap->n;
15034742e46bSJacob Faibussowitsch     PetscDeviceContext dctx;
15044742e46bSJacob Faibussowitsch 
15054742e46bSJacob Faibussowitsch     PetscCall(GetHandles_(&dctx));
15064742e46bSJacob Faibussowitsch     {
15074742e46bSJacob Faibussowitsch       const auto a = DeviceArrayWrite(dctx, A);
15084742e46bSJacob Faibussowitsch       PetscInt   lda;
15094742e46bSJacob Faibussowitsch 
15104742e46bSJacob Faibussowitsch       PetscCall(MatDenseGetLDA(A, &lda));
15114742e46bSJacob Faibussowitsch       if (lda > m) {
15124742e46bSJacob Faibussowitsch         for (PetscInt i = 0; i < n; i++) PetscCall(PetscRandomGetValues(rng, m, a.data() + i * lda));
15134742e46bSJacob Faibussowitsch       } else {
15144742e46bSJacob Faibussowitsch         PetscInt mn;
15154742e46bSJacob Faibussowitsch 
15164742e46bSJacob Faibussowitsch         PetscCall(PetscIntMultError(m, n, &mn));
15174742e46bSJacob Faibussowitsch         PetscCall(PetscRandomGetValues(rng, mn, a));
15184742e46bSJacob Faibussowitsch       }
15194742e46bSJacob Faibussowitsch     }
15204742e46bSJacob Faibussowitsch   } else {
15214742e46bSJacob Faibussowitsch     PetscCall(MatSetRandom_SeqDense(A, rng));
15224742e46bSJacob Faibussowitsch   }
15234742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15244742e46bSJacob Faibussowitsch }
15254742e46bSJacob Faibussowitsch 
15264742e46bSJacob Faibussowitsch // ==========================================================================================
15274742e46bSJacob Faibussowitsch 
15284742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
15294742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::GetColumnVector(Mat A, Vec v, PetscInt col) noexcept
15304742e46bSJacob Faibussowitsch {
15314742e46bSJacob Faibussowitsch   const auto         offloadmask = A->offloadmask;
15324742e46bSJacob Faibussowitsch   const auto         n           = A->rmap->n;
15334742e46bSJacob Faibussowitsch   const auto         col_offset  = [&](const PetscScalar *ptr) { return ptr + col * MatIMPLCast(A)->lda; };
15344742e46bSJacob Faibussowitsch   PetscBool          viscupm;
15354742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
15364742e46bSJacob Faibussowitsch   cupmStream_t       stream;
15374742e46bSJacob Faibussowitsch 
15384742e46bSJacob Faibussowitsch   PetscFunctionBegin;
15394742e46bSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(v), &viscupm, VecSeq_CUPM::VECSEQCUPM(), VecSeq_CUPM::VECMPICUPM(), VecSeq_CUPM::VECCUPM(), ""));
15404742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx, &stream));
15414742e46bSJacob Faibussowitsch   if (viscupm && !v->boundtocpu) {
15424742e46bSJacob Faibussowitsch     const auto x = VecSeq_CUPM::DeviceArrayWrite(dctx, v);
15434742e46bSJacob Faibussowitsch 
15444742e46bSJacob Faibussowitsch     // update device data
15454742e46bSJacob Faibussowitsch     if (PetscOffloadDevice(offloadmask)) {
15464742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMemcpyAsync(x.data(), col_offset(DeviceArrayRead(dctx, A)), n, cupmMemcpyDeviceToDevice, stream));
15474742e46bSJacob Faibussowitsch     } else {
15484742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMemcpyAsync(x.data(), col_offset(HostArrayRead(dctx, A)), n, cupmMemcpyHostToDevice, stream));
15494742e46bSJacob Faibussowitsch     }
15504742e46bSJacob Faibussowitsch   } else {
15514742e46bSJacob Faibussowitsch     PetscScalar *x;
15524742e46bSJacob Faibussowitsch 
15534742e46bSJacob Faibussowitsch     // update host data
15544742e46bSJacob Faibussowitsch     PetscCall(VecGetArrayWrite(v, &x));
15554742e46bSJacob Faibussowitsch     if (PetscOffloadUnallocated(offloadmask) || PetscOffloadHost(offloadmask)) {
15564742e46bSJacob Faibussowitsch       PetscCall(PetscArraycpy(x, col_offset(HostArrayRead(dctx, A)), n));
15574742e46bSJacob Faibussowitsch     } else if (PetscOffloadDevice(offloadmask)) {
15584742e46bSJacob Faibussowitsch       PetscCall(PetscCUPMMemcpyAsync(x, col_offset(DeviceArrayRead(dctx, A)), n, cupmMemcpyDeviceToHost, stream));
15594742e46bSJacob Faibussowitsch     }
15604742e46bSJacob Faibussowitsch     PetscCall(VecRestoreArrayWrite(v, &x));
15614742e46bSJacob Faibussowitsch   }
15624742e46bSJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15634742e46bSJacob Faibussowitsch }
15644742e46bSJacob Faibussowitsch 
15654742e46bSJacob Faibussowitsch template <device::cupm::DeviceType T>
15664742e46bSJacob Faibussowitsch template <PetscMemoryAccessMode access>
15674742e46bSJacob Faibussowitsch inline PetscErrorCode MatDense_Seq_CUPM<T>::GetColumnVec(Mat A, PetscInt col, Vec *v) noexcept
15684742e46bSJacob Faibussowitsch {
15694742e46bSJacob Faibussowitsch   using namespace vec::cupm;
15704742e46bSJacob Faibussowitsch   const auto         mimpl = MatIMPLCast(A);
15714742e46bSJacob Faibussowitsch   PetscDeviceContext dctx;
15724742e46bSJacob Faibussowitsch 
15734742e46bSJacob Faibussowitsch   PetscFunctionBegin;
15744742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
15754742e46bSJacob Faibussowitsch   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
15764742e46bSJacob Faibussowitsch   mimpl->vecinuse = col + 1;
15774742e46bSJacob Faibussowitsch   PetscCall(GetHandles_(&dctx));
15784742e46bSJacob Faibussowitsch   PetscCall(GetArray<PETSC_MEMTYPE_DEVICE, access>(A, const_cast<PetscScalar **>(&mimpl->ptrinuse), dctx));
15794742e46bSJacob Faibussowitsch   if (!mimpl->cvec) {
15804742e46bSJacob Faibussowitsch     // we pass the data of A, to prevent allocating needless GPU memory the first time
15814742e46bSJacob Faibussowitsch     // VecCUPMPlaceArray is called
15824742e46bSJacob Faibussowitsch     PetscCall(VecCreateSeqCUPMWithArraysAsync<T>(PetscObjectComm(PetscObjectCast(A)), A->rmap->bs, A->rmap->n, nullptr, mimpl->ptrinuse, &mimpl->cvec));
15834742e46bSJacob Faibussowitsch   }
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 {
16164742e46bSJacob Faibussowitsch   Mat                fact;
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
18294742e46bSJacob Faibussowitsch 
18304742e46bSJacob Faibussowitsch #endif // __cplusplus
18314742e46bSJacob Faibussowitsch 
18324742e46bSJacob Faibussowitsch #endif // PETSCMATSEQDENSECUPM_HPP
1833