Skip to content

Commit

Permalink
Use device to get native context
Browse files Browse the repository at this point in the history
  • Loading branch information
@hdelan
hdelan committed Mar 25, 2024
1 parent c968a5c commit 878b981
Show file tree
Hide file tree
Showing 5 changed files with 20 additions and 24 deletions.
10 changes: 5 additions & 5 deletions src/blas/backends/cublas/cublas_scope_handle.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,14 +41,14 @@ thread_local cublas_handle<pi_context> CublasScopedContextHandler::handle_helper
CublasScopedContextHandler::CublasScopedContextHandler(sycl::queue queue, sycl::interop_handle &ih)
: ih(ih),
needToRecover_(false) {
auto cudaDevice = sycl::get_native<sycl::backend::ext_oneapi_cuda>(queue.get_device());
placedContext_ = new sycl::context(queue.get_context());
auto cudaDevice = ih.get_native_device<sycl::backend::ext_oneapi_cuda>();
CUresult err;
CUcontext desired;
CUDA_ERROR_FUNC(cuDevicePrimaryCtxRetain, err, &desired, cudaDevice);
CUDA_ERROR_FUNC(cuCtxGetCurrent, err, &original_);
// Getting the primary context also sets it as the active context
CUDA_ERROR_FUNC(cuDevicePrimaryCtxRetain, err, &desired, cudaDevice);
if (original_ != desired) {
// Sets the desired context as the active one for the thread
CUDA_ERROR_FUNC(cuCtxSetCurrent, err, desired);
// No context is installed and the suggested context is primary
// This is the most common case. We can activate the context in the
// thread and leave it there until all the PI context referring to the
Expand Down Expand Up @@ -87,7 +87,7 @@ void ContextCallback(void *userData) {
}

cublasHandle_t CublasScopedContextHandler::get_handle(const sycl::queue &queue) {
auto cudaDevice = sycl::get_native<sycl::backend::ext_oneapi_cuda>(queue.get_device());
auto cudaDevice = ih.get_native_device<sycl::backend::ext_oneapi_cuda>();
CUresult cuErr;
CUcontext desired;
CUDA_ERROR_FUNC(cuDevicePrimaryCtxRetain, cuErr, &desired, cudaDevice);
Expand Down
10 changes: 5 additions & 5 deletions src/blas/backends/rocblas/rocblas_scope_handle.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,14 +57,14 @@ RocblasScopedContextHandler::RocblasScopedContextHandler(sycl::queue queue,
sycl::interop_handle &ih)
: interop_h(ih),
needToRecover_(false) {
auto hipDevice = sycl::get_native<sycl::backend::ext_oneapi_hip>(queue.get_device());
placedContext_ = new sycl::context(queue.get_context());
auto hipDevice = ih.get_native_device<sycl::backend::ext_oneapi_hip>();
hipError_t err;
hipCtx_t desired;
HIP_ERROR_FUNC(hipDevicePrimaryCtxRetain, err, &desired, hipDevice);
HIP_ERROR_FUNC(hipCtxGetCurrent, err, &original_);
// Getting the primary context also sets it as the active context
HIP_ERROR_FUNC(hipDevicePrimaryCtxRetain, err, &desired, hipDevice);
if (original_ != desired) {
// Sets the desired context as the active one for the thread
HIP_ERROR_FUNC(hipCtxSetCurrent, err, desired);
// No context is installed and the suggested context is primary
// This is the most common case. We can activate the context in the
// thread and leave it there until all the PI context referring to the
Expand Down Expand Up @@ -103,7 +103,7 @@ void ContextCallback(void *userData) {
}

rocblas_handle RocblasScopedContextHandler::get_handle(const sycl::queue &queue) {
auto hipDevice = sycl::get_native<sycl::backend::ext_oneapi_hip>(queue.get_device());
auto hipDevice = interop_h.get_native_device<sycl::backend::ext_oneapi_hip>();
hipError_t hipErr;
hipCtx_t desired;
HIP_ERROR_FUNC(hipDevicePrimaryCtxRetain, hipErr, &desired, hipDevice);
Expand Down
4 changes: 0 additions & 4 deletions src/dft/backends/cufft/commit.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -88,10 +88,6 @@ class cufft_commit final : public dft::detail::commit_impl<prec, dom> {
throw mkl::exception("dft/backends/cufft", __FUNCTION__,
"Failed to change cuda context.");
}
if (cuCtxSetCurrent(interopContext) != CUDA_SUCCESS) {
throw mkl::exception("dft/backends/cufft", __FUNCTION__,
"Failed to change cuda context.");
}
}
}

Expand Down
10 changes: 5 additions & 5 deletions src/lapack/backends/cusolver/cusolver_scope_handle.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,14 +42,14 @@ CusolverScopedContextHandler::CusolverScopedContextHandler(sycl::queue queue,
sycl::interop_handle &ih)
: ih(ih),
needToRecover_(false) {
auto cudaDevice = sycl::get_native<sycl::backend::ext_oneapi_cuda>(queue.get_device());
placedContext_ = new sycl::context(queue.get_context());
auto cudaDevice = ih.get_native_device<sycl::backend::ext_oneapi_cuda>();
CUresult err;
CUcontext desired;
CUDA_ERROR_FUNC(cuDevicePrimaryCtxRetain, err, &desired, cudaDevice);
CUDA_ERROR_FUNC(cuCtxGetCurrent, err, &original_);
// Getting the primary context also sets it as the active context
CUDA_ERROR_FUNC(cuDevicePrimaryCtxRetain, err, &desired, cudaDevice);
if (original_ != desired) {
// Sets the desired context as the active one for the thread
CUDA_ERROR_FUNC(cuCtxSetCurrent, err, desired);
// No context is installed and the suggested context is primary
// This is the most common case. We can activate the context in the
// thread and leave it there until all the PI context referring to the
Expand Down Expand Up @@ -88,7 +88,7 @@ void ContextCallback(void *userData) {
}

cusolverDnHandle_t CusolverScopedContextHandler::get_handle(const sycl::queue &queue) {
auto cudaDevice = sycl::get_native<sycl::backend::ext_oneapi_cuda>(queue.get_device());
auto cudaDevice = ih.get_native_device<sycl::backend::ext_oneapi_cuda>();
CUresult cuErr;
CUcontext desired;
CUDA_ERROR_FUNC(cuDevicePrimaryCtxRetain, cuErr, &desired, cudaDevice);
Expand Down
10 changes: 5 additions & 5 deletions src/lapack/backends/rocsolver/rocsolver_scope_handle.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,14 +44,14 @@ RocsolverScopedContextHandler::RocsolverScopedContextHandler(sycl::queue queue,
sycl::interop_handle &ih)
: ih(ih),
needToRecover_(false) {
auto hipDevice = sycl::get_native<sycl::backend::ext_oneapi_hip>(queue.get_device());
placedContext_ = new sycl::context(queue.get_context());
auto hipDevice = ih.get_native_device<sycl::backend::ext_oneapi_hip>();
hipError_t err;
hipCtx_t desired;
HIP_ERROR_FUNC(hipDevicePrimaryCtxRetain, err, &desired, hipDevice);
HIP_ERROR_FUNC(hipCtxGetCurrent, err, &original_);
// Getting the primary context also sets it as the active context
HIP_ERROR_FUNC(hipDevicePrimaryCtxRetain, err, &desired, hipDevice);
if (original_ != desired) {
// Sets the desired context as the active one for the thread
HIP_ERROR_FUNC(hipCtxSetCurrent, err, desired);
// No context is installed and the suggested context is primary
// This is the most common case. We can activate the context in the
// thread and leave it there until all the PI context referring to the
Expand Down Expand Up @@ -90,7 +90,7 @@ void ContextCallback(void *userData) {
}

rocblas_handle RocsolverScopedContextHandler::get_handle(const sycl::queue &queue) {
auto hipDevice = sycl::get_native<sycl::backend::ext_oneapi_hip>(queue.get_device());
auto hipDevice = ih.get_native_device<sycl::backend::ext_oneapi_hip>();
hipError_t hipErr;
hipCtx_t desired;
HIP_ERROR_FUNC(hipDevicePrimaryCtxRetain, hipErr, &desired, hipDevice);
Expand Down

0 comments on commit 878b981

Please sign in to comment.