UCT/CUDA: Removed unsafe usage of cuCtxGetId. (#10852)

Author: rakhmets

src/uct/cuda/base/cuda_iface.c

@@ -236,35 +236,11 @@ ucs_status_t uct_cuda_base_iface_flush(uct_iface_h tl_iface, unsigned flags,
     return UCS_OK;
 }
 
-static int uct_cuda_base_is_ctx_rsc_valid(const uct_cuda_ctx_rsc_t *ctx_rsc)
+void uct_cuda_base_stream_destroy(CUstream *stream)
 {
-#if CUDA_VERSION >= 12000
-    unsigned long long ctx_id;
-    CUresult result;
-
-    result = uct_cuda_base_ctx_get_id(ctx_rsc->ctx, &ctx_id);
-    if (result == CUDA_ERROR_CONTEXT_IS_DESTROYED) {
-        return 0;
-    } else if (result != CUDA_SUCCESS) {
-        UCT_CUDADRV_LOG(cuCtxGetId, UCS_LOG_LEVEL_WARN, result);
-        return 0;
-    }
-
-    return ctx_id == ctx_rsc->ctx_id;
-#else
-    /* Best effort check on older Cuda versions */
-    return uct_cuda_base_is_context_valid(ctx_rsc->ctx);
-#endif
-}
-
-void uct_cuda_base_stream_destroy(const uct_cuda_ctx_rsc_t *ctx_rsc,
-                                  CUstream *stream)
-{
-    if ((*stream == NULL) || !uct_cuda_base_is_ctx_rsc_valid(ctx_rsc)) {
-        return;
+    if (*stream != NULL) {
+        (void)UCT_CUDADRV_FUNC_LOG_WARN(cuStreamDestroy(*stream));
     }
-
-    UCT_CUDADRV_FUNC_LOG_WARN(cuStreamDestroy(*stream));
 }
 
 static void
@@ -279,12 +255,8 @@ uct_cuda_base_event_desc_init(ucs_mpool_t *mp, void *obj, void *chunk)
 static void uct_cuda_base_event_desc_cleanup(ucs_mpool_t *mp, void *obj)
 {
     uct_cuda_event_desc_t *event_desc = obj;
-    uct_cuda_ctx_rsc_t *ctx_rsc       = ucs_container_of(mp, uct_cuda_ctx_rsc_t,
-                                                         event_mp);
 
-    if (uct_cuda_base_is_ctx_rsc_valid(ctx_rsc)) {
-        UCT_CUDADRV_FUNC_LOG_WARN(cuEventDestroy(event_desc->event));
-    }
+    (void)UCT_CUDADRV_FUNC_LOG_WARN(cuEventDestroy(event_desc->event));
 }
 
 void uct_cuda_base_queue_desc_init(uct_cuda_queue_desc_t *qdesc)
@@ -302,7 +274,7 @@ void uct_cuda_base_queue_desc_destroy(const uct_cuda_ctx_rsc_t *ctx_rsc,
                  ucs_queue_length(&qdesc->event_queue));
     }
 
-    uct_cuda_base_stream_destroy(ctx_rsc, &qdesc->stream);
+    uct_cuda_base_stream_destroy(&qdesc->stream);
 }
 
 static ucs_mpool_ops_t uct_cuda_event_desc_mpool_ops = {

src/uct/cuda/base/cuda_iface.h

@@ -57,39 +57,15 @@ static UCS_F_ALWAYS_INLINE int uct_cuda_base_is_context_active()
 }
 
 
-static UCS_F_ALWAYS_INLINE int uct_cuda_base_is_context_valid(CUcontext ctx)
-{
-    unsigned version;
-    ucs_status_t status;
-
-    /* Check if CUDA context is valid by running a dummy operation on it */
-    status = UCT_CUDADRV_FUNC_LOG_DEBUG(cuCtxGetApiVersion(ctx, &version));
-    return (status == UCS_OK);
-}
-
-
-static UCS_F_ALWAYS_INLINE int uct_cuda_base_context_match(CUcontext ctx1,
-                                                           CUcontext ctx2)
-{
-    return ((ctx1 != NULL) && (ctx1 == ctx2) &&
-            uct_cuda_base_is_context_valid(ctx1));
-}
-
-
 static UCS_F_ALWAYS_INLINE CUresult
 uct_cuda_base_ctx_get_id(CUcontext ctx, unsigned long long *ctx_id_p)
 {
-    unsigned long long ctx_id = 0;
-
 #if CUDA_VERSION >= 12000
-    CUresult result = cuCtxGetId(ctx, &ctx_id);
-    if (ucs_unlikely(result != CUDA_SUCCESS)) {
-        return result;
-    }
-#endif
-
-    *ctx_id_p = ctx_id;
+    return cuCtxGetId(ctx, ctx_id_p);
+#else
+    *ctx_id_p = 0;
     return CUDA_SUCCESS;
+#endif
 }
 
 
@@ -192,8 +168,7 @@ void uct_cuda_base_queue_desc_init(uct_cuda_queue_desc_t *qdesc);
 void uct_cuda_base_queue_desc_destroy(const uct_cuda_ctx_rsc_t *ctx_rsc,
                                       uct_cuda_queue_desc_t *qdesc);
 
-void uct_cuda_base_stream_destroy(const uct_cuda_ctx_rsc_t *ctx_rsc,
-                                  CUstream *stream);
+void uct_cuda_base_stream_destroy(CUstream *stream);
 
 #if (__CUDACC_VER_MAJOR__ >= 100000)
 void CUDA_CB uct_cuda_base_iface_stream_cb_fxn(void *arg);

src/uct/cuda/cuda_copy/cuda_copy_iface.c

@@ -313,7 +313,7 @@ static void uct_cuda_copy_ctx_rsc_destroy(uct_iface_h tl_iface,
         }
     }
 
-    uct_cuda_base_stream_destroy(cuda_ctx_rsc, &ctx_rsc->short_stream);
+    uct_cuda_base_stream_destroy(&ctx_rsc->short_stream);
     ucs_free(ctx_rsc);
 }