UCT/IB/MLX5/GDAKI: Fixed retaining inactive primary context.

Author: rakhmets

src/uct/ib/mlx5/gdaki/gdaki.c

@@ -763,17 +763,11 @@ static UCS_CLASS_INIT_FUNC(uct_rc_gdaki_iface_t, uct_md_h tl_md,
         return status;
     }
 
-    status = UCT_CUDADRV_FUNC_LOG_ERR(
-            cuDevicePrimaryCtxRetain(&self->cuda_ctx, self->cuda_dev));
+    status = uct_cuda_ctx_primary_push(self->cuda_dev, 0, UCS_LOG_LEVEL_ERROR);
     if (status != UCS_OK) {
         return status;
     }
 
-    status = UCT_CUDADRV_FUNC_LOG_ERR(cuCtxPushCurrent(self->cuda_ctx));
-    if (status != UCS_OK) {
-        goto err_ctx_release;
-    }
-
     status = uct_rc_gdaki_alloc(sizeof(uint64_t), sizeof(uint64_t),
                                 (void**)&self->atomic_buff, &self->atomic_raw);
     if (status != UCS_OK) {
@@ -797,19 +791,17 @@ static UCS_CLASS_INIT_FUNC(uct_rc_gdaki_iface_t, uct_md_h tl_md,
 err_lock:
     ibv_dereg_mr(self->atomic_mr);
 err_atomic:
-    cuMemFree(self->atomic_raw);
+    (void)UCT_CUDADRV_FUNC_LOG_WARN(cuMemFree(self->atomic_raw));
 err_ctx:
-    (void)UCT_CUDADRV_FUNC_LOG_WARN(cuCtxPopCurrent(NULL));
-err_ctx_release:
-    (void)UCT_CUDADRV_FUNC_LOG_WARN(cuDevicePrimaryCtxRelease(self->cuda_dev));
+    uct_cuda_ctx_primary_pop_and_release(self->cuda_dev);
     return status;
 }
 
 static UCS_CLASS_CLEANUP_FUNC(uct_rc_gdaki_iface_t)
 {
     pthread_mutex_destroy(&self->ep_init_lock);
     ibv_dereg_mr(self->atomic_mr);
-    cuMemFree(self->atomic_raw);
+    (void)UCT_CUDADRV_FUNC_LOG_WARN(cuMemFree(self->atomic_raw));
     (void)UCT_CUDADRV_FUNC_LOG_WARN(cuDevicePrimaryCtxRelease(self->cuda_dev));
 }