117f48955SJacob Faibussowitsch #ifndef PETSCDEVICECONTEXTCUPM_HPP 2030f984aSJacob Faibussowitsch #define PETSCDEVICECONTEXTCUPM_HPP 3030f984aSJacob Faibussowitsch 4a4af0ceeSJacob Faibussowitsch #include <petsc/private/deviceimpl.h> 517f48955SJacob Faibussowitsch #include <petsc/private/cupmblasinterface.hpp> 67a101e5eSJacob Faibussowitsch #include <petsc/private/logimpl.h> 7030f984aSJacob Faibussowitsch 80e6b6b59SJacob Faibussowitsch #include <petsc/private/cpp/array.hpp> 9a4af0ceeSJacob Faibussowitsch 100e6b6b59SJacob Faibussowitsch #include "../segmentedmempool.hpp" 110e6b6b59SJacob Faibussowitsch #include "cupmallocator.hpp" 120e6b6b59SJacob Faibussowitsch #include "cupmstream.hpp" 130e6b6b59SJacob Faibussowitsch #include "cupmevent.hpp" 140e6b6b59SJacob Faibussowitsch 150e6b6b59SJacob Faibussowitsch #if defined(__cplusplus) 166797ed33SJacob Faibussowitsch 17*d71ae5a4SJacob Faibussowitsch namespace Petsc 18*d71ae5a4SJacob Faibussowitsch { 19a4af0ceeSJacob Faibussowitsch 20*d71ae5a4SJacob Faibussowitsch namespace device 21*d71ae5a4SJacob Faibussowitsch { 2217f48955SJacob Faibussowitsch 23*d71ae5a4SJacob Faibussowitsch namespace cupm 24*d71ae5a4SJacob Faibussowitsch { 2517f48955SJacob Faibussowitsch 26*d71ae5a4SJacob Faibussowitsch namespace impl 27*d71ae5a4SJacob Faibussowitsch { 28030f984aSJacob Faibussowitsch 2917f48955SJacob Faibussowitsch template <DeviceType T> 300e6b6b59SJacob Faibussowitsch class DeviceContext : BlasInterface<T> { 3117f48955SJacob Faibussowitsch public: 3217f48955SJacob Faibussowitsch PETSC_CUPMBLAS_INHERIT_INTERFACE_TYPEDEFS_USING(cupmBlasInterface_t, T); 3317f48955SJacob Faibussowitsch 3417f48955SJacob Faibussowitsch private: 359371c9d4SSatish Balay template <typename H, std::size_t> 369371c9d4SSatish Balay struct HandleTag { 379371c9d4SSatish Balay using type = H; 389371c9d4SSatish Balay }; 390e6b6b59SJacob Faibussowitsch 407a101e5eSJacob Faibussowitsch using stream_tag = HandleTag<cupmStream_t, 0>; 417a101e5eSJacob Faibussowitsch using blas_tag = HandleTag<cupmBlasHandle_t, 1>; 427a101e5eSJacob Faibussowitsch using solver_tag = HandleTag<cupmSolverHandle_t, 2>; 43a4af0ceeSJacob Faibussowitsch 440e6b6b59SJacob Faibussowitsch using stream_type = CUPMStream<T>; 450e6b6b59SJacob Faibussowitsch using event_type = CUPMEvent<T>; 460e6b6b59SJacob Faibussowitsch 47030f984aSJacob Faibussowitsch public: 48030f984aSJacob Faibussowitsch // This is the canonical PETSc "impls" struct that normally resides in a standalone impls 49030f984aSJacob Faibussowitsch // header, but since we are using the power of templates it must be declared part of 50030f984aSJacob Faibussowitsch // this class to have easy access the same typedefs. Technically one can make a 51030f984aSJacob Faibussowitsch // templated struct outside the class but it's more code for the same result. 520e6b6b59SJacob Faibussowitsch struct PetscDeviceContext_IMPLS : memory::PoolAllocated<PetscDeviceContext_IMPLS> { 530e6b6b59SJacob Faibussowitsch stream_type stream{}; 540e6b6b59SJacob Faibussowitsch cupmEvent_t event{}; 550e6b6b59SJacob Faibussowitsch cupmEvent_t begin{}; // timer-only 560e6b6b59SJacob Faibussowitsch cupmEvent_t end{}; // timer-only 57a4af0ceeSJacob Faibussowitsch #if PetscDefined(USE_DEBUG) 580e6b6b59SJacob Faibussowitsch PetscBool timerInUse{}; 59a4af0ceeSJacob Faibussowitsch #endif 600e6b6b59SJacob Faibussowitsch cupmBlasHandle_t blas{}; 610e6b6b59SJacob Faibussowitsch cupmSolverHandle_t solver{}; 62a4af0ceeSJacob Faibussowitsch 630e6b6b59SJacob Faibussowitsch constexpr PetscDeviceContext_IMPLS() noexcept = default; 640e6b6b59SJacob Faibussowitsch 65*d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD cupmStream_t get(stream_tag) const noexcept { return this->stream.get_stream(); } 660e6b6b59SJacob Faibussowitsch 67*d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD cupmBlasHandle_t get(blas_tag) const noexcept { return this->blas; } 680e6b6b59SJacob Faibussowitsch 69*d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD cupmSolverHandle_t get(solver_tag) const noexcept { return this->solver; } 70030f984aSJacob Faibussowitsch }; 71030f984aSJacob Faibussowitsch 72030f984aSJacob Faibussowitsch private: 7317f48955SJacob Faibussowitsch static bool initialized_; 7417f48955SJacob Faibussowitsch static std::array<cupmBlasHandle_t, PETSC_DEVICE_MAX_DEVICES> blashandles_; 7517f48955SJacob Faibussowitsch static std::array<cupmSolverHandle_t, PETSC_DEVICE_MAX_DEVICES> solverhandles_; 76030f984aSJacob Faibussowitsch 77*d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static constexpr PetscDeviceContext_IMPLS *impls_cast_(PetscDeviceContext ptr) noexcept { return static_cast<PetscDeviceContext_IMPLS *>(ptr->data); } 78a4af0ceeSJacob Faibussowitsch 79*d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static constexpr CUPMEvent<T> *event_cast_(PetscEvent event) noexcept { return static_cast<CUPMEvent<T> *>(event->data); } 800e6b6b59SJacob Faibussowitsch 81*d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static PetscLogEvent CUPMBLAS_HANDLE_CREATE() noexcept { return T == DeviceType::CUDA ? CUBLAS_HANDLE_CREATE : HIPBLAS_HANDLE_CREATE; } 827a101e5eSJacob Faibussowitsch 83*d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static PetscLogEvent CUPMSOLVER_HANDLE_CREATE() noexcept { return T == DeviceType::CUDA ? CUSOLVER_HANDLE_CREATE : HIPSOLVER_HANDLE_CREATE; } 847a101e5eSJacob Faibussowitsch 857a101e5eSJacob Faibussowitsch // this exists purely to satisfy the compiler so the tag-based dispatch works for the other 867a101e5eSJacob Faibussowitsch // handles 87*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode initialize_handle_(stream_tag, PetscDeviceContext)) { return 0; } 887a101e5eSJacob Faibussowitsch 89*d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static PetscErrorCode create_handle_(cupmBlasHandle_t &handle) noexcept 90*d71ae5a4SJacob Faibussowitsch { 917a101e5eSJacob Faibussowitsch PetscLogEvent event; 927a101e5eSJacob Faibussowitsch 93030f984aSJacob Faibussowitsch PetscFunctionBegin; 947a101e5eSJacob Faibussowitsch if (PetscLikely(handle)) PetscFunctionReturn(0); 957a101e5eSJacob Faibussowitsch PetscCall(PetscLogPauseCurrentEvent_Internal(&event)); 967a101e5eSJacob Faibussowitsch PetscCall(PetscLogEventBegin(CUPMBLAS_HANDLE_CREATE(), 0, 0, 0, 0)); 9717f48955SJacob Faibussowitsch for (auto i = 0; i < 3; ++i) { 9817f48955SJacob Faibussowitsch auto cberr = cupmBlasCreate(&handle); 9917f48955SJacob Faibussowitsch if (PetscLikely(cberr == CUPMBLAS_STATUS_SUCCESS)) break; 1009566063dSJacob Faibussowitsch if (PetscUnlikely(cberr != CUPMBLAS_STATUS_ALLOC_FAILED) && (cberr != CUPMBLAS_STATUS_NOT_INITIALIZED)) PetscCallCUPMBLAS(cberr); 10117f48955SJacob Faibussowitsch if (i != 2) { 1029566063dSJacob Faibussowitsch PetscCall(PetscSleep(3)); 10317f48955SJacob Faibussowitsch continue; 104a4af0ceeSJacob Faibussowitsch } 1055f80ce2aSJacob Faibussowitsch PetscCheck(cberr == CUPMBLAS_STATUS_SUCCESS, PETSC_COMM_SELF, PETSC_ERR_GPU_RESOURCE, "Unable to initialize %s", cupmBlasName()); 106a4af0ceeSJacob Faibussowitsch } 1077a101e5eSJacob Faibussowitsch PetscCall(PetscLogEventEnd(CUPMBLAS_HANDLE_CREATE(), 0, 0, 0, 0)); 1087a101e5eSJacob Faibussowitsch PetscCall(PetscLogEventResume_Internal(event)); 109030f984aSJacob Faibussowitsch PetscFunctionReturn(0); 110030f984aSJacob Faibussowitsch } 111030f984aSJacob Faibussowitsch 112*d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static PetscErrorCode initialize_handle_(blas_tag, PetscDeviceContext dctx) noexcept 113*d71ae5a4SJacob Faibussowitsch { 1147a101e5eSJacob Faibussowitsch const auto dci = impls_cast_(dctx); 1157a101e5eSJacob Faibussowitsch auto &handle = blashandles_[dctx->device->deviceId]; 11617f48955SJacob Faibussowitsch 11717f48955SJacob Faibussowitsch PetscFunctionBegin; 1187a101e5eSJacob Faibussowitsch PetscCall(create_handle_(handle)); 1190e6b6b59SJacob Faibussowitsch PetscCallCUPMBLAS(cupmBlasSetStream(handle, dci->stream.get_stream())); 1207a101e5eSJacob Faibussowitsch dci->blas = handle; 1217a101e5eSJacob Faibussowitsch PetscFunctionReturn(0); 1227a101e5eSJacob Faibussowitsch } 1237a101e5eSJacob Faibussowitsch 124*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode create_handle_(cupmSolverHandle_t &handle)) 125*d71ae5a4SJacob Faibussowitsch { 1267a101e5eSJacob Faibussowitsch PetscLogEvent event; 1277a101e5eSJacob Faibussowitsch 1287a101e5eSJacob Faibussowitsch PetscFunctionBegin; 1297a101e5eSJacob Faibussowitsch PetscCall(PetscLogPauseCurrentEvent_Internal(&event)); 1307a101e5eSJacob Faibussowitsch PetscCall(PetscLogEventBegin(CUPMSOLVER_HANDLE_CREATE(), 0, 0, 0, 0)); 1317a101e5eSJacob Faibussowitsch PetscCall(cupmBlasInterface_t::InitializeHandle(handle)); 1327a101e5eSJacob Faibussowitsch PetscCall(PetscLogEventEnd(CUPMSOLVER_HANDLE_CREATE(), 0, 0, 0, 0)); 1337a101e5eSJacob Faibussowitsch PetscCall(PetscLogEventResume_Internal(event)); 1347a101e5eSJacob Faibussowitsch PetscFunctionReturn(0); 1357a101e5eSJacob Faibussowitsch } 1367a101e5eSJacob Faibussowitsch 137*d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static PetscErrorCode initialize_handle_(solver_tag, PetscDeviceContext dctx) noexcept 138*d71ae5a4SJacob Faibussowitsch { 1397a101e5eSJacob Faibussowitsch const auto dci = impls_cast_(dctx); 1407a101e5eSJacob Faibussowitsch auto &handle = solverhandles_[dctx->device->deviceId]; 1417a101e5eSJacob Faibussowitsch 1427a101e5eSJacob Faibussowitsch PetscFunctionBegin; 1437a101e5eSJacob Faibussowitsch PetscCall(create_handle_(handle)); 1440e6b6b59SJacob Faibussowitsch PetscCall(cupmBlasInterface_t::SetHandleStream(handle, dci->stream.get_stream())); 1457a101e5eSJacob Faibussowitsch dci->solver = handle; 14617f48955SJacob Faibussowitsch PetscFunctionReturn(0); 14717f48955SJacob Faibussowitsch } 14817f48955SJacob Faibussowitsch 149*d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static PetscErrorCode check_current_device_(PetscDeviceContext dctxl, PetscDeviceContext dctxr) noexcept 150*d71ae5a4SJacob Faibussowitsch { 1510e6b6b59SJacob Faibussowitsch const auto devidl = dctxl->device->deviceId, devidr = dctxr->device->deviceId; 1520e6b6b59SJacob Faibussowitsch 1530e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 1540e6b6b59SJacob Faibussowitsch PetscCheck(devidl == devidr, PETSC_COMM_SELF, PETSC_ERR_GPU, "Device contexts must be on the same device; dctx A (id %" PetscInt64_FMT " device id %" PetscInt_FMT ") dctx B (id %" PetscInt64_FMT " device id %" PetscInt_FMT ")", 1550e6b6b59SJacob Faibussowitsch PetscObjectCast(dctxl)->id, devidl, PetscObjectCast(dctxr)->id, devidr); 1560e6b6b59SJacob Faibussowitsch PetscCall(PetscDeviceCheckDeviceCount_Internal(devidl)); 1570e6b6b59SJacob Faibussowitsch PetscCall(PetscDeviceCheckDeviceCount_Internal(devidr)); 1580e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmSetDevice(static_cast<int>(devidl))); 1590e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 1600e6b6b59SJacob Faibussowitsch } 1610e6b6b59SJacob Faibussowitsch 162*d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static PetscErrorCode check_current_device_(PetscDeviceContext dctx) noexcept { return check_current_device_(dctx, dctx); } 1630e6b6b59SJacob Faibussowitsch 164*d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static PetscErrorCode finalize_() noexcept 165*d71ae5a4SJacob Faibussowitsch { 16617f48955SJacob Faibussowitsch PetscFunctionBegin; 16717f48955SJacob Faibussowitsch for (auto &&handle : blashandles_) { 16817f48955SJacob Faibussowitsch if (handle) { 1699566063dSJacob Faibussowitsch PetscCallCUPMBLAS(cupmBlasDestroy(handle)); 17017f48955SJacob Faibussowitsch handle = nullptr; 17117f48955SJacob Faibussowitsch } 17217f48955SJacob Faibussowitsch } 17317f48955SJacob Faibussowitsch for (auto &&handle : solverhandles_) { 17417f48955SJacob Faibussowitsch if (handle) { 1759566063dSJacob Faibussowitsch PetscCall(cupmBlasInterface_t::DestroyHandle(handle)); 17617f48955SJacob Faibussowitsch handle = nullptr; 17717f48955SJacob Faibussowitsch } 17817f48955SJacob Faibussowitsch } 17917f48955SJacob Faibussowitsch initialized_ = false; 18017f48955SJacob Faibussowitsch PetscFunctionReturn(0); 18117f48955SJacob Faibussowitsch } 18217f48955SJacob Faibussowitsch 1830e6b6b59SJacob Faibussowitsch template <typename Allocator, typename PoolType = ::Petsc::memory::SegmentedMemoryPool<typename Allocator::value_type, stream_type, Allocator, 256 * sizeof(PetscScalar)>> 184*d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static PoolType &default_pool_() noexcept 185*d71ae5a4SJacob Faibussowitsch { 1860e6b6b59SJacob Faibussowitsch static PoolType pool; 1870e6b6b59SJacob Faibussowitsch return pool; 1880e6b6b59SJacob Faibussowitsch } 189030f984aSJacob Faibussowitsch 190*d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static PetscErrorCode check_memtype_(PetscMemType mtype, const char mess[]) noexcept 191*d71ae5a4SJacob Faibussowitsch { 1920e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 1930e6b6b59SJacob Faibussowitsch PetscCheck(PetscMemTypeHost(mtype) || (mtype == PETSC_MEMTYPE_DEVICE) || (mtype == PETSC_MEMTYPE_CUPM()), PETSC_COMM_SELF, PETSC_ERR_SUP, "%s device context can only handle %s (pinned) host or device memory", cupmName(), mess); 1940e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 1950e6b6b59SJacob Faibussowitsch } 1960e6b6b59SJacob Faibussowitsch 1970e6b6b59SJacob Faibussowitsch public: 198030f984aSJacob Faibussowitsch // All of these functions MUST be static in order to be callable from C, otherwise they 199030f984aSJacob Faibussowitsch // get the implicit 'this' pointer tacked on 20017f48955SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode destroy(PetscDeviceContext)); 20117f48955SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode changeStreamType(PetscDeviceContext, PetscStreamType)); 20217f48955SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode setUp(PetscDeviceContext)); 20317f48955SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode query(PetscDeviceContext, PetscBool *)); 20417f48955SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode waitForContext(PetscDeviceContext, PetscDeviceContext)); 20517f48955SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode synchronize(PetscDeviceContext)); 206a4af0ceeSJacob Faibussowitsch template <typename Handle_t> 20717f48955SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode getHandle(PetscDeviceContext, void *)); 20817f48955SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode beginTimer(PetscDeviceContext)); 20917f48955SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode endTimer(PetscDeviceContext, PetscLogDouble *)); 2106797ed33SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode memAlloc(PetscDeviceContext, PetscBool, PetscMemType, std::size_t, std::size_t, void **)); 2110e6b6b59SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode memFree(PetscDeviceContext, PetscMemType, void **)); 2120e6b6b59SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode memCopy(PetscDeviceContext, void *PETSC_RESTRICT, const void *PETSC_RESTRICT, std::size_t, PetscDeviceCopyMode)); 2130e6b6b59SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode memSet(PetscDeviceContext, PetscMemType, void *, PetscInt, std::size_t)); 2140e6b6b59SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode createEvent(PetscDeviceContext, PetscEvent)); 2150e6b6b59SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode recordEvent(PetscDeviceContext, PetscEvent)); 2160e6b6b59SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode waitForEvent(PetscDeviceContext, PetscEvent)); 2177a101e5eSJacob Faibussowitsch 2187a101e5eSJacob Faibussowitsch // not a PetscDeviceContext method, this registers the class 2197a101e5eSJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode initialize()); 2200e6b6b59SJacob Faibussowitsch 2210e6b6b59SJacob Faibussowitsch // clang-format off 2220e6b6b59SJacob Faibussowitsch const _DeviceContextOps ops = { 2230e6b6b59SJacob Faibussowitsch destroy, 2240e6b6b59SJacob Faibussowitsch changeStreamType, 2250e6b6b59SJacob Faibussowitsch setUp, 2260e6b6b59SJacob Faibussowitsch query, 2270e6b6b59SJacob Faibussowitsch waitForContext, 2280e6b6b59SJacob Faibussowitsch synchronize, 2290e6b6b59SJacob Faibussowitsch getHandle<blas_tag>, 2300e6b6b59SJacob Faibussowitsch getHandle<solver_tag>, 2310e6b6b59SJacob Faibussowitsch getHandle<stream_tag>, 2320e6b6b59SJacob Faibussowitsch beginTimer, 2330e6b6b59SJacob Faibussowitsch endTimer, 2340e6b6b59SJacob Faibussowitsch memAlloc, 2350e6b6b59SJacob Faibussowitsch memFree, 2360e6b6b59SJacob Faibussowitsch memCopy, 2370e6b6b59SJacob Faibussowitsch memSet, 2380e6b6b59SJacob Faibussowitsch createEvent, 2390e6b6b59SJacob Faibussowitsch recordEvent, 2400e6b6b59SJacob Faibussowitsch waitForEvent 2410e6b6b59SJacob Faibussowitsch }; 2420e6b6b59SJacob Faibussowitsch // clang-format on 243030f984aSJacob Faibussowitsch }; 244030f984aSJacob Faibussowitsch 2450e6b6b59SJacob Faibussowitsch // not a PetscDeviceContext method, this initializes the CLASS 24617f48955SJacob Faibussowitsch template <DeviceType T> 247*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::initialize()) 248*d71ae5a4SJacob Faibussowitsch { 2497a101e5eSJacob Faibussowitsch PetscFunctionBegin; 2507a101e5eSJacob Faibussowitsch if (PetscUnlikely(!initialized_)) { 2510e6b6b59SJacob Faibussowitsch cupmMemPool_t mempool; 2520e6b6b59SJacob Faibussowitsch uint64_t threshold = UINT64_MAX; 2530e6b6b59SJacob Faibussowitsch 2547a101e5eSJacob Faibussowitsch initialized_ = true; 2550e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmDeviceGetMemPool(&mempool, 0)); 2560e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmMemPoolSetAttribute(mempool, cupmMemPoolAttrReleaseThreshold, &threshold)); 2570e6b6b59SJacob Faibussowitsch blashandles_.fill(nullptr); 2580e6b6b59SJacob Faibussowitsch solverhandles_.fill(nullptr); 2597a101e5eSJacob Faibussowitsch PetscCall(PetscRegisterFinalize(finalize_)); 2607a101e5eSJacob Faibussowitsch } 2617a101e5eSJacob Faibussowitsch PetscFunctionReturn(0); 2627a101e5eSJacob Faibussowitsch } 2637a101e5eSJacob Faibussowitsch 2647a101e5eSJacob Faibussowitsch template <DeviceType T> 265*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::destroy(PetscDeviceContext dctx)) 266*d71ae5a4SJacob Faibussowitsch { 267030f984aSJacob Faibussowitsch PetscFunctionBegin; 2680e6b6b59SJacob Faibussowitsch if (const auto dci = impls_cast_(dctx)) { 2690e6b6b59SJacob Faibussowitsch PetscCall(dci->stream.destroy()); 2700e6b6b59SJacob Faibussowitsch if (dci->event) PetscCall(cupm_fast_event_pool<T>().deallocate(std::move(dci->event))); 2719566063dSJacob Faibussowitsch if (dci->begin) PetscCallCUPM(cupmEventDestroy(dci->begin)); 2729566063dSJacob Faibussowitsch if (dci->end) PetscCallCUPM(cupmEventDestroy(dci->end)); 2730e6b6b59SJacob Faibussowitsch delete dci; 2740e6b6b59SJacob Faibussowitsch dctx->data = nullptr; 2750e6b6b59SJacob Faibussowitsch } 276030f984aSJacob Faibussowitsch PetscFunctionReturn(0); 277030f984aSJacob Faibussowitsch } 278030f984aSJacob Faibussowitsch 27917f48955SJacob Faibussowitsch template <DeviceType T> 280*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::changeStreamType(PetscDeviceContext dctx, PETSC_UNUSED PetscStreamType stype)) 281*d71ae5a4SJacob Faibussowitsch { 2827a101e5eSJacob Faibussowitsch const auto dci = impls_cast_(dctx); 283030f984aSJacob Faibussowitsch 284030f984aSJacob Faibussowitsch PetscFunctionBegin; 2850e6b6b59SJacob Faibussowitsch PetscCall(dci->stream.destroy()); 286030f984aSJacob Faibussowitsch // set these to null so they aren't usable until setup is called again 287030f984aSJacob Faibussowitsch dci->blas = nullptr; 288030f984aSJacob Faibussowitsch dci->solver = nullptr; 289030f984aSJacob Faibussowitsch PetscFunctionReturn(0); 290030f984aSJacob Faibussowitsch } 291030f984aSJacob Faibussowitsch 29217f48955SJacob Faibussowitsch template <DeviceType T> 293*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::setUp(PetscDeviceContext dctx)) 294*d71ae5a4SJacob Faibussowitsch { 2957a101e5eSJacob Faibussowitsch const auto dci = impls_cast_(dctx); 2960e6b6b59SJacob Faibussowitsch auto &event = dci->event; 297030f984aSJacob Faibussowitsch 298030f984aSJacob Faibussowitsch PetscFunctionBegin; 2990e6b6b59SJacob Faibussowitsch PetscCall(check_current_device_(dctx)); 3000e6b6b59SJacob Faibussowitsch PetscCall(dci->stream.change_type(dctx->streamType)); 3010e6b6b59SJacob Faibussowitsch if (!event) PetscCall(cupm_fast_event_pool<T>().allocate(&event)); 302a4af0ceeSJacob Faibussowitsch #if PetscDefined(USE_DEBUG) 303a4af0ceeSJacob Faibussowitsch dci->timerInUse = PETSC_FALSE; 304a4af0ceeSJacob Faibussowitsch #endif 305030f984aSJacob Faibussowitsch PetscFunctionReturn(0); 306030f984aSJacob Faibussowitsch } 307030f984aSJacob Faibussowitsch 30817f48955SJacob Faibussowitsch template <DeviceType T> 309*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::query(PetscDeviceContext dctx, PetscBool *idle)) 310*d71ae5a4SJacob Faibussowitsch { 311030f984aSJacob Faibussowitsch PetscFunctionBegin; 3120e6b6b59SJacob Faibussowitsch PetscCall(check_current_device_(dctx)); 3130e6b6b59SJacob Faibussowitsch switch (const auto cerr = cupmStreamQuery(impls_cast_(dctx)->stream.get_stream())) { 314*d71ae5a4SJacob Faibussowitsch case cupmSuccess: 315*d71ae5a4SJacob Faibussowitsch *idle = PETSC_TRUE; 316*d71ae5a4SJacob Faibussowitsch break; 317*d71ae5a4SJacob Faibussowitsch case cupmErrorNotReady: 318*d71ae5a4SJacob Faibussowitsch *idle = PETSC_FALSE; 319*d71ae5a4SJacob Faibussowitsch break; 320*d71ae5a4SJacob Faibussowitsch default: 321*d71ae5a4SJacob Faibussowitsch PetscCallCUPM(cerr); 322*d71ae5a4SJacob Faibussowitsch PetscUnreachable(); 323030f984aSJacob Faibussowitsch } 324030f984aSJacob Faibussowitsch PetscFunctionReturn(0); 325030f984aSJacob Faibussowitsch } 326030f984aSJacob Faibussowitsch 32717f48955SJacob Faibussowitsch template <DeviceType T> 328*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::waitForContext(PetscDeviceContext dctxa, PetscDeviceContext dctxb)) 329*d71ae5a4SJacob Faibussowitsch { 3300e6b6b59SJacob Faibussowitsch const auto dcib = impls_cast_(dctxb); 3310e6b6b59SJacob Faibussowitsch const auto event = dcib->event; 332030f984aSJacob Faibussowitsch 333030f984aSJacob Faibussowitsch PetscFunctionBegin; 3340e6b6b59SJacob Faibussowitsch PetscCall(check_current_device_(dctxa, dctxb)); 3350e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmEventRecord(event, dcib->stream.get_stream())); 3360e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmStreamWaitEvent(impls_cast_(dctxa)->stream.get_stream(), event, 0)); 337030f984aSJacob Faibussowitsch PetscFunctionReturn(0); 338030f984aSJacob Faibussowitsch } 339030f984aSJacob Faibussowitsch 34017f48955SJacob Faibussowitsch template <DeviceType T> 341*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::synchronize(PetscDeviceContext dctx)) 342*d71ae5a4SJacob Faibussowitsch { 3430e6b6b59SJacob Faibussowitsch auto idle = PETSC_TRUE; 344030f984aSJacob Faibussowitsch 345030f984aSJacob Faibussowitsch PetscFunctionBegin; 3460e6b6b59SJacob Faibussowitsch PetscCall(query(dctx, &idle)); 3470e6b6b59SJacob Faibussowitsch if (!idle) PetscCallCUPM(cupmStreamSynchronize(impls_cast_(dctx)->stream.get_stream())); 348030f984aSJacob Faibussowitsch PetscFunctionReturn(0); 349030f984aSJacob Faibussowitsch } 350030f984aSJacob Faibussowitsch 35117f48955SJacob Faibussowitsch template <DeviceType T> 35217f48955SJacob Faibussowitsch template <typename handle_t> 353*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::getHandle(PetscDeviceContext dctx, void *handle)) 354*d71ae5a4SJacob Faibussowitsch { 355a4af0ceeSJacob Faibussowitsch PetscFunctionBegin; 3567a101e5eSJacob Faibussowitsch PetscCall(initialize_handle_(handle_t{}, dctx)); 3577a101e5eSJacob Faibussowitsch *static_cast<typename handle_t::type *>(handle) = impls_cast_(dctx)->get(handle_t{}); 358a4af0ceeSJacob Faibussowitsch PetscFunctionReturn(0); 359a4af0ceeSJacob Faibussowitsch } 360a4af0ceeSJacob Faibussowitsch 36117f48955SJacob Faibussowitsch template <DeviceType T> 362*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::beginTimer(PetscDeviceContext dctx)) 363*d71ae5a4SJacob Faibussowitsch { 3640e6b6b59SJacob Faibussowitsch const auto dci = impls_cast_(dctx); 365a4af0ceeSJacob Faibussowitsch 366a4af0ceeSJacob Faibussowitsch PetscFunctionBegin; 3670e6b6b59SJacob Faibussowitsch PetscCall(check_current_device_(dctx)); 368a4af0ceeSJacob Faibussowitsch #if PetscDefined(USE_DEBUG) 3695f80ce2aSJacob Faibussowitsch PetscCheck(!dci->timerInUse, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Forgot to call PetscLogGpuTimeEnd()?"); 370a4af0ceeSJacob Faibussowitsch dci->timerInUse = PETSC_TRUE; 371a4af0ceeSJacob Faibussowitsch #endif 37217f48955SJacob Faibussowitsch if (!dci->begin) { 3730e6b6b59SJacob Faibussowitsch PetscAssert(!dci->end, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Don't have a 'begin' event, but somehow have an end event"); 3749566063dSJacob Faibussowitsch PetscCallCUPM(cupmEventCreate(&dci->begin)); 3759566063dSJacob Faibussowitsch PetscCallCUPM(cupmEventCreate(&dci->end)); 37617f48955SJacob Faibussowitsch } 3770e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmEventRecord(dci->begin, dci->stream.get_stream())); 378a4af0ceeSJacob Faibussowitsch PetscFunctionReturn(0); 379a4af0ceeSJacob Faibussowitsch } 380a4af0ceeSJacob Faibussowitsch 38117f48955SJacob Faibussowitsch template <DeviceType T> 382*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::endTimer(PetscDeviceContext dctx, PetscLogDouble *elapsed)) 383*d71ae5a4SJacob Faibussowitsch { 384a4af0ceeSJacob Faibussowitsch float gtime; 3850e6b6b59SJacob Faibussowitsch const auto dci = impls_cast_(dctx); 3860e6b6b59SJacob Faibussowitsch const auto end = dci->end; 387a4af0ceeSJacob Faibussowitsch 388a4af0ceeSJacob Faibussowitsch PetscFunctionBegin; 3890e6b6b59SJacob Faibussowitsch PetscCall(check_current_device_(dctx)); 390a4af0ceeSJacob Faibussowitsch #if PetscDefined(USE_DEBUG) 3915f80ce2aSJacob Faibussowitsch PetscCheck(dci->timerInUse, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Forgot to call PetscLogGpuTimeBegin()?"); 392a4af0ceeSJacob Faibussowitsch dci->timerInUse = PETSC_FALSE; 393a4af0ceeSJacob Faibussowitsch #endif 3940e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmEventRecord(end, dci->stream.get_stream())); 3950e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmEventSynchronize(end)); 3960e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmEventElapsedTime(>ime, dci->begin, end)); 39717f48955SJacob Faibussowitsch *elapsed = static_cast<util::remove_pointer_t<decltype(elapsed)>>(gtime); 398a4af0ceeSJacob Faibussowitsch PetscFunctionReturn(0); 399a4af0ceeSJacob Faibussowitsch } 400a4af0ceeSJacob Faibussowitsch 4010e6b6b59SJacob Faibussowitsch template <DeviceType T> 402*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::memAlloc(PetscDeviceContext dctx, PetscBool clear, PetscMemType mtype, std::size_t n, std::size_t alignment, void **dest)) 403*d71ae5a4SJacob Faibussowitsch { 4040e6b6b59SJacob Faibussowitsch const auto &stream = impls_cast_(dctx)->stream; 4050e6b6b59SJacob Faibussowitsch 4060e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 4070e6b6b59SJacob Faibussowitsch PetscCall(check_current_device_(dctx)); 4080e6b6b59SJacob Faibussowitsch PetscCall(check_memtype_(mtype, "allocating")); 4090e6b6b59SJacob Faibussowitsch if (PetscMemTypeHost(mtype)) { 4106797ed33SJacob Faibussowitsch PetscCall(default_pool_<HostAllocator<T>>().allocate(n, reinterpret_cast<char **>(dest), &stream, alignment)); 4110e6b6b59SJacob Faibussowitsch } else { 4126797ed33SJacob Faibussowitsch PetscCall(default_pool_<DeviceAllocator<T>>().allocate(n, reinterpret_cast<char **>(dest), &stream, alignment)); 4130e6b6b59SJacob Faibussowitsch } 4146797ed33SJacob Faibussowitsch if (clear) PetscCallCUPM(cupmMemsetAsync(*dest, 0, n, stream.get_stream())); 4150e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 4160e6b6b59SJacob Faibussowitsch } 4170e6b6b59SJacob Faibussowitsch 4180e6b6b59SJacob Faibussowitsch template <DeviceType T> 419*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::memFree(PetscDeviceContext dctx, PetscMemType mtype, void **ptr)) 420*d71ae5a4SJacob Faibussowitsch { 4210e6b6b59SJacob Faibussowitsch const auto &stream = impls_cast_(dctx)->stream; 4220e6b6b59SJacob Faibussowitsch 4230e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 4240e6b6b59SJacob Faibussowitsch PetscCall(check_current_device_(dctx)); 4250e6b6b59SJacob Faibussowitsch PetscCall(check_memtype_(mtype, "freeing")); 4260e6b6b59SJacob Faibussowitsch if (!*ptr) PetscFunctionReturn(0); 4270e6b6b59SJacob Faibussowitsch if (PetscMemTypeHost(mtype)) { 4280e6b6b59SJacob Faibussowitsch PetscCall(default_pool_<HostAllocator<T>>().deallocate(reinterpret_cast<char **>(ptr), &stream)); 4290e6b6b59SJacob Faibussowitsch // if ptr exists still exists the pool didn't own it 4300e6b6b59SJacob Faibussowitsch if (*ptr) { 4310e6b6b59SJacob Faibussowitsch auto registered = PETSC_FALSE, managed = PETSC_FALSE; 4320e6b6b59SJacob Faibussowitsch 4330e6b6b59SJacob Faibussowitsch PetscCall(PetscCUPMGetMemType(*ptr, nullptr, ®istered, &managed)); 4340e6b6b59SJacob Faibussowitsch if (registered) { 4350e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmFreeHost(*ptr)); 4360e6b6b59SJacob Faibussowitsch } else if (managed) { 4370e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmFreeAsync(*ptr, stream.get_stream())); 4380e6b6b59SJacob Faibussowitsch } 4390e6b6b59SJacob Faibussowitsch } 4400e6b6b59SJacob Faibussowitsch } else { 4410e6b6b59SJacob Faibussowitsch PetscCall(default_pool_<DeviceAllocator<T>>().deallocate(reinterpret_cast<char **>(ptr), &stream)); 4420e6b6b59SJacob Faibussowitsch // if ptr exists still exists the pool didn't own it 4430e6b6b59SJacob Faibussowitsch if (*ptr) PetscCallCUPM(cupmFreeAsync(*ptr, stream.get_stream())); 4440e6b6b59SJacob Faibussowitsch } 4450e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 4460e6b6b59SJacob Faibussowitsch } 4470e6b6b59SJacob Faibussowitsch 4480e6b6b59SJacob Faibussowitsch template <DeviceType T> 449*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::memCopy(PetscDeviceContext dctx, void *PETSC_RESTRICT dest, const void *PETSC_RESTRICT src, std::size_t n, PetscDeviceCopyMode mode)) 450*d71ae5a4SJacob Faibussowitsch { 4510e6b6b59SJacob Faibussowitsch const auto stream = impls_cast_(dctx)->stream.get_stream(); 4520e6b6b59SJacob Faibussowitsch 4530e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 4540e6b6b59SJacob Faibussowitsch // can't use PetscCUPMMemcpyAsync here since we don't know sizeof(*src)... 4550e6b6b59SJacob Faibussowitsch if (mode == PETSC_DEVICE_COPY_HTOH) { 4560e6b6b59SJacob Faibussowitsch // yes this is faster 4570e6b6b59SJacob Faibussowitsch if (cupmStreamQuery(stream) == cupmSuccess) { 4580e6b6b59SJacob Faibussowitsch PetscCall(PetscMemcpy(dest, src, n)); 4590e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 4600e6b6b59SJacob Faibussowitsch } 4610e6b6b59SJacob Faibussowitsch // in case cupmStreamQuery() did not return cupmErrorNotReady 4620e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmGetLastError()); 4630e6b6b59SJacob Faibussowitsch } 4640e6b6b59SJacob Faibussowitsch PetscCall(cupmMemcpyAsync(dest, src, n, PetscDeviceCopyModeToCUPMMemcpyKind(mode), stream)); 4650e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 4660e6b6b59SJacob Faibussowitsch } 4670e6b6b59SJacob Faibussowitsch 4680e6b6b59SJacob Faibussowitsch template <DeviceType T> 469*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::memSet(PetscDeviceContext dctx, PetscMemType mtype, void *ptr, PetscInt v, std::size_t n)) 470*d71ae5a4SJacob Faibussowitsch { 4710e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 4720e6b6b59SJacob Faibussowitsch PetscCall(check_current_device_(dctx)); 4730e6b6b59SJacob Faibussowitsch PetscCall(check_memtype_(mtype, "zeroing")); 4746797ed33SJacob Faibussowitsch PetscCallCUPM(cupmMemsetAsync(ptr, static_cast<int>(v), n, impls_cast_(dctx)->stream.get_stream())); 4750e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 4760e6b6b59SJacob Faibussowitsch } 4770e6b6b59SJacob Faibussowitsch 4780e6b6b59SJacob Faibussowitsch template <DeviceType T> 479*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::createEvent(PetscDeviceContext dctx, PetscEvent event)) 480*d71ae5a4SJacob Faibussowitsch { 4810e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 4820e6b6b59SJacob Faibussowitsch PetscCallCXX(event->data = new event_type()); 4830e6b6b59SJacob Faibussowitsch event->destroy = [](PetscEvent event) { 4840e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 4850e6b6b59SJacob Faibussowitsch delete event_cast_(event); 4860e6b6b59SJacob Faibussowitsch event->data = nullptr; 4870e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 4880e6b6b59SJacob Faibussowitsch }; 4890e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 4900e6b6b59SJacob Faibussowitsch } 4910e6b6b59SJacob Faibussowitsch 4920e6b6b59SJacob Faibussowitsch template <DeviceType T> 493*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::recordEvent(PetscDeviceContext dctx, PetscEvent event)) 494*d71ae5a4SJacob Faibussowitsch { 4950e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 4960e6b6b59SJacob Faibussowitsch PetscCall(impls_cast_(dctx)->stream.record_event(*event_cast_(event))); 4970e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 4980e6b6b59SJacob Faibussowitsch } 4990e6b6b59SJacob Faibussowitsch 5000e6b6b59SJacob Faibussowitsch template <DeviceType T> 501*d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::waitForEvent(PetscDeviceContext dctx, PetscEvent event)) 502*d71ae5a4SJacob Faibussowitsch { 5030e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 5040e6b6b59SJacob Faibussowitsch PetscCall(impls_cast_(dctx)->stream.wait_for_event(*event_cast_(event))); 5050e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 5060e6b6b59SJacob Faibussowitsch } 5070e6b6b59SJacob Faibussowitsch 508030f984aSJacob Faibussowitsch // initialize the static member variables 5099371c9d4SSatish Balay template <DeviceType T> 5109371c9d4SSatish Balay bool DeviceContext<T>::initialized_ = false; 511030f984aSJacob Faibussowitsch 51217f48955SJacob Faibussowitsch template <DeviceType T> 51317f48955SJacob Faibussowitsch std::array<typename DeviceContext<T>::cupmBlasHandle_t, PETSC_DEVICE_MAX_DEVICES> DeviceContext<T>::blashandles_ = {}; 514030f984aSJacob Faibussowitsch 51517f48955SJacob Faibussowitsch template <DeviceType T> 51617f48955SJacob Faibussowitsch std::array<typename DeviceContext<T>::cupmSolverHandle_t, PETSC_DEVICE_MAX_DEVICES> DeviceContext<T>::solverhandles_ = {}; 51717f48955SJacob Faibussowitsch 5180e6b6b59SJacob Faibussowitsch } // namespace impl 519030f984aSJacob Faibussowitsch 520a4af0ceeSJacob Faibussowitsch // shorten this one up a bit (and instantiate the templates) 5210e6b6b59SJacob Faibussowitsch using CUPMContextCuda = impl::DeviceContext<DeviceType::CUDA>; 5220e6b6b59SJacob Faibussowitsch using CUPMContextHip = impl::DeviceContext<DeviceType::HIP>; 523030f984aSJacob Faibussowitsch 524030f984aSJacob Faibussowitsch // shorthand for what is an EXTREMELY long name 5250e6b6b59SJacob Faibussowitsch #define PetscDeviceContext_(IMPLS) ::Petsc::device::cupm::impl::DeviceContext<::Petsc::device::cupm::DeviceType::IMPLS>::PetscDeviceContext_IMPLS 526030f984aSJacob Faibussowitsch 5270e6b6b59SJacob Faibussowitsch } // namespace cupm 52817f48955SJacob Faibussowitsch 5290e6b6b59SJacob Faibussowitsch } // namespace device 53017f48955SJacob Faibussowitsch 53117f48955SJacob Faibussowitsch } // namespace Petsc 532030f984aSJacob Faibussowitsch 5330e6b6b59SJacob Faibussowitsch #endif // __cplusplus 5340e6b6b59SJacob Faibussowitsch 535a4af0ceeSJacob Faibussowitsch #endif // PETSCDEVICECONTEXTCUDA_HPP 536