1 #ifndef PETSCMATSEQDENSECUPM_HPP 2 #define PETSCMATSEQDENSECUPM_HPP 3 4 #include <petsc/private/matdensecupmimpl.h> /*I <petscmat.h> I*/ 5 #include <../src/mat/impls/dense/seq/dense.h> 6 7 #if defined(__cplusplus) 8 #include <petsc/private/deviceimpl.h> // PetscDeviceContextGetOptionalNullContext_Internal() 9 #include <petsc/private/randomimpl.h> // _p_PetscRandom 10 #include <petsc/private/vecimpl.h> // _p_Vec 11 #include <petsc/private/cupmobject.hpp> 12 #include <petsc/private/cupmsolverinterface.hpp> 13 14 #include <petsc/private/cpp/type_traits.hpp> // PetscObjectCast() 15 #include <petsc/private/cpp/utility.hpp> // util::exchange() 16 17 #include <../src/vec/vec/impls/seq/cupm/vecseqcupm.hpp> // for VecSeq_CUPM 18 19 namespace Petsc 20 { 21 22 namespace mat 23 { 24 25 namespace cupm 26 { 27 28 namespace impl 29 { 30 31 template <device::cupm::DeviceType T> 32 class MatDense_Seq_CUPM : MatDense_CUPM<T, MatDense_Seq_CUPM<T>> { 33 public: 34 MATDENSECUPM_HEADER(T, MatDense_Seq_CUPM<T>); 35 36 private: 37 struct Mat_SeqDenseCUPM { 38 PetscScalar *d_v; // pointer to the matrix on the GPU 39 PetscScalar *unplacedarray; // if one called MatCUPMDensePlaceArray(), this is where it stashed the original 40 bool d_user_alloc; 41 bool d_unplaced_user_alloc; 42 // factorization support 43 cupmBlasInt_t *d_fact_ipiv; // device pivots 44 cupmScalar_t *d_fact_tau; // device QR tau vector 45 cupmBlasInt_t *d_fact_info; // device info 46 cupmScalar_t *d_fact_work; // device workspace 47 cupmBlasInt_t d_fact_lwork; // size of device workspace 48 // workspace 49 Vec workvec; 50 }; 51 52 static PetscErrorCode SetPreallocation_(Mat, PetscDeviceContext, PetscScalar *) noexcept; 53 54 static PetscErrorCode HostToDevice_(Mat, PetscDeviceContext) noexcept; 55 static PetscErrorCode DeviceToHost_(Mat, PetscDeviceContext) noexcept; 56 57 static PetscErrorCode CheckCUPMSolverInfo_(const cupmBlasInt_t *, cupmStream_t) noexcept; 58 59 template <typename Derived> 60 struct SolveCommon; 61 struct SolveQR; 62 struct SolveCholesky; 63 struct SolveLU; 64 65 template <typename Solver, bool transpose> 66 static PetscErrorCode MatSolve_Factored_Dispatch_(Mat, Vec, Vec) noexcept; 67 template <typename Solver, bool transpose> 68 static PetscErrorCode MatMatSolve_Factored_Dispatch_(Mat, Mat, Mat) noexcept; 69 template <bool transpose> 70 static PetscErrorCode MatMultAdd_Dispatch_(Mat, Vec, Vec, Vec) noexcept; 71 72 template <bool to_host> 73 static PetscErrorCode Convert_Dispatch_(Mat, MatType, MatReuse, Mat *) noexcept; 74 75 PETSC_NODISCARD static constexpr MatType MATIMPLCUPM_() noexcept; 76 PETSC_NODISCARD static constexpr Mat_SeqDense *MatIMPLCast_(Mat) noexcept; 77 78 public: 79 PETSC_NODISCARD static constexpr Mat_SeqDenseCUPM *MatCUPMCast(Mat) noexcept; 80 81 // define these by hand since they don't fit the above mold 82 PETSC_NODISCARD static constexpr const char *MatConvert_seqdensecupm_seqdense_C() noexcept; 83 PETSC_NODISCARD static constexpr const char *MatProductSetFromOptions_seqaij_seqdensecupm_C() noexcept; 84 85 static PetscErrorCode Create(Mat) noexcept; 86 static PetscErrorCode Destroy(Mat) noexcept; 87 static PetscErrorCode SetUp(Mat) noexcept; 88 static PetscErrorCode Reset(Mat) noexcept; 89 90 static PetscErrorCode BindToCPU(Mat, PetscBool) noexcept; 91 static PetscErrorCode Convert_SeqDense_SeqDenseCUPM(Mat, MatType, MatReuse, Mat *) noexcept; 92 static PetscErrorCode Convert_SeqDenseCUPM_SeqDense(Mat, MatType, MatReuse, Mat *) noexcept; 93 94 template <PetscMemType, PetscMemoryAccessMode> 95 static PetscErrorCode GetArray(Mat, PetscScalar **, PetscDeviceContext) noexcept; 96 template <PetscMemType, PetscMemoryAccessMode> 97 static PetscErrorCode RestoreArray(Mat, PetscScalar **, PetscDeviceContext) noexcept; 98 template <PetscMemoryAccessMode> 99 static PetscErrorCode GetArrayAndMemType(Mat, PetscScalar **, PetscMemType *, PetscDeviceContext) noexcept; 100 template <PetscMemoryAccessMode> 101 static PetscErrorCode RestoreArrayAndMemType(Mat, PetscScalar **, PetscDeviceContext) noexcept; 102 103 private: 104 template <PetscMemType mtype, PetscMemoryAccessMode mode> 105 static PetscErrorCode GetArrayC_(Mat m, PetscScalar **p) noexcept 106 { 107 PetscDeviceContext dctx; 108 109 PetscFunctionBegin; 110 PetscCall(GetHandles_(&dctx)); 111 PetscCall(GetArray<mtype, mode>(m, p, dctx)); 112 PetscFunctionReturn(PETSC_SUCCESS); 113 } 114 115 template <PetscMemType mtype, PetscMemoryAccessMode mode> 116 static PetscErrorCode RestoreArrayC_(Mat m, PetscScalar **p) noexcept 117 { 118 PetscDeviceContext dctx; 119 120 PetscFunctionBegin; 121 PetscCall(GetHandles_(&dctx)); 122 PetscCall(RestoreArray<mtype, mode>(m, p, dctx)); 123 PetscFunctionReturn(PETSC_SUCCESS); 124 } 125 126 template <PetscMemoryAccessMode mode> 127 static PetscErrorCode GetArrayAndMemTypeC_(Mat m, PetscScalar **p, PetscMemType *tp) noexcept 128 { 129 PetscDeviceContext dctx; 130 131 PetscFunctionBegin; 132 PetscCall(GetHandles_(&dctx)); 133 PetscCall(GetArrayAndMemType<mode>(m, p, tp, dctx)); 134 PetscFunctionReturn(PETSC_SUCCESS); 135 } 136 137 template <PetscMemoryAccessMode mode> 138 static PetscErrorCode RestoreArrayAndMemTypeC_(Mat m, PetscScalar **p) noexcept 139 { 140 PetscDeviceContext dctx; 141 142 PetscFunctionBegin; 143 PetscCall(GetHandles_(&dctx)); 144 PetscCall(RestoreArrayAndMemType<mode>(m, p, dctx)); 145 PetscFunctionReturn(PETSC_SUCCESS); 146 } 147 148 public: 149 static PetscErrorCode PlaceArray(Mat, const PetscScalar *) noexcept; 150 static PetscErrorCode ReplaceArray(Mat, const PetscScalar *) noexcept; 151 static PetscErrorCode ResetArray(Mat) noexcept; 152 153 template <bool transpose_A, bool transpose_B> 154 static PetscErrorCode MatMatMult_Numeric_Dispatch(Mat, Mat, Mat) noexcept; 155 static PetscErrorCode Copy(Mat, Mat, MatStructure) noexcept; 156 static PetscErrorCode ZeroEntries(Mat) noexcept; 157 static PetscErrorCode Scale(Mat, PetscScalar) noexcept; 158 static PetscErrorCode Shift(Mat, PetscScalar) noexcept; 159 static PetscErrorCode AXPY(Mat, PetscScalar, Mat, MatStructure) noexcept; 160 static PetscErrorCode Duplicate(Mat, MatDuplicateOption, Mat *) noexcept; 161 static PetscErrorCode SetRandom(Mat, PetscRandom) noexcept; 162 163 static PetscErrorCode GetColumnVector(Mat, Vec, PetscInt) noexcept; 164 template <PetscMemoryAccessMode> 165 static PetscErrorCode GetColumnVec(Mat, PetscInt, Vec *) noexcept; 166 template <PetscMemoryAccessMode> 167 static PetscErrorCode RestoreColumnVec(Mat, PetscInt, Vec *) noexcept; 168 169 static PetscErrorCode GetFactor(Mat, MatFactorType, Mat *) noexcept; 170 static PetscErrorCode InvertFactors(Mat) noexcept; 171 172 static PetscErrorCode GetSubMatrix(Mat, PetscInt, PetscInt, PetscInt, PetscInt, Mat *) noexcept; 173 static PetscErrorCode RestoreSubMatrix(Mat, Mat *) noexcept; 174 }; 175 176 } // namespace impl 177 178 namespace 179 { 180 181 // Declare this here so that the functions below can make use of it 182 template <device::cupm::DeviceType T> 183 inline PetscErrorCode MatCreateSeqDenseCUPM(MPI_Comm comm, PetscInt m, PetscInt n, PetscScalar *data, Mat *A, PetscDeviceContext dctx = nullptr, bool preallocate = true) noexcept 184 { 185 PetscFunctionBegin; 186 PetscCall(impl::MatDense_Seq_CUPM<T>::CreateIMPLDenseCUPM(comm, m, n, m, n, data, A, dctx, preallocate)); 187 PetscFunctionReturn(PETSC_SUCCESS); 188 } 189 190 } // anonymous namespace 191 192 namespace impl 193 { 194 195 // ========================================================================================== 196 // MatDense_Seq_CUPM - Private API - Utility 197 // ========================================================================================== 198 199 template <device::cupm::DeviceType T> 200 inline PetscErrorCode MatDense_Seq_CUPM<T>::SetPreallocation_(Mat m, PetscDeviceContext dctx, PetscScalar *user_device_array) noexcept 201 { 202 const auto mcu = MatCUPMCast(m); 203 const auto nrows = m->rmap->n; 204 const auto ncols = m->cmap->n; 205 auto &lda = MatIMPLCast(m)->lda; 206 cupmStream_t stream; 207 208 PetscFunctionBegin; 209 PetscCheckTypeName(m, MATSEQDENSECUPM()); 210 PetscValidDeviceContext(dctx, 2); 211 PetscCall(checkCupmBlasIntCast(nrows)); 212 PetscCall(checkCupmBlasIntCast(ncols)); 213 PetscCall(GetHandlesFrom_(dctx, &stream)); 214 if (lda <= 0) lda = nrows; 215 if (!mcu->d_user_alloc) PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream)); 216 if (user_device_array) { 217 mcu->d_user_alloc = PETSC_TRUE; 218 mcu->d_v = user_device_array; 219 } else { 220 PetscInt size; 221 222 mcu->d_user_alloc = PETSC_FALSE; 223 PetscCall(PetscIntMultError(lda, ncols, &size)); 224 PetscCall(PetscCUPMMallocAsync(&mcu->d_v, size, stream)); 225 PetscCall(PetscCUPMMemsetAsync(mcu->d_v, 0, size, stream)); 226 } 227 m->offloadmask = PETSC_OFFLOAD_GPU; 228 PetscFunctionReturn(PETSC_SUCCESS); 229 } 230 231 template <device::cupm::DeviceType T> 232 inline PetscErrorCode MatDense_Seq_CUPM<T>::HostToDevice_(Mat m, PetscDeviceContext dctx) noexcept 233 { 234 const auto nrows = m->rmap->n; 235 const auto ncols = m->cmap->n; 236 const auto copy = m->offloadmask == PETSC_OFFLOAD_CPU || m->offloadmask == PETSC_OFFLOAD_UNALLOCATED; 237 238 PetscFunctionBegin; 239 PetscCheckTypeName(m, MATSEQDENSECUPM()); 240 if (m->boundtocpu) PetscFunctionReturn(PETSC_SUCCESS); 241 PetscCall(PetscInfo(m, "%s matrix %" PetscInt_FMT " x %" PetscInt_FMT "\n", copy ? "Copy" : "Reusing", nrows, ncols)); 242 if (copy) { 243 const auto mcu = MatCUPMCast(m); 244 cupmStream_t stream; 245 246 // Allocate GPU memory if not present 247 if (!mcu->d_v) PetscCall(SetPreallocation(m, dctx, nullptr)); 248 PetscCall(GetHandlesFrom_(dctx, &stream)); 249 PetscCall(PetscLogEventBegin(MAT_DenseCopyToGPU, m, 0, 0, 0)); 250 { 251 const auto mimpl = MatIMPLCast(m); 252 const auto lda = mimpl->lda; 253 const auto src = mimpl->v; 254 const auto dest = mcu->d_v; 255 256 if (lda > nrows) { 257 PetscCall(PetscCUPMMemcpy2DAsync(dest, lda, src, lda, nrows, ncols, cupmMemcpyHostToDevice, stream)); 258 } else { 259 PetscCall(PetscCUPMMemcpyAsync(dest, src, lda * ncols, cupmMemcpyHostToDevice, stream)); 260 } 261 } 262 PetscCall(PetscLogEventEnd(MAT_DenseCopyToGPU, m, 0, 0, 0)); 263 // order important, ensure that offloadmask is PETSC_OFFLOAD_BOTH 264 m->offloadmask = PETSC_OFFLOAD_BOTH; 265 } 266 PetscFunctionReturn(PETSC_SUCCESS); 267 } 268 269 template <device::cupm::DeviceType T> 270 inline PetscErrorCode MatDense_Seq_CUPM<T>::DeviceToHost_(Mat m, PetscDeviceContext dctx) noexcept 271 { 272 const auto nrows = m->rmap->n; 273 const auto ncols = m->cmap->n; 274 const auto copy = m->offloadmask == PETSC_OFFLOAD_GPU; 275 276 PetscFunctionBegin; 277 PetscCheckTypeName(m, MATSEQDENSECUPM()); 278 PetscCall(PetscInfo(m, "%s matrix %" PetscInt_FMT " x %" PetscInt_FMT "\n", copy ? "Copy" : "Reusing", nrows, ncols)); 279 if (copy) { 280 const auto mimpl = MatIMPLCast(m); 281 cupmStream_t stream; 282 283 // MatCreateSeqDenseCUPM may not allocate CPU memory. Allocate if needed 284 if (!mimpl->v) PetscCall(MatSeqDenseSetPreallocation(m, nullptr)); 285 PetscCall(GetHandlesFrom_(dctx, &stream)); 286 PetscCall(PetscLogEventBegin(MAT_DenseCopyFromGPU, m, 0, 0, 0)); 287 { 288 const auto lda = mimpl->lda; 289 const auto dest = mimpl->v; 290 const auto src = MatCUPMCast(m)->d_v; 291 292 if (lda > nrows) { 293 PetscCall(PetscCUPMMemcpy2DAsync(dest, lda, src, lda, nrows, ncols, cupmMemcpyDeviceToHost, stream)); 294 } else { 295 PetscCall(PetscCUPMMemcpyAsync(dest, src, lda * ncols, cupmMemcpyDeviceToHost, stream)); 296 } 297 } 298 PetscCall(PetscLogEventEnd(MAT_DenseCopyFromGPU, m, 0, 0, 0)); 299 // order is important, MatSeqDenseSetPreallocation() might set offloadmask 300 m->offloadmask = PETSC_OFFLOAD_BOTH; 301 } 302 PetscFunctionReturn(PETSC_SUCCESS); 303 } 304 305 template <device::cupm::DeviceType T> 306 inline PetscErrorCode MatDense_Seq_CUPM<T>::CheckCUPMSolverInfo_(const cupmBlasInt_t *fact_info, cupmStream_t stream) noexcept 307 { 308 PetscFunctionBegin; 309 if (PetscDefined(USE_DEBUG)) { 310 cupmBlasInt_t info = 0; 311 312 PetscCall(PetscCUPMMemcpyAsync(&info, fact_info, 1, cupmMemcpyDeviceToHost, stream)); 313 if (stream) PetscCallCUPM(cupmStreamSynchronize(stream)); 314 static_assert(std::is_same<decltype(info), int>::value, ""); 315 PetscCheck(info <= 0, PETSC_COMM_SELF, PETSC_ERR_MAT_CH_ZRPVT, "Bad factorization: zero pivot in row %d", info - 1); 316 PetscCheck(info >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Wrong argument to cupmSolver %d", -info); 317 } 318 PetscFunctionReturn(PETSC_SUCCESS); 319 } 320 321 // ========================================================================================== 322 // MatDense_Seq_CUPM - Private API - Solver Dispatch 323 // ========================================================================================== 324 325 // specific solvers called through the dispatch_() family of functions 326 template <device::cupm::DeviceType T> 327 template <typename Derived> 328 struct MatDense_Seq_CUPM<T>::SolveCommon { 329 using derived_type = Derived; 330 331 template <typename F> 332 static PetscErrorCode ResizeFactLwork(Mat_SeqDenseCUPM *mcu, cupmStream_t stream, F &&cupmSolverComputeFactLwork) noexcept 333 { 334 cupmBlasInt_t lwork; 335 336 PetscFunctionBegin; 337 PetscCallCUPMSOLVER(cupmSolverComputeFactLwork(&lwork)); 338 if (lwork > mcu->d_fact_lwork) { 339 mcu->d_fact_lwork = lwork; 340 PetscCallCUPM(cupmFreeAsync(mcu->d_fact_work, stream)); 341 PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_work, lwork, stream)); 342 } 343 PetscFunctionReturn(PETSC_SUCCESS); 344 } 345 346 static PetscErrorCode FactorPrepare(Mat A, cupmStream_t stream) noexcept 347 { 348 const auto mcu = MatCUPMCast(A); 349 350 PetscFunctionBegin; 351 PetscCall(PetscInfo(A, "%s factor %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", derived_type::NAME(), A->rmap->n, A->cmap->n)); 352 A->factortype = derived_type::MATFACTORTYPE(); 353 A->ops->solve = MatSolve_Factored_Dispatch_<derived_type, false>; 354 A->ops->solvetranspose = MatSolve_Factored_Dispatch_<derived_type, true>; 355 A->ops->matsolve = MatMatSolve_Factored_Dispatch_<derived_type, false>; 356 A->ops->matsolvetranspose = MatMatSolve_Factored_Dispatch_<derived_type, true>; 357 358 PetscCall(PetscStrFreeAllocpy(MATSOLVERCUPM(), &A->solvertype)); 359 if (!mcu->d_fact_info) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_info, 1, stream)); 360 PetscFunctionReturn(PETSC_SUCCESS); 361 } 362 }; 363 364 template <device::cupm::DeviceType T> 365 struct MatDense_Seq_CUPM<T>::SolveLU : SolveCommon<SolveLU> { 366 using base_type = SolveCommon<SolveLU>; 367 368 static constexpr const char *NAME() noexcept { return "LU"; } 369 static constexpr MatFactorType MATFACTORTYPE() noexcept { return MAT_FACTOR_LU; } 370 371 static PetscErrorCode Factor(Mat A, IS, IS, const MatFactorInfo *) noexcept 372 { 373 const auto m = static_cast<cupmBlasInt_t>(A->rmap->n); 374 const auto n = static_cast<cupmBlasInt_t>(A->cmap->n); 375 cupmStream_t stream; 376 cupmSolverHandle_t handle; 377 PetscDeviceContext dctx; 378 379 PetscFunctionBegin; 380 if (!m || !n) PetscFunctionReturn(PETSC_SUCCESS); 381 PetscCall(GetHandles_(&dctx, &handle, &stream)); 382 PetscCall(base_type::FactorPrepare(A, stream)); 383 { 384 const auto mcu = MatCUPMCast(A); 385 const auto da = DeviceArrayReadWrite(dctx, A); 386 const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda); 387 388 // clang-format off 389 PetscCall( 390 base_type::ResizeFactLwork( 391 mcu, stream, 392 [&](cupmBlasInt_t *fact_lwork) 393 { 394 return cupmSolverXgetrf_bufferSize(handle, m, n, da.cupmdata(), lda, fact_lwork); 395 } 396 ) 397 ); 398 // clang-format on 399 if (!mcu->d_fact_ipiv) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_ipiv, n, stream)); 400 401 PetscCall(PetscLogGpuTimeBegin()); 402 PetscCallCUPMSOLVER(cupmSolverXgetrf(handle, m, n, da.cupmdata(), lda, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_ipiv, mcu->d_fact_info)); 403 PetscCall(PetscLogGpuTimeEnd()); 404 PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream)); 405 } 406 PetscCall(PetscLogGpuFlops(2.0 * n * n * m / 3.0)); 407 PetscFunctionReturn(PETSC_SUCCESS); 408 } 409 410 template <bool transpose> 411 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 412 { 413 const auto mcu = MatCUPMCast(A); 414 const auto fact_info = mcu->d_fact_info; 415 const auto fact_ipiv = mcu->d_fact_ipiv; 416 cupmSolverHandle_t handle; 417 418 PetscFunctionBegin; 419 PetscCall(GetHandlesFrom_(dctx, &handle)); 420 PetscCall(PetscInfo(A, "%s solve %d x %d on backend\n", NAME(), m, k)); 421 PetscCall(PetscLogGpuTimeBegin()); 422 { 423 constexpr auto op = transpose ? CUPMSOLVER_OP_T : CUPMSOLVER_OP_N; 424 const auto da = DeviceArrayRead(dctx, A); 425 const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda); 426 427 // clang-format off 428 PetscCall( 429 base_type::ResizeFactLwork( 430 mcu, stream, 431 [&](cupmBlasInt_t *lwork) 432 { 433 return cupmSolverXgetrs_bufferSize( 434 handle, op, m, nrhs, da.cupmdata(), lda, fact_ipiv, x, ldx, lwork 435 ); 436 } 437 ) 438 ); 439 // clang-format on 440 PetscCallCUPMSOLVER(cupmSolverXgetrs(handle, op, m, nrhs, da.cupmdata(), lda, fact_ipiv, x, ldx, mcu->d_fact_work, mcu->d_fact_lwork, fact_info)); 441 PetscCall(CheckCUPMSolverInfo_(fact_info, stream)); 442 } 443 PetscCall(PetscLogGpuTimeEnd()); 444 PetscCall(PetscLogGpuFlops(nrhs * (2.0 * m * m - m))); 445 PetscFunctionReturn(PETSC_SUCCESS); 446 } 447 }; 448 449 template <device::cupm::DeviceType T> 450 struct MatDense_Seq_CUPM<T>::SolveCholesky : SolveCommon<SolveCholesky> { 451 using base_type = SolveCommon<SolveCholesky>; 452 453 static constexpr const char *NAME() noexcept { return "Cholesky"; } 454 static constexpr MatFactorType MATFACTORTYPE() noexcept { return MAT_FACTOR_CHOLESKY; } 455 456 static PetscErrorCode Factor(Mat A, IS, const MatFactorInfo *) noexcept 457 { 458 const auto n = static_cast<cupmBlasInt_t>(A->rmap->n); 459 PetscDeviceContext dctx; 460 cupmSolverHandle_t handle; 461 cupmStream_t stream; 462 463 PetscFunctionBegin; 464 if (!n || !A->cmap->n) PetscFunctionReturn(PETSC_SUCCESS); 465 PetscCheck(A->spd == PETSC_BOOL3_TRUE, PETSC_COMM_SELF, PETSC_ERR_SUP, "%ssytrs unavailable. Use MAT_FACTOR_LU", cupmSolverName()); 466 PetscCall(GetHandles_(&dctx, &handle, &stream)); 467 PetscCall(base_type::FactorPrepare(A, stream)); 468 { 469 const auto mcu = MatCUPMCast(A); 470 const auto da = DeviceArrayReadWrite(dctx, A); 471 const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda); 472 473 // clang-format off 474 PetscCall( 475 base_type::ResizeFactLwork( 476 mcu, stream, 477 [&](cupmBlasInt_t *fact_lwork) 478 { 479 return cupmSolverXpotrf_bufferSize( 480 handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, fact_lwork 481 ); 482 } 483 ) 484 ); 485 // clang-format on 486 PetscCall(PetscLogGpuTimeBegin()); 487 PetscCallCUPMSOLVER(cupmSolverXpotrf(handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_info)); 488 PetscCall(PetscLogGpuTimeEnd()); 489 PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream)); 490 } 491 PetscCall(PetscLogGpuFlops(1.0 * n * n * n / 3.0)); 492 493 #if 0 494 // At the time of writing this interface (cuda 10.0), cusolverDn does not implement *sytrs 495 // and *hetr* routines. The code below should work, and it can be activated when *sytrs 496 // routines will be available 497 if (!mcu->d_fact_ipiv) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_ipiv, n, stream)); 498 if (!mcu->d_fact_lwork) { 499 PetscCallCUPMSOLVER(cupmSolverDnXsytrf_bufferSize(handle, n, da.cupmdata(), lda, &mcu->d_fact_lwork)); 500 PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_work, mcu->d_fact_lwork, stream)); 501 } 502 if (mcu->d_fact_info) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_info, 1, stream)); 503 PetscCall(PetscLogGpuTimeBegin()); 504 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)); 505 PetscCall(PetscLogGpuTimeEnd()); 506 #endif 507 PetscFunctionReturn(PETSC_SUCCESS); 508 } 509 510 template <bool transpose> 511 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 512 { 513 const auto mcu = MatCUPMCast(A); 514 const auto fact_info = mcu->d_fact_info; 515 cupmSolverHandle_t handle; 516 517 PetscFunctionBegin; 518 PetscAssert(!mcu->d_fact_ipiv, PETSC_COMM_SELF, PETSC_ERR_LIB, "%ssytrs not implemented", cupmSolverName()); 519 PetscCall(GetHandlesFrom_(dctx, &handle)); 520 PetscCall(PetscInfo(A, "%s solve %d x %d on backend\n", NAME(), m, k)); 521 PetscCall(PetscLogGpuTimeBegin()); 522 { 523 const auto da = DeviceArrayRead(dctx, A); 524 const auto lda = static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda); 525 526 // clang-format off 527 PetscCall( 528 base_type::ResizeFactLwork( 529 mcu, stream, 530 [&](cupmBlasInt_t *lwork) 531 { 532 return cupmSolverXpotrs_bufferSize( 533 handle, CUPMSOLVER_FILL_MODE_LOWER, m, nrhs, da.cupmdata(), lda, x, ldx, lwork 534 ); 535 } 536 ) 537 ); 538 // clang-format on 539 PetscCallCUPMSOLVER(cupmSolverXpotrs(handle, CUPMSOLVER_FILL_MODE_LOWER, m, nrhs, da.cupmdata(), lda, x, ldx, mcu->d_fact_work, mcu->d_fact_lwork, fact_info)); 540 PetscCall(CheckCUPMSolverInfo_(fact_info, stream)); 541 } 542 PetscCall(PetscLogGpuTimeEnd()); 543 PetscCall(PetscLogGpuFlops(nrhs * (2.0 * m * m - m))); 544 PetscFunctionReturn(PETSC_SUCCESS); 545 } 546 }; 547 548 template <device::cupm::DeviceType T> 549 struct MatDense_Seq_CUPM<T>::SolveQR : SolveCommon<SolveQR> { 550 using base_type = SolveCommon<SolveQR>; 551 552 static constexpr const char *NAME() noexcept { return "QR"; } 553 static constexpr MatFactorType MATFACTORTYPE() noexcept { return MAT_FACTOR_QR; } 554 555 static PetscErrorCode Factor(Mat A, IS, const MatFactorInfo *) noexcept 556 { 557 const auto m = static_cast<cupmBlasInt_t>(A->rmap->n); 558 const auto n = static_cast<cupmBlasInt_t>(A->cmap->n); 559 const auto min = std::min(m, n); 560 const auto mimpl = MatIMPLCast(A); 561 cupmStream_t stream; 562 cupmSolverHandle_t handle; 563 PetscDeviceContext dctx; 564 565 PetscFunctionBegin; 566 if (!m || !n) PetscFunctionReturn(PETSC_SUCCESS); 567 PetscCall(GetHandles_(&dctx, &handle, &stream)); 568 PetscCall(base_type::FactorPrepare(A, stream)); 569 mimpl->rank = min; 570 { 571 const auto mcu = MatCUPMCast(A); 572 const auto da = DeviceArrayReadWrite(dctx, A); 573 const auto lda = static_cast<cupmBlasInt_t>(mimpl->lda); 574 575 if (!mcu->workvec) PetscCall(vec::cupm::VecCreateSeqCUPMAsync<T>(PetscObjectComm(PetscObjectCast(A)), m, &mcu->workvec)); 576 if (!mcu->d_fact_tau) PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_tau, min, stream)); 577 // clang-format off 578 PetscCall( 579 base_type::ResizeFactLwork( 580 mcu, stream, 581 [&](cupmBlasInt_t *fact_lwork) 582 { 583 return cupmSolverXgeqrf_bufferSize(handle, m, n, da.cupmdata(), lda, fact_lwork); 584 } 585 ) 586 ); 587 // clang-format on 588 PetscCall(PetscLogGpuTimeBegin()); 589 PetscCallCUPMSOLVER(cupmSolverXgeqrf(handle, m, n, da.cupmdata(), lda, mcu->d_fact_tau, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_info)); 590 PetscCall(PetscLogGpuTimeEnd()); 591 PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream)); 592 } 593 PetscCall(PetscLogGpuFlops(2.0 * min * min * (std::max(m, n) - min / 3.0))); 594 PetscFunctionReturn(PETSC_SUCCESS); 595 } 596 597 template <bool transpose> 598 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 599 { 600 const auto mimpl = MatIMPLCast(A); 601 const auto rank = static_cast<cupmBlasInt_t>(mimpl->rank); 602 const auto mcu = MatCUPMCast(A); 603 const auto fact_info = mcu->d_fact_info; 604 const auto fact_tau = mcu->d_fact_tau; 605 const auto fact_work = mcu->d_fact_work; 606 const auto fact_lwork = mcu->d_fact_lwork; 607 cupmSolverHandle_t solver_handle; 608 cupmBlasHandle_t blas_handle; 609 610 PetscFunctionBegin; 611 PetscCall(GetHandlesFrom_(dctx, &blas_handle, &solver_handle)); 612 PetscCall(PetscInfo(A, "%s solve %d x %d on backend\n", NAME(), m, k)); 613 PetscCall(PetscLogGpuTimeBegin()); 614 { 615 const auto da = DeviceArrayRead(dctx, A); 616 const auto one = cupmScalarCast(1.0); 617 const auto lda = static_cast<cupmBlasInt_t>(mimpl->lda); 618 619 if (transpose) { 620 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)); 621 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)); 622 PetscCall(CheckCUPMSolverInfo_(fact_info, stream)); 623 } else { 624 constexpr auto op = PetscDefined(USE_COMPLEX) ? CUPMSOLVER_OP_C : CUPMSOLVER_OP_T; 625 626 PetscCallCUPMSOLVER(cupmSolverXormqr(solver_handle, CUPMSOLVER_SIDE_LEFT, op, m, nrhs, rank, da.cupmdata(), lda, fact_tau, x, ldx, fact_work, fact_lwork, fact_info)); 627 PetscCall(CheckCUPMSolverInfo_(fact_info, stream)); 628 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)); 629 } 630 } 631 PetscCall(PetscLogGpuTimeEnd()); 632 PetscCall(PetscLogFlops(nrhs * (4.0 * m * rank - (rank * rank)))); 633 PetscFunctionReturn(PETSC_SUCCESS); 634 } 635 }; 636 637 template <device::cupm::DeviceType T> 638 template <typename Solver, bool transpose> 639 inline PetscErrorCode MatDense_Seq_CUPM<T>::MatSolve_Factored_Dispatch_(Mat A, Vec x, Vec y) noexcept 640 { 641 using namespace vec::cupm; 642 const auto pobj_A = PetscObjectCast(A); 643 const auto m = static_cast<cupmBlasInt_t>(A->rmap->n); 644 const auto k = static_cast<cupmBlasInt_t>(A->cmap->n); 645 auto &workvec = MatCUPMCast(A)->workvec; 646 PetscScalar *y_array = nullptr; 647 PetscDeviceContext dctx; 648 PetscBool xiscupm, yiscupm, aiscupm; 649 bool use_y_array_directly; 650 cupmStream_t stream; 651 652 PetscFunctionBegin; 653 PetscCheck(A->factortype != MAT_FACTOR_NONE, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix must be factored to solve"); 654 PetscCall(PetscObjectTypeCompare(PetscObjectCast(x), VecSeq_CUPM::VECSEQCUPM(), &xiscupm)); 655 PetscCall(PetscObjectTypeCompare(PetscObjectCast(y), VecSeq_CUPM::VECSEQCUPM(), &yiscupm)); 656 PetscCall(PetscObjectTypeCompare(pobj_A, MATSEQDENSECUPM(), &aiscupm)); 657 PetscAssert(aiscupm, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Matrix A is somehow not CUPM?????????????????????????????"); 658 PetscCall(GetHandles_(&dctx, &stream)); 659 use_y_array_directly = yiscupm && (k >= m); 660 { 661 const PetscScalar *x_array; 662 const auto xisdevice = xiscupm && PetscOffloadDevice(x->offloadmask); 663 const auto copy_mode = xisdevice ? cupmMemcpyDeviceToDevice : cupmMemcpyHostToDevice; 664 665 if (!use_y_array_directly && !workvec) PetscCall(VecCreateSeqCUPMAsync<T>(PetscObjectComm(pobj_A), m, &workvec)); 666 // The logic here is to try to minimize the amount of memory copying: 667 // 668 // If we call VecCUPMGetArrayRead(X, &x) every time xiscupm and the data is not offloaded 669 // to the GPU yet, then the data is copied to the GPU. But we are only trying to get the 670 // data in order to copy it into the y array. So the array x will be wherever the data 671 // already is so that only one memcpy is performed 672 if (xisdevice) { 673 PetscCall(VecCUPMGetArrayReadAsync<T>(x, &x_array, dctx)); 674 } else { 675 PetscCall(VecGetArrayRead(x, &x_array)); 676 } 677 PetscCall(VecCUPMGetArrayWriteAsync<T>(use_y_array_directly ? y : workvec, &y_array, dctx)); 678 PetscCall(PetscCUPMMemcpyAsync(y_array, x_array, m, copy_mode, stream)); 679 if (xisdevice) { 680 PetscCall(VecCUPMRestoreArrayReadAsync<T>(x, &x_array, dctx)); 681 } else { 682 PetscCall(VecRestoreArrayRead(x, &x_array)); 683 } 684 } 685 686 if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A)); 687 PetscCall(Solver{}.template Solve<transpose>(A, cupmScalarPtrCast(y_array), m, m, 1, k, dctx, stream)); 688 if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A)); 689 690 if (use_y_array_directly) { 691 PetscCall(VecCUPMRestoreArrayWriteAsync<T>(y, &y_array, dctx)); 692 } else { 693 const auto copy_mode = yiscupm ? cupmMemcpyDeviceToDevice : cupmMemcpyDeviceToHost; 694 PetscScalar *yv; 695 696 // The logic here is that the data is not yet in either y's GPU array or its CPU array. 697 // There is nothing in the interface to say where the user would like it to end up. So we 698 // choose the GPU, because it is the faster option 699 if (yiscupm) { 700 PetscCall(VecCUPMGetArrayWriteAsync<T>(y, &yv, dctx)); 701 } else { 702 PetscCall(VecGetArray(y, &yv)); 703 } 704 PetscCall(PetscCUPMMemcpyAsync(yv, y_array, k, copy_mode, stream)); 705 if (yiscupm) { 706 PetscCall(VecCUPMRestoreArrayWriteAsync<T>(y, &yv, dctx)); 707 } else { 708 PetscCall(VecRestoreArray(y, &yv)); 709 } 710 PetscCall(VecCUPMRestoreArrayWriteAsync<T>(workvec, &y_array)); 711 } 712 PetscFunctionReturn(PETSC_SUCCESS); 713 } 714 715 template <device::cupm::DeviceType T> 716 template <typename Solver, bool transpose> 717 inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMatSolve_Factored_Dispatch_(Mat A, Mat B, Mat X) noexcept 718 { 719 const auto m = static_cast<cupmBlasInt_t>(A->rmap->n); 720 const auto k = static_cast<cupmBlasInt_t>(A->cmap->n); 721 cupmBlasInt_t nrhs, ldb, ldx, ldy; 722 PetscScalar *y; 723 PetscBool biscupm, xiscupm, aiscupm; 724 PetscDeviceContext dctx; 725 cupmStream_t stream; 726 727 PetscFunctionBegin; 728 PetscCheck(A->factortype != MAT_FACTOR_NONE, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix must be factored to solve"); 729 PetscCall(PetscObjectTypeCompare(PetscObjectCast(B), MATSEQDENSECUPM(), &biscupm)); 730 PetscCall(PetscObjectTypeCompare(PetscObjectCast(X), MATSEQDENSECUPM(), &xiscupm)); 731 PetscCall(PetscObjectTypeCompare(PetscObjectCast(A), MATSEQDENSECUPM(), &aiscupm)); 732 PetscCall(GetHandles_(&dctx, &stream)); 733 { 734 PetscInt n; 735 736 PetscCall(MatGetSize(B, nullptr, &n)); 737 PetscCall(PetscCUPMBlasIntCast(n, &nrhs)); 738 PetscCall(MatDenseGetLDA(B, &n)); 739 PetscCall(PetscCUPMBlasIntCast(n, &ldb)); 740 PetscCall(MatDenseGetLDA(X, &n)); 741 PetscCall(PetscCUPMBlasIntCast(n, &ldx)); 742 } 743 { 744 // The logic here is to try to minimize the amount of memory copying: 745 // 746 // If we call MatDenseCUPMGetArrayRead(B, &b) every time biscupm and the data is not 747 // offloaded to the GPU yet, then the data is copied to the GPU. But we are only trying to 748 // get the data in order to copy it into the y array. So the array b will be wherever the 749 // data already is so that only one memcpy is performed 750 const auto bisdevice = biscupm && PetscOffloadDevice(B->offloadmask); 751 const auto copy_mode = bisdevice ? cupmMemcpyDeviceToDevice : cupmMemcpyHostToDevice; 752 const PetscScalar *b; 753 754 if (bisdevice) { 755 b = DeviceArrayRead(dctx, B); 756 } else if (biscupm) { 757 b = HostArrayRead(dctx, B); 758 } else { 759 PetscCall(MatDenseGetArrayRead(B, &b)); 760 } 761 762 if (ldx < m || !xiscupm) { 763 // X's array cannot serve as the array (too small or not on device), B's array cannot 764 // serve as the array (const), so allocate a new array 765 ldy = m; 766 PetscCall(PetscCUPMMallocAsync(&y, nrhs * m)); 767 } else { 768 // X's array should serve as the array 769 ldy = ldx; 770 y = DeviceArrayWrite(dctx, X); 771 } 772 PetscCall(PetscCUPMMemcpy2DAsync(y, ldy, b, ldb, m, nrhs, copy_mode, stream)); 773 if (!bisdevice && !biscupm) PetscCall(MatDenseRestoreArrayRead(B, &b)); 774 } 775 776 // convert to CUPM twice?????????????????????????????????? 777 // but A should already be CUPM?????????????????????????????????????? 778 if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A)); 779 PetscCall(Solver{}.template Solve<transpose>(A, cupmScalarPtrCast(y), ldy, m, nrhs, k, dctx, stream)); 780 if (!aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A)); 781 782 if (ldx < m || !xiscupm) { 783 const auto copy_mode = xiscupm ? cupmMemcpyDeviceToDevice : cupmMemcpyDeviceToHost; 784 PetscScalar *x; 785 786 // The logic here is that the data is not yet in either X's GPU array or its CPU 787 // array. There is nothing in the interface to say where the user would like it to end up. 788 // So we choose the GPU, because it is the faster option 789 if (xiscupm) { 790 x = DeviceArrayWrite(dctx, X); 791 } else { 792 PetscCall(MatDenseGetArray(X, &x)); 793 } 794 PetscCall(PetscCUPMMemcpy2DAsync(x, ldx, y, ldy, k, nrhs, copy_mode, stream)); 795 if (!xiscupm) PetscCall(MatDenseRestoreArray(X, &x)); 796 PetscCallCUPM(cupmFreeAsync(y, stream)); 797 } 798 PetscFunctionReturn(PETSC_SUCCESS); 799 } 800 801 template <device::cupm::DeviceType T> 802 template <bool transpose> 803 inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMultAdd_Dispatch_(Mat A, Vec xx, Vec yy, Vec zz) noexcept 804 { 805 const auto m = static_cast<cupmBlasInt_t>(A->rmap->n); 806 const auto n = static_cast<cupmBlasInt_t>(A->cmap->n); 807 cupmBlasHandle_t handle; 808 PetscDeviceContext dctx; 809 810 PetscFunctionBegin; 811 if (yy && yy != zz) PetscCall(VecSeq_CUPM::Copy(yy, zz)); // mult add 812 if (!m || !n) { 813 // mult only 814 if (!yy) PetscCall(VecSeq_CUPM::Set(zz, 0.0)); 815 PetscFunctionReturn(PETSC_SUCCESS); 816 } 817 PetscCall(PetscInfo(A, "Matrix-vector product %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " on backend\n", m, n)); 818 PetscCall(GetHandles_(&dctx, &handle)); 819 { 820 constexpr auto op = transpose ? CUPMBLAS_OP_T : CUPMBLAS_OP_N; 821 const auto one = cupmScalarCast(1.0); 822 const auto zero = cupmScalarCast(0.0); 823 const auto da = DeviceArrayRead(dctx, A); 824 const auto dxx = VecSeq_CUPM::DeviceArrayRead(dctx, xx); 825 const auto dzz = VecSeq_CUPM::DeviceArrayReadWrite(dctx, zz); 826 827 PetscCall(PetscLogGpuTimeBegin()); 828 PetscCallCUPMBLAS(cupmBlasXgemv(handle, op, m, n, &one, da.cupmdata(), static_cast<cupmBlasInt_t>(MatIMPLCast(A)->lda), dxx.cupmdata(), 1, (yy ? &one : &zero), dzz.cupmdata(), 1)); 829 PetscCall(PetscLogGpuTimeEnd()); 830 } 831 PetscCall(PetscLogGpuFlops(2.0 * m * n - (yy ? 0 : m))); 832 PetscFunctionReturn(PETSC_SUCCESS); 833 } 834 835 // ========================================================================================== 836 // MatDense_Seq_CUPM - Private API - Conversion Dispatch 837 // ========================================================================================== 838 839 template <device::cupm::DeviceType T> 840 template <bool to_host> 841 inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_Dispatch_(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept 842 { 843 PetscFunctionBegin; 844 if (reuse == MAT_REUSE_MATRIX || reuse == MAT_INITIAL_MATRIX) { 845 // TODO these cases should be optimized 846 PetscCall(MatConvert_Basic(M, type, reuse, newmat)); 847 } else { 848 const auto B = *newmat; 849 const auto pobj = PetscObjectCast(B); 850 851 if (to_host) { 852 PetscCall(BindToCPU(B, PETSC_TRUE)); 853 PetscCall(Reset(B)); 854 } else { 855 PetscCall(PetscDeviceInitialize(PETSC_DEVICE_CUPM())); 856 } 857 858 PetscCall(PetscStrFreeAllocpy(to_host ? VECSTANDARD : VecSeq_CUPM::VECCUPM(), &B->defaultvectype)); 859 PetscCall(PetscObjectChangeTypeName(pobj, to_host ? MATSEQDENSE : MATSEQDENSECUPM())); 860 // cvec might be the wrong VecType, destroy and rebuild it if necessary 861 // REVIEW ME: this is possibly very inefficient 862 PetscCall(VecDestroy(&MatIMPLCast(B)->cvec)); 863 864 MatComposeOp_CUPM(to_host, pobj, MatConvert_seqdensecupm_seqdense_C(), nullptr, Convert_SeqDenseCUPM_SeqDense); 865 MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArray_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ_WRITE>); 866 MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArrayRead_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ>); 867 MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMGetArrayWrite_C(), nullptr, GetArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_WRITE>); 868 MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArray_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ_WRITE>); 869 MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArrayRead_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_READ>); 870 MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMRestoreArrayWrite_C(), nullptr, RestoreArrayC_<PETSC_MEMTYPE_DEVICE, PETSC_MEMORY_ACCESS_WRITE>); 871 MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMPlaceArray_C(), nullptr, PlaceArray); 872 MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMResetArray_C(), nullptr, ResetArray); 873 MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMReplaceArray_C(), nullptr, ReplaceArray); 874 MatComposeOp_CUPM(to_host, pobj, MatProductSetFromOptions_seqaij_seqdensecupm_C(), nullptr, MatProductSetFromOptions_SeqAIJ_SeqDense); 875 MatComposeOp_CUPM(to_host, pobj, MatDenseCUPMSetPreallocation_C(), nullptr, SetPreallocation); 876 877 if (to_host) { 878 B->offloadmask = PETSC_OFFLOAD_CPU; 879 } else { 880 Mat_SeqDenseCUPM *mcu; 881 882 PetscCall(PetscNew(&mcu)); 883 B->spptr = mcu; 884 B->offloadmask = PETSC_OFFLOAD_UNALLOCATED; // REVIEW ME: why not offload host?? 885 PetscCall(BindToCPU(B, PETSC_FALSE)); 886 } 887 888 MatSetOp_CUPM(to_host, B, bindtocpu, nullptr, BindToCPU); 889 MatSetOp_CUPM(to_host, B, destroy, MatDestroy_SeqDense, Destroy); 890 } 891 PetscFunctionReturn(PETSC_SUCCESS); 892 } 893 894 // ========================================================================================== 895 // MatDense_Seq_CUPM - Public API 896 // ========================================================================================== 897 898 template <device::cupm::DeviceType T> 899 inline constexpr MatType MatDense_Seq_CUPM<T>::MATIMPLCUPM_() noexcept 900 { 901 return MATSEQDENSECUPM(); 902 } 903 904 template <device::cupm::DeviceType T> 905 inline constexpr typename MatDense_Seq_CUPM<T>::Mat_SeqDenseCUPM *MatDense_Seq_CUPM<T>::MatCUPMCast(Mat m) noexcept 906 { 907 return static_cast<Mat_SeqDenseCUPM *>(m->spptr); 908 } 909 910 template <device::cupm::DeviceType T> 911 inline constexpr Mat_SeqDense *MatDense_Seq_CUPM<T>::MatIMPLCast_(Mat m) noexcept 912 { 913 return static_cast<Mat_SeqDense *>(m->data); 914 } 915 916 template <device::cupm::DeviceType T> 917 inline constexpr const char *MatDense_Seq_CUPM<T>::MatConvert_seqdensecupm_seqdense_C() noexcept 918 { 919 return T == device::cupm::DeviceType::CUDA ? "MatConvert_seqdensecuda_seqdense_C" : "MatConvert_seqdensehip_seqdense_C"; 920 } 921 922 template <device::cupm::DeviceType T> 923 inline constexpr const char *MatDense_Seq_CUPM<T>::MatProductSetFromOptions_seqaij_seqdensecupm_C() noexcept 924 { 925 return T == device::cupm::DeviceType::CUDA ? "MatProductSetFromOptions_seqaij_seqdensecuda_C" : "MatProductSetFromOptions_seqaij_seqdensehip_C"; 926 } 927 928 // ========================================================================================== 929 930 // MatCreate_SeqDenseCUPM() 931 template <device::cupm::DeviceType T> 932 inline PetscErrorCode MatDense_Seq_CUPM<T>::Create(Mat A) noexcept 933 { 934 PetscFunctionBegin; 935 PetscCall(PetscDeviceInitialize(PETSC_DEVICE_CUPM())); 936 PetscCall(MatCreate_SeqDense(A)); 937 PetscCall(Convert_SeqDense_SeqDenseCUPM(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A)); 938 PetscFunctionReturn(PETSC_SUCCESS); 939 } 940 941 template <device::cupm::DeviceType T> 942 inline PetscErrorCode MatDense_Seq_CUPM<T>::Destroy(Mat A) noexcept 943 { 944 PetscFunctionBegin; 945 // prevent copying back data if we own the data pointer 946 if (!MatIMPLCast(A)->user_alloc) A->offloadmask = PETSC_OFFLOAD_CPU; 947 PetscCall(Convert_SeqDenseCUPM_SeqDense(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A)); 948 PetscCall(MatDestroy_SeqDense(A)); 949 PetscFunctionReturn(PETSC_SUCCESS); 950 } 951 952 // obj->ops->setup() 953 template <device::cupm::DeviceType T> 954 inline PetscErrorCode MatDense_Seq_CUPM<T>::SetUp(Mat A) noexcept 955 { 956 PetscFunctionBegin; 957 PetscCall(PetscLayoutSetUp(A->rmap)); 958 PetscCall(PetscLayoutSetUp(A->cmap)); 959 if (!A->preallocated) { 960 PetscDeviceContext dctx; 961 962 PetscCall(GetHandles_(&dctx)); 963 PetscCall(SetPreallocation(A, dctx, nullptr)); 964 } 965 PetscFunctionReturn(PETSC_SUCCESS); 966 } 967 968 template <device::cupm::DeviceType T> 969 inline PetscErrorCode MatDense_Seq_CUPM<T>::Reset(Mat A) noexcept 970 { 971 PetscFunctionBegin; 972 if (const auto mcu = MatCUPMCast(A)) { 973 cupmStream_t stream; 974 975 PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME()); 976 PetscCall(GetHandles_(&stream)); 977 if (!mcu->d_user_alloc) PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream)); 978 PetscCallCUPM(cupmFreeAsync(mcu->d_fact_tau, stream)); 979 PetscCallCUPM(cupmFreeAsync(mcu->d_fact_ipiv, stream)); 980 PetscCallCUPM(cupmFreeAsync(mcu->d_fact_info, stream)); 981 PetscCallCUPM(cupmFreeAsync(mcu->d_fact_work, stream)); 982 PetscCall(VecDestroy(&mcu->workvec)); 983 PetscCall(PetscFree(A->spptr /* mcu */)); 984 } 985 PetscFunctionReturn(PETSC_SUCCESS); 986 } 987 988 // ========================================================================================== 989 990 template <device::cupm::DeviceType T> 991 inline PetscErrorCode MatDense_Seq_CUPM<T>::BindToCPU(Mat A, PetscBool to_host) noexcept 992 { 993 const auto mimpl = MatIMPLCast(A); 994 const auto pobj = PetscObjectCast(A); 995 996 PetscFunctionBegin; 997 PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first"); 998 PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 999 A->boundtocpu = to_host; 1000 PetscCall(PetscStrFreeAllocpy(to_host ? PETSCRANDER48 : PETSCDEVICERAND(), &A->defaultrandtype)); 1001 if (to_host) { 1002 PetscDeviceContext dctx; 1003 1004 // make sure we have an up-to-date copy on the CPU 1005 PetscCall(GetHandles_(&dctx)); 1006 PetscCall(DeviceToHost_(A, dctx)); 1007 } else { 1008 PetscBool iscupm; 1009 1010 if (auto &cvec = mimpl->cvec) { 1011 PetscCall(PetscObjectTypeCompare(PetscObjectCast(cvec), VecSeq_CUPM::VECSEQCUPM(), &iscupm)); 1012 if (!iscupm) PetscCall(VecDestroy(&cvec)); 1013 } 1014 if (auto &cmat = mimpl->cmat) { 1015 PetscCall(PetscObjectTypeCompare(PetscObjectCast(cmat), MATSEQDENSECUPM(), &iscupm)); 1016 if (!iscupm) PetscCall(MatDestroy(&cmat)); 1017 } 1018 } 1019 1020 // ============================================================ 1021 // Composed ops 1022 // ============================================================ 1023 MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArray_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ_WRITE>); 1024 MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayRead_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_READ>); 1025 MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayWrite_C", MatDenseGetArray_SeqDense, GetArrayC_<PETSC_MEMTYPE_HOST, PETSC_MEMORY_ACCESS_WRITE>); 1026 MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ_WRITE>); 1027 MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ_WRITE>); 1028 MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayReadAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ>); 1029 MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayReadAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_READ>); 1030 MatComposeOp_CUPM(to_host, pobj, "MatDenseGetArrayWriteAndMemType_C", nullptr, GetArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_WRITE>); 1031 MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreArrayWriteAndMemType_C", nullptr, RestoreArrayAndMemTypeC_<PETSC_MEMORY_ACCESS_WRITE>); 1032 MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVec_C", MatDenseGetColumnVec_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_READ_WRITE>); 1033 MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVec_C", MatDenseRestoreColumnVec_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_READ_WRITE>); 1034 MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVecRead_C", MatDenseGetColumnVecRead_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_READ>); 1035 MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVecRead_C", MatDenseRestoreColumnVecRead_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_READ>); 1036 MatComposeOp_CUPM(to_host, pobj, "MatDenseGetColumnVecWrite_C", MatDenseGetColumnVecWrite_SeqDense, GetColumnVec<PETSC_MEMORY_ACCESS_WRITE>); 1037 MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreColumnVecWrite_C", MatDenseRestoreColumnVecWrite_SeqDense, RestoreColumnVec<PETSC_MEMORY_ACCESS_WRITE>); 1038 MatComposeOp_CUPM(to_host, pobj, "MatDenseGetSubMatrix_C", MatDenseGetSubMatrix_SeqDense, GetSubMatrix); 1039 MatComposeOp_CUPM(to_host, pobj, "MatDenseRestoreSubMatrix_C", MatDenseRestoreSubMatrix_SeqDense, RestoreSubMatrix); 1040 MatComposeOp_CUPM(to_host, pobj, "MatQRFactor_C", MatQRFactor_SeqDense, SolveQR::Factor); 1041 // always the same 1042 PetscCall(PetscObjectComposeFunction(pobj, "MatDenseSetLDA_C", MatDenseSetLDA_SeqDense)); 1043 1044 // ============================================================ 1045 // Function pointer ops 1046 // ============================================================ 1047 MatSetOp_CUPM(to_host, A, duplicate, MatDuplicate_SeqDense, Duplicate); 1048 MatSetOp_CUPM(to_host, A, mult, MatMult_SeqDense, [](Mat A, Vec xx, Vec yy) { return MatMultAdd_Dispatch_</* transpose */ false>(A, xx, nullptr, yy); }); 1049 MatSetOp_CUPM(to_host, A, multtranspose, MatMultTranspose_SeqDense, [](Mat A, Vec xx, Vec yy) { return MatMultAdd_Dispatch_</* transpose */ true>(A, xx, nullptr, yy); }); 1050 MatSetOp_CUPM(to_host, A, multadd, MatMultAdd_SeqDense, MatMultAdd_Dispatch_</* transpose */ false>); 1051 MatSetOp_CUPM(to_host, A, multtransposeadd, MatMultTransposeAdd_SeqDense, MatMultAdd_Dispatch_</* transpose */ true>); 1052 MatSetOp_CUPM(to_host, A, matmultnumeric, MatMatMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ false, /* transpose_B */ false>); 1053 MatSetOp_CUPM(to_host, A, mattransposemultnumeric, MatMatTransposeMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ false, /* transpose_B */ true>); 1054 MatSetOp_CUPM(to_host, A, transposematmultnumeric, MatTransposeMatMultNumeric_SeqDense_SeqDense, MatMatMult_Numeric_Dispatch</* transpose_A */ true, /* transpose_B */ false>); 1055 MatSetOp_CUPM(to_host, A, axpy, MatAXPY_SeqDense, AXPY); 1056 MatSetOp_CUPM(to_host, A, choleskyfactor, MatCholeskyFactor_SeqDense, SolveCholesky::Factor); 1057 MatSetOp_CUPM(to_host, A, lufactor, MatLUFactor_SeqDense, SolveLU::Factor); 1058 MatSetOp_CUPM(to_host, A, getcolumnvector, MatGetColumnVector_SeqDense, GetColumnVector); 1059 MatSetOp_CUPM(to_host, A, scale, MatScale_SeqDense, Scale); 1060 MatSetOp_CUPM(to_host, A, shift, MatShift_SeqDense, Shift); 1061 MatSetOp_CUPM(to_host, A, copy, MatCopy_SeqDense, Copy); 1062 MatSetOp_CUPM(to_host, A, zeroentries, MatZeroEntries_SeqDense, ZeroEntries); 1063 MatSetOp_CUPM(to_host, A, setup, MatSetUp_SeqDense, SetUp); 1064 MatSetOp_CUPM(to_host, A, setrandom, MatSetRandom_SeqDense, SetRandom); 1065 // seemingly always the same 1066 A->ops->productsetfromoptions = MatProductSetFromOptions_SeqDense; 1067 1068 if (const auto cmat = mimpl->cmat) PetscCall(MatBindToCPU(cmat, to_host)); 1069 PetscFunctionReturn(PETSC_SUCCESS); 1070 } 1071 1072 template <device::cupm::DeviceType T> 1073 inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_SeqDenseCUPM_SeqDense(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept 1074 { 1075 PetscFunctionBegin; 1076 PetscCall(Convert_Dispatch_</* to host */ true>(M, type, reuse, newmat)); 1077 PetscFunctionReturn(PETSC_SUCCESS); 1078 } 1079 1080 template <device::cupm::DeviceType T> 1081 inline PetscErrorCode MatDense_Seq_CUPM<T>::Convert_SeqDense_SeqDenseCUPM(Mat M, MatType type, MatReuse reuse, Mat *newmat) noexcept 1082 { 1083 PetscFunctionBegin; 1084 PetscCall(Convert_Dispatch_</* to host */ false>(M, type, reuse, newmat)); 1085 PetscFunctionReturn(PETSC_SUCCESS); 1086 } 1087 1088 // ========================================================================================== 1089 1090 template <device::cupm::DeviceType T> 1091 template <PetscMemType mtype, PetscMemoryAccessMode access> 1092 inline PetscErrorCode MatDense_Seq_CUPM<T>::GetArray(Mat m, PetscScalar **array, PetscDeviceContext dctx) noexcept 1093 { 1094 constexpr auto hostmem = PetscMemTypeHost(mtype); 1095 constexpr auto read_access = PetscMemoryAccessRead(access); 1096 1097 PetscFunctionBegin; 1098 static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), ""); 1099 PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx)); 1100 if (hostmem) { 1101 if (read_access) { 1102 PetscCall(DeviceToHost_(m, dctx)); 1103 } else if (!MatIMPLCast(m)->v) { 1104 // MatCreateSeqDenseCUPM may not allocate CPU memory. Allocate if needed 1105 PetscCall(MatSeqDenseSetPreallocation(m, nullptr)); 1106 } 1107 *array = MatIMPLCast(m)->v; 1108 } else { 1109 if (read_access) { 1110 PetscCall(HostToDevice_(m, dctx)); 1111 } else if (!MatCUPMCast(m)->d_v) { 1112 // write-only 1113 PetscCall(SetPreallocation(m, dctx, nullptr)); 1114 } 1115 *array = MatCUPMCast(m)->d_v; 1116 } 1117 if (PetscMemoryAccessWrite(access)) { 1118 m->offloadmask = hostmem ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU; 1119 PetscCall(PetscObjectStateIncrease(PetscObjectCast(m))); 1120 } 1121 PetscFunctionReturn(PETSC_SUCCESS); 1122 } 1123 1124 template <device::cupm::DeviceType T> 1125 template <PetscMemType mtype, PetscMemoryAccessMode access> 1126 inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreArray(Mat m, PetscScalar **array, PetscDeviceContext) noexcept 1127 { 1128 PetscFunctionBegin; 1129 static_assert((mtype == PETSC_MEMTYPE_HOST) || (mtype == PETSC_MEMTYPE_DEVICE), ""); 1130 if (PetscMemoryAccessWrite(access)) { 1131 // WRITE or READ_WRITE 1132 m->offloadmask = PetscMemTypeHost(mtype) ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU; 1133 PetscCall(PetscObjectStateIncrease(PetscObjectCast(m))); 1134 } 1135 if (array) { 1136 PetscCall(CheckPointerMatchesMemType_(*array, mtype)); 1137 *array = nullptr; 1138 } 1139 PetscFunctionReturn(PETSC_SUCCESS); 1140 } 1141 1142 template <device::cupm::DeviceType T> 1143 template <PetscMemoryAccessMode access> 1144 inline PetscErrorCode MatDense_Seq_CUPM<T>::GetArrayAndMemType(Mat m, PetscScalar **array, PetscMemType *mtype, PetscDeviceContext dctx) noexcept 1145 { 1146 PetscFunctionBegin; 1147 PetscCall(GetArray<PETSC_MEMTYPE_DEVICE, access>(m, array, dctx)); 1148 if (mtype) *mtype = PETSC_MEMTYPE_CUPM(); 1149 PetscFunctionReturn(PETSC_SUCCESS); 1150 } 1151 1152 template <device::cupm::DeviceType T> 1153 template <PetscMemoryAccessMode access> 1154 inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreArrayAndMemType(Mat m, PetscScalar **array, PetscDeviceContext dctx) noexcept 1155 { 1156 PetscFunctionBegin; 1157 PetscCall(RestoreArray<PETSC_MEMTYPE_DEVICE, access>(m, array, dctx)); 1158 PetscFunctionReturn(PETSC_SUCCESS); 1159 } 1160 1161 // ========================================================================================== 1162 1163 template <device::cupm::DeviceType T> 1164 inline PetscErrorCode MatDense_Seq_CUPM<T>::PlaceArray(Mat A, const PetscScalar *array) noexcept 1165 { 1166 const auto mimpl = MatIMPLCast(A); 1167 const auto mcu = MatCUPMCast(A); 1168 1169 PetscFunctionBegin; 1170 PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first"); 1171 PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 1172 PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME()); 1173 if (mimpl->v) { 1174 PetscDeviceContext dctx; 1175 1176 PetscCall(GetHandles_(&dctx)); 1177 PetscCall(HostToDevice_(A, dctx)); 1178 } 1179 mcu->unplacedarray = util::exchange(mcu->d_v, const_cast<PetscScalar *>(array)); 1180 mcu->d_unplaced_user_alloc = util::exchange(mcu->d_user_alloc, PETSC_TRUE); 1181 PetscFunctionReturn(PETSC_SUCCESS); 1182 } 1183 1184 template <device::cupm::DeviceType T> 1185 inline PetscErrorCode MatDense_Seq_CUPM<T>::ReplaceArray(Mat A, const PetscScalar *array) noexcept 1186 { 1187 const auto mimpl = MatIMPLCast(A); 1188 const auto mcu = MatCUPMCast(A); 1189 1190 PetscFunctionBegin; 1191 PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first"); 1192 PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 1193 PetscCheck(!mcu->unplacedarray, PETSC_COMM_SELF, PETSC_ERR_ORDER, "MatDense%sResetArray() must be called first", cupmNAME()); 1194 if (!mcu->d_user_alloc) { 1195 cupmStream_t stream; 1196 1197 PetscCall(GetHandles_(&stream)); 1198 PetscCallCUPM(cupmFreeAsync(mcu->d_v, stream)); 1199 } 1200 mcu->d_v = const_cast<PetscScalar *>(array); 1201 mcu->d_user_alloc = PETSC_FALSE; 1202 PetscFunctionReturn(PETSC_SUCCESS); 1203 } 1204 1205 template <device::cupm::DeviceType T> 1206 inline PetscErrorCode MatDense_Seq_CUPM<T>::ResetArray(Mat A) noexcept 1207 { 1208 const auto mimpl = MatIMPLCast(A); 1209 const auto mcu = MatCUPMCast(A); 1210 1211 PetscFunctionBegin; 1212 PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first"); 1213 PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 1214 if (mimpl->v) { 1215 PetscDeviceContext dctx; 1216 1217 PetscCall(GetHandles_(&dctx)); 1218 PetscCall(HostToDevice_(A, dctx)); 1219 } 1220 mcu->d_v = util::exchange(mcu->unplacedarray, nullptr); 1221 mcu->d_user_alloc = mcu->d_unplaced_user_alloc; 1222 PetscFunctionReturn(PETSC_SUCCESS); 1223 } 1224 1225 // ========================================================================================== 1226 1227 template <device::cupm::DeviceType T> 1228 template <bool transpose_A, bool transpose_B> 1229 inline PetscErrorCode MatDense_Seq_CUPM<T>::MatMatMult_Numeric_Dispatch(Mat A, Mat B, Mat C) noexcept 1230 { 1231 cupmBlasInt_t m, n, k; 1232 PetscBool Aiscupm, Biscupm; 1233 PetscDeviceContext dctx; 1234 cupmBlasHandle_t handle; 1235 1236 PetscFunctionBegin; 1237 PetscCall(PetscCUPMBlasIntCast(C->rmap->n, &m)); 1238 PetscCall(PetscCUPMBlasIntCast(C->cmap->n, &n)); 1239 PetscCall(PetscCUPMBlasIntCast(transpose_A ? A->rmap->n : A->cmap->n, &k)); 1240 if (!m || !n || !k) PetscFunctionReturn(PETSC_SUCCESS); 1241 1242 // we may end up with SEQDENSE as one of the arguments 1243 // REVIEW ME: how? and why is it not B and C???????? 1244 PetscCall(PetscObjectTypeCompare(PetscObjectCast(A), MATSEQDENSECUPM(), &Aiscupm)); 1245 PetscCall(PetscObjectTypeCompare(PetscObjectCast(B), MATSEQDENSECUPM(), &Biscupm)); 1246 if (!Aiscupm) PetscCall(MatConvert(A, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &A)); 1247 if (!Biscupm) PetscCall(MatConvert(B, MATSEQDENSECUPM(), MAT_INPLACE_MATRIX, &B)); 1248 PetscCall(PetscInfo(C, "Matrix-Matrix product %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " x %" PetscBLASInt_FMT " on backend\n", m, k, n)); 1249 PetscCall(GetHandles_(&dctx, &handle)); 1250 1251 PetscCall(PetscLogGpuTimeBegin()); 1252 { 1253 const auto one = cupmScalarCast(1.0); 1254 const auto zero = cupmScalarCast(0.0); 1255 const auto da = DeviceArrayRead(dctx, A); 1256 const auto db = DeviceArrayRead(dctx, B); 1257 const auto dc = DeviceArrayWrite(dctx, C); 1258 PetscInt alda, blda, clda; 1259 1260 PetscCall(MatDenseGetLDA(A, &alda)); 1261 PetscCall(MatDenseGetLDA(B, &blda)); 1262 PetscCall(MatDenseGetLDA(C, &clda)); 1263 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)); 1264 } 1265 PetscCall(PetscLogGpuTimeEnd()); 1266 1267 PetscCall(PetscLogGpuFlops(1.0 * m * n * k + 1.0 * m * n * (k - 1))); 1268 if (!Aiscupm) PetscCall(MatConvert(A, MATSEQDENSE, MAT_INPLACE_MATRIX, &A)); 1269 if (!Biscupm) PetscCall(MatConvert(B, MATSEQDENSE, MAT_INPLACE_MATRIX, &B)); 1270 PetscFunctionReturn(PETSC_SUCCESS); 1271 } 1272 1273 template <device::cupm::DeviceType T> 1274 inline PetscErrorCode MatDense_Seq_CUPM<T>::Copy(Mat A, Mat B, MatStructure str) noexcept 1275 { 1276 const auto m = A->rmap->n; 1277 const auto n = A->cmap->n; 1278 1279 PetscFunctionBegin; 1280 PetscAssert(m == B->rmap->n && n == B->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "size(B) != size(A)"); 1281 // The two matrices must have the same copy implementation to be eligible for fast copy 1282 if (A->ops->copy == B->ops->copy) { 1283 PetscDeviceContext dctx; 1284 cupmStream_t stream; 1285 1286 PetscCall(GetHandles_(&dctx, &stream)); 1287 PetscCall(PetscLogGpuTimeBegin()); 1288 { 1289 const auto va = DeviceArrayRead(dctx, A); 1290 const auto vb = DeviceArrayWrite(dctx, B); 1291 // order is important, DeviceArrayRead/Write() might call SetPreallocation() which sets 1292 // lda! 1293 const auto lda_a = MatIMPLCast(A)->lda; 1294 const auto lda_b = MatIMPLCast(B)->lda; 1295 1296 if (lda_a > m || lda_b > m) { 1297 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()); 1298 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()); 1299 PetscCall(PetscCUPMMemcpy2DAsync(vb.data(), lda_b, va.data(), lda_a, m, n, cupmMemcpyDeviceToDevice, stream)); 1300 } else { 1301 PetscCall(PetscCUPMMemcpyAsync(vb.data(), va.data(), m * n, cupmMemcpyDeviceToDevice, stream)); 1302 } 1303 } 1304 PetscCall(PetscLogGpuTimeEnd()); 1305 } else { 1306 PetscCall(MatCopy_Basic(A, B, str)); 1307 } 1308 PetscFunctionReturn(PETSC_SUCCESS); 1309 } 1310 1311 template <device::cupm::DeviceType T> 1312 inline PetscErrorCode MatDense_Seq_CUPM<T>::ZeroEntries(Mat m) noexcept 1313 { 1314 PetscDeviceContext dctx; 1315 cupmStream_t stream; 1316 1317 PetscFunctionBegin; 1318 PetscCall(GetHandles_(&dctx, &stream)); 1319 PetscCall(PetscLogGpuTimeBegin()); 1320 { 1321 const auto va = DeviceArrayWrite(dctx, m); 1322 const auto lda = MatIMPLCast(m)->lda; 1323 const auto ma = m->rmap->n; 1324 const auto na = m->cmap->n; 1325 1326 if (lda > ma) { 1327 PetscCall(PetscCUPMMemset2DAsync(va.data(), lda, 0, ma, na, stream)); 1328 } else { 1329 PetscCall(PetscCUPMMemsetAsync(va.data(), 0, ma * na, stream)); 1330 } 1331 } 1332 PetscCall(PetscLogGpuTimeEnd()); 1333 PetscFunctionReturn(PETSC_SUCCESS); 1334 } 1335 1336 namespace detail 1337 { 1338 1339 // ========================================================================================== 1340 // SubMatIndexFunctor 1341 // 1342 // Iterator which permutes a linear index range into matrix indices for am nrows x ncols 1343 // submat with leading dimension lda. Essentially SubMatIndexFunctor(i) returns the index for 1344 // the i'th sequential entry in the matrix. 1345 // ========================================================================================== 1346 template <typename T> 1347 struct SubMatIndexFunctor { 1348 PETSC_HOSTDEVICE_INLINE_DECL T operator()(T x) const noexcept { return ((x / nrows) * lda) + (x % nrows); } 1349 1350 PetscInt nrows; 1351 PetscInt ncols; 1352 PetscInt lda; 1353 }; 1354 1355 template <typename Iterator> 1356 struct SubMatrixIterator : MatrixIteratorBase<Iterator, SubMatIndexFunctor<typename thrust::iterator_difference<Iterator>::type>> { 1357 using base_type = MatrixIteratorBase<Iterator, SubMatIndexFunctor<typename thrust::iterator_difference<Iterator>::type>>; 1358 1359 using iterator = typename base_type::iterator; 1360 1361 constexpr SubMatrixIterator(Iterator first, Iterator last, PetscInt nrows, PetscInt ncols, PetscInt lda) noexcept : 1362 base_type{ 1363 std::move(first), std::move(last), {nrows, ncols, lda} 1364 } 1365 { 1366 } 1367 1368 PETSC_NODISCARD iterator end() const noexcept { return this->begin() + (this->func.nrows * this->func.ncols); } 1369 }; 1370 1371 namespace 1372 { 1373 1374 template <typename T> 1375 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 1376 { 1377 const auto nrows = rend - rstart; 1378 const auto ncols = cend - cstart; 1379 const auto dptr = thrust::device_pointer_cast(ptr); 1380 1381 return {dptr + (rstart * lda) + cstart, dptr + ((rstart + nrows) * lda) + cstart, nrows, ncols, lda}; 1382 } 1383 1384 } // namespace 1385 1386 } // namespace detail 1387 1388 template <device::cupm::DeviceType T> 1389 inline PetscErrorCode MatDense_Seq_CUPM<T>::Scale(Mat A, PetscScalar alpha) noexcept 1390 { 1391 const auto m = A->rmap->n; 1392 const auto n = A->cmap->n; 1393 const auto N = m * n; 1394 PetscDeviceContext dctx; 1395 1396 PetscFunctionBegin; 1397 PetscCall(PetscInfo(A, "Performing Scale %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m, n)); 1398 PetscCall(GetHandles_(&dctx)); 1399 { 1400 const auto da = DeviceArrayReadWrite(dctx, A); 1401 const auto lda = MatIMPLCast(A)->lda; 1402 1403 if (lda > m) { 1404 cupmStream_t stream; 1405 1406 PetscCall(GetHandlesFrom_(dctx, &stream)); 1407 // clang-format off 1408 PetscCallThrust( 1409 const auto sub_mat = detail::make_submat_iterator(0, m, 0, n, lda, da.data()); 1410 1411 THRUST_CALL( 1412 thrust::transform, 1413 stream, 1414 sub_mat.begin(), sub_mat.end(), sub_mat.begin(), 1415 device::cupm::functors::make_times_equals(alpha) 1416 ) 1417 ); 1418 // clang-format on 1419 } else { 1420 const auto cu_alpha = cupmScalarCast(alpha); 1421 cupmBlasHandle_t handle; 1422 1423 PetscCall(GetHandlesFrom_(dctx, &handle)); 1424 PetscCall(PetscLogGpuTimeBegin()); 1425 PetscCallCUPMBLAS(cupmBlasXscal(handle, N, &cu_alpha, da.cupmdata(), 1)); 1426 PetscCall(PetscLogGpuTimeEnd()); 1427 } 1428 } 1429 PetscCall(PetscLogGpuFlops(N)); 1430 PetscFunctionReturn(PETSC_SUCCESS); 1431 } 1432 1433 template <device::cupm::DeviceType T> 1434 inline PetscErrorCode MatDense_Seq_CUPM<T>::Shift(Mat A, PetscScalar alpha) noexcept 1435 { 1436 const auto m = A->rmap->n; 1437 const auto n = A->cmap->n; 1438 PetscDeviceContext dctx; 1439 1440 PetscFunctionBegin; 1441 PetscCall(GetHandles_(&dctx)); 1442 PetscCall(PetscInfo(A, "Performing Shift %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m, n)); 1443 PetscCall(DiagonalUnaryTransform(A, 0, m, n, dctx, device::cupm::functors::make_plus_equals(alpha))); 1444 PetscFunctionReturn(PETSC_SUCCESS); 1445 } 1446 1447 template <device::cupm::DeviceType T> 1448 inline PetscErrorCode MatDense_Seq_CUPM<T>::AXPY(Mat Y, PetscScalar alpha, Mat X, MatStructure) noexcept 1449 { 1450 const auto m_x = X->rmap->n, m_y = Y->rmap->n; 1451 const auto n_x = X->cmap->n, n_y = Y->cmap->n; 1452 const auto N = m_x * n_x; 1453 PetscDeviceContext dctx; 1454 1455 PetscFunctionBegin; 1456 if (!m_x || !n_x || alpha == (PetscScalar)0.0) PetscFunctionReturn(PETSC_SUCCESS); 1457 PetscCall(PetscInfo(Y, "Performing AXPY %" PetscInt_FMT " x %" PetscInt_FMT " on backend\n", m_y, n_y)); 1458 PetscCall(GetHandles_(&dctx)); 1459 { 1460 const auto dx = DeviceArrayRead(dctx, X); 1461 const auto dy = DeviceArrayReadWrite(dctx, Y); 1462 const auto lda_x = MatIMPLCast(X)->lda; 1463 const auto lda_y = MatIMPLCast(Y)->lda; 1464 1465 if (lda_x > m_x || lda_y > m_x) { 1466 cupmStream_t stream; 1467 1468 PetscCall(GetHandlesFrom_(dctx, &stream)); 1469 // clang-format off 1470 PetscCallThrust( 1471 const auto sub_mat_y = detail::make_submat_iterator(0, m_y, 0, n_y, lda_y, dy.data()); 1472 const auto sub_mat_x = detail::make_submat_iterator(0, m_x, 0, n_x, lda_x, dx.data()); 1473 1474 THRUST_CALL( 1475 thrust::transform, 1476 stream, 1477 sub_mat_x.begin(), sub_mat_x.end(), sub_mat_y.begin(), sub_mat_y.begin(), 1478 device::cupm::functors::make_axpy(alpha) 1479 ); 1480 ); 1481 // clang-format on 1482 } else { 1483 const auto cu_alpha = cupmScalarCast(alpha); 1484 cupmBlasHandle_t handle; 1485 1486 PetscCall(GetHandlesFrom_(dctx, &handle)); 1487 PetscCall(PetscLogGpuTimeBegin()); 1488 PetscCallCUPMBLAS(cupmBlasXaxpy(handle, N, &cu_alpha, dx.cupmdata(), 1, dy.cupmdata(), 1)); 1489 PetscCall(PetscLogGpuTimeEnd()); 1490 } 1491 } 1492 PetscCall(PetscLogGpuFlops(PetscMax(2 * N - 1, 0))); 1493 PetscFunctionReturn(PETSC_SUCCESS); 1494 } 1495 1496 template <device::cupm::DeviceType T> 1497 inline PetscErrorCode MatDense_Seq_CUPM<T>::Duplicate(Mat A, MatDuplicateOption opt, Mat *B) noexcept 1498 { 1499 const auto hopt = (opt == MAT_COPY_VALUES && A->offloadmask != PETSC_OFFLOAD_CPU) ? MAT_DO_NOT_COPY_VALUES : opt; 1500 PetscDeviceContext dctx; 1501 1502 PetscFunctionBegin; 1503 PetscCall(GetHandles_(&dctx)); 1504 // do not call SetPreallocation() yet, we call it afterwards?? 1505 PetscCall(MatCreateSeqDenseCUPM<T>(PetscObjectComm(PetscObjectCast(A)), A->rmap->n, A->cmap->n, nullptr, B, dctx, /* preallocate */ false)); 1506 PetscCall(MatDuplicateNoCreate_SeqDense(*B, A, hopt)); 1507 if (opt == MAT_COPY_VALUES && hopt != MAT_COPY_VALUES) PetscCall(Copy(A, *B, SAME_NONZERO_PATTERN)); 1508 // allocate memory if needed 1509 if (opt != MAT_COPY_VALUES && !MatCUPMCast(*B)->d_v) PetscCall(SetPreallocation(*B, dctx, nullptr)); 1510 PetscFunctionReturn(PETSC_SUCCESS); 1511 } 1512 1513 template <device::cupm::DeviceType T> 1514 inline PetscErrorCode MatDense_Seq_CUPM<T>::SetRandom(Mat A, PetscRandom rng) noexcept 1515 { 1516 PetscBool device; 1517 1518 PetscFunctionBegin; 1519 PetscCall(PetscObjectTypeCompare(PetscObjectCast(rng), PETSCDEVICERAND(), &device)); 1520 if (device) { 1521 const auto m = A->rmap->n; 1522 const auto n = A->cmap->n; 1523 PetscDeviceContext dctx; 1524 1525 PetscCall(GetHandles_(&dctx)); 1526 { 1527 const auto a = DeviceArrayWrite(dctx, A); 1528 PetscInt lda; 1529 1530 PetscCall(MatDenseGetLDA(A, &lda)); 1531 if (lda > m) { 1532 for (PetscInt i = 0; i < n; i++) PetscCall(PetscRandomGetValues(rng, m, a.data() + i * lda)); 1533 } else { 1534 PetscInt mn; 1535 1536 PetscCall(PetscIntMultError(m, n, &mn)); 1537 PetscCall(PetscRandomGetValues(rng, mn, a)); 1538 } 1539 } 1540 } else { 1541 PetscCall(MatSetRandom_SeqDense(A, rng)); 1542 } 1543 PetscFunctionReturn(PETSC_SUCCESS); 1544 } 1545 1546 // ========================================================================================== 1547 1548 template <device::cupm::DeviceType T> 1549 inline PetscErrorCode MatDense_Seq_CUPM<T>::GetColumnVector(Mat A, Vec v, PetscInt col) noexcept 1550 { 1551 const auto offloadmask = A->offloadmask; 1552 const auto n = A->rmap->n; 1553 const auto col_offset = [&](const PetscScalar *ptr) { return ptr + col * MatIMPLCast(A)->lda; }; 1554 PetscBool viscupm; 1555 PetscDeviceContext dctx; 1556 cupmStream_t stream; 1557 1558 PetscFunctionBegin; 1559 PetscCall(PetscObjectTypeCompareAny(PetscObjectCast(v), &viscupm, VecSeq_CUPM::VECSEQCUPM(), VecSeq_CUPM::VECMPICUPM(), VecSeq_CUPM::VECCUPM(), "")); 1560 PetscCall(GetHandles_(&dctx, &stream)); 1561 if (viscupm && !v->boundtocpu) { 1562 const auto x = VecSeq_CUPM::DeviceArrayWrite(dctx, v); 1563 1564 // update device data 1565 if (PetscOffloadDevice(offloadmask)) { 1566 PetscCall(PetscCUPMMemcpyAsync(x.data(), col_offset(DeviceArrayRead(dctx, A)), n, cupmMemcpyDeviceToDevice, stream)); 1567 } else { 1568 PetscCall(PetscCUPMMemcpyAsync(x.data(), col_offset(HostArrayRead(dctx, A)), n, cupmMemcpyHostToDevice, stream)); 1569 } 1570 } else { 1571 PetscScalar *x; 1572 1573 // update host data 1574 PetscCall(VecGetArrayWrite(v, &x)); 1575 if (PetscOffloadUnallocated(offloadmask) || PetscOffloadHost(offloadmask)) { 1576 PetscCall(PetscArraycpy(x, col_offset(HostArrayRead(dctx, A)), n)); 1577 } else if (PetscOffloadDevice(offloadmask)) { 1578 PetscCall(PetscCUPMMemcpyAsync(x, col_offset(DeviceArrayRead(dctx, A)), n, cupmMemcpyDeviceToHost, stream)); 1579 } 1580 PetscCall(VecRestoreArrayWrite(v, &x)); 1581 } 1582 PetscFunctionReturn(PETSC_SUCCESS); 1583 } 1584 1585 template <device::cupm::DeviceType T> 1586 template <PetscMemoryAccessMode access> 1587 inline PetscErrorCode MatDense_Seq_CUPM<T>::GetColumnVec(Mat A, PetscInt col, Vec *v) noexcept 1588 { 1589 using namespace vec::cupm; 1590 const auto mimpl = MatIMPLCast(A); 1591 PetscDeviceContext dctx; 1592 1593 PetscFunctionBegin; 1594 PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first"); 1595 PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 1596 mimpl->vecinuse = col + 1; 1597 if (!mimpl->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &mimpl->cvec)); 1598 PetscCall(GetHandles_(&dctx)); 1599 PetscCall(GetArray<PETSC_MEMTYPE_DEVICE, access>(A, const_cast<PetscScalar **>(&mimpl->ptrinuse), dctx)); 1600 PetscCall(VecCUPMPlaceArrayAsync<T>(mimpl->cvec, mimpl->ptrinuse + static_cast<std::size_t>(col) * static_cast<std::size_t>(mimpl->lda))); 1601 if (access == PETSC_MEMORY_ACCESS_READ) PetscCall(VecLockReadPush(mimpl->cvec)); 1602 *v = mimpl->cvec; 1603 PetscFunctionReturn(PETSC_SUCCESS); 1604 } 1605 1606 template <device::cupm::DeviceType T> 1607 template <PetscMemoryAccessMode access> 1608 inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreColumnVec(Mat A, PetscInt, Vec *v) noexcept 1609 { 1610 using namespace vec::cupm; 1611 const auto mimpl = MatIMPLCast(A); 1612 const auto cvec = mimpl->cvec; 1613 PetscDeviceContext dctx; 1614 1615 PetscFunctionBegin; 1616 PetscCheck(mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first"); 1617 PetscCheck(cvec, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column vector"); 1618 mimpl->vecinuse = 0; 1619 if (access == PETSC_MEMORY_ACCESS_READ) PetscCall(VecLockReadPop(cvec)); 1620 PetscCall(VecCUPMResetArrayAsync<T>(cvec)); 1621 PetscCall(GetHandles_(&dctx)); 1622 PetscCall(RestoreArray<PETSC_MEMTYPE_DEVICE, access>(A, const_cast<PetscScalar **>(&mimpl->ptrinuse), dctx)); 1623 if (v) *v = nullptr; 1624 PetscFunctionReturn(PETSC_SUCCESS); 1625 } 1626 1627 // ========================================================================================== 1628 1629 template <device::cupm::DeviceType T> 1630 inline PetscErrorCode MatDense_Seq_CUPM<T>::GetFactor(Mat A, MatFactorType ftype, Mat *fact_out) noexcept 1631 { 1632 Mat fact = nullptr; 1633 PetscDeviceContext dctx; 1634 1635 PetscFunctionBegin; 1636 PetscCall(GetHandles_(&dctx)); 1637 PetscCall(MatCreateSeqDenseCUPM<T>(PetscObjectComm(PetscObjectCast(A)), A->rmap->n, A->cmap->n, nullptr, &fact, dctx, /* preallocate */ false)); 1638 fact->factortype = ftype; 1639 switch (ftype) { 1640 case MAT_FACTOR_LU: 1641 case MAT_FACTOR_ILU: // fall-through 1642 fact->ops->lufactorsymbolic = MatLUFactorSymbolic_SeqDense; 1643 fact->ops->ilufactorsymbolic = MatLUFactorSymbolic_SeqDense; 1644 break; 1645 case MAT_FACTOR_CHOLESKY: 1646 case MAT_FACTOR_ICC: // fall-through 1647 fact->ops->choleskyfactorsymbolic = MatCholeskyFactorSymbolic_SeqDense; 1648 break; 1649 case MAT_FACTOR_QR: { 1650 const auto pobj = PetscObjectCast(fact); 1651 1652 PetscCall(PetscObjectComposeFunction(pobj, "MatQRFactor_C", MatQRFactor_SeqDense)); 1653 PetscCall(PetscObjectComposeFunction(pobj, "MatQRFactorSymbolic_C", MatQRFactorSymbolic_SeqDense)); 1654 } break; 1655 case MAT_FACTOR_NONE: 1656 case MAT_FACTOR_ILUDT: // fall-through 1657 case MAT_FACTOR_NUM_TYPES: // fall-through 1658 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "MatFactorType %s not supported", MatFactorTypes[ftype]); 1659 } 1660 PetscCall(PetscStrFreeAllocpy(MATSOLVERCUPM(), &fact->solvertype)); 1661 PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_LU)); 1662 PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_ILU)); 1663 PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_CHOLESKY)); 1664 PetscCall(PetscStrallocpy(MATORDERINGEXTERNAL, const_cast<char **>(fact->preferredordering) + MAT_FACTOR_ICC)); 1665 *fact_out = fact; 1666 PetscFunctionReturn(PETSC_SUCCESS); 1667 } 1668 1669 template <device::cupm::DeviceType T> 1670 inline PetscErrorCode MatDense_Seq_CUPM<T>::InvertFactors(Mat A) noexcept 1671 { 1672 const auto mimpl = MatIMPLCast(A); 1673 const auto mcu = MatCUPMCast(A); 1674 const auto n = static_cast<cupmBlasInt_t>(A->cmap->n); 1675 cupmSolverHandle_t handle; 1676 PetscDeviceContext dctx; 1677 cupmStream_t stream; 1678 1679 PetscFunctionBegin; 1680 #if PetscDefined(HAVE_CUDA) && PetscDefined(USING_NVCC) 1681 // HIP appears to have this by default?? 1682 PetscCheck(PETSC_PKG_CUDA_VERSION_GE(10, 1, 0), PETSC_COMM_SELF, PETSC_ERR_SUP, "Upgrade to CUDA version 10.1.0 or higher"); 1683 #endif 1684 if (!n || !A->rmap->n) PetscFunctionReturn(PETSC_SUCCESS); 1685 PetscCheck(A->factortype == MAT_FACTOR_CHOLESKY, PETSC_COMM_SELF, PETSC_ERR_LIB, "Factor type %s not implemented", MatFactorTypes[A->factortype]); 1686 // spd 1687 PetscCheck(!mcu->d_fact_ipiv, PETSC_COMM_SELF, PETSC_ERR_LIB, "%sDnsytri not implemented", cupmSolverName()); 1688 1689 PetscCall(GetHandles_(&dctx, &handle, &stream)); 1690 { 1691 const auto da = DeviceArrayReadWrite(dctx, A); 1692 const auto lda = static_cast<cupmBlasInt_t>(mimpl->lda); 1693 cupmBlasInt_t il; 1694 1695 PetscCallCUPMSOLVER(cupmSolverXpotri_bufferSize(handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, &il)); 1696 if (il > mcu->d_fact_lwork) { 1697 mcu->d_fact_lwork = il; 1698 PetscCallCUPM(cupmFreeAsync(mcu->d_fact_work, stream)); 1699 PetscCall(PetscCUPMMallocAsync(&mcu->d_fact_work, il, stream)); 1700 } 1701 PetscCall(PetscLogGpuTimeBegin()); 1702 PetscCallCUPMSOLVER(cupmSolverXpotri(handle, CUPMSOLVER_FILL_MODE_LOWER, n, da.cupmdata(), lda, mcu->d_fact_work, mcu->d_fact_lwork, mcu->d_fact_info)); 1703 PetscCall(PetscLogGpuTimeEnd()); 1704 } 1705 PetscCall(CheckCUPMSolverInfo_(mcu->d_fact_info, stream)); 1706 // TODO (write cuda kernel) 1707 PetscCall(MatSeqDenseSymmetrize_Private(A, PETSC_TRUE)); 1708 PetscCall(PetscLogGpuFlops(1.0 * n * n * n / 3.0)); 1709 1710 A->ops->solve = nullptr; 1711 A->ops->solvetranspose = nullptr; 1712 A->ops->matsolve = nullptr; 1713 A->factortype = MAT_FACTOR_NONE; 1714 1715 PetscCall(PetscFree(A->solvertype)); 1716 PetscFunctionReturn(PETSC_SUCCESS); 1717 } 1718 1719 // ========================================================================================== 1720 1721 template <device::cupm::DeviceType T> 1722 inline PetscErrorCode MatDense_Seq_CUPM<T>::GetSubMatrix(Mat A, PetscInt rbegin, PetscInt rend, PetscInt cbegin, PetscInt cend, Mat *mat) noexcept 1723 { 1724 const auto mimpl = MatIMPLCast(A); 1725 const auto array_offset = [&](PetscScalar *ptr) { return ptr + rbegin + static_cast<std::size_t>(cbegin) * mimpl->lda; }; 1726 const auto n = rend - rbegin; 1727 const auto m = cend - cbegin; 1728 auto &cmat = mimpl->cmat; 1729 PetscDeviceContext dctx; 1730 1731 PetscFunctionBegin; 1732 PetscCheck(!mimpl->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first"); 1733 PetscCheck(!mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 1734 mimpl->matinuse = cbegin + 1; 1735 1736 PetscCall(GetHandles_(&dctx)); 1737 PetscCall(HostToDevice_(A, dctx)); 1738 1739 if (cmat && ((m != cmat->cmap->N) || (n != cmat->rmap->N))) PetscCall(MatDestroy(&cmat)); 1740 { 1741 const auto device_array = array_offset(MatCUPMCast(A)->d_v); 1742 1743 if (cmat) { 1744 PetscCall(PlaceArray(cmat, device_array)); 1745 } else { 1746 PetscCall(MatCreateSeqDenseCUPM<T>(PetscObjectComm(PetscObjectCast(A)), n, m, device_array, &cmat, dctx)); 1747 } 1748 } 1749 PetscCall(MatDenseSetLDA(cmat, mimpl->lda)); 1750 // place CPU array if present but do not copy any data 1751 if (const auto host_array = mimpl->v) { 1752 cmat->offloadmask = PETSC_OFFLOAD_GPU; 1753 PetscCall(MatDensePlaceArray(cmat, array_offset(host_array))); 1754 } 1755 1756 cmat->offloadmask = A->offloadmask; 1757 *mat = cmat; 1758 PetscFunctionReturn(PETSC_SUCCESS); 1759 } 1760 1761 template <device::cupm::DeviceType T> 1762 inline PetscErrorCode MatDense_Seq_CUPM<T>::RestoreSubMatrix(Mat A, Mat *m) noexcept 1763 { 1764 const auto mimpl = MatIMPLCast(A); 1765 const auto cmat = mimpl->cmat; 1766 const auto reset = static_cast<bool>(mimpl->v); 1767 bool copy, was_offload_host; 1768 1769 PetscFunctionBegin; 1770 PetscCheck(mimpl->matinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetSubMatrix() first"); 1771 PetscCheck(cmat, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column matrix"); 1772 PetscCheck(*m == cmat, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Not the matrix obtained from MatDenseGetSubMatrix()"); 1773 mimpl->matinuse = 0; 1774 1775 // calls to ResetArray may change it, so save it here 1776 was_offload_host = cmat->offloadmask == PETSC_OFFLOAD_CPU; 1777 if (was_offload_host && !reset) { 1778 copy = true; 1779 PetscCall(MatSeqDenseSetPreallocation(A, nullptr)); 1780 } else { 1781 copy = false; 1782 } 1783 1784 PetscCall(ResetArray(cmat)); 1785 if (reset) PetscCall(MatDenseResetArray(cmat)); 1786 if (copy) { 1787 PetscDeviceContext dctx; 1788 1789 PetscCall(GetHandles_(&dctx)); 1790 PetscCall(DeviceToHost_(A, dctx)); 1791 } else { 1792 A->offloadmask = was_offload_host ? PETSC_OFFLOAD_CPU : PETSC_OFFLOAD_GPU; 1793 } 1794 1795 cmat->offloadmask = PETSC_OFFLOAD_UNALLOCATED; 1796 *m = nullptr; 1797 PetscFunctionReturn(PETSC_SUCCESS); 1798 } 1799 1800 // ========================================================================================== 1801 1802 namespace 1803 { 1804 1805 template <device::cupm::DeviceType T> 1806 inline PetscErrorCode MatMatMultNumeric_SeqDenseCUPM_SeqDenseCUPM(Mat A, Mat B, Mat C, PetscBool TA, PetscBool TB) noexcept 1807 { 1808 PetscFunctionBegin; 1809 if (TA) { 1810 if (TB) { 1811 PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<true, true>(A, B, C)); 1812 } else { 1813 PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<true, false>(A, B, C)); 1814 } 1815 } else { 1816 if (TB) { 1817 PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<false, true>(A, B, C)); 1818 } else { 1819 PetscCall(MatDense_Seq_CUPM<T>::template MatMatMult_Numeric_Dispatch<false, false>(A, B, C)); 1820 } 1821 } 1822 PetscFunctionReturn(PETSC_SUCCESS); 1823 } 1824 1825 template <device::cupm::DeviceType T> 1826 inline PetscErrorCode MatSolverTypeRegister_DENSECUPM() noexcept 1827 { 1828 PetscFunctionBegin; 1829 for (auto ftype : util::make_array(MAT_FACTOR_LU, MAT_FACTOR_CHOLESKY, MAT_FACTOR_QR)) { 1830 PetscCall(MatSolverTypeRegister(MatDense_Seq_CUPM<T>::MATSOLVERCUPM(), MATSEQDENSE, ftype, MatDense_Seq_CUPM<T>::GetFactor)); 1831 PetscCall(MatSolverTypeRegister(MatDense_Seq_CUPM<T>::MATSOLVERCUPM(), MatDense_Seq_CUPM<T>::MATSEQDENSECUPM(), ftype, MatDense_Seq_CUPM<T>::GetFactor)); 1832 } 1833 PetscFunctionReturn(PETSC_SUCCESS); 1834 } 1835 1836 } // anonymous namespace 1837 1838 } // namespace impl 1839 1840 } // namespace cupm 1841 1842 } // namespace mat 1843 1844 } // namespace Petsc 1845 1846 #endif // __cplusplus 1847 1848 #endif // PETSCMATSEQDENSECUPM_HPP 1849