Actual source code: matseqdensecupm.hpp

      1: #pragma once

      3: #include <petsc/private/matdensecupmimpl.h>
      4: #include <../src/mat/impls/dense/seq/dense.h>

      6: #include <petsc/private/deviceimpl.h>
      7: #include <petsc/private/randomimpl.h>
      8: #include <petsc/private/vecimpl.h>
      9: #include <petsc/private/cupmobject.hpp>
     10: #include <petsc/private/cupmsolverinterface.hpp>

     12: #include <petsc/private/cpp/type_traits.hpp>
     13: #include <petsc/private/cpp/utility.hpp>

     15: #include <../src/vec/vec/impls/seq/cupm/vecseqcupm.hpp>

     17: namespace Petsc
     18: {

     20: namespace mat
     21: {

     23: namespace cupm
     24: {

     26: namespace impl
     27: {

     29: template <device::cupm::DeviceType T>
     30: class PETSC_SINGLE_LIBRARY_VISIBILITY_INTERNAL MatDense_Seq_CUPM : MatDense_CUPM<T, MatDense_Seq_CUPM<T>> {
     31: public:
     32:   MATDENSECUPM_HEADER(T, MatDense_Seq_CUPM<T>);

     34: private:
     35:   struct Mat_SeqDenseCUPM {
     36:     PetscScalar *d_v;           // pointer to the matrix on the GPU
     37:     PetscScalar *unplacedarray; // if one called MatCUPMDensePlaceArray(), this is where it stashed the original
     38:     bool         d_user_alloc;
     39:     bool         d_unplaced_user_alloc;
     40:     // factorization support
     41:     cupmBlasInt_t *d_fact_ipiv;  // device pivots
     42:     cupmScalar_t  *d_fact_tau;   // device QR tau vector
     43:     cupmBlasInt_t *d_fact_info;  // device info
     44:     cupmScalar_t  *d_fact_work;  // device workspace
     45:     cupmBlasInt_t  d_fact_lwork; // size of device workspace
     46:     // workspace
     47:     Vec workvec;
     48:   };

     50:   static PetscErrorCode SetPreallocation_(Mat, PetscDeviceContext, PetscScalar *) noexcept;

     52:   static PetscErrorCode HostToDevice_(Mat, PetscDeviceContext) noexcept;
     53:   static PetscErrorCode DeviceToHost_(Mat, PetscDeviceContext) noexcept;

     55:   static PetscErrorCode CheckCUPMSolverInfo_(const cupmBlasInt_t *, cupmStream_t) noexcept;

     57:   template <typename Derived>
     58:   struct SolveCommon;
     59:   struct SolveQR;
     60:   struct SolveCholesky;
     61:   struct SolveLU;

     63:   template <typename Solver, bool transpose>
     64:   static PetscErrorCode MatSolve_Factored_Dispatch_(Mat, Vec, Vec) noexcept;
     65:   template <typename Solver, bool transpose>
     66:   static PetscErrorCode MatMatSolve_Factored_Dispatch_(Mat, Mat, Mat) noexcept;
     67:   template <bool transpose, bool hermitian>
     68:   static PetscErrorCode MatMultAddColumnRange_Dispatch_(Mat, Vec, Vec, Vec, PetscInt, PetscInt) noexcept;
     69:   template <bool transpose, bool hermitian>
     70:   static PetscErrorCode MatMultColumnRange_Dispatch_(Mat, Vec, Vec, PetscInt, PetscInt) noexcept;
     71:   template <bool transpose, bool hermitian>
     72:   static PetscErrorCode MatMultAdd_Dispatch_(Mat, Vec, Vec, Vec) noexcept;

     74:   template <bool to_host>
     75:   static PetscErrorCode Convert_Dispatch_(Mat, MatType, MatReuse, Mat *) noexcept;

     77:   PETSC_NODISCARD static constexpr MatType       MATIMPLCUPM_() noexcept;
     78:   PETSC_NODISCARD static constexpr Mat_SeqDense *MatIMPLCast_(Mat) noexcept;

     80: public:
     81:   PETSC_NODISCARD static constexpr Mat_SeqDenseCUPM *MatCUPMCast(Mat) noexcept;

     83:   // define these by hand since they don't fit the above mold
     84:   PETSC_NODISCARD static constexpr const char *MatConvert_seqdensecupm_seqdense_C() noexcept;
     85:   PETSC_NODISCARD static constexpr const char *MatProductSetFromOptions_seqaij_seqdensecupm_C() noexcept;

     87:   static PetscErrorCode Create(Mat) noexcept;
     88:   static PetscErrorCode Destroy(Mat) noexcept;
     89:   static PetscErrorCode SetUp(Mat) noexcept;
     90:   static PetscErrorCode Reset(Mat) noexcept;

     92:   static PetscErrorCode BindToCPU(Mat, PetscBool) noexcept;
     93:   static PetscErrorCode Convert_SeqDense_SeqDenseCUPM(Mat, MatType, MatReuse, Mat *) noexcept;
     94:   static PetscErrorCode Convert_SeqDenseCUPM_SeqDense(Mat, MatType, MatReuse, Mat *) noexcept;

     96:   template <PetscMemType, PetscMemoryAccessMode>
     97:   static PetscErrorCode GetArray(Mat, PetscScalar **, PetscDeviceContext) noexcept;
     98:   template <PetscMemType, PetscMemoryAccessMode>
     99:   static PetscErrorCode RestoreArray(Mat, PetscScalar **, PetscDeviceContext) noexcept;
    100:   template <PetscMemoryAccessMode>
    101:   static PetscErrorCode GetArrayAndMemType(Mat, PetscScalar **, PetscMemType *, PetscDeviceContext) noexcept;
    102:   template <PetscMemoryAccessMode>
    103:   static PetscErrorCode RestoreArrayAndMemType(Mat, PetscScalar **, PetscDeviceContext) noexcept;

    105: private:
    106:   template <PetscMemType mtype, PetscMemoryAccessMode mode>
    107:   static PetscErrorCode GetArrayC_(Mat m, PetscScalar **p) noexcept
    108:   {
    109:     PetscDeviceContext dctx;

    111:     PetscFunctionBegin;
    112:     PetscCall(GetHandles_(&dctx));
    113:     PetscCall(GetArray<mtype, mode>(m, p, dctx));
    114:     PetscFunctionReturn(PETSC_SUCCESS);
    115:   }

    117:   template <PetscMemType mtype, PetscMemoryAccessMode mode>
    118:   static PetscErrorCode RestoreArrayC_(Mat m, PetscScalar **p) noexcept
    119:   {
    120:     PetscDeviceContext dctx;

    122:     PetscFunctionBegin;
    123:     PetscCall(GetHandles_(&dctx));
    124:     PetscCall(RestoreArray<mtype, mode>(m, p, dctx));
    125:     PetscFunctionReturn(PETSC_SUCCESS);
    126:   }

    128:   template <PetscMemoryAccessMode mode>
    129:   static PetscErrorCode GetArrayAndMemTypeC_(Mat m, PetscScalar **p, PetscMemType *tp) noexcept
    130:   {
    131:     PetscDeviceContext dctx;

    133:     PetscFunctionBegin;
    134:     PetscCall(GetHandles_(&dctx));
    135:     PetscCall(GetArrayAndMemType<mode>(m, p, tp, dctx));
    136:     PetscFunctionReturn(PETSC_SUCCESS);
    137:   }

    139:   template <PetscMemoryAccessMode mode>
    140:   static PetscErrorCode RestoreArrayAndMemTypeC_(Mat m, PetscScalar **p) noexcept
    141:   {
    142:     PetscDeviceContext dctx;

    144:     PetscFunctionBegin;
    145:     PetscCall(GetHandles_(&dctx));
    146:     PetscCall(RestoreArrayAndMemType<mode>(m, p, dctx));
    147:     PetscFunctionReturn(PETSC_SUCCESS);
    148:   }

    150: public:
    151:   static PetscErrorCode PlaceArray(Mat, const PetscScalar *) noexcept;
    152:   static PetscErrorCode ReplaceArray(Mat, const PetscScalar *) noexcept;
    153:   static PetscErrorCode ResetArray(Mat) noexcept;

    155:   template <bool transpose_A, bool transpose_B>
    156:   static PetscErrorCode MatMatMult_Numeric_Dispatch(Mat, Mat, Mat) noexcept;
    157:   static PetscErrorCode Copy(Mat, Mat, MatStructure) noexcept;
    158:   static PetscErrorCode ZeroEntries(Mat) noexcept;
    159:   static PetscErrorCode Conjugate(Mat) noexcept;
    160:   static PetscErrorCode Scale(Mat, PetscScalar) noexcept;
    161:   static PetscErrorCode AXPY(Mat, PetscScalar, Mat, MatStructure) noexcept;
    162:   static PetscErrorCode Duplicate(Mat, MatDuplicateOption, Mat *) noexcept;
    163:   static PetscErrorCode SetRandom(Mat, PetscRandom) noexcept;

    165:   static PetscErrorCode GetColumnVector(Mat, Vec, PetscInt) noexcept;
    166:   template <PetscMemoryAccessMode>
    167:   static PetscErrorCode GetColumnVec(Mat, PetscInt, Vec *) noexcept;
    168:   template <PetscMemoryAccessMode>
    169:   static PetscErrorCode RestoreColumnVec(Mat, PetscInt, Vec *) noexcept;

    171:   static PetscErrorCode GetFactor(Mat, MatFactorType, Mat *) noexcept;
    172:   static PetscErrorCode InvertFactors(Mat) noexcept;

    174:   static PetscErrorCode GetSubMatrix(Mat, PetscInt, PetscInt, PetscInt, PetscInt, Mat *) noexcept;
    175:   static PetscErrorCode RestoreSubMatrix(Mat, Mat *) noexcept;
    176: };

    178: } // namespace impl

    180: namespace
    181: {

    183: // Declare this here so that the functions below can make use of it
    184: template <device::cupm::DeviceType T>
    185: inline PetscErrorCode MatCreateSeqDenseCUPM(MPI_Comm comm, PetscInt m, PetscInt n, PetscScalar *data, Mat *A, PetscDeviceContext dctx = nullptr, bool preallocate = true) noexcept
    186: {
    187:   PetscFunctionBegin;
    188:   PetscCall(impl::MatDense_Seq_CUPM<T>::CreateIMPLDenseCUPM(comm, m, n, m, n, data, A, dctx, preallocate));
    189:   PetscFunctionReturn(PETSC_SUCCESS);
    190: }

    192: } // anonymous namespace

    194: namespace impl
    195: {

    197: // ==========================================================================================
    198: // MatDense_Seq_CUPM - Private API - Utility
    199: // ==========================================================================================

    201: template <device::cupm::DeviceType T>
    202: inline PetscErrorCode MatDense_Seq_CUPM<T>::SetPreallocation_(Mat m, PetscDeviceContext dctx, PetscScalar *user_device_array) noexcept
    203: {
    204:   const auto   mcu   = MatCUPMCast(m);
    205:   const auto   nrows = m->rmap->n;
    206:   const auto   ncols = m->cmap->n;
    207:   auto        &lda   = MatIMPLCast(m)->lda;
    208:   cupmStream_t stream;

    210:   PetscFunctionBegin;
    211:   PetscCheckTypeName(m, MATSEQDENSECUPM());
    213:   PetscCall(checkCupmBlasIntCast(nrows));
    214:   PetscCall(checkCupmBlasIntCast(ncols));
    215:   PetscCall(GetHandlesFrom_(dctx, &stream));
    216:   if (lda <= 0) lda = nrows;
    217:   if (!mcu->d_user_alloc) PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream));
    218:   if (user_device_array) {
    219:     mcu->d_user_alloc = PETSC_TRUE;
    220:     mcu->d_v          = user_device_array;
    221:   } else {
    222:     std::size_t size;

    224:     mcu->d_user_alloc = PETSC_FALSE;
    225:     size              = lda * ncols;
    226:     PetscCall(PetscCUPMMallocAsync(&mcu->d_v, size, stream));
    227:     PetscCall(PetscCUPMMemsetAsync(mcu->d_v, 0, size, stream));
    228:   }
    229:   m->offloadmask = PETSC_OFFLOAD_GPU;
    230:   PetscFunctionReturn(PETSC_SUCCESS);
    231: }

    233: template <device::cupm::DeviceType T>
    234: inline PetscErrorCode MatDense_Seq_CUPM<T>::HostToDevice_(Mat m, PetscDeviceContext dctx) noexcept
    235: {
    236:   const auto nrows = m->rmap->n;
    237:   const auto ncols = m->cmap->n;
    238:   const auto copy  = m->offloadmask == PETSC_OFFLOAD_CPU || m->offloadmask == PETSC_OFFLOAD_UNALLOCATED;

    240:   PetscFunctionBegin;
    241:   PetscCheckTypeName(m, MATSEQDENSECUPM());
    242:   if (m->boundtocpu) PetscFunctionReturn(PETSC_SUCCESS);
    243:   PetscCall(PetscInfo(m, "%s matrix %" PetscInt_FMT " x %" PetscInt_FMT "\n", copy ? "Copy" : "Reusing", nrows, ncols));
    244:   if (copy) {
    245:     const auto   mcu = MatCUPMCast(m);
    246:     cupmStream_t stream;

    248:     // Allocate GPU memory if not present
    249:     if (!mcu->d_v) PetscCall(SetPreallocation(m, dctx, nullptr));
    250:     PetscCall(GetHandlesFrom_(dctx, &stream));
    251:     PetscCall(PetscLogEventBegin(MAT_DenseCopyToGPU, m, 0, 0, 0));
    252:     {
    253:       const auto mimpl = MatIMPLCast(m);
    254:       const auto lda   = mimpl->lda;
    255:       const auto src   = mimpl->v;
    256:       const auto dest  = mcu->d_v;

    258:       if (lda > nrows) {
    259:         PetscCall(PetscCUPMMemcpy2DAsync(dest, lda, src, lda, nrows, ncols, cupmMemcpyHostToDevice, stream));
    260:       } else {
    261:         PetscCall(PetscCUPMMemcpyAsync(dest, src, lda * ncols, cupmMemcpyHostToDevice, stream));
    262:       }
    263:     }
    264:     PetscCall(PetscLogEventEnd(MAT_DenseCopyToGPU, m, 0, 0, 0));
    265:     // order important, ensure that offloadmask is PETSC_OFFLOAD_BOTH
    266:     m->offloadmask = PETSC_OFFLOAD_BOTH;
    267:   }
    268:   PetscFunctionReturn(PETSC_SUCCESS);
    269: }

    271: template <device::cupm::DeviceType T>
    272: inline PetscErrorCode MatDense_Seq_CUPM<T>::DeviceToHost_(Mat m, PetscDeviceContext dctx) noexcept
    273: {
    274:   const auto nrows = m->rmap->n;
    275:   const auto ncols = m->cmap->n;
    276:   const auto copy  = m->offloadmask == PETSC_OFFLOAD_GPU;

    278:   PetscFunctionBegin;
    279:   PetscCheckTypeName(m, MATSEQDENSECUPM());
    280:   PetscCall(PetscInfo(m, "%s matrix %" PetscInt_FMT " x %" PetscInt_FMT "\n", copy ? "Copy" : "Reusing", nrows, ncols));
    281:   if (copy) {
    282:     const auto   mimpl = MatIMPLCast(m);
    283:     cupmStream_t stream;

    285:     // MatCreateSeqDenseCUPM may not allocate CPU memory. Allocate if needed
    286:     if (!mimpl->v) PetscCall(MatSeqDenseSetPreallocation(m, nullptr));
    287:     PetscCall(GetHandlesFrom_(dctx, &stream));
    288:     PetscCall(PetscLogEventBegin(MAT_DenseCopyFromGPU, m, 0, 0, 0));
    289:     {
    290:       const auto lda  = mimpl->lda;
    291:       const auto dest = mimpl->v;
    292:       const auto src  = MatCUPMCast(m)->d_v;

    294:       if (lda > nrows) {
    295:         PetscCall(PetscCUPMMemcpy2DAsync(dest, lda, src, lda, nrows, ncols, cupmMemcpyDeviceToHost, stream));
    296:       } else {
    297:         PetscCall(PetscCUPMMemcpyAsync(dest, src, lda * ncols, cupmMemcpyDeviceToHost, stream));
    298:       }
    299:     }
    300:     PetscCall(PetscLogEventEnd(MAT_DenseCopyFromGPU, m, 0, 0, 0));
    301:     // order is important, MatSeqDenseSetPreallocation() might set offloadmask
    302:     m->offloadmask = PETSC_OFFLOAD_BOTH;
    303:   }
    304:   PetscFunctionReturn(PETSC_SUCCESS);
    305: }

    307: template <device::cupm::DeviceType T>
    308: inline PetscErrorCode MatDense_Seq_CUPM<T>::CheckCUPMSolverInfo_(const cupmBlasInt_t *fact_info, cupmStream_t stream) noexcept
    309: {
    310:   PetscFunctionBegin;
    311:   if (PetscDefined(USE_DEBUG)) {
    312:     cupmBlasInt_t info = 0;

    314:     PetscCall(PetscCUPMMemcpyAsync(&info, fact_info, 1, cupmMemcpyDeviceToHost, stream));
    315:     if (stream) PetscCallCUPM(cupmStreamSynchronize(stream));
    316:     static_assert(std::is_same<decltype(info), int>::value, "");
    317:     PetscCheck(info <= 0, PETSC_COMM_SELF, PETSC_ERR_MAT_CH_ZRPVT, "Bad factorization: zero pivot in row %d", info - 1);
    318:     PetscCheck(info >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Wrong argument to cupmSolver %d", -info);
    319:   }
    320:   PetscFunctionReturn(PETSC_SUCCESS);
    321: }

    323: // ==========================================================================================
    324: // MatDense_Seq_CUPM - Private API - Solver Dispatch
    325: // ==========================================================================================

    327: // specific solvers called through the dispatch_() family of functions
    328: template <device::cupm::DeviceType T>
    329: template <typename Derived>
    330: struct MatDense_Seq_CUPM<T>::SolveCommon {
    331:   using derived_type = Derived;

    333:   template <typename F>
    334:   static PetscErrorCode ResizeFactLwork(Mat_SeqDenseCUPM *mcu, cupmStream_t stream, F &&cupmSolverComputeFactLwork) noexcept
    335:   {
    336:     cupmBlasInt_t lwork;

    338:     PetscFunctionBegin;
    339:     PetscCallCUPMSOLVER(cupmSolverComputeFactLwork(&lwork));
    340:     if (lwork > mcu->d_fact_lwork) {
    341:       mcu->d_fact_lwork = lwork;
    342:       PetscCallCUPM(cupmFreeAsync(mcu->d_fact_work, stream));
    343:       PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_work, lwork, stream));
    344:     }
    345:     PetscFunctionReturn(PETSC_SUCCESS);
    346:   }

    348:   static PetscErrorCode FactorPrepare(Mat A, cupmStream_t stream) noexcept
    349:   {
    350:     const auto mcu = MatCUPMCast(A);

    352:     PetscFunctionBegin;
    353:     PetscCall(PetscInfo(A, "%s factor %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", derived_type::NAME(), A->rmap->n, A->cmap->n));
    354:     A->factortype             = derived_type::MATFACTORTYPE();
    355:     A->ops->solve             = MatSolve_Factored_Dispatch_<derived_type, false>;
    356:     A->ops->solvetranspose    = MatSolve_Factored_Dispatch_<derived_type, true>;
    357:     A->ops->matsolve          = MatMatSolve_Factored_Dispatch_<derived_type, false>;
    358:     A->ops->matsolvetranspose = MatMatSolve_Factored_Dispatch_<derived_type, true>;

    360:     PetscCall(PetscStrFreeAllocpy(MATSOLVERCUPM(), &A->solvertype));
    361:     if (!mcu->d_fact_info) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_info, 1, stream));
    362:     PetscFunctionReturn(PETSC_SUCCESS);
    363:   }
    364: };

    366: template <device::cupm::DeviceType T>
    367: struct MatDense_Seq_CUPM<T>::SolveLU : SolveCommon<SolveLU> {
    368:   using base_type = SolveCommon<SolveLU>;

    370:   static constexpr const char   *NAME() noexcept { return "LU"; }
    371:   static constexpr MatFactorType MATFACTORTYPE() noexcept { return MAT_FACTOR_LU; }

    373:   static PetscErrorCode Factor(Mat A, IS, IS, const MatFactorInfo *) noexcept
    374:   {
    375:     const auto         m = static_cast<cupmBlasInt_t>(A->rmap->n);
    376:     const auto         n = static_cast<cupmBlasInt_t>(A->cmap->n);
    377:     cupmStream_t       stream;
    378:     cupmSolverHandle_t handle;
    379:     PetscDeviceContext dctx;

    381:     PetscFunctionBegin;
    382:     if (!m || !n) PetscFunctionReturn(PETSC_SUCCESS);
    383:     PetscCall(GetHandles_(&dctx, &handle, &stream));
    384:     PetscCall(base_type::FactorPrepare(A, stream));
    385:     {
    386:       const auto mcu = MatCUPMCast(A);
    387:       const auto da  = DeviceArrayReadWrite(dctx, A);
    388:       const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);

    390:       // clang-format off
    391:       PetscCall(
    392:         base_type::ResizeFactLwork(
    393:           mcu, stream,
    394:           [&](cupmBlasInt_t *fact_lwork)
    395:           {
    396:             return cupmSolverXgetrf_bufferSize(handle, m, n, da.cupmdata(), lda, fact_lwork);
    397:           }
    398:         )
    399:       );
    400:       // clang-format on
    401:       if (!mcu->d_fact_ipiv) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_ipiv, n, stream));

    403:       PetscCall(PetscLogGpuTimeBegin());
    404:       PetscCallCUPMSOLVER(cupmSolverXgetrf(handle, m, n, da.cupmdata(), lda, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_ipiv, mcu->d_fact_info));
    405:       PetscCall(PetscLogGpuTimeEnd());
    406:       PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream));
    407:     }
    408:     PetscCall(PetscLogGpuFlops(2.0 * n * n * m / 3.0));
    409:     PetscFunctionReturn(PETSC_SUCCESS);
    410:   }

    412:   template <bool transpose>
    413:   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
    414:   {
    415:     const auto         mcu       = MatCUPMCast(A);
    416:     const auto         fact_info = mcu->d_fact_info;
    417:     const auto         fact_ipiv = mcu->d_fact_ipiv;
    418:     cupmSolverHandle_t handle;

    420:     PetscFunctionBegin;
    421:     PetscCall(GetHandlesFrom_(dctx, &handle));
    422:     PetscCall(PetscInfo(A, "%s solve %d x %d on backend\n", NAME(), m, k));
    423:     PetscCall(PetscLogGpuTimeBegin());
    424:     {
    425:       constexpr auto op  = transpose ? CUPMSOLVER_OP_T : CUPMSOLVER_OP_N;
    426:       const auto     da  = DeviceArrayRead(dctx, A);
    427:       const auto     lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);

    429:       // clang-format off
    430:       PetscCall(
    431:         base_type::ResizeFactLwork(
    432:           mcu, stream,
    433:           [&](cupmBlasInt_t *lwork)
    434:           {
    435:             return cupmSolverXgetrs_bufferSize(
    436:               handle, op, m, nrhs, da.cupmdata(), lda, fact_ipiv, x, ldx, lwork
    437:             );
    438:           }
    439:         )
    440:       );
    441:       // clang-format on
    442:       PetscCallCUPMSOLVER(cupmSolverXgetrs(handle, op, m, nrhs, da.cupmdata(), lda, fact_ipiv, x, ldx, mcu->d_fact_work, mcu->d_fact_lwork, fact_info));
    443:       PetscCall(CheckCUPMSolverInfo_(fact_info, stream));
    444:     }
    445:     PetscCall(PetscLogGpuTimeEnd());
    446:     PetscCall(PetscLogGpuFlops(nrhs * (2.0 * m * m - m)));
    447:     PetscFunctionReturn(PETSC_SUCCESS);
    448:   }
    449: };

    451: template <device::cupm::DeviceType T>
    452: struct MatDense_Seq_CUPM<T>::SolveCholesky : SolveCommon<SolveCholesky> {
    453:   using base_type = SolveCommon<SolveCholesky>;

    455:   static constexpr const char   *NAME() noexcept { return "Cholesky"; }
    456:   static constexpr MatFactorType MATFACTORTYPE() noexcept { return MAT_FACTOR_CHOLESKY; }

    458:   static PetscErrorCode Factor(Mat A, IS, const MatFactorInfo *) noexcept
    459:   {
    460:     const auto         n = static_cast<cupmBlasInt_t>(A->rmap->n);
    461:     PetscDeviceContext dctx;
    462:     cupmSolverHandle_t handle;
    463:     cupmStream_t       stream;

    465:     PetscFunctionBegin;
    466:     if (!n || !A->cmap->n) PetscFunctionReturn(PETSC_SUCCESS);
    467:     PetscCheck(A->spd == PETSC_BOOL3_TRUE, PETSC_COMM_SELF, PETSC_ERR_SUP, "%ssytrs unavailable. Use MAT_FACTOR_LU", cupmSolverName());
    468:     PetscCall(GetHandles_(&dctx, &handle, &stream));
    469:     PetscCall(base_type::FactorPrepare(A, stream));
    470:     {
    471:       const auto mcu = MatCUPMCast(A);
    472:       const auto da  = DeviceArrayReadWrite(dctx, A);
    473:       const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);

    475:       // clang-format off
    476:       PetscCall(
    477:         base_type::ResizeFactLwork(
    478:           mcu, stream,
    479:           [&](cupmBlasInt_t *fact_lwork)
    480:           {
    481:             return cupmSolverXpotrf_bufferSize(
    482:               handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, fact_lwork
    483:             );
    484:           }
    485:         )
    486:       );
    487:       // clang-format on
    488:       PetscCall(PetscLogGpuTimeBegin());
    489:       PetscCallCUPMSOLVER(cupmSolverXpotrf(handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_info));
    490:       PetscCall(PetscLogGpuTimeEnd());
    491:       PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream));
    492:     }
    493:     PetscCall(PetscLogGpuFlops(1.0 * n * n * n / 3.0));

    495: #if 0
    496:     // At the time of writing this interface (cuda 10.0), cusolverDn does not implement *sytrs
    497:     // and *hetr* routines. The code below should work, and it can be activated when *sytrs
    498:     // routines will be available
    499:     if (!mcu->d_fact_ipiv) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_ipiv, n, stream));
    500:     if (!mcu->d_fact_lwork) {
    501:       PetscCallCUPMSOLVER(cupmSolverDnXsytrf_bufferSize(handle, n, da.cupmdata(), lda, &mcu->d_fact_lwork));
    502:       PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_work, mcu->d_fact_lwork, stream));
    503:     }
    504:     if (mcu->d_fact_info) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_info, 1, stream));
    505:     PetscCall(PetscLogGpuTimeBegin());
    506:     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));
    507:     PetscCall(PetscLogGpuTimeEnd());
    508: #endif
    509:     PetscFunctionReturn(PETSC_SUCCESS);
    510:   }

    512:   template <bool transpose>
    513:   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
    514:   {
    515:     const auto         mcu       = MatCUPMCast(A);
    516:     const auto         fact_info = mcu->d_fact_info;
    517:     cupmSolverHandle_t handle;

    519:     PetscFunctionBegin;
    520:     PetscAssert(!mcu->d_fact_ipiv, PETSC_COMM_SELF, PETSC_ERR_LIB, "%ssytrs not implemented", cupmSolverName());
    521:     PetscCall(GetHandlesFrom_(dctx, &handle));
    522:     PetscCall(PetscInfo(A, "%s solve %d x %d on backend\n", NAME(), m, k));
    523:     PetscCall(PetscLogGpuTimeBegin());
    524:     {
    525:       const auto da  = DeviceArrayRead(dctx, A);
    526:       const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);

    528:       // clang-format off
    529:       PetscCall(
    530:         base_type::ResizeFactLwork(
    531:           mcu, stream,
    532:           [&](cupmBlasInt_t *lwork)
    533:           {
    534:             return cupmSolverXpotrs_bufferSize(
    535:               handle, CUPMSOLVER_FILL_MODE_LOWER, m, nrhs, da.cupmdata(), lda, x, ldx, lwork
    536:             );
    537:           }
    538:         )
    539:       );
    540:       // clang-format on
    541:       PetscCallCUPMSOLVER(cupmSolverXpotrs(handle, CUPMSOLVER_FILL_MODE_LOWER, m, nrhs, da.cupmdata(), lda, x, ldx, mcu->d_fact_work, mcu->d_fact_lwork, fact_info));
    542:       PetscCall(CheckCUPMSolverInfo_(fact_info, stream));
    543:     }
    544:     PetscCall(PetscLogGpuTimeEnd());
    545:     PetscCall(PetscLogGpuFlops(nrhs * (2.0 * m * m - m)));
    546:     PetscFunctionReturn(PETSC_SUCCESS);
    547:   }
    548: };

    550: template <device::cupm::DeviceType T>
    551: struct MatDense_Seq_CUPM<T>::SolveQR : SolveCommon<SolveQR> {
    552:   using base_type = SolveCommon<SolveQR>;

    554:   static constexpr const char   *NAME() noexcept { return "QR"; }
    555:   static constexpr MatFactorType MATFACTORTYPE() noexcept { return MAT_FACTOR_QR; }

    557:   static PetscErrorCode Factor(Mat A, IS, const MatFactorInfo *) noexcept
    558:   {
    559:     const auto         m     = static_cast<cupmBlasInt_t>(A->rmap->n);
    560:     const auto         n     = static_cast<cupmBlasInt_t>(A->cmap->n);
    561:     const auto         min   = std::min(m, n);
    562:     const auto         mimpl = MatIMPLCast(A);
    563:     cupmStream_t       stream;
    564:     cupmSolverHandle_t handle;
    565:     PetscDeviceContext dctx;

    567:     PetscFunctionBegin;
    568:     if (!m || !n) PetscFunctionReturn(PETSC_SUCCESS);
    569:     PetscCall(GetHandles_(&dctx, &handle, &stream));
    570:     PetscCall(base_type::FactorPrepare(A, stream));
    571:     mimpl->rank = min;
    572:     {
    573:       const auto mcu = MatCUPMCast(A);
    574:       const auto da  = DeviceArrayReadWrite(dctx, A);
    575:       const auto lda = static_cast<cupmBlasInt_t>(mimpl->lda);

    577:       if (!mcu->workvec) PetscCall(vec::cupm::VecCreateSeqCUPMAsync<T>(PetscObjectComm(PetscObjectCast(A)), m, &mcu->workvec));
    578:       if (!mcu->d_fact_tau) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_tau, min, stream));
    579:       // clang-format off
    580:       PetscCall(
    581:         base_type::ResizeFactLwork(
    582:           mcu, stream,
    583:           [&](cupmBlasInt_t *fact_lwork)
    584:           {
    585:             return cupmSolverXgeqrf_bufferSize(handle, m, n, da.cupmdata(), lda, fact_lwork);
    586:           }
    587:         )
    588:       );
    589:       // clang-format on
    590:       PetscCall(PetscLogGpuTimeBegin());
    591:       PetscCallCUPMSOLVER(cupmSolverXgeqrf(handle, m, n, da.cupmdata(), lda, mcu->d_fact_tau, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_info));
    592:       PetscCall(PetscLogGpuTimeEnd());
    593:       PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream));
    594:     }
    595:     PetscCall(PetscLogGpuFlops(2.0 * min * min * (std::max(m, n) - min / 3.0)));
    596:     PetscFunctionReturn(PETSC_SUCCESS);
    597:   }

    599:   template <bool transpose>
    600:   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
    601:   {
    602:     const auto         mimpl      = MatIMPLCast(A);
    603:     const auto         rank       = static_cast<cupmBlasInt_t>(mimpl->rank);
    604:     const auto         mcu        = MatCUPMCast(A);
    605:     const auto         fact_info  = mcu->d_fact_info;
    606:     const auto         fact_tau   = mcu->d_fact_tau;
    607:     const auto         fact_work  = mcu->d_fact_work;
    608:     const auto         fact_lwork = mcu->d_fact_lwork;
    609:     cupmSolverHandle_t solver_handle;
    610:     cupmBlasHandle_t   blas_handle;

    612:     PetscFunctionBegin;
    613:     PetscCall(GetHandlesFrom_(dctx, &blas_handle, &solver_handle));
    614:     PetscCall(PetscInfo(A, "%s solve %d x %d on backend\n", NAME(), m, k));
    615:     PetscCall(PetscLogGpuTimeBegin());
    616:     {
    617:       const auto da  = DeviceArrayRead(dctx, A);
    618:       const auto one = cupmScalarCast(1.0);
    619:       const auto lda = static_cast<cupmBlasInt_t>(mimpl->lda);

    621:       if (transpose) {
    622:         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));
    623:         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));
    624:         PetscCall(CheckCUPMSolverInfo_(fact_info, stream));
    625:       } else {
    626:         constexpr auto op = PetscDefined(USE_COMPLEX) ? CUPMSOLVER_OP_C : CUPMSOLVER_OP_T;

    628:         PetscCallCUPMSOLVER(cupmSolverXormqr(solver_handle, CUPMSOLVER_SIDE_LEFT, op, m, nrhs, rank, da.cupmdata(), lda, fact_tau, x, ldx, fact_work, fact_lwork, fact_info));
    629:         PetscCall(CheckCUPMSolverInfo_(fact_info, stream));
    630:         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));
    631:       }
    632:     }
    633:     PetscCall(PetscLogGpuTimeEnd());
    634:     PetscCall(PetscLogFlops(nrhs * (4.0 * m * rank - (rank * rank))));
    635:     PetscFunctionReturn(PETSC_SUCCESS);
    636:   }
    637: };

    639: template <device::cupm::DeviceType T>
    640: template <typename Solver, bool transpose>
    641: inline PetscErrorCode MatDense_Seq_CUPM<T>::MatSolve_Factored_Dispatch_(Mat A, Vec x, Vec y) noexcept
    642: {
    643:   using namespace vec::cupm;
    644:   const auto         pobj_A  = PetscObjectCast(A);
    645:   const auto         m       = static_cast<cupmBlasInt_t>(A->rmap->n);
    646:   const auto         k       = static_cast<cupmBlasInt_t>(A->cmap->n);
    647:   auto              &workvec = MatCUPMCast(A)->workvec;
    648:   PetscScalar       *y_array = nullptr;
    649:   PetscDeviceContext dctx;
    650:   PetscBool          xiscupm, yiscupm, aiscupm;
    651:   bool               use_y_array_directly;
    652:   cupmStream_t       stream;

    654:   PetscFunctionBegin;
    655:   PetscCheck(A->factortype != MAT_FACTOR_NONE, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix must be factored to solve");
    656:   PetscCall(PetscObjectTypeCompare(PetscObjectCast(x), VecSeq_CUPM::VECSEQCUPM(), &xiscupm));
    657:   PetscCall(PetscObjectTypeCompare(PetscObjectCast(y), VecSeq_CUPM::VECSEQCUPM(), &yiscupm));
    658:   PetscCall(PetscObjectTypeCompare(pobj_A, MATSEQDENSECUPM(), &aiscupm));
    659:   PetscAssert(aiscupm, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Matrix A is somehow not CUPM?????????????????????????????");
    660:   PetscCall(GetHandles_(&dctx, &stream));
    661:   use_y_array_directly = yiscupm && (k >= m);
    662:   {
    663:     const PetscScalar *x_array;
    664:     const auto         xisdevice = xiscupm && PetscOffloadDevice(x->offloadmask);
    665:     const auto         copy_mode = xisdevice ? cupmMemcpyDeviceToDevice : cupmMemcpyHostToDevice;

    667:     if (!use_y_array_directly && !workvec) PetscCall(VecCreateSeqCUPMAsync<T>(PetscObjectComm(pobj_A), m, &workvec));
    668:     // The logic here is to try to minimize the amount of memory copying:
    669:     //
    670:     // If we call VecCUPMGetArrayRead(X, &x) every time xiscupm and the data is not offloaded
    671:     // to the GPU yet, then the data is copied to the GPU. But we are only trying to get the
    672:     // data in order to copy it into the y array. So the array x will be wherever the data
    673:     // already is so that only one memcpy is performed
    674:     if (xisdevice) {
    675:       PetscCall(VecCUPMGetArrayReadAsync<T>(x, &x_array, dctx));
    676:     } else {
    677:       PetscCall(VecGetArrayRead(x, &x_array));
    678:     }
    679:     PetscCall(VecCUPMGetArrayWriteAsync<T>(use_y_array_directly ? y : workvec, &y_array, dctx));
    680:     PetscCall(PetscCUPMMemcpyAsync(y_array, x_array, m, copy_mode, stream));
    681:     if (xisdevice) {
    682:       PetscCall(VecCUPMRestoreArrayReadAsync<T>(x, &x_array, dctx));
    683:     } else {
    684:       PetscCall(VecRestoreArrayRead(x, &x_array));
    685:     }
    686:   }

    688:   if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
    689:   PetscCall(Solver{}.template Solve<transpose>(A, cupmScalarPtrCast(y_array), m, m, 1, k, dctx, stream));
    690:   if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A));

    692:   if (use_y_array_directly) {
    693:     PetscCall(VecCUPMRestoreArrayWriteAsync<T>(y, &y_array, dctx));
    694:   } else {
    695:     const auto   copy_mode = yiscupm ? cupmMemcpyDeviceToDevice : cupmMemcpyDeviceToHost;
    696:     PetscScalar *yv;

    698:     // The logic here is that the data is not yet in either y's GPU array or its CPU array.
    699:     // There is nothing in the interface to say where the user would like it to end up. So we
    700:     // choose the GPU, because it is the faster option
    701:     if (yiscupm) {
    702:       PetscCall(VecCUPMGetArrayWriteAsync<T>(y, &yv, dctx));
    703:     } else {
    704:       PetscCall(VecGetArray(y, &yv));
    705:     }
    706:     PetscCall(PetscCUPMMemcpyAsync(yv, y_array, k, copy_mode, stream));
    707:     if (yiscupm) {
    708:       PetscCall(VecCUPMRestoreArrayWriteAsync<T>(y, &yv, dctx));
    709:     } else {
    710:       PetscCall(VecRestoreArray(y, &yv));
    711:     }
    712:     PetscCall(VecCUPMRestoreArrayWriteAsync<T>(workvec, &y_array));
    713:   }
    714:   PetscFunctionReturn(PETSC_SUCCESS);
    715: }

    717: template <device::cupm::DeviceType T>
    718: template <typename Solver, bool transpose>
    719: inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMatSolve_Factored_Dispatch_(Mat A, Mat B, Mat X) noexcept
    720: {
    721:   const auto         m = static_cast<cupmBlasInt_t>(A->rmap->n);
    722:   const auto         k = static_cast<cupmBlasInt_t>(A->cmap->n);
    723:   cupmBlasInt_t      nrhs, ldb, ldx, ldy;
    724:   PetscScalar       *y;
    725:   PetscBool          biscupm, xiscupm, aiscupm;
    726:   PetscDeviceContext dctx;
    727:   cupmStream_t       stream;

    729:   PetscFunctionBegin;
    730:   PetscCheck(A->factortype != MAT_FACTOR_NONE, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix must be factored to solve");
    731:   PetscCall(PetscObjectTypeCompare(PetscObjectCast(B), MATSEQDENSECUPM(), &biscupm));
    732:   PetscCall(PetscObjectTypeCompare(PetscObjectCast(X), MATSEQDENSECUPM(), &xiscupm));
    733:   PetscCall(PetscObjectTypeCompare(PetscObjectCast(A), MATSEQDENSECUPM(), &aiscupm));
    734:   PetscCall(GetHandles_(&dctx, &stream));
    735:   {
    736:     PetscInt n;

    738:     PetscCall(MatGetSize(B, nullptr, &n));
    739:     PetscCall(PetscCUPMBlasIntCast(n, &nrhs));
    740:     PetscCall(MatDenseGetLDA(B, &n));
    741:     PetscCall(PetscCUPMBlasIntCast(n, &ldb));
    742:     PetscCall(MatDenseGetLDA(X, &n));
    743:     PetscCall(PetscCUPMBlasIntCast(n, &ldx));
    744:   }
    745:   {
    746:     // The logic here is to try to minimize the amount of memory copying:
    747:     //
    748:     // If we call MatDenseCUPMGetArrayRead(B, &b) every time biscupm and the data is not
    749:     // offloaded to the GPU yet, then the data is copied to the GPU. But we are only trying to
    750:     // get the data in order to copy it into the y array. So the array b will be wherever the
    751:     // data already is so that only one memcpy is performed
    752:     const auto         bisdevice = biscupm && PetscOffloadDevice(B->offloadmask);
    753:     const auto         copy_mode = bisdevice ? cupmMemcpyDeviceToDevice : cupmMemcpyHostToDevice;
    754:     const PetscScalar *b;

    756:     if (bisdevice) {
    757:       b = DeviceArrayRead(dctx, B);
    758:     } else if (biscupm) {
    759:       b = HostArrayRead(dctx, B);
    760:     } else {
    761:       PetscCall(MatDenseGetArrayRead(B, &b));
    762:     }

    764:     if (ldx < m || !xiscupm) {
    765:       // X's array cannot serve as the array (too small or not on device), B's array cannot
    766:       // serve as the array (const), so allocate a new array
    767:       ldy = m;
    768:       PetscCall(PetscCUPMMallocAsync(&y, nrhs * m));
    769:     } else {
    770:       // X's array should serve as the array
    771:       ldy = ldx;
    772:       y   = DeviceArrayWrite(dctx, X);
    773:     }
    774:     PetscCall(PetscCUPMMemcpy2DAsync(y, ldy, b, ldb, m, nrhs, copy_mode, stream));
    775:     if (!bisdevice && !biscupm) PetscCall(MatDenseRestoreArrayRead(B, &b));
    776:   }

    778:   // convert to CUPM twice??????????????????????????????????
    779:   // but A should already be CUPM??????????????????????????????????????
    780:   if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
    781:   PetscCall(Solver{}.template Solve<transpose>(A, cupmScalarPtrCast(y), ldy, m, nrhs, k, dctx, stream));
    782:   if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));

    784:   if (ldx < m || !xiscupm) {
    785:     const auto   copy_mode = xiscupm ? cupmMemcpyDeviceToDevice : cupmMemcpyDeviceToHost;
    786:     PetscScalar *x;

    788:     // The logic here is that the data is not yet in either X's GPU array or its CPU
    789:     // array. There is nothing in the interface to say where the user would like it to end up.
    790:     // So we choose the GPU, because it is the faster option
    791:     if (xiscupm) {
    792:       x = DeviceArrayWrite(dctx, X);
    793:     } else {
    794:       PetscCall(MatDenseGetArray(X, &x));
    795:     }
    796:     PetscCall(PetscCUPMMemcpy2DAsync(x, ldx, y, ldy, k, nrhs, copy_mode, stream));
    797:     if (!xiscupm) PetscCall(MatDenseRestoreArray(X, &x));
    798:     PetscCallCUPM(cupmFreeAsync(y, stream));
    799:   }
    800:   PetscFunctionReturn(PETSC_SUCCESS);
    801: }

    803: template <device::cupm::DeviceType T>
    804: template <bool transpose, bool hermitian>
    805: inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMultAddColumnRange_Dispatch_(Mat A, Vec xx, Vec yy, Vec zz, PetscInt c_start, PetscInt c_end) noexcept
    806: {
    807:   const auto         m   = static_cast<cupmBlasInt_t>(A->rmap->n);
    808:   const auto         n   = static_cast<cupmBlasInt_t>(c_end - c_start);
    809:   const auto         lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda);
    810:   PetscBool          xiscupm, yiscupm, ziscupm;
    811:   cupmBlasHandle_t   handle;
    812:   Vec                x = xx, y = yy, z = zz;
    813:   PetscDeviceContext dctx;

    815:   PetscFunctionBegin;
    816:   PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(xx), &xiscupm, VecSeq_CUPM::VECSEQCUPM(), VecSeq_CUPM::VECMPICUPM(), VecSeq_CUPM::VECCUPM(), ""));
    817:   if (!xiscupm || xx->boundtocpu) {
    818:     PetscCall(VecCreate(PetscObjectComm(PetscObjectCast(xx)), &x));
    819:     PetscCall(VecSetLayout(x, xx->map));
    820:     PetscCall(VecSetType(x, VecSeq_CUPM::VECCUPM()));
    821:     PetscCall(VecCopy(xx, x));
    822:   }

    824:   if (yy) {
    825:     PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(yy), &yiscupm, VecSeq_CUPM::VECSEQCUPM(), VecSeq_CUPM::VECMPICUPM(), VecSeq_CUPM::VECCUPM(), ""));
    826:     if (!yiscupm || yy->boundtocpu) {
    827:       PetscCall(VecCreate(PetscObjectComm(PetscObjectCast(yy)), &y));
    828:       PetscCall(VecSetLayout(y, yy->map));
    829:       PetscCall(VecSetType(y, VecSeq_CUPM::VECCUPM()));
    830:       PetscCall(VecCopy(yy, y));
    831:     }
    832:   }

    834:   if (zz != yy) {
    835:     PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(zz), &ziscupm, VecSeq_CUPM::VECSEQCUPM(), VecSeq_CUPM::VECMPICUPM(), VecSeq_CUPM::VECCUPM(), ""));
    836:     if (!ziscupm || zz->boundtocpu) {
    837:       PetscCall(VecCreate(PetscObjectComm(PetscObjectCast(zz)), &z));
    838:       PetscCall(VecSetLayout(z, zz->map));
    839:       PetscCall(VecSetType(z, VecSeq_CUPM::VECCUPM()));
    840:     }
    841:   } else {
    842:     z = y;
    843:   }

    845:   if (y && y != z) PetscCall(VecSeq_CUPM::Copy(y, z)); // mult add
    846:   if (!m || !n) {
    847:     // mult only
    848:     if (!y) PetscCall(VecSeq_CUPM::Set(z, 0.0));
    849:     PetscFunctionReturn(PETSC_SUCCESS);
    850:   }
    851:   PetscCall(PetscInfo(A, "Matrix-vector product %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " on backend\n", m, n));
    852:   PetscCall(GetHandles_(&dctx, &handle));
    853:   {
    854:     constexpr auto op   = transpose ? (hermitian ? CUPMBLAS_OP_C : CUPMBLAS_OP_T) : CUPMBLAS_OP_N;
    855:     const auto     one  = cupmScalarCast(1.0);
    856:     const auto     zero = cupmScalarCast(0.0);
    857:     const auto     da   = DeviceArrayRead(dctx, A);
    858:     const auto     dxx  = VecSeq_CUPM::DeviceArrayRead(dctx, x);
    859:     const auto     dzz  = VecSeq_CUPM::DeviceArrayReadWrite(dctx, z);

    861:     PetscCall(PetscLogGpuTimeBegin());
    862:     PetscCallCUPMBLAS(cupmBlasXgemv(handle, op, m, n, &one, da.cupmdata() + c_start * lda, lda, dxx.cupmdata() + (transpose ? 0 : c_start), 1, y ? &one : &zero, dzz.cupmdata() + (transpose ? c_start : 0), 1));
    863:     PetscCall(PetscLogGpuTimeEnd());
    864:   }
    865:   PetscCall(PetscLogGpuFlops(2.0 * m * n - (yy ? 0 : m)));
    866:   if (z != zz) {
    867:     PetscCall(VecCopy(z, zz));
    868:     if (z != y) PetscCall(VecDestroy(&z));
    869:   }
    870:   if (y != yy) PetscCall(VecDestroy(&y));
    871:   if (x != xx) PetscCall(VecDestroy(&x));
    872:   PetscFunctionReturn(PETSC_SUCCESS);
    873: }

    875: template <device::cupm::DeviceType T>
    876: template <bool transpose, bool hermitian>
    877: inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMultColumnRange_Dispatch_(Mat A, Vec xx, Vec yy, PetscInt c_start, PetscInt c_end) noexcept
    878: {
    879:   PetscFunctionBegin;
    880:   PetscCall(MatMultAddColumnRange_Dispatch_<transpose, hermitian>(A, xx, nullptr, yy, c_start, c_end));
    881:   PetscFunctionReturn(PETSC_SUCCESS);
    882: }

    884: template <device::cupm::DeviceType T>
    885: template <bool transpose, bool hermitian>
    886: inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMultAdd_Dispatch_(Mat A, Vec xx, Vec yy, Vec zz) noexcept
    887: {
    888:   PetscFunctionBegin;
    889:   PetscCall(MatMultAddColumnRange_Dispatch_<transpose, hermitian>(A, xx, yy, zz, 0, A->cmap->n));
    890:   PetscFunctionReturn(PETSC_SUCCESS);
    891: }

    893: // ==========================================================================================
    894: // MatDense_Seq_CUPM - Private API - Conversion Dispatch
    895: // ==========================================================================================

    897: template <device::cupm::DeviceType T>
    898: template <bool to_host>
    899: inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_Dispatch_(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept
    900: {
    901:   PetscFunctionBegin;
    902:   if (reuse == MAT_REUSE_MATRIX || reuse == MAT_INITIAL_MATRIX) {
    903:     // TODO these cases should be optimized
    904:     PetscCall(MatConvert_Basic(M, type, reuse, newmat));
    905:   } else {
    906:     const auto B    = *newmat;
    907:     const auto pobj = PetscObjectCast(B);

    909:     if (to_host) {
    910:       PetscCall(BindToCPU(B, PETSC_TRUE));
    911:       PetscCall(Reset(B));
    912:     } else {
    913:       PetscCall(PetscDeviceInitialize(PETSC_DEVICE_CUPM()));
    914:     }

    916:     PetscCall(PetscStrFreeAllocpy(to_host ? VECSTANDARD : VecSeq_CUPM::VECCUPM(), &B->defaultvectype));
    917:     PetscCall(PetscObjectChangeTypeName(pobj, to_host ? MATSEQDENSE : MATSEQDENSECUPM()));
    918:     // cvec might be the wrong VecType, destroy and rebuild it if necessary
    919:     // REVIEW ME: this is possibly very inefficient
    920:     PetscCall(VecDestroy(&MatIMPLCast(B)->cvec));

    922:     MatComposeOp_CUPM(to_host, pobj, MatConvert_seqdensecupm_seqdense_C(), nullptr, Convert_SeqDenseCUPM_SeqDense);
    923:     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArray_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ_WRITE>);
    924:     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArrayRead_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ>);
    925:     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArrayWrite_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_WRITE>);
    926:     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArray_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ_WRITE>);
    927:     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArrayRead_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ>);
    928:     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArrayWrite_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_WRITE>);
    929:     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMPlaceArray_C(), nullptr, PlaceArray);
    930:     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMResetArray_C(), nullptr, ResetArray);
    931:     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMReplaceArray_C(), nullptr, ReplaceArray);
    932:     MatComposeOp_CUPM(to_host, pobj, MatProductSetFromOptions_seqaij_seqdensecupm_C(), nullptr, MatProductSetFromOptions_SeqAIJ_SeqDense);
    933:     MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMSetPreallocation_C(), nullptr, SetPreallocation);

    935:     if (to_host) {
    936:       B->offloadmask = PETSC_OFFLOAD_CPU;
    937:     } else {
    938:       Mat_SeqDenseCUPM *mcu;

    940:       PetscCall(PetscNew(&mcu));
    941:       B->spptr       = mcu;
    942:       B->offloadmask = PETSC_OFFLOAD_UNALLOCATED; // REVIEW ME: why not offload host??
    943:       PetscCall(BindToCPU(B, PETSC_FALSE));
    944:     }

    946:     MatSetOp_CUPM(to_host, B, bindtocpu, nullptr, BindToCPU);
    947:     MatSetOp_CUPM(to_host, B, destroy, MatDestroy_SeqDense, Destroy);
    948:   }
    949:   PetscFunctionReturn(PETSC_SUCCESS);
    950: }

    952: // ==========================================================================================
    953: // MatDense_Seq_CUPM - Public API
    954: // ==========================================================================================

    956: template <device::cupm::DeviceType T>
    957: inline constexpr MatType MatDense_Seq_CUPM<T>::MATIMPLCUPM_() noexcept
    958: {
    959:   return MATSEQDENSECUPM();
    960: }

    962: template <device::cupm::DeviceType T>
    963: inline constexpr typename MatDense_Seq_CUPM<T>::Mat_SeqDenseCUPM *MatDense_Seq_CUPM<T>::MatCUPMCast(Mat m) noexcept
    964: {
    965:   return static_cast<Mat_SeqDenseCUPM *>(m->spptr);
    966: }

    968: template <device::cupm::DeviceType T>
    969: inline constexpr Mat_SeqDense *MatDense_Seq_CUPM<T>::MatIMPLCast_(Mat m) noexcept
    970: {
    971:   return static_cast<Mat_SeqDense *>(m->data);
    972: }

    974: template <device::cupm::DeviceType T>
    975: inline constexpr const char *MatDense_Seq_CUPM<T>::MatConvert_seqdensecupm_seqdense_C() noexcept
    976: {
    977:   return T == device::cupm::DeviceType::CUDA ? "MatConvert_seqdensecuda_seqdense_C" : "MatConvert_seqdensehip_seqdense_C";
    978: }

    980: template <device::cupm::DeviceType T>
    981: inline constexpr const char *MatDense_Seq_CUPM<T>::MatProductSetFromOptions_seqaij_seqdensecupm_C() noexcept
    982: {
    983:   return T == device::cupm::DeviceType::CUDA ? "MatProductSetFromOptions_seqaij_seqdensecuda_C" : "MatProductSetFromOptions_seqaij_seqdensehip_C";
    984: }

    986: // ==========================================================================================

    988: // MatCreate_SeqDenseCUPM()
    989: template <device::cupm::DeviceType T>
    990: inline PetscErrorCode MatDense_Seq_CUPM<T>::Create(Mat A) noexcept
    991: {
    992:   PetscFunctionBegin;
    993:   PetscCall(PetscDeviceInitialize(PETSC_DEVICE_CUPM()));
    994:   PetscCall(MatCreate_SeqDense(A));
    995:   PetscCall(Convert_SeqDense_SeqDenseCUPM(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
    996:   PetscFunctionReturn(PETSC_SUCCESS);
    997: }

    999: template <device::cupm::DeviceType T>
    1000: inline PetscErrorCode MatDense_Seq_CUPM<T>::Destroy(Mat A) noexcept
    1001: {
    1002:   PetscFunctionBegin;
    1003:   // prevent copying back data if we own the data pointer
    1004:   if (!MatIMPLCast(A)->user_alloc) A->offloadmask = PETSC_OFFLOAD_CPU;
    1005:   PetscCall(Convert_SeqDenseCUPM_SeqDense(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A));
    1006:   PetscCall(MatDestroy_SeqDense(A));
    1007:   PetscFunctionReturn(PETSC_SUCCESS);
    1008: }

    1010: // obj->ops->setup()
    1011: template <device::cupm::DeviceType T>
    1012: inline PetscErrorCode MatDense_Seq_CUPM<T>::SetUp(Mat A) noexcept
    1013: {
    1014:   PetscFunctionBegin;
    1015:   PetscCall(PetscLayoutSetUp(A->rmap));
    1016:   PetscCall(PetscLayoutSetUp(A->cmap));
    1017:   if (!A->preallocated) {
    1018:     PetscDeviceContext dctx;

    1020:     PetscCall(GetHandles_(&dctx));
    1021:     PetscCall(SetPreallocation(A, dctx, nullptr));
    1022:   }
    1023:   PetscFunctionReturn(PETSC_SUCCESS);
    1024: }

    1026: template <device::cupm::DeviceType T>
    1027: inline PetscErrorCode MatDense_Seq_CUPM<T>::Reset(Mat A) noexcept
    1028: {
    1029:   PetscFunctionBegin;
    1030:   if (const auto mcu = MatCUPMCast(A)) {
    1031:     cupmStream_t stream;

    1033:     PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME());
    1034:     PetscCall(GetHandles_(&stream));
    1035:     if (!mcu->d_user_alloc) PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream));
    1036:     PetscCallCUPM(cupmFreeAsync(mcu->d_fact_tau, stream));
    1037:     PetscCallCUPM(cupmFreeAsync(mcu->d_fact_ipiv, stream));
    1038:     PetscCallCUPM(cupmFreeAsync(mcu->d_fact_info, stream));
    1039:     PetscCallCUPM(cupmFreeAsync(mcu->d_fact_work, stream));
    1040:     PetscCall(VecDestroy(&mcu->workvec));
    1041:     PetscCall(PetscFree(A->spptr /* mcu */));
    1042:   }
    1043:   PetscFunctionReturn(PETSC_SUCCESS);
    1044: }

    1046: // ==========================================================================================

    1048: template <device::cupm::DeviceType T>
    1049: inline PetscErrorCode MatDense_Seq_CUPM<T>::BindToCPU(Mat A, PetscBool to_host) noexcept
    1050: {
    1051:   const auto mimpl = MatIMPLCast(A);
    1052:   const auto pobj  = PetscObjectCast(A);

    1054:   PetscFunctionBegin;
    1055:   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
    1056:   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
    1057:   A->boundtocpu = to_host;
    1058:   PetscCall(PetscStrFreeAllocpy(to_host ? PETSCRANDER48 : PETSCDEVICERAND(), &A->defaultrandtype));
    1059:   if (to_host) {
    1060:     PetscDeviceContext dctx;

    1062:     // make sure we have an up-to-date copy on the CPU
    1063:     PetscCall(GetHandles_(&dctx));
    1064:     PetscCall(DeviceToHost_(A, dctx));
    1065:   } else {
    1066:     PetscBool iscupm;

    1068:     if (auto &cvec = mimpl->cvec) {
    1069:       PetscCall(PetscObjectTypeCompare(PetscObjectCast(cvec), VecSeq_CUPM::VECSEQCUPM(), &iscupm));
    1070:       if (!iscupm) PetscCall(VecDestroy(&cvec));
    1071:     }
    1072:     if (auto &cmat = mimpl->cmat) {
    1073:       PetscCall(PetscObjectTypeCompare(PetscObjectCast(cmat), MATSEQDENSECUPM(), &iscupm));
    1074:       if (!iscupm) PetscCall(MatDestroy(&cmat));
    1075:     }
    1076:   }

    1078:   // ============================================================
    1079:   // Composed ops
    1080:   // ============================================================
    1081:   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArray_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ_WRITE>);
    1082:   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayRead_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ>);
    1083:   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayWrite_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_WRITE>);
    1084:   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ_WRITE>);
    1085:   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ_WRITE>);
    1086:   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayReadAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ>);
    1087:   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayReadAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ>);
    1088:   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayWriteAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_WRITE>);
    1089:   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayWriteAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_WRITE>);
    1090:   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVec_C", MatDenseGetColumnVec_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_READ_WRITE>);
    1091:   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVec_C", MatDenseRestoreColumnVec_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_READ_WRITE>);
    1092:   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVecRead_C", MatDenseGetColumnVecRead_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_READ>);
    1093:   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVecRead_C", MatDenseRestoreColumnVecRead_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_READ>);
    1094:   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVecWrite_C", MatDenseGetColumnVecWrite_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_WRITE>);
    1095:   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVecWrite_C", MatDenseRestoreColumnVecWrite_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_WRITE>);
    1096:   MatComposeOp_CUPM(to_host, pobj, "MatDenseGetSubMatrix_C", MatDenseGetSubMatrix_SeqDense, GetSubMatrix);
    1097:   MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreSubMatrix_C", MatDenseRestoreSubMatrix_SeqDense, RestoreSubMatrix);
    1098:   MatComposeOp_CUPM(to_host, pobj, "MatQRFactor_C", MatQRFactor_SeqDense, SolveQR::Factor);
    1099:   MatComposeOp_CUPM(to_host, pobj, "MatMultColumnRange_C", MatMultColumnRange_SeqDense, MatMultColumnRange_Dispatch_</* transpose */ false, /* hermitian */ false>);
    1100:   MatComposeOp_CUPM(to_host, pobj, "MatMultAddColumnRange_C", MatMultAddColumnRange_SeqDense, MatMultAddColumnRange_Dispatch_</* transpose */ false, /* hermitian */ false>);
    1101:   MatComposeOp_CUPM(to_host, pobj, "MatMultHermitianTransposeColumnRange_C", MatMultHermitianTransposeColumnRange_SeqDense, MatMultColumnRange_Dispatch_</* transpose */ true, /* hermitian */ true>);
    1102:   MatComposeOp_CUPM(to_host, pobj, "MatMultHermitianTransposeAddColumnRange_C", MatMultHermitianTransposeAddColumnRange_SeqDense, MatMultAddColumnRange_Dispatch_</* transpose */ true, /* hermitian */ true>);
    1103:   // always the same
    1104:   PetscCall(PetscObjectComposeFunction(pobj, "MatDenseSetLDA_C", MatDenseSetLDA_SeqDense));

    1106:   // ============================================================
    1107:   // Function pointer ops
    1108:   // ============================================================
    1109:   MatSetOp_CUPM(to_host, A, duplicate, MatDuplicate_SeqDense, Duplicate);
    1110:   MatSetOp_CUPM(to_host, A, mult, MatMult_SeqDense, [](Mat A, Vec xx, Vec yy) { return MatMultAdd_Dispatch_</* transpose */ false, /* hermitian */ false>(A, xx, nullptr, yy); });
    1111:   MatSetOp_CUPM(to_host, A, multtranspose, MatMultTranspose_SeqDense, [](Mat A, Vec xx, Vec yy) { return MatMultAdd_Dispatch_</* transpose */ true, /* hermitian */ false>(A, xx, nullptr, yy); });
    1112:   MatSetOp_CUPM(to_host, A, multhermitiantranspose, MatMultTranspose_SeqDense, [](Mat A, Vec xx, Vec yy) { return MatMultAdd_Dispatch_</* transpose */ true, /* hermitian */ true>(A, xx, nullptr, yy); });
    1113:   MatSetOp_CUPM(to_host, A, multadd, MatMultAdd_SeqDense, MatMultAdd_Dispatch_</* transpose */ false, /* hermitian */ false>);
    1114:   MatSetOp_CUPM(to_host, A, multtransposeadd, MatMultTransposeAdd_SeqDense, MatMultAdd_Dispatch_</* transpose */ true, /* hermitian */ false>);
    1115:   MatSetOp_CUPM(to_host, A, multhermitiantransposeadd, MatMultHermitianTransposeAdd_SeqDense, MatMultAdd_Dispatch_</* transpose */ true, /* hermitian */ true>);
    1116:   MatSetOp_CUPM(to_host, A, matmultnumeric, MatMatMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ false, /* transpose_B */ false>);
    1117:   MatSetOp_CUPM(to_host, A, mattransposemultnumeric, MatMatTransposeMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ false, /* transpose_B */ true>);
    1118:   MatSetOp_CUPM(to_host, A, transposematmultnumeric, MatTransposeMatMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ true, /* transpose_B */ false>);
    1119:   MatSetOp_CUPM(to_host, A, axpy, MatAXPY_SeqDense, AXPY);
    1120:   MatSetOp_CUPM(to_host, A, choleskyfactor, MatCholeskyFactor_SeqDense, SolveCholesky::Factor);
    1121:   MatSetOp_CUPM(to_host, A, lufactor, MatLUFactor_SeqDense, SolveLU::Factor);
    1122:   MatSetOp_CUPM(to_host, A, getcolumnvector, MatGetColumnVector_SeqDense, GetColumnVector);
    1123:   MatSetOp_CUPM(to_host, A, conjugate, MatConjugate_SeqDense, Conjugate);
    1124:   MatSetOp_CUPM(to_host, A, scale, MatScale_SeqDense, Scale);
    1125:   MatSetOp_CUPM(to_host, A, shift, MatShift_SeqDense, Shift);
    1126:   MatSetOp_CUPM(to_host, A, copy, MatCopy_SeqDense, Copy);
    1127:   MatSetOp_CUPM(to_host, A, zeroentries, MatZeroEntries_SeqDense, ZeroEntries);
    1128:   MatSetOp_CUPM(to_host, A, setup, MatSetUp_SeqDense, SetUp);
    1129:   MatSetOp_CUPM(to_host, A, setrandom, MatSetRandom_SeqDense, SetRandom);
    1130:   MatSetOp_CUPM(to_host, A, getdiagonal, MatGetDiagonal_SeqDense, GetDiagonal);
    1131:   // seemingly always the same
    1132:   A->ops->productsetfromoptions = MatProductSetFromOptions_SeqDense;

    1134:   if (const auto cmat = mimpl->cmat) PetscCall(MatBindToCPU(cmat, to_host));
    1135:   PetscFunctionReturn(PETSC_SUCCESS);
    1136: }

    1138: template <device::cupm::DeviceType T>
    1139: inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_SeqDenseCUPM_SeqDense(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept
    1140: {
    1141:   PetscFunctionBegin;
    1142:   PetscCall(Convert_Dispatch_</* to host */ true>(M, type, reuse, newmat));
    1143:   PetscFunctionReturn(PETSC_SUCCESS);
    1144: }

    1146: template <device::cupm::DeviceType T>
    1147: inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_SeqDense_SeqDenseCUPM(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept
    1148: {
    1149:   PetscFunctionBegin;
    1150:   PetscCall(Convert_Dispatch_</* to host */ false>(M, type, reuse, newmat));
    1151:   PetscFunctionReturn(PETSC_SUCCESS);
    1152: }

    1154: // ==========================================================================================

    1156: template <device::cupm::DeviceType T>
    1157: template <PetscMemType mtype, PetscMemoryAccessMode access>
    1158: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetArray(Mat m, PetscScalar **array, PetscDeviceContext dctx) noexcept
    1159: {
    1160:   constexpr auto hostmem     = PetscMemTypeHost(mtype);
    1161:   constexpr auto read_access = PetscMemoryAccessRead(access);

    1163:   PetscFunctionBegin;
    1164:   static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), "");
    1165:   PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
    1166:   if (hostmem) {
    1167:     if (read_access) {
    1168:       PetscCall(DeviceToHost_(m, dctx));
    1169:     } else if (!MatIMPLCast(m)->v) {
    1170:       // MatCreateSeqDenseCUPM may not allocate CPU memory. Allocate if needed
    1171:       PetscCall(MatSeqDenseSetPreallocation(m, nullptr));
    1172:     }
    1173:     *array = MatIMPLCast(m)->v;
    1174:   } else {
    1175:     if (read_access) {
    1176:       PetscCall(HostToDevice_(m, dctx));
    1177:     } else if (!MatCUPMCast(m)->d_v) {
    1178:       // write-only
    1179:       PetscCall(SetPreallocation(m, dctx, nullptr));
    1180:     }
    1181:     *array = MatCUPMCast(m)->d_v;
    1182:   }
    1183:   if (PetscMemoryAccessWrite(access)) {
    1184:     m->offloadmask = hostmem ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
    1185:     PetscCall(PetscObjectStateIncrease(PetscObjectCast(m)));
    1186:   }
    1187:   PetscFunctionReturn(PETSC_SUCCESS);
    1188: }

    1190: template <device::cupm::DeviceType T>
    1191: template <PetscMemType mtype, PetscMemoryAccessMode access>
    1192: inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreArray(Mat m, PetscScalar **array, PetscDeviceContext) noexcept
    1193: {
    1194:   PetscFunctionBegin;
    1195:   static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), "");
    1196:   if (PetscMemoryAccessWrite(access)) {
    1197:     // WRITE or READ_WRITE
    1198:     m->offloadmask = PetscMemTypeHost(mtype) ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
    1199:     PetscCall(PetscObjectStateIncrease(PetscObjectCast(m)));
    1200:   }
    1201:   if (array) {
    1202:     PetscCall(CheckPointerMatchesMemType_(*array, mtype));
    1203:     *array = nullptr;
    1204:   }
    1205:   PetscFunctionReturn(PETSC_SUCCESS);
    1206: }

    1208: template <device::cupm::DeviceType T>
    1209: template <PetscMemoryAccessMode access>
    1210: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetArrayAndMemType(Mat m, PetscScalar **array, PetscMemType *mtype, PetscDeviceContext dctx) noexcept
    1211: {
    1212:   PetscFunctionBegin;
    1213:   PetscCall(GetArray<PETSC_MEMTYPE_DEVICE, access>(m, array, dctx));
    1214:   if (mtype) *mtype = PETSC_MEMTYPE_CUPM();
    1215:   PetscFunctionReturn(PETSC_SUCCESS);
    1216: }

    1218: template <device::cupm::DeviceType T>
    1219: template <PetscMemoryAccessMode access>
    1220: inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreArrayAndMemType(Mat m, PetscScalar **array, PetscDeviceContext dctx) noexcept
    1221: {
    1222:   PetscFunctionBegin;
    1223:   PetscCall(RestoreArray<PETSC_MEMTYPE_DEVICE, access>(m, array, dctx));
    1224:   PetscFunctionReturn(PETSC_SUCCESS);
    1225: }

    1227: // ==========================================================================================

    1229: template <device::cupm::DeviceType T>
    1230: inline PetscErrorCode MatDense_Seq_CUPM<T>::PlaceArray(Mat A, const PetscScalar *array) noexcept
    1231: {
    1232:   const auto mimpl = MatIMPLCast(A);
    1233:   const auto mcu   = MatCUPMCast(A);

    1235:   PetscFunctionBegin;
    1236:   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
    1237:   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
    1238:   PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME());
    1239:   if (mimpl->v) {
    1240:     PetscDeviceContext dctx;

    1242:     PetscCall(GetHandles_(&dctx));
    1243:     PetscCall(HostToDevice_(A, dctx));
    1244:   }
    1245:   mcu->unplacedarray         = util::exchange(mcu->d_v, const_cast<PetscScalar *>(array));
    1246:   mcu->d_unplaced_user_alloc = util::exchange(mcu->d_user_alloc, PETSC_TRUE);
    1247:   PetscFunctionReturn(PETSC_SUCCESS);
    1248: }

    1250: template <device::cupm::DeviceType T>
    1251: inline PetscErrorCode MatDense_Seq_CUPM<T>::ReplaceArray(Mat A, const PetscScalar *array) noexcept
    1252: {
    1253:   const auto mimpl = MatIMPLCast(A);
    1254:   const auto mcu   = MatCUPMCast(A);

    1256:   PetscFunctionBegin;
    1257:   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
    1258:   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
    1259:   PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME());
    1260:   if (!mcu->d_user_alloc) {
    1261:     cupmStream_t stream;

    1263:     PetscCall(GetHandles_(&stream));
    1264:     PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream));
    1265:   }
    1266:   mcu->d_v          = const_cast<PetscScalar *>(array);
    1267:   mcu->d_user_alloc = PETSC_FALSE;
    1268:   PetscFunctionReturn(PETSC_SUCCESS);
    1269: }

    1271: template <device::cupm::DeviceType T>
    1272: inline PetscErrorCode MatDense_Seq_CUPM<T>::ResetArray(Mat A) noexcept
    1273: {
    1274:   const auto mimpl = MatIMPLCast(A);
    1275:   const auto mcu   = MatCUPMCast(A);

    1277:   PetscFunctionBegin;
    1278:   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
    1279:   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
    1280:   if (mimpl->v) {
    1281:     PetscDeviceContext dctx;

    1283:     PetscCall(GetHandles_(&dctx));
    1284:     PetscCall(HostToDevice_(A, dctx));
    1285:   }
    1286:   mcu->d_v          = util::exchange(mcu->unplacedarray, nullptr);
    1287:   mcu->d_user_alloc = mcu->d_unplaced_user_alloc;
    1288:   PetscFunctionReturn(PETSC_SUCCESS);
    1289: }

    1291: // ==========================================================================================

    1293: template <device::cupm::DeviceType T>
    1294: template <bool transpose_A, bool transpose_B>
    1295: inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMatMult_Numeric_Dispatch(Mat A, Mat B, Mat C) noexcept
    1296: {
    1297:   cupmBlasInt_t      m, n, k;
    1298:   PetscBool          Aiscupm, Biscupm;
    1299:   PetscDeviceContext dctx;
    1300:   cupmBlasHandle_t   handle;

    1302:   PetscFunctionBegin;
    1303:   PetscCall(PetscCUPMBlasIntCast(C->rmap->n, &m));
    1304:   PetscCall(PetscCUPMBlasIntCast(C->cmap->n, &n));
    1305:   PetscCall(PetscCUPMBlasIntCast(transpose_A ? A->rmap->n : A->cmap->n, &k));
    1306:   if (!m || !n || !k) PetscFunctionReturn(PETSC_SUCCESS);

    1308:   // we may end up with SEQDENSE as one of the arguments
    1309:   // REVIEW ME: how? and why is it not B and C????????
    1310:   PetscCall(PetscObjectTypeCompare(PetscObjectCast(A), MATSEQDENSECUPM(), &Aiscupm));
    1311:   PetscCall(PetscObjectTypeCompare(PetscObjectCast(B), MATSEQDENSECUPM(), &Biscupm));
    1312:   if (!Aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A));
    1313:   if (!Biscupm) PetscCall(MatConvert(B, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &B));
    1314:   PetscCall(PetscInfo(C, "Matrix-Matrix product %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " on backend\n", m, k, n));
    1315:   PetscCall(GetHandles_(&dctx, &handle));

    1317:   PetscCall(PetscLogGpuTimeBegin());
    1318:   {
    1319:     const auto one  = cupmScalarCast(1.0);
    1320:     const auto zero = cupmScalarCast(0.0);
    1321:     const auto da   = DeviceArrayRead(dctx, A);
    1322:     const auto db   = DeviceArrayRead(dctx, B);
    1323:     const auto dc   = DeviceArrayWrite(dctx, C);
    1324:     PetscInt   alda, blda, clda;

    1326:     PetscCall(MatDenseGetLDA(A, &alda));
    1327:     PetscCall(MatDenseGetLDA(B, &blda));
    1328:     PetscCall(MatDenseGetLDA(C, &clda));
    1329:     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));
    1330:   }
    1331:   PetscCall(PetscLogGpuTimeEnd());

    1333:   PetscCall(PetscLogGpuFlops(1.0 * m * n * k + 1.0 * m * n * (k - 1)));
    1334:   if (!Aiscupm) PetscCall(MatConvert(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A));
    1335:   if (!Biscupm) PetscCall(MatConvert(B, MATSEQDENSE, MAT_INPLACE_MATRIX, &B));
    1336:   PetscFunctionReturn(PETSC_SUCCESS);
    1337: }

    1339: template <device::cupm::DeviceType T>
    1340: inline PetscErrorCode MatDense_Seq_CUPM<T>::Copy(Mat A, Mat B, MatStructure str) noexcept
    1341: {
    1342:   const auto m = A->rmap->n;
    1343:   const auto n = A->cmap->n;

    1345:   PetscFunctionBegin;
    1346:   PetscAssert(m == B->rmap->n && n == B->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "size(B) != size(A)");
    1347:   // The two matrices must have the same copy implementation to be eligible for fast copy
    1348:   if (A->ops->copy == B->ops->copy) {
    1349:     PetscDeviceContext dctx;
    1350:     cupmStream_t       stream;

    1352:     PetscCall(GetHandles_(&dctx, &stream));
    1353:     PetscCall(PetscLogGpuTimeBegin());
    1354:     {
    1355:       const auto va = DeviceArrayRead(dctx, A);
    1356:       const auto vb = DeviceArrayWrite(dctx, B);
    1357:       // order is important, DeviceArrayRead/Write() might call SetPreallocation() which sets
    1358:       // lda!
    1359:       const auto lda_a = MatIMPLCast(A)->lda;
    1360:       const auto lda_b = MatIMPLCast(B)->lda;

    1362:       if (lda_a > m || lda_b > m) {
    1363:         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());
    1364:         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());
    1365:         PetscCall(PetscCUPMMemcpy2DAsync(vb.data(), lda_b, va.data(), lda_a, m, n, cupmMemcpyDeviceToDevice, stream));
    1366:       } else {
    1367:         PetscCall(PetscCUPMMemcpyAsync(vb.data(), va.data(), m * n, cupmMemcpyDeviceToDevice, stream));
    1368:       }
    1369:     }
    1370:     PetscCall(PetscLogGpuTimeEnd());
    1371:   } else {
    1372:     PetscCall(MatCopy_Basic(A, B, str));
    1373:   }
    1374:   PetscFunctionReturn(PETSC_SUCCESS);
    1375: }

    1377: template <device::cupm::DeviceType T>
    1378: inline PetscErrorCode MatDense_Seq_CUPM<T>::ZeroEntries(Mat m) noexcept
    1379: {
    1380:   PetscDeviceContext dctx;
    1381:   cupmStream_t       stream;

    1383:   PetscFunctionBegin;
    1384:   PetscCall(GetHandles_(&dctx, &stream));
    1385:   PetscCall(PetscLogGpuTimeBegin());
    1386:   {
    1387:     const auto va  = DeviceArrayWrite(dctx, m);
    1388:     const auto lda = MatIMPLCast(m)->lda;
    1389:     const auto ma  = m->rmap->n;
    1390:     const auto na  = m->cmap->n;

    1392:     if (lda > ma) {
    1393:       PetscCall(PetscCUPMMemset2DAsync(va.data(), lda, 0, ma, na, stream));
    1394:     } else {
    1395:       PetscCall(PetscCUPMMemsetAsync(va.data(), 0, ma * na, stream));
    1396:     }
    1397:   }
    1398:   PetscCall(PetscLogGpuTimeEnd());
    1399:   PetscFunctionReturn(PETSC_SUCCESS);
    1400: }

    1402: namespace detail
    1403: {

    1405: // ==========================================================================================
    1406: // SubMatIndexFunctor
    1407: //
    1408: // Iterator which permutes a linear index range into matrix indices for an nrows x ncols
    1409: // submat with leading dimension lda. Essentially SubMatIndexFunctor(i) returns the index for
    1410: // the i'th sequential entry in the matrix.
    1411: // ==========================================================================================
    1412: template <typename T>
    1413: struct SubMatIndexFunctor {
    1414:   PETSC_HOSTDEVICE_INLINE_DECL T operator()(T x) const noexcept { return ((x / nrows) * lda) + (x % nrows); }

    1416:   PetscInt nrows;
    1417:   PetscInt ncols;
    1418:   PetscInt lda;
    1419: };

    1421: template <typename Iterator>
    1422: struct SubMatrixIterator : MatrixIteratorBase<Iterator, SubMatIndexFunctor<iter_difference_t<Iterator>>> {
    1423:   using base_type = MatrixIteratorBase<Iterator, SubMatIndexFunctor<iter_difference_t<Iterator>>>;

    1425:   using iterator = typename base_type::iterator;

    1427:   constexpr SubMatrixIterator(Iterator first, Iterator last, PetscInt nrows, PetscInt ncols, PetscInt lda) noexcept :
    1428:     base_type{
    1429:       std::move(first), std::move(last), {nrows, ncols, lda}
    1430:   }
    1431:   {
    1432:   }

    1434:   PETSC_NODISCARD iterator end() const noexcept { return this->begin() + (this->func.nrows * this->func.ncols); }
    1435: };

    1437: namespace
    1438: {

    1440: template <typename T>
    1441: 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
    1442: {
    1443:   const auto nrows = rend - rstart;
    1444:   const auto ncols = cend - cstart;
    1445:   const auto dptr  = thrust::device_pointer_cast(ptr);

    1447:   return {dptr + (rstart * lda) + cstart, dptr + ((rstart + nrows) * lda) + cstart, nrows, ncols, lda};
    1448: }

    1450: } // namespace

    1452: struct conjugate {
    1453:   PETSC_NODISCARD PETSC_HOSTDEVICE_INLINE_DECL PetscScalar operator()(const PetscScalar &x) const noexcept { return PetscConj(x); }
    1454: };

    1456: } // namespace detail

    1458: template <device::cupm::DeviceType T>
    1459: inline PetscErrorCode MatDense_Seq_CUPM<T>::Conjugate(Mat A) noexcept
    1460: {
    1461:   const auto         m = A->rmap->n;
    1462:   const auto         n = A->cmap->n;
    1463:   const auto         N = m * n;
    1464:   PetscDeviceContext dctx;
    1465:   cupmStream_t       stream;

    1467:   PetscFunctionBegin;
    1468:   if (PetscDefined(USE_COMPLEX)) {
    1469:     PetscCall(GetHandles_(&dctx, &stream));
    1470:     PetscCall(PetscLogGpuTimeBegin());
    1471:     {
    1472:       const auto   da  = DeviceArrayReadWrite(dctx, A);
    1473:       const auto   lda = MatIMPLCast(A)->lda;
    1474:       cupmStream_t stream;
    1475:       PetscCall(GetHandlesFrom_(dctx, &stream));

    1477:       if (lda > m) {
    1478:         // clang-format off
    1479:         PetscCallThrust(
    1480:           const auto sub_mat = detail::make_submat_iterator(0, m, 0, n, lda, da.data());

    1482:           THRUST_CALL(
    1483:             thrust::transform,
    1484:             stream,
    1485:             sub_mat.begin(), sub_mat.end(), sub_mat.begin(),
    1486:             detail::conjugate{}
    1487:           )
    1488:         );
    1489:         // clang-format on
    1490:       } else {
    1491:         // clang-format off
    1492:         PetscCallThrust(
    1493:           const auto aptr = thrust::device_pointer_cast(da.data());

    1495:           THRUST_CALL(
    1496:             thrust::transform,
    1497:             stream,
    1498:             aptr, aptr + N, aptr,
    1499:             detail::conjugate{}
    1500:           )
    1501:         );
    1502:         // clang-format on
    1503:       }
    1504:     }
    1505:     PetscCall(PetscLogGpuTimeEnd());
    1506:   }
    1507:   PetscFunctionReturn(PETSC_SUCCESS);
    1508: }

    1510: template <device::cupm::DeviceType T>
    1511: inline PetscErrorCode MatDense_Seq_CUPM<T>::Scale(Mat A, PetscScalar alpha) noexcept
    1512: {
    1513:   const auto         m = A->rmap->n;
    1514:   const auto         n = A->cmap->n;
    1515:   const auto         N = m * n;
    1516:   PetscDeviceContext dctx;

    1518:   PetscFunctionBegin;
    1519:   PetscCall(PetscInfo(A, "Performing Scale %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m, n));
    1520:   PetscCall(GetHandles_(&dctx));
    1521:   {
    1522:     const auto da  = DeviceArrayReadWrite(dctx, A);
    1523:     const auto lda = MatIMPLCast(A)->lda;

    1525:     if (lda > m) {
    1526:       cupmStream_t stream;

    1528:       PetscCall(GetHandlesFrom_(dctx, &stream));
    1529:       // clang-format off
    1530:       PetscCallThrust(
    1531:         const auto sub_mat = detail::make_submat_iterator(0, m, 0, n, lda, da.data());

    1533:         THRUST_CALL(
    1534:           thrust::transform,
    1535:           stream,
    1536:           sub_mat.begin(), sub_mat.end(), sub_mat.begin(),
    1537:           device::cupm::functors::make_times_equals(alpha)
    1538:         )
    1539:       );
    1540:       // clang-format on
    1541:     } else {
    1542:       const auto       cu_alpha = cupmScalarCast(alpha);
    1543:       cupmBlasHandle_t handle;

    1545:       PetscCall(GetHandlesFrom_(dctx, &handle));
    1546:       PetscCall(PetscLogGpuTimeBegin());
    1547:       PetscCallCUPMBLAS(cupmBlasXscal(handle, N, &cu_alpha, da.cupmdata(), 1));
    1548:       PetscCall(PetscLogGpuTimeEnd());
    1549:     }
    1550:   }
    1551:   PetscCall(PetscLogGpuFlops(N));
    1552:   PetscFunctionReturn(PETSC_SUCCESS);
    1553: }

    1555: template <device::cupm::DeviceType T>
    1556: inline PetscErrorCode MatDense_Seq_CUPM<T>::AXPY(Mat Y, PetscScalar alpha, Mat X, MatStructure) noexcept
    1557: {
    1558:   const auto         m_x = X->rmap->n, m_y = Y->rmap->n;
    1559:   const auto         n_x = X->cmap->n, n_y = Y->cmap->n;
    1560:   const auto         N = m_x * n_x;
    1561:   PetscDeviceContext dctx;

    1563:   PetscFunctionBegin;
    1564:   if (!m_x || !n_x || alpha == (PetscScalar)0.0) PetscFunctionReturn(PETSC_SUCCESS);
    1565:   PetscCall(PetscInfo(Y, "Performing AXPY %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m_y, n_y));
    1566:   PetscCall(GetHandles_(&dctx));
    1567:   {
    1568:     const auto dx    = DeviceArrayRead(dctx, X);
    1569:     const auto dy    = DeviceArrayReadWrite(dctx, Y);
    1570:     const auto lda_x = MatIMPLCast(X)->lda;
    1571:     const auto lda_y = MatIMPLCast(Y)->lda;

    1573:     if (lda_x > m_x || lda_y > m_x) {
    1574:       cupmStream_t stream;

    1576:       PetscCall(GetHandlesFrom_(dctx, &stream));
    1577:       // clang-format off
    1578:       PetscCallThrust(
    1579:         const auto sub_mat_y = detail::make_submat_iterator(0, m_y, 0, n_y, lda_y, dy.data());
    1580:         const auto sub_mat_x = detail::make_submat_iterator(0, m_x, 0, n_x, lda_x, dx.data());

    1582:         THRUST_CALL(
    1583:           thrust::transform,
    1584:           stream,
    1585:           sub_mat_x.begin(), sub_mat_x.end(), sub_mat_y.begin(), sub_mat_y.begin(),
    1586:           device::cupm::functors::make_axpy(alpha)
    1587:         );
    1588:       );
    1589:       // clang-format on
    1590:     } else {
    1591:       const auto       cu_alpha = cupmScalarCast(alpha);
    1592:       cupmBlasHandle_t handle;

    1594:       PetscCall(GetHandlesFrom_(dctx, &handle));
    1595:       PetscCall(PetscLogGpuTimeBegin());
    1596:       PetscCallCUPMBLAS(cupmBlasXaxpy(handle, N, &cu_alpha, dx.cupmdata(), 1, dy.cupmdata(), 1));
    1597:       PetscCall(PetscLogGpuTimeEnd());
    1598:     }
    1599:   }
    1600:   PetscCall(PetscLogGpuFlops(PetscMax(2 * N - 1, 0)));
    1601:   PetscFunctionReturn(PETSC_SUCCESS);
    1602: }

    1604: template <device::cupm::DeviceType T>
    1605: inline PetscErrorCode MatDense_Seq_CUPM<T>::Duplicate(Mat A, MatDuplicateOption opt, Mat *B) noexcept
    1606: {
    1607:   const auto         hopt = (opt == MAT_COPY_VALUES && A->offloadmask != PETSC_OFFLOAD_CPU) ? MAT_DO_NOT_COPY_VALUES : opt;
    1608:   PetscDeviceContext dctx;

    1610:   PetscFunctionBegin;
    1611:   PetscCall(GetHandles_(&dctx));
    1612:   // do not call SetPreallocation() yet, we call it afterwards??
    1613:   PetscCall(MatCreateSeqDenseCUPM<T>(PetscObjectComm(PetscObjectCast(A)), A->rmap->n, A->cmap->n, nullptr, B, dctx, /* preallocate */ false));
    1614:   PetscCall(MatDuplicateNoCreate_SeqDense(*B, A, hopt));
    1615:   if (opt == MAT_COPY_VALUES && hopt != MAT_COPY_VALUES) PetscCall(Copy(A, *B, SAME_NONZERO_PATTERN));
    1616:   // allocate memory if needed
    1617:   if (opt != MAT_COPY_VALUES && !MatCUPMCast(*B)->d_v) PetscCall(SetPreallocation(*B, dctx, nullptr));
    1618:   PetscFunctionReturn(PETSC_SUCCESS);
    1619: }

    1621: template <device::cupm::DeviceType T>
    1622: inline PetscErrorCode MatDense_Seq_CUPM<T>::SetRandom(Mat A, PetscRandom rng) noexcept
    1623: {
    1624:   PetscBool device_rand_is_rander48;
    1625:   PetscBool device = PETSC_FALSE;

    1627:   PetscFunctionBegin;
    1628:   // CUPMObject<hip>::PETSCDEVICERAD() is PETSCRANDER48 until PetscRandom is implemented for hiprand
    1629:   PetscCall(PetscStrncmp(PETSCDEVICERAND(), PETSCRANDER48, sizeof(PETSCRANDER48), &device_rand_is_rander48));
    1630:   if (!device_rand_is_rander48) PetscCall(PetscObjectTypeCompare(PetscObjectCast(rng), PETSCDEVICERAND(), &device));
    1631:   if (device) {
    1632:     const auto         m = A->rmap->n;
    1633:     const auto         n = A->cmap->n;
    1634:     PetscDeviceContext dctx;

    1636:     PetscCall(GetHandles_(&dctx));
    1637:     {
    1638:       const auto a = DeviceArrayWrite(dctx, A);
    1639:       PetscInt   lda;

    1641:       PetscCall(MatDenseGetLDA(A, &lda));
    1642:       if (lda > m) {
    1643:         for (PetscInt i = 0; i < n; i++) PetscCall(PetscRandomGetValues(rng, m, a.data() + i * lda));
    1644:       } else {
    1645:         PetscInt mn;

    1647:         PetscCall(PetscIntMultError(m, n, &mn));
    1648:         PetscCall(PetscRandomGetValues(rng, mn, a));
    1649:       }
    1650:     }
    1651:   } else {
    1652:     PetscCall(MatSetRandom_SeqDense(A, rng));
    1653:   }
    1654:   PetscFunctionReturn(PETSC_SUCCESS);
    1655: }

    1657: // ==========================================================================================

    1659: template <device::cupm::DeviceType T>
    1660: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetColumnVector(Mat A, Vec v, PetscInt col) noexcept
    1661: {
    1662:   const auto         offloadmask = A->offloadmask;
    1663:   const auto         n           = A->rmap->n;
    1664:   const auto         col_offset  = [&](const PetscScalar *ptr) { return ptr + col * MatIMPLCast(A)->lda; };
    1665:   PetscBool          viscupm;
    1666:   PetscDeviceContext dctx;
    1667:   cupmStream_t       stream;

    1669:   PetscFunctionBegin;
    1670:   PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(v), &viscupm, VecSeq_CUPM::VECSEQCUPM(), VecSeq_CUPM::VECMPICUPM(), VecSeq_CUPM::VECCUPM(), ""));
    1671:   PetscCall(GetHandles_(&dctx, &stream));
    1672:   if (viscupm && !v->boundtocpu) {
    1673:     const auto x = VecSeq_CUPM::DeviceArrayWrite(dctx, v);

    1675:     // update device data
    1676:     if (PetscOffloadDevice(offloadmask)) {
    1677:       PetscCall(PetscCUPMMemcpyAsync(x.data(), col_offset(DeviceArrayRead(dctx, A)), n, cupmMemcpyDeviceToDevice, stream));
    1678:     } else {
    1679:       PetscCall(PetscCUPMMemcpyAsync(x.data(), col_offset(HostArrayRead(dctx, A)), n, cupmMemcpyHostToDevice, stream));
    1680:     }
    1681:   } else {
    1682:     PetscScalar *x;

    1684:     // update host data
    1685:     PetscCall(VecGetArrayWrite(v, &x));
    1686:     if (PetscOffloadUnallocated(offloadmask) || PetscOffloadHost(offloadmask)) {
    1687:       PetscCall(PetscArraycpy(x, col_offset(HostArrayRead(dctx, A)), n));
    1688:     } else if (PetscOffloadDevice(offloadmask)) {
    1689:       PetscCall(PetscCUPMMemcpyAsync(x, col_offset(DeviceArrayRead(dctx, A)), n, cupmMemcpyDeviceToHost, stream));
    1690:     }
    1691:     PetscCall(VecRestoreArrayWrite(v, &x));
    1692:   }
    1693:   PetscFunctionReturn(PETSC_SUCCESS);
    1694: }

    1696: template <device::cupm::DeviceType T>
    1697: template <PetscMemoryAccessMode access>
    1698: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetColumnVec(Mat A, PetscInt col, Vec *v) noexcept
    1699: {
    1700:   using namespace vec::cupm;
    1701:   const auto         mimpl = MatIMPLCast(A);
    1702:   PetscDeviceContext dctx;

    1704:   PetscFunctionBegin;
    1705:   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
    1706:   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
    1707:   mimpl->vecinuse = col + 1;
    1708:   if (!mimpl->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &mimpl->cvec));
    1709:   PetscCall(GetHandles_(&dctx));
    1710:   PetscCall(GetArray<PETSC_MEMTYPE_DEVICE, access>(A, const_cast<PetscScalar **>(&mimpl->ptrinuse), dctx));
    1711:   PetscCall(VecCUPMPlaceArrayAsync<T>(mimpl->cvec, mimpl->ptrinuse + static_cast<std::size_t>(col) * static_cast<std::size_t>(mimpl->lda)));
    1712:   if (access == PETSC_MEMORY_ACCESS_READ) PetscCall(VecLockReadPush(mimpl->cvec));
    1713:   *v = mimpl->cvec;
    1714:   PetscFunctionReturn(PETSC_SUCCESS);
    1715: }

    1717: template <device::cupm::DeviceType T>
    1718: template <PetscMemoryAccessMode access>
    1719: inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreColumnVec(Mat A, PetscInt, Vec *v) noexcept
    1720: {
    1721:   using namespace vec::cupm;
    1722:   const auto         mimpl = MatIMPLCast(A);
    1723:   const auto         cvec  = mimpl->cvec;
    1724:   PetscDeviceContext dctx;

    1726:   PetscFunctionBegin;
    1727:   PetscCheck(mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
    1728:   PetscCheck(cvec, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column vector");
    1729:   mimpl->vecinuse = 0;
    1730:   if (access == PETSC_MEMORY_ACCESS_READ) PetscCall(VecLockReadPop(cvec));
    1731:   PetscCall(VecCUPMResetArrayAsync<T>(cvec));
    1732:   PetscCall(GetHandles_(&dctx));
    1733:   PetscCall(RestoreArray<PETSC_MEMTYPE_DEVICE, access>(A, const_cast<PetscScalar **>(&mimpl->ptrinuse), dctx));
    1734:   if (v) *v = nullptr;
    1735:   PetscFunctionReturn(PETSC_SUCCESS);
    1736: }

    1738: // ==========================================================================================

    1740: template <device::cupm::DeviceType T>
    1741: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetFactor(Mat A, MatFactorType ftype, Mat *fact_out) noexcept
    1742: {
    1743:   Mat                fact = nullptr;
    1744:   PetscDeviceContext dctx;

    1746:   PetscFunctionBegin;
    1747:   PetscCall(GetHandles_(&dctx));
    1748:   PetscCall(MatCreateSeqDenseCUPM<T>(PetscObjectComm(PetscObjectCast(A)), A->rmap->n, A->cmap->n, nullptr, &fact, dctx, /* preallocate */ false));
    1749:   fact->factortype = ftype;
    1750:   switch (ftype) {
    1751:   case MAT_FACTOR_LU:
    1752:   case MAT_FACTOR_ILU: // fall-through
    1753:     fact->ops->lufactorsymbolic  = MatLUFactorSymbolic_SeqDense;
    1754:     fact->ops->ilufactorsymbolic = MatLUFactorSymbolic_SeqDense;
    1755:     break;
    1756:   case MAT_FACTOR_CHOLESKY:
    1757:   case MAT_FACTOR_ICC: // fall-through
    1758:     fact->ops->choleskyfactorsymbolic = MatCholeskyFactorSymbolic_SeqDense;
    1759:     break;
    1760:   case MAT_FACTOR_QR: {
    1761:     const auto pobj = PetscObjectCast(fact);

    1763:     PetscCall(PetscObjectComposeFunction(pobj, "MatQRFactor_C", MatQRFactor_SeqDense));
    1764:     PetscCall(PetscObjectComposeFunction(pobj, "MatQRFactorSymbolic_C", MatQRFactorSymbolic_SeqDense));
    1765:   } break;
    1766:   case MAT_FACTOR_NONE:
    1767:   case MAT_FACTOR_ILUDT:     // fall-through
    1768:   case MAT_FACTOR_NUM_TYPES: // fall-through
    1769:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "MatFactorType %s not supported", MatFactorTypes[ftype]);
    1770:   }
    1771:   PetscCall(PetscStrFreeAllocpy(MATSOLVERCUPM(), &fact->solvertype));
    1772:   PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_LU));
    1773:   PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_ILU));
    1774:   PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_CHOLESKY));
    1775:   PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_ICC));
    1776:   *fact_out = fact;
    1777:   PetscFunctionReturn(PETSC_SUCCESS);
    1778: }

    1780: template <device::cupm::DeviceType T>
    1781: inline PetscErrorCode MatDense_Seq_CUPM<T>::InvertFactors(Mat A) noexcept
    1782: {
    1783:   const auto         mimpl = MatIMPLCast(A);
    1784:   const auto         mcu   = MatCUPMCast(A);
    1785:   const auto         n     = static_cast<cupmBlasInt_t>(A->cmap->n);
    1786:   cupmSolverHandle_t handle;
    1787:   PetscDeviceContext dctx;
    1788:   cupmStream_t       stream;

    1790:   PetscFunctionBegin;
    1791: #if PetscDefined(HAVE_CUDA) && PetscDefined(USING_NVCC)
    1792:   // HIP appears to have this by default??
    1793:   PetscCheck(PETSC_PKG_CUDA_VERSION_GE(10, 1, 0), PETSC_COMM_SELF, PETSC_ERR_SUP, "Upgrade to CUDA version 10.1.0 or higher");
    1794: #endif
    1795:   if (!n || !A->rmap->n) PetscFunctionReturn(PETSC_SUCCESS);
    1796:   PetscCheck(A->factortype == MAT_FACTOR_CHOLESKY, PETSC_COMM_SELF, PETSC_ERR_LIB, "Factor type %s not implemented", MatFactorTypes[A->factortype]);
    1797:   // spd
    1798:   PetscCheck(!mcu->d_fact_ipiv, PETSC_COMM_SELF, PETSC_ERR_LIB, "%sDnsytri not implemented", cupmSolverName());

    1800:   PetscCall(GetHandles_(&dctx, &handle, &stream));
    1801:   {
    1802:     const auto    da  = DeviceArrayReadWrite(dctx, A);
    1803:     const auto    lda = static_cast<cupmBlasInt_t>(mimpl->lda);
    1804:     cupmBlasInt_t il;

    1806:     PetscCallCUPMSOLVER(cupmSolverXpotri_bufferSize(handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, &il));
    1807:     if (il > mcu->d_fact_lwork) {
    1808:       mcu->d_fact_lwork = il;
    1809:       PetscCallCUPM(cupmFreeAsync(mcu->d_fact_work, stream));
    1810:       PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_work, il, stream));
    1811:     }
    1812:     PetscCall(PetscLogGpuTimeBegin());
    1813:     PetscCallCUPMSOLVER(cupmSolverXpotri(handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_info));
    1814:     PetscCall(PetscLogGpuTimeEnd());
    1815:   }
    1816:   PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream));
    1817:   // TODO (write cuda kernel)
    1818:   PetscCall(MatSeqDenseSymmetrize_Private(A, PETSC_TRUE));
    1819:   PetscCall(PetscLogGpuFlops(1.0 * n * n * n / 3.0));

    1821:   A->ops->solve          = nullptr;
    1822:   A->ops->solvetranspose = nullptr;
    1823:   A->ops->matsolve       = nullptr;
    1824:   A->factortype          = MAT_FACTOR_NONE;

    1826:   PetscCall(PetscFree(A->solvertype));
    1827:   PetscFunctionReturn(PETSC_SUCCESS);
    1828: }

    1830: // ==========================================================================================

    1832: template <device::cupm::DeviceType T>
    1833: inline PetscErrorCode MatDense_Seq_CUPM<T>::GetSubMatrix(Mat A, PetscInt rbegin, PetscInt rend, PetscInt cbegin, PetscInt cend, Mat *mat) noexcept
    1834: {
    1835:   const auto         mimpl        = MatIMPLCast(A);
    1836:   const auto         array_offset = [&](PetscScalar *ptr) { return ptr + rbegin + static_cast<std::size_t>(cbegin) * mimpl->lda; };
    1837:   const auto         n            = rend - rbegin;
    1838:   const auto         m            = cend - cbegin;
    1839:   auto              &cmat         = mimpl->cmat;
    1840:   PetscDeviceContext dctx;

    1842:   PetscFunctionBegin;
    1843:   PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
    1844:   PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
    1845:   mimpl->matinuse = cbegin + 1;

    1847:   PetscCall(GetHandles_(&dctx));
    1848:   PetscCall(HostToDevice_(A, dctx));

    1850:   if (cmat && ((m != cmat->cmap->N) || (n != cmat->rmap->N))) PetscCall(MatDestroy(&cmat));
    1851:   {
    1852:     const auto device_array = array_offset(MatCUPMCast(A)->d_v);

    1854:     if (cmat) {
    1855:       PetscCall(PlaceArray(cmat, device_array));
    1856:     } else {
    1857:       PetscCall(MatCreateSeqDenseCUPM<T>(PetscObjectComm(PetscObjectCast(A)), n, m, device_array, &cmat, dctx));
    1858:     }
    1859:   }
    1860:   PetscCall(MatDenseSetLDA(cmat, mimpl->lda));
    1861:   // place CPU array if present but do not copy any data
    1862:   if (const auto host_array = mimpl->v) {
    1863:     cmat->offloadmask = PETSC_OFFLOAD_GPU;
    1864:     PetscCall(MatDensePlaceArray(cmat, array_offset(host_array)));
    1865:   }

    1867:   cmat->offloadmask = A->offloadmask;
    1868:   *mat              = cmat;
    1869:   PetscFunctionReturn(PETSC_SUCCESS);
    1870: }

    1872: template <device::cupm::DeviceType T>
    1873: inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreSubMatrix(Mat A, Mat *m) noexcept
    1874: {
    1875:   const auto mimpl = MatIMPLCast(A);
    1876:   const auto cmat  = mimpl->cmat;
    1877:   const auto reset = static_cast<bool>(mimpl->v);
    1878:   bool       copy, was_offload_host;

    1880:   PetscFunctionBegin;
    1881:   PetscCheck(mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetSubMatrix() first");
    1882:   PetscCheck(cmat, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column matrix");
    1883:   PetscCheck(*m == cmat, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Not the matrix obtained from MatDenseGetSubMatrix()");
    1884:   mimpl->matinuse = 0;

    1886:   // calls to ResetArray may change it, so save it here
    1887:   was_offload_host = cmat->offloadmask == PETSC_OFFLOAD_CPU;
    1888:   if (was_offload_host && !reset) {
    1889:     copy = true;
    1890:     PetscCall(MatSeqDenseSetPreallocation(A, nullptr));
    1891:   } else {
    1892:     copy = false;
    1893:   }

    1895:   PetscCall(ResetArray(cmat));
    1896:   if (reset) PetscCall(MatDenseResetArray(cmat));
    1897:   if (copy) {
    1898:     PetscDeviceContext dctx;

    1900:     PetscCall(GetHandles_(&dctx));
    1901:     PetscCall(DeviceToHost_(A, dctx));
    1902:   } else {
    1903:     A->offloadmask = was_offload_host ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU;
    1904:   }

    1906:   cmat->offloadmask = PETSC_OFFLOAD_UNALLOCATED;
    1907:   *m                = nullptr;
    1908:   PetscFunctionReturn(PETSC_SUCCESS);
    1909: }

    1911: // ==========================================================================================

    1913: namespace
    1914: {

    1916: template <device::cupm::DeviceType T>
    1917: inline PetscErrorCode MatMatMultNumeric_SeqDenseCUPM_SeqDenseCUPM(Mat A, Mat B, Mat C, PetscBool TA, PetscBool TB) noexcept
    1918: {
    1919:   PetscFunctionBegin;
    1920:   if (TA) {
    1921:     if (TB) {
    1922:       PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<true, true>(A, B, C));
    1923:     } else {
    1924:       PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<true, false>(A, B, C));
    1925:     }
    1926:   } else {
    1927:     if (TB) {
    1928:       PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<false, true>(A, B, C));
    1929:     } else {
    1930:       PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<false, false>(A, B, C));
    1931:     }
    1932:   }
    1933:   PetscFunctionReturn(PETSC_SUCCESS);
    1934: }

    1936: template <device::cupm::DeviceType T>
    1937: inline PetscErrorCode MatSolverTypeRegister_DENSECUPM() noexcept
    1938: {
    1939:   PetscFunctionBegin;
    1940:   for (auto ftype : util::make_array(MAT_FACTOR_LU, MAT_FACTOR_CHOLESKY, MAT_FACTOR_QR)) {
    1941:     PetscCall(MatSolverTypeRegister(MatDense_Seq_CUPM<T>::MATSOLVERCUPM(), MATSEQDENSE, ftype, MatDense_Seq_CUPM<T>::GetFactor));
    1942:     PetscCall(MatSolverTypeRegister(MatDense_Seq_CUPM<T>::MATSOLVERCUPM(), MatDense_Seq_CUPM<T>::MATSEQDENSECUPM(), ftype, MatDense_Seq_CUPM<T>::GetFactor));
    1943:   }
    1944:   PetscFunctionReturn(PETSC_SUCCESS);
    1945: }

    1947: } // anonymous namespace

    1949: } // namespace impl

    1951: } // namespace cupm

    1953: } // namespace mat

    1955: } // namespace Petsc