[gcc/devel/omp/gcc-14] libgomp/plugin/plugin-nvptx.c: Fix device used for stream creation

[Tobias Burnus via Gcc-cvs]

https://gcc.gnu.org/g:ea561defcb718dc6f3f03237be661f0697ef3095

commit ea561defcb718dc6f3f03237be661f0697ef3095
Author: Tobias Burnus <tbur...@baylibre.com>
Date:   Mon Mar 24 16:08:20 2025 +0100

    libgomp/plugin/plugin-nvptx.c: Fix device used for stream creation
    
    libgomp/ChangeLog:
    
            * plugin/plugin-nvptx.c (GOMP_OFFLOAD_interop): Set context for
            stream creation to use the specified device.
    
    (cherry picked from commit 1c5a375c21a262eb636449f88e671a09e140404e)

Diff:
---
 libgomp/ChangeLog.omp         |  8 ++++++++
 libgomp/plugin/plugin-nvptx.c | 18 ++++++++++++++++--
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/libgomp/ChangeLog.omp b/libgomp/ChangeLog.omp
index 931c55777fda..d0a2d53881dd 100644
--- a/libgomp/ChangeLog.omp
+++ b/libgomp/ChangeLog.omp
@@ -1,3 +1,11 @@
+2025-03-24  Tobias Burnus  <tbur...@baylibre.com>
+
+       Backported from master:
+       2025-03-24  Tobias Burnus  <tbur...@baylibre.com>
+
+       * plugin/plugin-nvptx.c (GOMP_OFFLOAD_interop): Set context for
+       stream creation to use the specified device.
+
 2025-03-21  Tobias Burnus  <tbur...@baylibre.com>
 
        Backported from master:
diff --git a/libgomp/plugin/plugin-nvptx.c b/libgomp/plugin/plugin-nvptx.c
index f49fef397203..75552fcc8af0 100644
--- a/libgomp/plugin/plugin-nvptx.c
+++ b/libgomp/plugin/plugin-nvptx.c
@@ -2508,12 +2508,26 @@ GOMP_OFFLOAD_interop (struct interop_obj_t *obj, int 
ord,
          break;
       }
 
-  obj->device_data = ptx_devices[ord];
+  struct ptx_device *ptx_dev = obj->device_data = ptx_devices[ord];
 
   if (targetsync)
     {
       CUstream stream = NULL;
-      CUDA_CALL_ASSERT (cuStreamCreate, &stream, CU_STREAM_DEFAULT);
+      CUdevice cur_ctx_dev;
+      CUresult res = CUDA_CALL_NOCHECK (cuCtxGetDevice, &cur_ctx_dev);
+      if (res != CUDA_SUCCESS && res != CUDA_ERROR_INVALID_CONTEXT)
+       GOMP_PLUGIN_fatal ("cuCtxGetDevice error: %s", cuda_error (res));
+      if (res != CUDA_ERROR_INVALID_CONTEXT && ptx_dev->dev == cur_ctx_dev)
+       CUDA_CALL_ASSERT (cuStreamCreate, &stream, CU_STREAM_DEFAULT);
+      else
+       {
+         CUcontext old_ctx;
+         assert (ptx_dev->ctx);
+         CUDA_CALL_ASSERT (cuCtxPushCurrent, ptx_dev->ctx);
+         CUDA_CALL_ASSERT (cuStreamCreate, &stream, CU_STREAM_DEFAULT);
+         if (res != CUDA_ERROR_INVALID_CONTEXT)
+           CUDA_CALL_ASSERT (cuCtxPopCurrent, &old_ctx);
+       }
       obj->stream = stream;
     }
 }