From e7e0165b81e7ed2f1dbfaf2ba9a7d7e3d923b0da Mon Sep 17 00:00:00 2001
From: Nicolas Miller <nicolas.miller@codeplay.com>
Date: Fri, 27 Jan 2023 10:55:57 +0000
Subject: [PATCH 1/2] [SYCL][CUDA] Fix context setup for device infos

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, 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.
---
 sycl/plugins/cuda/pi_cuda.cpp | 41 +++++++++++++++++------------------
 1 file changed, 20 insertions(+), 21 deletions(-)

diff --git a/sycl/plugins/cuda/pi_cuda.cpp b/sycl/plugins/cuda/pi_cuda.cpp
index b02462259ea9d..29ca5dc0b733f 100644
--- a/sycl/plugins/cuda/pi_cuda.cpp
+++ b/sycl/plugins/cuda/pi_cuda.cpp
@@ -188,7 +188,7 @@ 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;
     }
@@ -196,9 +196,22 @@ class ScopedContext {
     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
@@ -1168,6 +1183,8 @@ pi_result cuda_piDeviceGetInfo(pi_device device, pi_device_info param_name,
 
   assert(device != nullptr);
 
+  ScopedContext active(device);
+
   switch (param_name) {
   case PI_DEVICE_INFO_TYPE: {
     return getInfo(param_value_size, param_value, param_value_size_ret,
@@ -1946,29 +1963,11 @@ 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));
-    }
     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);
   }

From 1744d31e19333118fb339b2d42a6fe91879cbe4a Mon Sep 17 00:00:00 2001
From: Nicolas Miller <nicolas.miller@codeplay.com>
Date: Wed, 1 Feb 2023 10:37:06 +0000
Subject: [PATCH 2/2] [SYCL][CUDA] Only use scoped context for free memory

---
 sycl/plugins/cuda/pi_cuda.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/sycl/plugins/cuda/pi_cuda.cpp b/sycl/plugins/cuda/pi_cuda.cpp
index 29ca5dc0b733f..b7c64ef9f9a58 100644
--- a/sycl/plugins/cuda/pi_cuda.cpp
+++ b/sycl/plugins/cuda/pi_cuda.cpp
@@ -1183,8 +1183,6 @@ pi_result cuda_piDeviceGetInfo(pi_device device, pi_device_info param_name,
 
   assert(device != nullptr);
 
-  ScopedContext active(device);
-
   switch (param_name) {
   case PI_DEVICE_INFO_TYPE: {
     return getInfo(param_value_size, param_value, param_value_size_ret,
@@ -1963,6 +1961,7 @@ pi_result cuda_piDeviceGetInfo(pi_device device, pi_device_info param_name,
   }
 
   case PI_EXT_INTEL_DEVICE_INFO_FREE_MEMORY: {
+    ScopedContext active(device);
     size_t FreeMemory = 0;
     size_t TotalMemory = 0;
     sycl::detail::pi::assertion(cuMemGetInfo(&FreeMemory, &TotalMemory) ==