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 17d71ae5a4SJacob Faibussowitsch namespace Petsc 18d71ae5a4SJacob Faibussowitsch { 19a4af0ceeSJacob Faibussowitsch 20d71ae5a4SJacob Faibussowitsch namespace device 21d71ae5a4SJacob Faibussowitsch { 2217f48955SJacob Faibussowitsch 23d71ae5a4SJacob Faibussowitsch namespace cupm 24d71ae5a4SJacob Faibussowitsch { 2517f48955SJacob Faibussowitsch 26d71ae5a4SJacob Faibussowitsch namespace impl 27d71ae5a4SJacob 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 65d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD cupmStream_t get(stream_tag) const noexcept { return this->stream.get_stream(); } 660e6b6b59SJacob Faibussowitsch 67d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD cupmBlasHandle_t get(blas_tag) const noexcept { return this->blas; } 680e6b6b59SJacob Faibussowitsch 69d71ae5a4SJacob 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 77d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static constexpr PetscDeviceContext_IMPLS *impls_cast_(PetscDeviceContext ptr) noexcept { return static_cast<PetscDeviceContext_IMPLS *>(ptr->data); } 78a4af0ceeSJacob Faibussowitsch 79d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static constexpr CUPMEvent<T> *event_cast_(PetscEvent event) noexcept { return static_cast<CUPMEvent<T> *>(event->data); } 800e6b6b59SJacob Faibussowitsch 81d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static PetscLogEvent CUPMBLAS_HANDLE_CREATE() noexcept { return T == DeviceType::CUDA ? CUBLAS_HANDLE_CREATE : HIPBLAS_HANDLE_CREATE; } 827a101e5eSJacob Faibussowitsch 83d71ae5a4SJacob 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 87d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DECL(PetscErrorCode initialize_handle_(stream_tag, PetscDeviceContext)) { return 0; } 887a101e5eSJacob Faibussowitsch 8947d993e7Ssuyashtn PETSC_NODISCARD static PetscErrorCode create_handle_(blas_tag, cupmBlasHandle_t &handle) noexcept 90d71ae5a4SJacob 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 11247d993e7Ssuyashtn PETSC_NODISCARD static PetscErrorCode initialize_handle_(blas_tag tag, PetscDeviceContext dctx) noexcept 113d71ae5a4SJacob Faibussowitsch { 1147a101e5eSJacob Faibussowitsch const auto dci = impls_cast_(dctx); 1157a101e5eSJacob Faibussowitsch auto &handle = blashandles_[dctx->device->deviceId]; 11617f48955SJacob Faibussowitsch 11717f48955SJacob Faibussowitsch PetscFunctionBegin; 11847d993e7Ssuyashtn PetscCall(create_handle_(tag, handle)); 1190e6b6b59SJacob Faibussowitsch PetscCallCUPMBLAS(cupmBlasSetStream(handle, dci->stream.get_stream())); 1207a101e5eSJacob Faibussowitsch dci->blas = handle; 1217a101e5eSJacob Faibussowitsch PetscFunctionReturn(0); 1227a101e5eSJacob Faibussowitsch } 1237a101e5eSJacob Faibussowitsch 12447d993e7Ssuyashtn PETSC_CXX_COMPAT_DECL(PetscErrorCode create_handle_(solver_tag, cupmSolverHandle_t &handle)) 125d71ae5a4SJacob 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 13747d993e7Ssuyashtn PETSC_NODISCARD static PetscErrorCode initialize_handle_(solver_tag tag, PetscDeviceContext dctx) noexcept 138d71ae5a4SJacob Faibussowitsch { 1397a101e5eSJacob Faibussowitsch const auto dci = impls_cast_(dctx); 1407a101e5eSJacob Faibussowitsch auto &handle = solverhandles_[dctx->device->deviceId]; 1417a101e5eSJacob Faibussowitsch 1427a101e5eSJacob Faibussowitsch PetscFunctionBegin; 14347d993e7Ssuyashtn PetscCall(create_handle_(tag, 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 149d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static PetscErrorCode check_current_device_(PetscDeviceContext dctxl, PetscDeviceContext dctxr) noexcept 150d71ae5a4SJacob 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 162d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static PetscErrorCode check_current_device_(PetscDeviceContext dctx) noexcept { return check_current_device_(dctx, dctx); } 1630e6b6b59SJacob Faibussowitsch 164d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static PetscErrorCode finalize_() noexcept 165d71ae5a4SJacob 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)>> 184d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static PoolType &default_pool_() noexcept 185d71ae5a4SJacob Faibussowitsch { 1860e6b6b59SJacob Faibussowitsch static PoolType pool; 1870e6b6b59SJacob Faibussowitsch return pool; 1880e6b6b59SJacob Faibussowitsch } 189030f984aSJacob Faibussowitsch 190d71ae5a4SJacob Faibussowitsch PETSC_NODISCARD static PetscErrorCode check_memtype_(PetscMemType mtype, const char mess[]) noexcept 191d71ae5a4SJacob 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> 247d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::initialize()) 248d71ae5a4SJacob 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> 265d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::destroy(PetscDeviceContext dctx)) 266d71ae5a4SJacob 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> 280d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::changeStreamType(PetscDeviceContext dctx, PETSC_UNUSED PetscStreamType stype)) 281d71ae5a4SJacob 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> 293d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::setUp(PetscDeviceContext dctx)) 294d71ae5a4SJacob 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> 309d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::query(PetscDeviceContext dctx, PetscBool *idle)) 310d71ae5a4SJacob Faibussowitsch { 311030f984aSJacob Faibussowitsch PetscFunctionBegin; 3120e6b6b59SJacob Faibussowitsch PetscCall(check_current_device_(dctx)); 313*4b955ea4SJacob Faibussowitsch switch (auto cerr = cupmStreamQuery(impls_cast_(dctx)->stream.get_stream())) { 314d71ae5a4SJacob Faibussowitsch case cupmSuccess: 315d71ae5a4SJacob Faibussowitsch *idle = PETSC_TRUE; 316d71ae5a4SJacob Faibussowitsch break; 317d71ae5a4SJacob Faibussowitsch case cupmErrorNotReady: 318d71ae5a4SJacob Faibussowitsch *idle = PETSC_FALSE; 319*4b955ea4SJacob Faibussowitsch // reset the error 320*4b955ea4SJacob Faibussowitsch cerr = cupmGetLastError(); 321*4b955ea4SJacob Faibussowitsch static_cast<void>(cerr); 322d71ae5a4SJacob Faibussowitsch break; 323d71ae5a4SJacob Faibussowitsch default: 324d71ae5a4SJacob Faibussowitsch PetscCallCUPM(cerr); 325d71ae5a4SJacob Faibussowitsch PetscUnreachable(); 326030f984aSJacob Faibussowitsch } 327030f984aSJacob Faibussowitsch PetscFunctionReturn(0); 328030f984aSJacob Faibussowitsch } 329030f984aSJacob Faibussowitsch 33017f48955SJacob Faibussowitsch template <DeviceType T> 331d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::waitForContext(PetscDeviceContext dctxa, PetscDeviceContext dctxb)) 332d71ae5a4SJacob Faibussowitsch { 3330e6b6b59SJacob Faibussowitsch const auto dcib = impls_cast_(dctxb); 3340e6b6b59SJacob Faibussowitsch const auto event = dcib->event; 335030f984aSJacob Faibussowitsch 336030f984aSJacob Faibussowitsch PetscFunctionBegin; 3370e6b6b59SJacob Faibussowitsch PetscCall(check_current_device_(dctxa, dctxb)); 3380e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmEventRecord(event, dcib->stream.get_stream())); 3390e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmStreamWaitEvent(impls_cast_(dctxa)->stream.get_stream(), event, 0)); 340030f984aSJacob Faibussowitsch PetscFunctionReturn(0); 341030f984aSJacob Faibussowitsch } 342030f984aSJacob Faibussowitsch 34317f48955SJacob Faibussowitsch template <DeviceType T> 344d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::synchronize(PetscDeviceContext dctx)) 345d71ae5a4SJacob Faibussowitsch { 3460e6b6b59SJacob Faibussowitsch auto idle = PETSC_TRUE; 347030f984aSJacob Faibussowitsch 348030f984aSJacob Faibussowitsch PetscFunctionBegin; 3490e6b6b59SJacob Faibussowitsch PetscCall(query(dctx, &idle)); 3500e6b6b59SJacob Faibussowitsch if (!idle) PetscCallCUPM(cupmStreamSynchronize(impls_cast_(dctx)->stream.get_stream())); 351030f984aSJacob Faibussowitsch PetscFunctionReturn(0); 352030f984aSJacob Faibussowitsch } 353030f984aSJacob Faibussowitsch 35417f48955SJacob Faibussowitsch template <DeviceType T> 35517f48955SJacob Faibussowitsch template <typename handle_t> 356d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::getHandle(PetscDeviceContext dctx, void *handle)) 357d71ae5a4SJacob Faibussowitsch { 358a4af0ceeSJacob Faibussowitsch PetscFunctionBegin; 3597a101e5eSJacob Faibussowitsch PetscCall(initialize_handle_(handle_t{}, dctx)); 3607a101e5eSJacob Faibussowitsch *static_cast<typename handle_t::type *>(handle) = impls_cast_(dctx)->get(handle_t{}); 361a4af0ceeSJacob Faibussowitsch PetscFunctionReturn(0); 362a4af0ceeSJacob Faibussowitsch } 363a4af0ceeSJacob Faibussowitsch 36417f48955SJacob Faibussowitsch template <DeviceType T> 365d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::beginTimer(PetscDeviceContext dctx)) 366d71ae5a4SJacob Faibussowitsch { 3670e6b6b59SJacob Faibussowitsch const auto dci = impls_cast_(dctx); 368a4af0ceeSJacob Faibussowitsch 369a4af0ceeSJacob Faibussowitsch PetscFunctionBegin; 3700e6b6b59SJacob Faibussowitsch PetscCall(check_current_device_(dctx)); 371a4af0ceeSJacob Faibussowitsch #if PetscDefined(USE_DEBUG) 3725f80ce2aSJacob Faibussowitsch PetscCheck(!dci->timerInUse, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Forgot to call PetscLogGpuTimeEnd()?"); 373a4af0ceeSJacob Faibussowitsch dci->timerInUse = PETSC_TRUE; 374a4af0ceeSJacob Faibussowitsch #endif 37517f48955SJacob Faibussowitsch if (!dci->begin) { 3760e6b6b59SJacob Faibussowitsch PetscAssert(!dci->end, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Don't have a 'begin' event, but somehow have an end event"); 3779566063dSJacob Faibussowitsch PetscCallCUPM(cupmEventCreate(&dci->begin)); 3789566063dSJacob Faibussowitsch PetscCallCUPM(cupmEventCreate(&dci->end)); 37917f48955SJacob Faibussowitsch } 3800e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmEventRecord(dci->begin, dci->stream.get_stream())); 381a4af0ceeSJacob Faibussowitsch PetscFunctionReturn(0); 382a4af0ceeSJacob Faibussowitsch } 383a4af0ceeSJacob Faibussowitsch 38417f48955SJacob Faibussowitsch template <DeviceType T> 385d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::endTimer(PetscDeviceContext dctx, PetscLogDouble *elapsed)) 386d71ae5a4SJacob Faibussowitsch { 387a4af0ceeSJacob Faibussowitsch float gtime; 3880e6b6b59SJacob Faibussowitsch const auto dci = impls_cast_(dctx); 3890e6b6b59SJacob Faibussowitsch const auto end = dci->end; 390a4af0ceeSJacob Faibussowitsch 391a4af0ceeSJacob Faibussowitsch PetscFunctionBegin; 3920e6b6b59SJacob Faibussowitsch PetscCall(check_current_device_(dctx)); 393a4af0ceeSJacob Faibussowitsch #if PetscDefined(USE_DEBUG) 3945f80ce2aSJacob Faibussowitsch PetscCheck(dci->timerInUse, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Forgot to call PetscLogGpuTimeBegin()?"); 395a4af0ceeSJacob Faibussowitsch dci->timerInUse = PETSC_FALSE; 396a4af0ceeSJacob Faibussowitsch #endif 3970e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmEventRecord(end, dci->stream.get_stream())); 3980e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmEventSynchronize(end)); 3990e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmEventElapsedTime(>ime, dci->begin, end)); 40017f48955SJacob Faibussowitsch *elapsed = static_cast<util::remove_pointer_t<decltype(elapsed)>>(gtime); 401a4af0ceeSJacob Faibussowitsch PetscFunctionReturn(0); 402a4af0ceeSJacob Faibussowitsch } 403a4af0ceeSJacob Faibussowitsch 4040e6b6b59SJacob Faibussowitsch template <DeviceType T> 405d71ae5a4SJacob 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)) 406d71ae5a4SJacob Faibussowitsch { 4070e6b6b59SJacob Faibussowitsch const auto &stream = impls_cast_(dctx)->stream; 4080e6b6b59SJacob Faibussowitsch 4090e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 4100e6b6b59SJacob Faibussowitsch PetscCall(check_current_device_(dctx)); 4110e6b6b59SJacob Faibussowitsch PetscCall(check_memtype_(mtype, "allocating")); 4120e6b6b59SJacob Faibussowitsch if (PetscMemTypeHost(mtype)) { 4136797ed33SJacob Faibussowitsch PetscCall(default_pool_<HostAllocator<T>>().allocate(n, reinterpret_cast<char **>(dest), &stream, alignment)); 4140e6b6b59SJacob Faibussowitsch } else { 4156797ed33SJacob Faibussowitsch PetscCall(default_pool_<DeviceAllocator<T>>().allocate(n, reinterpret_cast<char **>(dest), &stream, alignment)); 4160e6b6b59SJacob Faibussowitsch } 4176797ed33SJacob Faibussowitsch if (clear) PetscCallCUPM(cupmMemsetAsync(*dest, 0, n, stream.get_stream())); 4180e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 4190e6b6b59SJacob Faibussowitsch } 4200e6b6b59SJacob Faibussowitsch 4210e6b6b59SJacob Faibussowitsch template <DeviceType T> 422d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::memFree(PetscDeviceContext dctx, PetscMemType mtype, void **ptr)) 423d71ae5a4SJacob Faibussowitsch { 4240e6b6b59SJacob Faibussowitsch const auto &stream = impls_cast_(dctx)->stream; 4250e6b6b59SJacob Faibussowitsch 4260e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 4270e6b6b59SJacob Faibussowitsch PetscCall(check_current_device_(dctx)); 4280e6b6b59SJacob Faibussowitsch PetscCall(check_memtype_(mtype, "freeing")); 4290e6b6b59SJacob Faibussowitsch if (!*ptr) PetscFunctionReturn(0); 4300e6b6b59SJacob Faibussowitsch if (PetscMemTypeHost(mtype)) { 4310e6b6b59SJacob Faibussowitsch PetscCall(default_pool_<HostAllocator<T>>().deallocate(reinterpret_cast<char **>(ptr), &stream)); 4320e6b6b59SJacob Faibussowitsch // if ptr exists still exists the pool didn't own it 4330e6b6b59SJacob Faibussowitsch if (*ptr) { 4340e6b6b59SJacob Faibussowitsch auto registered = PETSC_FALSE, managed = PETSC_FALSE; 4350e6b6b59SJacob Faibussowitsch 4360e6b6b59SJacob Faibussowitsch PetscCall(PetscCUPMGetMemType(*ptr, nullptr, ®istered, &managed)); 4370e6b6b59SJacob Faibussowitsch if (registered) { 4380e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmFreeHost(*ptr)); 4390e6b6b59SJacob Faibussowitsch } else if (managed) { 4400e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmFreeAsync(*ptr, stream.get_stream())); 4410e6b6b59SJacob Faibussowitsch } 4420e6b6b59SJacob Faibussowitsch } 4430e6b6b59SJacob Faibussowitsch } else { 4440e6b6b59SJacob Faibussowitsch PetscCall(default_pool_<DeviceAllocator<T>>().deallocate(reinterpret_cast<char **>(ptr), &stream)); 4450e6b6b59SJacob Faibussowitsch // if ptr exists still exists the pool didn't own it 4460e6b6b59SJacob Faibussowitsch if (*ptr) PetscCallCUPM(cupmFreeAsync(*ptr, stream.get_stream())); 4470e6b6b59SJacob Faibussowitsch } 4480e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 4490e6b6b59SJacob Faibussowitsch } 4500e6b6b59SJacob Faibussowitsch 4510e6b6b59SJacob Faibussowitsch template <DeviceType T> 452d71ae5a4SJacob 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)) 453d71ae5a4SJacob Faibussowitsch { 4540e6b6b59SJacob Faibussowitsch const auto stream = impls_cast_(dctx)->stream.get_stream(); 4550e6b6b59SJacob Faibussowitsch 4560e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 4570e6b6b59SJacob Faibussowitsch // can't use PetscCUPMMemcpyAsync here since we don't know sizeof(*src)... 4580e6b6b59SJacob Faibussowitsch if (mode == PETSC_DEVICE_COPY_HTOH) { 4590e6b6b59SJacob Faibussowitsch // yes this is faster 4600e6b6b59SJacob Faibussowitsch if (cupmStreamQuery(stream) == cupmSuccess) { 4610e6b6b59SJacob Faibussowitsch PetscCall(PetscMemcpy(dest, src, n)); 4620e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 4630e6b6b59SJacob Faibussowitsch } 4640e6b6b59SJacob Faibussowitsch // in case cupmStreamQuery() did not return cupmErrorNotReady 4650e6b6b59SJacob Faibussowitsch PetscCallCUPM(cupmGetLastError()); 4660e6b6b59SJacob Faibussowitsch } 4670e6b6b59SJacob Faibussowitsch PetscCall(cupmMemcpyAsync(dest, src, n, PetscDeviceCopyModeToCUPMMemcpyKind(mode), stream)); 4680e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 4690e6b6b59SJacob Faibussowitsch } 4700e6b6b59SJacob Faibussowitsch 4710e6b6b59SJacob Faibussowitsch template <DeviceType T> 472d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::memSet(PetscDeviceContext dctx, PetscMemType mtype, void *ptr, PetscInt v, std::size_t n)) 473d71ae5a4SJacob Faibussowitsch { 4740e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 4750e6b6b59SJacob Faibussowitsch PetscCall(check_current_device_(dctx)); 4760e6b6b59SJacob Faibussowitsch PetscCall(check_memtype_(mtype, "zeroing")); 4776797ed33SJacob Faibussowitsch PetscCallCUPM(cupmMemsetAsync(ptr, static_cast<int>(v), n, impls_cast_(dctx)->stream.get_stream())); 4780e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 4790e6b6b59SJacob Faibussowitsch } 4800e6b6b59SJacob Faibussowitsch 4810e6b6b59SJacob Faibussowitsch template <DeviceType T> 482d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::createEvent(PetscDeviceContext dctx, PetscEvent event)) 483d71ae5a4SJacob Faibussowitsch { 4840e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 4850e6b6b59SJacob Faibussowitsch PetscCallCXX(event->data = new event_type()); 4860e6b6b59SJacob Faibussowitsch event->destroy = [](PetscEvent event) { 4870e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 4880e6b6b59SJacob Faibussowitsch delete event_cast_(event); 4890e6b6b59SJacob Faibussowitsch event->data = nullptr; 4900e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 4910e6b6b59SJacob Faibussowitsch }; 4920e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 4930e6b6b59SJacob Faibussowitsch } 4940e6b6b59SJacob Faibussowitsch 4950e6b6b59SJacob Faibussowitsch template <DeviceType T> 496d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::recordEvent(PetscDeviceContext dctx, PetscEvent event)) 497d71ae5a4SJacob Faibussowitsch { 4980e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 4990e6b6b59SJacob Faibussowitsch PetscCall(impls_cast_(dctx)->stream.record_event(*event_cast_(event))); 5000e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 5010e6b6b59SJacob Faibussowitsch } 5020e6b6b59SJacob Faibussowitsch 5030e6b6b59SJacob Faibussowitsch template <DeviceType T> 504d71ae5a4SJacob Faibussowitsch PETSC_CXX_COMPAT_DEFN(PetscErrorCode DeviceContext<T>::waitForEvent(PetscDeviceContext dctx, PetscEvent event)) 505d71ae5a4SJacob Faibussowitsch { 5060e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 5070e6b6b59SJacob Faibussowitsch PetscCall(impls_cast_(dctx)->stream.wait_for_event(*event_cast_(event))); 5080e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 5090e6b6b59SJacob Faibussowitsch } 5100e6b6b59SJacob Faibussowitsch 511030f984aSJacob Faibussowitsch // initialize the static member variables 5129371c9d4SSatish Balay template <DeviceType T> 5139371c9d4SSatish Balay bool DeviceContext<T>::initialized_ = false; 514030f984aSJacob Faibussowitsch 51517f48955SJacob Faibussowitsch template <DeviceType T> 51617f48955SJacob Faibussowitsch std::array<typename DeviceContext<T>::cupmBlasHandle_t, PETSC_DEVICE_MAX_DEVICES> DeviceContext<T>::blashandles_ = {}; 517030f984aSJacob Faibussowitsch 51817f48955SJacob Faibussowitsch template <DeviceType T> 51917f48955SJacob Faibussowitsch std::array<typename DeviceContext<T>::cupmSolverHandle_t, PETSC_DEVICE_MAX_DEVICES> DeviceContext<T>::solverhandles_ = {}; 52017f48955SJacob Faibussowitsch 5210e6b6b59SJacob Faibussowitsch } // namespace impl 522030f984aSJacob Faibussowitsch 523a4af0ceeSJacob Faibussowitsch // shorten this one up a bit (and instantiate the templates) 5240e6b6b59SJacob Faibussowitsch using CUPMContextCuda = impl::DeviceContext<DeviceType::CUDA>; 5250e6b6b59SJacob Faibussowitsch using CUPMContextHip = impl::DeviceContext<DeviceType::HIP>; 526030f984aSJacob Faibussowitsch 527030f984aSJacob Faibussowitsch // shorthand for what is an EXTREMELY long name 5280e6b6b59SJacob Faibussowitsch #define PetscDeviceContext_(IMPLS) ::Petsc::device::cupm::impl::DeviceContext<::Petsc::device::cupm::DeviceType::IMPLS>::PetscDeviceContext_IMPLS 529030f984aSJacob Faibussowitsch 5300e6b6b59SJacob Faibussowitsch } // namespace cupm 53117f48955SJacob Faibussowitsch 5320e6b6b59SJacob Faibussowitsch } // namespace device 53317f48955SJacob Faibussowitsch 53417f48955SJacob Faibussowitsch } // namespace Petsc 535030f984aSJacob Faibussowitsch 5360e6b6b59SJacob Faibussowitsch #endif // __cplusplus 5370e6b6b59SJacob Faibussowitsch 538a4af0ceeSJacob Faibussowitsch #endif // PETSCDEVICECONTEXTCUDA_HPP 539