Split continue

Author: KFilipek

src/provider/provider_cuda.c

@@ -377,45 +377,46 @@ static umf_result_t cu_memory_provider_initialize(const void *params,
     cu_provider->context = cu_params->cuda_context_handle;
     cu_provider->device = cu_params->cuda_device_handle;
     cu_provider->memory_type = cu_params->memory_type;
-    cu_provider->min_alignment = min_alignment;
+    cu_provider->alloc_flags = cu_params->alloc_flags;
 
-    // If the memory type is shared (CUDA managed), the allocation flags must
-    // be set. NOTE: we do not check here if the flags are valid -
-    // this will be done by CUDA runtime.
-    if (cu_provider->memory_type == UMF_MEMORY_TYPE_SHARED &&
-        cu_params->alloc_flags == 0) {
-        // the default setting is CU_MEM_ATTACH_GLOBAL
-        cu_provider->alloc_flags = CU_MEM_ATTACH_GLOBAL;
-    } else {
-        cu_provider->alloc_flags = cu_params->alloc_flags;
-    }
+    *provider = cu_provider;
+
+    return UMF_RESULT_SUCCESS;
+}
+
+static umf_result_t cu_memory_provider_finalize(void *provider) {
+    umf_ba_global_free(provider);
+    return UMF_RESULT_SUCCESS;
+}
+
+static umf_result_t cu_memory_provider_post_initialize(void *provider) {
+    cu_memory_provider_t *cu_provider = (cu_memory_provider_t *)provider;
 
     // CUDA alloc functions doesn't allow to provide user alignment - get the
     // minimum one from the driver
     size_t min_alignment = 0;
     CUmemAllocationProp allocProps = {0};
     allocProps.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
     allocProps.type = CU_MEM_ALLOCATION_TYPE_PINNED;
-    allocProps.location.id = cu_params->cuda_device_handle;
+    allocProps.location.id = cu_provider->device;
     CUresult cu_result = g_cu_ops.cuMemGetAllocationGranularity(
         &min_alignment, &allocProps, CU_MEM_ALLOC_GRANULARITY_MINIMUM);
     if (cu_result != CUDA_SUCCESS) {
         umf_ba_global_free(cu_provider);
         return cu2umf_result(cu_result);
     }
 
-    *provider = cu_provider;
-
-    return UMF_RESULT_SUCCESS;
-}
 
-static umf_result_t cu_memory_provider_finalize(void *provider) {
-    umf_ba_global_free(provider);
-    return UMF_RESULT_SUCCESS;
-}
+    cu_provider->min_alignment = min_alignment;
 
-static umf_result_t cu_memory_provider_post_initialize(void *provider) {
-    (void)provider;
+    // If the memory type is shared (CUDA managed), the allocation flags must
+    // be set. NOTE: we do not check here if the flags are valid -
+    // this will be done by CUDA runtime.
+    if (cu_provider->memory_type == UMF_MEMORY_TYPE_SHARED &&
+        cu_provider->alloc_flags == 0) {
+        // the default setting is CU_MEM_ATTACH_GLOBAL
+        cu_provider->alloc_flags = CU_MEM_ATTACH_GLOBAL;
+    }
     // For initial version, just return success
     return UMF_RESULT_SUCCESS;
 }