use current ctx and dev by default in CUDA prov

Author: bratpiorka

include/umf/providers/provider_cuda.h

@@ -34,13 +34,17 @@ umf_result_t umfCUDAMemoryProviderParamsDestroy(
 /// @brief  Set the CUDA context handle in the parameters struct.
 /// @param  hParams handle to the parameters of the CUDA Memory Provider.
 /// @param  hContext handle to the CUDA context.
+/// @note   This function is optional - if no context is set in the params
+///         struct, the current context will be used.
 /// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
 umf_result_t umfCUDAMemoryProviderParamsSetContext(
     umf_cuda_memory_provider_params_handle_t hParams, void *hContext);
 
 /// @brief  Set the CUDA device handle in the parameters struct.
 /// @param  hParams handle to the parameters of the CUDA Memory Provider.
 /// @param  hDevice handle to the CUDA device.
+/// @note   This function is optional - if no device is set in the params
+///         struct, the current context's device will be used.
 /// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
 umf_result_t umfCUDAMemoryProviderParamsSetDevice(
     umf_cuda_memory_provider_params_handle_t hParams, int hDevice);

src/provider/provider_cuda.c

@@ -260,8 +260,23 @@ umf_result_t umfCUDAMemoryProviderParamsCreate(
         return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
     }
 
-    params_data->cuda_context_handle = NULL;
-    params_data->cuda_device_handle = -1;
+    // 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 = 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;
 

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);

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

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,41 @@ 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(&params);
+    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);
+
+    umfMemoryProviderDestroy(provider);
+    umfCUDAMemoryProviderParamsDestroy(defaultParams);
+}
+
 TEST_P(umfCUDAProviderTest, cudaProviderNullParams) {
     umf_result_t res = umfCUDAMemoryProviderParamsCreate(nullptr);
     EXPECT_EQ(res, UMF_RESULT_ERROR_INVALID_ARGUMENT);