add CUDA multi context test

Author: bratpiorka

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
+    cuCtxSetCurrent(ctx2);
+    // free memory from ctx1
+    umf_result = umfMemoryProviderFree(provider1, ptr1, size);
+    ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
+
+    cuCtxSetCurrent(ctx1);
+    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;