a

Author: bratpiorka

benchmark/ubench.c

@@ -454,8 +454,9 @@ static void do_umf_mem_props_benchmark(ze_context_handle_t context,
                 assert(res == UMF_RESULT_SUCCESS);
 
                 umf_usm_memory_type_t type = UMF_MEMORY_TYPE_UNKNOWN;
-                res = umfGetMemoryProperty(
-                    props_handle, UMF_MEMORY_PROPERTY_POINTER_TYPE, &type);
+                res = umfGetMemoryProperty(props_handle,
+                                           UMF_MEMORY_PROPERTY_POINTER_TYPE,
+                                           sizeof(type), &type);
                 assert(res == UMF_RESULT_SUCCESS);
                 if (type != UMF_MEMORY_TYPE_DEVICE) {
                     fprintf(stderr,

include/umf/memory_props.h

@@ -27,13 +27,15 @@ umfGetMemoryPropertiesHandle(const void *ptr,
 /// @brief Get a specific memory property from the properties handle
 /// @param props_handle handle to the memory properties
 /// @param memory_property_id ID of the memory property to get
-/// @param value [out] pointer to the value of the memory property which will
-///        be filled. NOTE: the type and size of the value depends on the
-///        memory property ID and should be checked in the documentation
+/// @param max_property_size size of the property value buffer
+/// @param property_value [out] pointer to the value of the memory property
+///         which will be filled
+// TODO check return type
 /// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure
 umf_result_t umfGetMemoryProperty(umf_memory_properties_handle_t props_handle,
                                   umf_memory_property_id_t memory_property_id,
-                                  void *value);
+                                  size_t max_property_size,
+                                  void *property_value);
 
 #ifdef __cplusplus
 }

include/umf/memory_provider.h

@@ -266,14 +266,18 @@ umfMemoryProviderAllocationMerge(umf_memory_provider_handle_t hProvider,
 /// @param hProvider pointer to the memory provider
 /// @param ptr pointer to the allocated memory
 /// @param propertyId identifier of the memory property to retrieve
-/// @param value [out] pointer to the preallocated variable where the value will be stored
+/// @param max_property_size size of the property value buffer
+/// @param property_value [out] pointer to the value of the memory property
+///         which will be filled
+// TODO check
 /// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure,
 ///         UMF_RESULT_ERROR_INVALID_ARGUMENT if propertyId is invalid or value is NULL,
 ///         UMF_RESULT_ERROR_NOT_SUPPORTED if the property is not supported by this provider.
 ///
 umf_result_t umfMemoryProviderGetAllocationProperties(
     umf_memory_provider_handle_t hProvider, const void *ptr,
-    umf_memory_property_id_t propertyId, void *value);
+    umf_memory_property_id_t propertyId, size_t max_property_size,
+    void *property_value);
 
 #ifdef __cplusplus
 }

include/umf/memory_provider_ops.h

@@ -274,14 +274,18 @@ typedef struct umf_memory_provider_ops_t {
     /// @param provider pointer to the memory provider
     /// @param ptr pointer to the allocated memory
     /// @param memory_property_id identifier of the memory property to retrieve
-    /// @param value [out] pointer to the preallocated variable where the value will be stored
+    /// @param max_property_size size of the property value buffer
+    /// @param property_value [out] pointer to the value of the memory property
+    ///         which will be filled
+    // TODO check
     /// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure,
     ///         UMF_RESULT_ERROR_INVALID_ARGUMENT if memory_property_id is invalid or value is NULL,
     ///         UMF_RESULT_ERROR_NOT_SUPPORTED if the property is not supported by this provider.
     ///
     umf_result_t (*ext_get_allocation_properties)(
         void *provider, const void *ptr,
-        umf_memory_property_id_t memory_property_id, void *value);
+        umf_memory_property_id_t memory_property_id, size_t max_property_size,
+        void *value);
 
 } umf_memory_provider_ops_t;
 

src/memory_props.c

@@ -7,6 +7,8 @@
  *
  */
 
+#include <inttypes.h>
+
 #include <umf/memory_props.h>
 #include <umf/providers/provider_cuda.h>
 #include <umf/providers/provider_level_zero.h>
@@ -19,6 +21,9 @@ umf_result_t
 umfGetMemoryPropertiesHandle(const void *ptr,
                              umf_memory_properties_handle_t *props_handle) {
 
+    LOG_DEBUG("umfGetMemoryPropertiesHandle: ptr=%p, props_handle=%p", ptr,
+              props_handle);
+
     if (props_handle == NULL) {
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
     }
@@ -31,39 +36,67 @@ umfGetMemoryPropertiesHandle(const void *ptr,
 
     *props_handle = &info->props;
 
+    LOG_DEBUG("umfGetMemoryPropertiesHandle: props_handle=%p, id=%" PRIu64,
+              *props_handle, (*props_handle)->id);
+
     return UMF_RESULT_SUCCESS;
 }
 
 umf_result_t umfGetMemoryProperty(umf_memory_properties_handle_t props_handle,
                                   umf_memory_property_id_t memory_property_id,
-                                  void *value) {
-    if ((value == NULL) || (props_handle == NULL)) {
+                                  size_t max_property_size, void *value) {
+
+    LOG_DEBUG("umfGetMemoryProperty: props_handle=%p, memory_property_id=%d, "
+              "max_property_size=%zu, value=%p",
+              props_handle, memory_property_id, max_property_size, value);
+
+    if ((value == NULL) || (props_handle == NULL) || (max_property_size == 0)) {
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
     }
 
     umf_memory_provider_t *provider = props_handle->provider;
 
+    LOG_DEBUG("umfGetMemoryProperty: provider=%p", provider);
+    LOG_DEBUG("dereferencing value...");
+
+    LOG_DEBUG("value: %zu", *(size_t *)value);
+
     switch (memory_property_id) {
     case UMF_MEMORY_PROPERTY_INVALID:
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
 
     case UMF_MEMORY_PROPERTY_POOL_HANDLE:
+        if (max_property_size < sizeof(umf_memory_pool_handle_t)) {
+            return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        }
         *(umf_memory_pool_handle_t *)value = props_handle->pool;
         return UMF_RESULT_SUCCESS;
 
     case UMF_MEMORY_PROPERTY_PROVIDER_HANDLE:
+        if (max_property_size < sizeof(umf_memory_provider_handle_t)) {
+            return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        }
         *(umf_memory_provider_handle_t *)value = provider;
         return UMF_RESULT_SUCCESS;
 
     case UMF_MEMORY_PROPERTY_BUFFER_ID:
+        if (max_property_size < sizeof(uint64_t)) {
+            return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        }
         *(uint64_t *)value = props_handle->id;
         return UMF_RESULT_SUCCESS;
 
     case UMF_MEMORY_PROPERTY_BASE_ADDRESS:
+        if (max_property_size < sizeof(uintptr_t)) {
+            return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        }
         *(uintptr_t *)value = (uintptr_t)props_handle->base;
         return UMF_RESULT_SUCCESS;
 
     case UMF_MEMORY_PROPERTY_BASE_SIZE:
+        if (max_property_size < sizeof(size_t)) {
+            return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        }
         *(size_t *)value = props_handle->base_size;
         return UMF_RESULT_SUCCESS;
 
@@ -73,7 +106,7 @@ umf_result_t umfGetMemoryProperty(umf_memory_properties_handle_t props_handle,
     case UMF_MEMORY_PROPERTY_DEVICE:
         return provider->ops.ext_get_allocation_properties(
             provider->provider_priv, props_handle->ptr, memory_property_id,
-            value);
+            max_property_size, value);
 
     default:
         break;

src/memory_provider.c

@@ -424,9 +424,11 @@ umfMemoryProviderCloseIPCHandle(umf_memory_provider_handle_t hProvider,
 
 umf_result_t umfMemoryProviderGetAllocationProperties(
     umf_memory_provider_handle_t hProvider, const void *ptr,
-    umf_memory_property_id_t propertyId, void *value) {
+    umf_memory_property_id_t propertyId, size_t max_property_size,
+    void *property_value) {
     UMF_CHECK((hProvider != NULL), UMF_RESULT_ERROR_INVALID_ARGUMENT);
-    UMF_CHECK((value != NULL), UMF_RESULT_ERROR_INVALID_ARGUMENT);
+    UMF_CHECK((property_value != NULL), UMF_RESULT_ERROR_INVALID_ARGUMENT);
     return hProvider->ops.ext_get_allocation_properties(
-        hProvider->provider_priv, ptr, propertyId, value);
+        hProvider->provider_priv, ptr, propertyId, max_property_size,
+        property_value);
 }

src/provider/provider_cuda.c

@@ -683,7 +683,8 @@ cu_memory_provider_close_ipc_handle(void *provider, void *ptr, size_t size) {
 
 static umf_result_t cu_memory_provider_get_allocation_properties(
     void *provider, const void *ptr,
-    umf_memory_property_id_t memory_property_id, void *value) {
+    umf_memory_property_id_t memory_property_id, size_t max_property_size,
+    void *value) {
 
     // unused
     (void)ptr;
@@ -692,14 +693,23 @@ static umf_result_t cu_memory_provider_get_allocation_properties(
 
     switch (memory_property_id) {
     case UMF_MEMORY_PROPERTY_POINTER_TYPE:
+        if (max_property_size < sizeof(umf_usm_memory_type_t)) {
+            return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        }
         *(umf_usm_memory_type_t *)value = cuda_provider->memory_type;
         return UMF_RESULT_SUCCESS;
 
     case UMF_MEMORY_PROPERTY_CONTEXT:
+        if (max_property_size < sizeof(CUcontext)) {
+            return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        }
         *(CUcontext *)value = cuda_provider->context;
         return UMF_RESULT_SUCCESS;
 
     case UMF_MEMORY_PROPERTY_DEVICE:
+        if (max_property_size < sizeof(CUdevice)) {
+            return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        }
         *(CUdevice *)value = cuda_provider->device;
         return UMF_RESULT_SUCCESS;
 

src/provider/provider_level_zero.c

@@ -793,23 +793,33 @@ ze_memory_provider_close_ipc_handle(void *provider, void *ptr, size_t size) {
 
 static umf_result_t ze_memory_provider_get_allocation_properties(
     void *provider, const void *ptr,
-    umf_memory_property_id_t memory_property_id, void *value) {
+    umf_memory_property_id_t memory_property_id, size_t max_property_size,
+    void *value) {
     (void)ptr;
 
     struct ze_memory_provider_t *ze_provider =
         (struct ze_memory_provider_t *)provider;
 
     switch (memory_property_id) {
     case UMF_MEMORY_PROPERTY_POINTER_TYPE:
+        if (max_property_size < sizeof(umf_usm_memory_type_t)) {
+            return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        }
         *(umf_usm_memory_type_t *)value =
             ze2umf_memory_type(ze_provider->memory_type);
         return UMF_RESULT_SUCCESS;
 
     case UMF_MEMORY_PROPERTY_CONTEXT:
+        if (max_property_size < sizeof(ze_context_handle_t)) {
+            return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        }
         *(ze_context_handle_t *)value = ze_provider->context;
         return UMF_RESULT_SUCCESS;
 
     case UMF_MEMORY_PROPERTY_DEVICE:
+        if (max_property_size < sizeof(ze_device_handle_t)) {
+            return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        }
         *(ze_device_handle_t *)value = ze_provider->device;
         return UMF_RESULT_SUCCESS;
 

src/provider/provider_tracking.c

@@ -1291,11 +1291,11 @@ static umf_result_t trackingCloseIpcHandle(void *provider, void *ptr,
 static umf_result_t
 trackingGetAllocationProperties(void *provider, const void *ptr,
                                 umf_memory_property_id_t memory_property_id,
-                                void *value) {
+                                size_t max_property_size, void *value) {
     umf_tracking_memory_provider_t *p =
         (umf_tracking_memory_provider_t *)provider;
-    return umfMemoryProviderGetAllocationProperties(p->hUpstream, ptr,
-                                                    memory_property_id, value);
+    return umfMemoryProviderGetAllocationProperties(
+        p->hUpstream, ptr, memory_property_id, max_property_size, value);
 }
 
 umf_memory_provider_ops_t UMF_TRACKING_MEMORY_PROVIDER_OPS = {

test/common/provider_null.c

@@ -135,10 +135,12 @@ static umf_result_t nullCloseIpcHandle(void *provider, void *ptr, size_t size) {
 
 static umf_result_t
 nullGetAllocationProperties(void *provider, const void *ptr,
-                            umf_memory_property_id_t propertyId, void *value) {
+                            umf_memory_property_id_t propertyId,
+                            size_t max_property_size, void *value) {
     (void)provider;
     (void)ptr;
     (void)propertyId;
+    (void)max_property_size;
     (void)value;
     return UMF_RESULT_SUCCESS;
 }