diff --git a/src/provider/provider_cuda.c b/src/provider/provider_cuda.c index 350bd016f..7a7b0a467 100644 --- a/src/provider/provider_cuda.c +++ b/src/provider/provider_cuda.c @@ -433,6 +433,14 @@ static umf_result_t cu_memory_provider_free(void *provider, void *ptr, cu_memory_provider_t *cu_provider = (cu_memory_provider_t *)provider; + // Remember current context and set the one from the provider + CUcontext restore_ctx = NULL; + umf_result_t umf_result = set_context(cu_provider->context, &restore_ctx); + if (umf_result != UMF_RESULT_SUCCESS) { + LOG_ERR("Failed to set CUDA context, ret = %d", umf_result); + return umf_result; + } + CUresult cu_result = CUDA_SUCCESS; switch (cu_provider->memory_type) { case UMF_MEMORY_TYPE_HOST: { @@ -451,6 +459,11 @@ static umf_result_t cu_memory_provider_free(void *provider, void *ptr, return UMF_RESULT_ERROR_UNKNOWN; } + umf_result = set_context(restore_ctx, &restore_ctx); + if (umf_result != UMF_RESULT_SUCCESS) { + LOG_ERR("Failed to restore CUDA context, ret = %d", umf_result); + } + return cu2umf_result(cu_result); } diff --git a/test/providers/cuda_helpers.cpp b/test/providers/cuda_helpers.cpp index 9c41d9382..bed9906c0 100644 --- a/test/providers/cuda_helpers.cpp +++ b/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,7 +251,7 @@ 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(¤t_ctx); if (cu_result != CUDA_SUCCESS) { @@ -259,7 +259,10 @@ static CUresult set_context(CUcontext required_ctx, CUcontext *restore_ctx) { 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) { diff --git a/test/providers/cuda_helpers.h b/test/providers/cuda_helpers.h index fc06c1fcf..65f4fbbf5 100644 --- a/test/providers/cuda_helpers.h +++ b/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, diff --git a/test/providers/provider_cuda.cpp b/test/providers/provider_cuda.cpp index 4f1d35911..8a7fdd28a 100644 --- a/test/providers/provider_cuda.cpp +++ b/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,71 @@ 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 + ret = set_context(ctx2, NULL); + ASSERT_EQ(ret, 0); + // free memory from ctx1 + umf_result = umfMemoryProviderFree(provider1, ptr1, size); + ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS); + + ret = set_context(ctx1, NULL); + ASSERT_EQ(ret, 0); + 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;