use current ctx and dev by default in CUDA prov

bratpiorka · bratpiorka · commit 3a08d6c76516 · 2025-03-12T17:18:56.000Z
diff --git a/include/umf/providers/provider_cuda.h b/include/umf/providers/provider_cuda.h
@@ -20,7 +20,8 @@ typedef struct umf_cuda_memory_provider_params_t
     *umf_cuda_memory_provider_params_handle_t;
 
 /// @brief  Create a struct to store parameters of the CUDA Memory Provider.
-/// @param  hParams [out] handle to the newly created parameters struct.
+/// @param  hParams [out] handle to the newly created parameters structure,
+///         initialized with the default (current) context and device ID.
 /// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
 umf_result_t umfCUDAMemoryProviderParamsCreate(
     umf_cuda_memory_provider_params_handle_t *hParams);
diff --git a/src/provider/provider_cuda.c b/src/provider/provider_cuda.c
@@ -139,6 +139,7 @@ typedef struct cu_ops_t {
     CUresult (*cuGetErrorName)(CUresult error, const char **pStr);
     CUresult (*cuGetErrorString)(CUresult error, const char **pStr);
     CUresult (*cuCtxGetCurrent)(CUcontext *pctx);
+    CUresult (*cuCtxGetDevice)(CUdevice *device);
     CUresult (*cuCtxSetCurrent)(CUcontext ctx);
     CUresult (*cuIpcGetMemHandle)(CUipcMemHandle *pHandle, CUdeviceptr dptr);
     CUresult (*cuIpcOpenMemHandle)(CUdeviceptr *pdptr, CUipcMemHandle handle,
@@ -221,6 +222,8 @@ static void init_cu_global_state(void) {
         utils_get_symbol_addr(lib_handle, "cuGetErrorString", lib_name);
     *(void **)&g_cu_ops.cuCtxGetCurrent =
         utils_get_symbol_addr(lib_handle, "cuCtxGetCurrent", lib_name);
+    *(void **)&g_cu_ops.cuCtxGetDevice =
+        utils_get_symbol_addr(lib_handle, "cuCtxGetDevice", lib_name);
     *(void **)&g_cu_ops.cuCtxSetCurrent =
         utils_get_symbol_addr(lib_handle, "cuCtxSetCurrent", lib_name);
     *(void **)&g_cu_ops.cuIpcGetMemHandle =
@@ -234,9 +237,9 @@ static void init_cu_global_state(void) {
         !g_cu_ops.cuMemHostAlloc || !g_cu_ops.cuMemAllocManaged ||
         !g_cu_ops.cuMemFree || !g_cu_ops.cuMemFreeHost ||
         !g_cu_ops.cuGetErrorName || !g_cu_ops.cuGetErrorString ||
-        !g_cu_ops.cuCtxGetCurrent || !g_cu_ops.cuCtxSetCurrent ||
-        !g_cu_ops.cuIpcGetMemHandle || !g_cu_ops.cuIpcOpenMemHandle ||
-        !g_cu_ops.cuIpcCloseMemHandle) {
+        !g_cu_ops.cuCtxGetCurrent || !g_cu_ops.cuCtxGetDevice ||
+        !g_cu_ops.cuCtxSetCurrent || !g_cu_ops.cuIpcGetMemHandle ||
+        !g_cu_ops.cuIpcOpenMemHandle || !g_cu_ops.cuIpcCloseMemHandle) {
         LOG_FATAL("Required CUDA symbols not found.");
         Init_cu_global_state_failed = true;
         utils_close_library(lib_handle);
@@ -260,8 +263,29 @@ umf_result_t umfCUDAMemoryProviderParamsCreate(
         return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
     }
 
-    params_data->cuda_context_handle = NULL;
-    params_data->cuda_device_handle = -1;
+    utils_init_once(&cu_is_initialized, init_cu_global_state);
+    if (Init_cu_global_state_failed) {
+        LOG_FATAL("Loading CUDA symbols failed");
+        return UMF_RESULT_ERROR_DEPENDENCY_UNAVAILABLE;
+    }
+
+    // initialize context and device to the current ones
+    CUcontext current_ctx = NULL;
+    CUresult cu_result = g_cu_ops.cuCtxGetCurrent(&current_ctx);
+    if (cu_result == CUDA_SUCCESS) {
+        params_data->cuda_context_handle = current_ctx;
+    } else {
+        params_data->cuda_context_handle = NULL;
+    }
+
+    CUdevice current_device = -1;
+    cu_result = g_cu_ops.cuCtxGetDevice(&current_device);
+    if (cu_result == CUDA_SUCCESS) {
+        params_data->cuda_device_handle = current_device;
+    } else {
+        params_data->cuda_device_handle = -1;
+    }
+
     params_data->memory_type = UMF_MEMORY_TYPE_UNKNOWN;
     params_data->alloc_flags = 0;
 
@@ -342,6 +366,12 @@ static umf_result_t cu_memory_provider_initialize(void *params,
     }
 
     if (cu_params->cuda_context_handle == NULL) {
+        LOG_ERR("Invalid context handle");
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    if (cu_params->cuda_device_handle < 0) {
+        LOG_ERR("Invalid device handle");
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
     }
 
diff --git a/test/providers/cuda_helpers.cpp b/test/providers/cuda_helpers.cpp
@@ -412,6 +412,18 @@ CUcontext get_mem_context(void *ptr) {
     return context;
 }
 
+int get_mem_device(void *ptr) {
+    int device;
+    CUresult res = libcu_ops.cuPointerGetAttribute(
+        &device, CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL, (CUdeviceptr)ptr);
+    if (res != CUDA_SUCCESS) {
+        fprintf(stderr, "cuPointerGetAttribute() failed!\n");
+        return -1;
+    }
+
+    return device;
+}
+
 CUcontext get_current_context() {
     CUcontext context;
     CUresult res = libcu_ops.cuCtxGetCurrent(&context);
diff --git a/test/providers/cuda_helpers.h b/test/providers/cuda_helpers.h
@@ -48,6 +48,8 @@ unsigned int get_mem_host_alloc_flags(void *ptr);
 
 CUcontext get_mem_context(void *ptr);
 
+int get_mem_device(void *ptr);
+
 CUcontext get_current_context();
 
 #ifdef __cplusplus
diff --git a/test/providers/provider_cuda.cpp b/test/providers/provider_cuda.cpp
@@ -142,14 +142,15 @@ struct umfCUDAProviderTest
 
         memAccessor = nullptr;
         expected_context = cudaTestHelper.get_test_context();
+        expected_device = cudaTestHelper.get_test_device();
         params = create_cuda_prov_params(cudaTestHelper.get_test_context(),
                                          cudaTestHelper.get_test_device(),
                                          memory_type, 0 /* alloc flags */);
         ASSERT_NE(expected_context, nullptr);
+        ASSERT_GE(expected_device, 0);
 
         switch (memory_type) {
         case UMF_MEMORY_TYPE_DEVICE:
-
             memAccessor = std::make_unique<CUDAMemoryAccessor>(
                 cudaTestHelper.get_test_context(),
                 cudaTestHelper.get_test_device());
@@ -178,6 +179,7 @@ struct umfCUDAProviderTest
 
     std::unique_ptr<MemoryAccessor> memAccessor = nullptr;
     CUcontext expected_context = nullptr;
+    int expected_device = -1;
     umf_usm_memory_type_t expected_memory_type;
 };
 
@@ -328,6 +330,44 @@ TEST_P(umfCUDAProviderTest, getPageSizeInvalidArgs) {
     umfMemoryProviderDestroy(provider);
 }
 
+TEST_P(umfCUDAProviderTest, cudaProviderDefaultParams) {
+    umf_cuda_memory_provider_params_handle_t defaultParams = nullptr;
+    umf_result_t umf_result = umfCUDAMemoryProviderParamsCreate(&defaultParams);
+    ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
+
+    umf_result = umfCUDAMemoryProviderParamsSetMemoryType(defaultParams,
+                                                          expected_memory_type);
+    ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
+
+    // NOTE: we intentionally do not set any context and device params
+
+    umf_memory_provider_handle_t provider = nullptr;
+    umf_result = umfMemoryProviderCreate(umfCUDAMemoryProviderOps(),
+                                         defaultParams, &provider);
+    ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
+    ASSERT_NE(provider, nullptr);
+
+    // do single alloc and check if the context and device id of allocated
+    // memory are correct
+
+    void *ptr = nullptr;
+    umf_result = umfMemoryProviderAlloc(provider, 128, 0, &ptr);
+    ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
+    ASSERT_NE(ptr, nullptr);
+
+    CUcontext actual_mem_context = get_mem_context(ptr);
+    ASSERT_EQ(actual_mem_context, expected_context);
+
+    int actual_device = get_mem_device(ptr);
+    ASSERT_EQ(actual_device, expected_device);
+
+    umf_result = umfMemoryProviderFree(provider, ptr, 128);
+    ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
+
+    umfMemoryProviderDestroy(provider);
+    umfCUDAMemoryProviderParamsDestroy(defaultParams);
+}
+
 TEST_P(umfCUDAProviderTest, cudaProviderNullParams) {
     umf_result_t res = umfCUDAMemoryProviderParamsCreate(nullptr);
     EXPECT_EQ(res, UMF_RESULT_ERROR_INVALID_ARGUMENT);
