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

Extend the `ScopedContext` to work with just a device, in that case it
will simply use the primary context.

This is helpful for entry points that only have a `pi_device` and no
`pi_context` but that still need some cuda calls that require an active
context, such as for the device infos.

This addresses a bug where getting the amount of free memory before
creating any queues or context, would simply crash.

This was partially solved in a previous PR
(#7906), however the previous PR was
releasing the primary context, but leaving it active on the current
thread, so getting the device info twice in a row would end up crashing
again since it would just use the active but released primary context.

This should address: #8117

Author: npmiller

sycl/plugins/cuda/pi_cuda.cpp

@@ -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) {
@@ -212,6 +225,8 @@ class ScopedContext {
       PI_CHECK_ERROR(cuCtxSetCurrent(desired));
     }
   }
+
+  pi_device device;
 };
 
 /// \cond NODOXY
@@ -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);
   }