add CUDA multi context test

Author: bratpiorka

test/providers/cuda_helpers.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2024 Intel Corporation
+ * Copyright (C) 2024-2025 Intel Corporation
  *
  * Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
  * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
@@ -251,15 +251,18 @@ int InitCUDAOps() {
 }
 #endif // USE_DLOPEN
 
-static CUresult set_context(CUcontext required_ctx, CUcontext *restore_ctx) {
+CUresult set_context(CUcontext required_ctx, CUcontext *restore_ctx) {
     CUcontext current_ctx = NULL;
     CUresult cu_result = libcu_ops.cuCtxGetCurrent(&current_ctx);
     if (cu_result != CUDA_SUCCESS) {
         fprintf(stderr, "cuCtxGetCurrent() failed.\n");
         return cu_result;
     }
 
-    *restore_ctx = current_ctx;
+    if (restore_ctx != NULL) {
+        *restore_ctx = current_ctx;
+    }
+
     if (current_ctx != required_ctx) {
         cu_result = libcu_ops.cuCtxSetCurrent(required_ctx);
         if (cu_result != CUDA_SUCCESS) {

test/providers/cuda_helpers.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2024 Intel Corporation
+ * Copyright (C) 2024-2025 Intel Corporation
  *
  * Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
  * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
@@ -30,6 +30,8 @@ int get_cuda_device(CUdevice *device);
 
 int create_context(CUdevice device, CUcontext *context);
 
+CUresult set_context(CUcontext required_ctx, CUcontext *restore_ctx);
+
 int destroy_context(CUcontext context);
 
 int cuda_fill(CUcontext context, CUdevice device, void *ptr, size_t size,

test/providers/provider_cuda.cpp

@@ -1,4 +1,4 @@
-// Copyright (C) 2024 Intel Corporation
+// Copyright (C) 2024-2025 Intel Corporation
 // Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
@@ -315,6 +315,69 @@ TEST_P(umfCUDAProviderTest, cudaProviderNullParams) {
     EXPECT_EQ(res, UMF_RESULT_ERROR_INVALID_ARGUMENT);
 }
 
+TEST_P(umfCUDAProviderTest, multiContext) {
+    CUdevice device;
+    int ret = get_cuda_device(&device);
+    ASSERT_EQ(ret, 0);
+
+    // create two CUDA contexts and two providers
+    CUcontext ctx1, ctx2;
+    ret = create_context(device, &ctx1);
+    ASSERT_EQ(ret, 0);
+    ret = create_context(device, &ctx2);
+    ASSERT_EQ(ret, 0);
+
+    cuda_params_unique_handle_t params1 =
+        create_cuda_prov_params(ctx1, device, UMF_MEMORY_TYPE_HOST);
+    ASSERT_NE(params1, nullptr);
+    umf_memory_provider_handle_t provider1;
+    umf_result_t umf_result = umfMemoryProviderCreate(
+        umfCUDAMemoryProviderOps(), params1.get(), &provider1);
+    ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
+    ASSERT_NE(provider1, nullptr);
+
+    cuda_params_unique_handle_t params2 =
+        create_cuda_prov_params(ctx2, device, UMF_MEMORY_TYPE_HOST);
+    ASSERT_NE(params2, nullptr);
+    umf_memory_provider_handle_t provider2;
+    umf_result = umfMemoryProviderCreate(umfCUDAMemoryProviderOps(),
+                                         params2.get(), &provider2);
+    ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
+    ASSERT_NE(provider2, nullptr);
+
+    // use the providers
+    // allocate from 1, then from 2, then free 1, then free 2
+    void *ptr1, *ptr2;
+    const int size = 128;
+    // NOTE: we use ctx1 here
+    umf_result = umfMemoryProviderAlloc(provider1, size, 0, &ptr1);
+    ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
+    ASSERT_NE(ptr1, nullptr);
+
+    // NOTE: we use ctx2 here
+    umf_result = umfMemoryProviderAlloc(provider2, size, 0, &ptr2);
+    ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
+    ASSERT_NE(ptr2, nullptr);
+
+    // even if we change the context, we should be able to free the memory
+    set_context(ctx2, NULL);
+    // free memory from ctx1
+    umf_result = umfMemoryProviderFree(provider1, ptr1, size);
+    ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
+
+    set_context(ctx1, NULL);
+    umf_result = umfMemoryProviderFree(provider2, ptr2, size);
+    ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
+
+    // cleanup
+    umfMemoryProviderDestroy(provider2);
+    umfMemoryProviderDestroy(provider1);
+    ret = destroy_context(ctx1);
+    ASSERT_EQ(ret, 0);
+    ret = destroy_context(ctx2);
+    ASSERT_EQ(ret, 0);
+}
+
 // TODO add tests that mixes CUDA Memory Provider and Disjoint Pool
 
 CUDATestHelper cudaTestHelper;