Skip to content

[SYCL][CUDA] Fix context setup for device infos #8124

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Feb 2, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 19 additions & 21 deletions sycl/plugins/cuda/pi_cuda.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -188,17 +188,30 @@ pi_result check_error(CUresult result, const char *function, int line,
/// contexts to be restored by SYCL.
class ScopedContext {
public:
ScopedContext(pi_context ctxt) {
ScopedContext(pi_context ctxt) : device(nullptr) {
if (!ctxt) {
throw PI_ERROR_INVALID_CONTEXT;
}

set_context(ctxt->get());
}

ScopedContext(CUcontext ctxt) { set_context(ctxt); }
ScopedContext(CUcontext ctxt) : device(nullptr) { set_context(ctxt); }

~ScopedContext() {}
// Creating a scoped context from a device will simply use the primary
// context, this should be used when there is no other appropriate context,
// such as for the device infos.
ScopedContext(pi_device device) : device(device) {
CUcontext ctxt;
cuDevicePrimaryCtxRetain(&ctxt, device->get());

set_context(ctxt);
}

~ScopedContext() {
if (device)
cuDevicePrimaryCtxRelease(device->get());
}

private:
void set_context(CUcontext desired) {
Expand All @@ -212,6 +225,8 @@ class ScopedContext {
PI_CHECK_ERROR(cuCtxSetCurrent(desired));
}
}

pi_device device;
};

/// \cond NODOXY
Expand Down Expand Up @@ -1946,29 +1961,12 @@ pi_result cuda_piDeviceGetInfo(pi_device device, pi_device_info param_name,
}

case PI_EXT_INTEL_DEVICE_INFO_FREE_MEMORY: {
// Check the device of the currently set context uses the same device.
// CUDA_ERROR_INVALID_CONTEXT signifies the absence of an active context.
CUdevice current_ctx_device;
CUresult current_ctx_device_ret = cuCtxGetDevice(&current_ctx_device);
if (current_ctx_device_ret != CUDA_ERROR_INVALID_CONTEXT)
PI_CHECK_ERROR(current_ctx_device_ret);
bool need_primary_ctx =
current_ctx_device_ret == CUDA_ERROR_INVALID_CONTEXT ||
current_ctx_device != device->get();
if (need_primary_ctx) {
// Use the primary context for the device if no context with the device is
// set.
CUcontext primary_context;
PI_CHECK_ERROR(cuDevicePrimaryCtxRetain(&primary_context, device->get()));
PI_CHECK_ERROR(cuCtxSetCurrent(primary_context));
}
ScopedContext active(device);
size_t FreeMemory = 0;
size_t TotalMemory = 0;
sycl::detail::pi::assertion(cuMemGetInfo(&FreeMemory, &TotalMemory) ==
CUDA_SUCCESS,
"failed cuMemGetInfo() API.");
if (need_primary_ctx)
PI_CHECK_ERROR(cuDevicePrimaryCtxRelease(device->get()));
return getInfo(param_value_size, param_value, param_value_size_ret,
FreeMemory);
}
Expand Down