Skip to content

Commit f5908e0

Browse files
bratpiorkabratpiorka
committed
x
1 parent 9283934 commit f5908e0

File tree

11 files changed

+102
-36
lines changed

11 files changed

+102
-36
lines changed

‎src/memory_pool.c‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -290,6 +290,7 @@ umf_result_t umfPoolByPtr(const void *ptr, umf_memory_pool_handle_t *pool) {
290290
umf_memory_properties_handle_t props = NULL;
291291
umf_result_t ret = umfGetMemoryPropertiesHandle(ptr, &props);
292292
if (ret != UMF_RESULT_SUCCESS || props == NULL || props->pool == NULL) {
293+
*pool = NULL;
293294
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
294295
}
295296

‎src/memory_props.c‎

Lines changed: 50 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -17,27 +17,39 @@
1717
#include "memory_provider_internal.h"
1818
#include "provider/provider_tracking.h"
1919

20+
// TODO flag - IPC?
2021
umf_result_t
2122
umfGetMemoryPropertiesHandle(const void *ptr,
2223
umf_memory_properties_handle_t *props_handle) {
2324

24-
LOG_DEBUG("umfGetMemoryPropertiesHandle: ptr=%p, props_handle=%p", ptr,
25-
props_handle);
25+
//LOG_DEBUG("umfGetMemoryPropertiesHandle: ptr=%p, props_handle=%p", ptr,
26+
// props_handle);
2627

27-
if (props_handle == NULL) {
28+
if (UNLIKELY(props_handle == NULL)) {
2829
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
2930
}
3031

3132
tracker_alloc_info_t *info = NULL;
3233
umf_result_t ret = umfMemoryTrackerGetAllocInfo(ptr, &info);
33-
if (ret != UMF_RESULT_SUCCESS) {
34-
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
34+
if (ret == UMF_RESULT_SUCCESS) {
35+
*props_handle = &info->props;
36+
//LOG_DEBUG("umfGetMemoryPropertiesHandle: props_handle=%p, id=%" PRIu64,
37+
// *props_handle, (*props_handle)->id);
38+
return UMF_RESULT_SUCCESS;
3539
}
3640

37-
*props_handle = &info->props;
41+
// try to get IPC info
42+
umf_ipc_info_t ipc_info;
43+
ret = umfMemoryTrackerGetIpcInfo(ptr, &ipc_info);
44+
if (ret != UMF_RESULT_SUCCESS) {
45+
LOG_ERR("Failed to get memory properties handle for ptr=%p", ptr);
46+
return ret;
47+
}
3848

39-
LOG_DEBUG("umfGetMemoryPropertiesHandle: props_handle=%p, id=%" PRIu64,
40-
*props_handle, (*props_handle)->id);
49+
*props_handle = ipc_info.props;
50+
//LOG_DEBUG(
51+
// "umfGetMemoryPropertiesHandle (IPC info): props_handle=%p, id=%" PRIu64,
52+
// *props_handle, (*props_handle)->id);
4153

4254
return UMF_RESULT_SUCCESS;
4355
}
@@ -46,62 +58,85 @@ umf_result_t umfGetMemoryProperty(umf_memory_properties_handle_t props_handle,
4658
umf_memory_property_id_t memory_property_id,
4759
size_t max_property_size, void *value) {
4860

61+
/*
4962
LOG_DEBUG("umfGetMemoryProperty: props_handle=%p, memory_property_id=%d, "
5063
"max_property_size=%zu, value=%p",
5164
props_handle, memory_property_id, max_property_size, value);
65+
*/
5266

53-
if ((value == NULL) || (props_handle == NULL) || (max_property_size == 0)) {
67+
if (UNLIKELY((value == NULL) || (props_handle == NULL) ||
68+
(max_property_size == 0))) {
5469
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
5570
}
5671

5772
umf_memory_provider_t *provider = props_handle->provider;
5873

74+
/*
5975
LOG_DEBUG("umfGetMemoryProperty: provider=%p", provider);
6076
LOG_DEBUG("dereferencing value...");
6177
6278
LOG_DEBUG("value: %zu", *(size_t *)value);
79+
*/
80+
81+
/*
82+
int x = (int)memory_property_id;
83+
for (int i = 0; i < 1000; i++) {
84+
x += (double)i / 10000000.0;
85+
memory_property_id = x;
86+
}
87+
*/
6388

6489
switch (memory_property_id) {
6590
case UMF_MEMORY_PROPERTY_INVALID:
6691
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
6792

6893
case UMF_MEMORY_PROPERTY_POOL_HANDLE:
69-
if (max_property_size < sizeof(umf_memory_pool_handle_t)) {
94+
if (UNLIKELY(max_property_size < sizeof(umf_memory_pool_handle_t))) {
7095
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
7196
}
7297
*(umf_memory_pool_handle_t *)value = props_handle->pool;
7398
return UMF_RESULT_SUCCESS;
7499

75100
case UMF_MEMORY_PROPERTY_PROVIDER_HANDLE:
76-
if (max_property_size < sizeof(umf_memory_provider_handle_t)) {
101+
if (UNLIKELY(max_property_size <
102+
sizeof(umf_memory_provider_handle_t))) {
77103
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
78104
}
79105
*(umf_memory_provider_handle_t *)value = provider;
80106
return UMF_RESULT_SUCCESS;
81107

82108
case UMF_MEMORY_PROPERTY_BUFFER_ID:
83-
if (max_property_size < sizeof(uint64_t)) {
109+
if (UNLIKELY(max_property_size < sizeof(uint64_t))) {
84110
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
85111
}
86112
*(uint64_t *)value = props_handle->id;
87113
return UMF_RESULT_SUCCESS;
88114

89115
case UMF_MEMORY_PROPERTY_BASE_ADDRESS:
90-
if (max_property_size < sizeof(uintptr_t)) {
116+
if (UNLIKELY(max_property_size < sizeof(uintptr_t))) {
91117
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
92118
}
93119
*(uintptr_t *)value = (uintptr_t)props_handle->base;
94120
return UMF_RESULT_SUCCESS;
95121

96122
case UMF_MEMORY_PROPERTY_BASE_SIZE:
97-
if (max_property_size < sizeof(size_t)) {
123+
if (UNLIKELY(max_property_size < sizeof(size_t))) {
98124
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
99125
}
100126
*(size_t *)value = props_handle->base_size;
101127
return UMF_RESULT_SUCCESS;
102128

103-
// GPU Memory Provider specific properties
104129
case UMF_MEMORY_PROPERTY_POINTER_TYPE:
130+
// NOTE: this property is "cached" in the props_handle but the value is
131+
// determined by the memory provider and set during addition to the
132+
// tracker.
133+
if (UNLIKELY(max_property_size < sizeof(umf_usm_memory_type_t))) {
134+
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
135+
}
136+
*(umf_usm_memory_type_t *)value = props_handle->memory_type;
137+
return UMF_RESULT_SUCCESS;
138+
139+
// GPU Memory Provider specific properties
105140
case UMF_MEMORY_PROPERTY_CONTEXT:
106141
case UMF_MEMORY_PROPERTY_DEVICE:
107142
return provider->ops.ext_get_allocation_properties(

‎src/memory_props_internal.h‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,27 +12,23 @@
1212

1313
#include <stdbool.h>
1414

15-
#include <umf/base.h>
1615
#include <umf/memory_pool.h>
1716
#include <umf/memory_props.h>
1817
#include <umf/memory_provider.h>
19-
20-
#if UMF_BUILD_LEVEL_ZERO_PROVIDER
21-
#include "ze_api.h"
22-
#endif
18+
#include <umf/memory_provider_gpu.h>
2319

2420
#ifdef __cplusplus
2521
extern "C" {
2622
#endif
2723

2824
typedef struct umf_memory_properties_t {
29-
// TODO alloc_info_t
3025
void *ptr;
3126
umf_memory_pool_handle_t pool;
3227
umf_memory_provider_handle_t provider;
3328
uint64_t id;
3429
void *base;
3530
size_t base_size;
31+
umf_usm_memory_type_t memory_type;
3632
} umf_memory_properties_t;
3733

3834
#ifdef __cplusplus

‎src/pool/pool_disjoint.c‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -898,16 +898,17 @@ umf_result_t disjoint_pool_malloc_usable_size(void *pool, const void *ptr,
898898
umf_memory_properties_handle_t props = NULL;
899899
umf_result_t umf_result = umfGetMemoryPropertiesHandle(ptr, &props);
900900
if (umf_result != UMF_RESULT_SUCCESS) {
901-
return 0;
901+
return umf_result;
902902
}
903903

904904
if (props == NULL) {
905905
TLS_last_allocation_error = UMF_RESULT_ERROR_UNKNOWN;
906906
LOG_ERR("failed to get allocation info from the memory tracker");
907-
return 0;
907+
return UMF_RESULT_ERROR_UNKNOWN;
908908
}
909909

910-
return props->base_size;
910+
*size = props->base_size;
911+
return UMF_RESULT_SUCCESS;
911912
}
912913

913914
// Get the unaligned pointer

‎src/provider/provider_cuda.c‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -699,21 +699,21 @@ static umf_result_t cu_memory_provider_get_allocation_properties(
699699

700700
switch (memory_property_id) {
701701
case UMF_MEMORY_PROPERTY_POINTER_TYPE:
702-
if (max_property_size < sizeof(umf_usm_memory_type_t)) {
702+
if (UNLIKELY(max_property_size < sizeof(umf_usm_memory_type_t))) {
703703
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
704704
}
705705
*(umf_usm_memory_type_t *)value = cuda_provider->memory_type;
706706
return UMF_RESULT_SUCCESS;
707707

708708
case UMF_MEMORY_PROPERTY_CONTEXT:
709-
if (max_property_size < sizeof(CUcontext)) {
709+
if (UNLIKELY(max_property_size < sizeof(CUcontext))) {
710710
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
711711
}
712712
*(CUcontext *)value = cuda_provider->context;
713713
return UMF_RESULT_SUCCESS;
714714

715715
case UMF_MEMORY_PROPERTY_DEVICE:
716-
if (max_property_size < sizeof(CUdevice)) {
716+
if (UNLIKELY(max_property_size < sizeof(CUdevice))) {
717717
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
718718
}
719719
*(CUdevice *)value = cuda_provider->device;

‎src/provider/provider_level_zero.c‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -805,22 +805,22 @@ static umf_result_t ze_memory_provider_get_allocation_properties(
805805

806806
switch (memory_property_id) {
807807
case UMF_MEMORY_PROPERTY_POINTER_TYPE:
808-
if (max_property_size < sizeof(umf_usm_memory_type_t)) {
808+
if (UNLIKELY(max_property_size < sizeof(umf_usm_memory_type_t))) {
809809
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
810810
}
811811
*(umf_usm_memory_type_t *)value =
812812
ze2umf_memory_type(ze_provider->memory_type);
813813
return UMF_RESULT_SUCCESS;
814814

815815
case UMF_MEMORY_PROPERTY_CONTEXT:
816-
if (max_property_size < sizeof(ze_context_handle_t)) {
816+
if (UNLIKELY(max_property_size < sizeof(ze_context_handle_t))) {
817817
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
818818
}
819819
*(ze_context_handle_t *)value = ze_provider->context;
820820
return UMF_RESULT_SUCCESS;
821821

822822
case UMF_MEMORY_PROPERTY_DEVICE:
823-
if (max_property_size < sizeof(ze_device_handle_t)) {
823+
if (UNLIKELY(max_property_size < sizeof(ze_device_handle_t))) {
824824
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
825825
}
826826
*(ze_device_handle_t *)value = ze_provider->device;

‎src/provider/provider_tracking.c‎

Lines changed: 28 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515

1616
#include <umf/memory_pool.h>
1717
#include <umf/memory_provider.h>
18+
#include <umf/memory_provider_gpu.h>
1819
#include <umf/memory_provider_ops.h>
1920

2021
#include "base_alloc_global.h"
@@ -23,6 +24,7 @@
2324
#include "ipc_internal.h"
2425
#include "memory_pool_internal.h"
2526
#include "memory_props_internal.h"
27+
#include "memory_provider_internal.h"
2628
#include "provider_tracking.h"
2729
#include "utils_common.h"
2830
#include "utils_concurrency.h"
@@ -47,6 +49,7 @@ struct umf_memory_tracker_t {
4749
};
4850

4951
typedef struct tracker_ipc_info_t {
52+
umf_memory_properties_t props;
5053
size_t size;
5154
umf_memory_provider_handle_t provider;
5255
ipc_opened_cache_value_t *ipc_cache_value;
@@ -184,13 +187,24 @@ umfMemoryTrackerAddAtLevel(umf_memory_tracker_handle_t hTracker, int level,
184187
return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
185188
}
186189

190+
// get provider
191+
umf_memory_provider_handle_t provider = NULL;
192+
umfPoolGetMemoryProvider(pool, &provider);
193+
194+
// get memory type
195+
umf_usm_memory_type_t memory_type = UMF_MEMORY_TYPE_UNKNOWN;
196+
provider->ops.ext_get_allocation_properties(
197+
provider->provider_priv, ptr, UMF_MEMORY_PROPERTY_POINTER_TYPE,
198+
sizeof(umf_usm_memory_type_t), &memory_type);
199+
187200
memset(&value->props, 0, sizeof(umf_memory_properties_t));
188-
umfPoolGetMemoryProvider(pool, &value->props.provider);
189201
value->props.id = utils_atomic_increment_u64(&unique_alloc_id);
190202
value->props.base = (void *)ptr;
191203
value->props.base_size = size;
192204
value->props.pool = pool;
205+
value->props.provider = provider;
193206
value->props.ptr = (void *)ptr;
207+
value->props.memory_type = memory_type;
194208

195209
value->n_children = 0;
196210
#if !defined(NDEBUG) && defined(UMF_DEVELOPER_MODE)
@@ -410,6 +424,14 @@ umfMemoryTrackerAddIpcSegment(umf_memory_tracker_handle_t hTracker,
410424
value->provider = provider;
411425
value->ipc_cache_value = cache_entry;
412426

427+
memset(&value->props, 0, sizeof(umf_memory_properties_t));
428+
value->props.id = utils_atomic_increment_u64(&unique_alloc_id);
429+
value->props.base = (void *)ptr;
430+
value->props.base_size = size;
431+
value->props.pool = NULL; // unknown
432+
value->props.provider = provider;
433+
value->props.ptr = (void *)ptr;
434+
413435
int ret =
414436
critnib_insert(hTracker->ipc_segments_map, (uintptr_t)ptr, value, 0);
415437
if (ret == 0) {
@@ -463,16 +485,16 @@ umf_result_t umfMemoryTrackerGetAllocInfo(const void *ptr,
463485
tracker_alloc_info_t **info) {
464486
assert(info);
465487

466-
if (ptr == NULL) {
488+
if (UNLIKELY(ptr == NULL)) {
467489
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
468490
}
469491

470-
if (TRACKER == NULL) {
492+
if (UNLIKELY(TRACKER == NULL)) {
471493
LOG_ERR("tracker does not exist");
472494
return UMF_RESULT_ERROR_NOT_SUPPORTED;
473495
}
474496

475-
if (TRACKER->alloc_segments_map[0] == NULL) {
497+
if (UNLIKELY(TRACKER->alloc_segments_map[0] == NULL)) {
476498
LOG_ERR("tracker's alloc_segments_map does not exist");
477499
return UMF_RESULT_ERROR_NOT_SUPPORTED;
478500
}
@@ -593,6 +615,8 @@ umf_result_t umfMemoryTrackerGetIpcInfo(const void *ptr,
593615
pIpcInfo->baseSize = rvalue->size;
594616
pIpcInfo->provider = rvalue->provider;
595617

618+
pIpcInfo->props = &rvalue->props;
619+
596620
if (ref_value) {
597621
critnib_release(TRACKER->ipc_segments_map, ref_value);
598622
}

‎src/provider/provider_tracking.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,8 @@ umf_result_t umfMemoryTrackerGetAllocInfo(const void *ptr,
5050
tracker_alloc_info_t **info);
5151

5252
typedef struct umf_ipc_info_t {
53+
umf_memory_properties_handle_t props;
54+
5355
void *base;
5456
size_t baseSize;
5557
umf_memory_provider_handle_t provider;

‎src/utils/utils_common.h‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,11 +57,17 @@ typedef enum umf_purge_advise_t {
5757

5858
#define __TLS __declspec(thread)
5959

60+
#define LIKELY(x) (x)
61+
#define UNLIKELY(x) (x)
62+
6063
#else /* Linux */
6164

6265
#define __TLS __thread
6366

64-
#endif /* _WIN32 */
67+
#define LIKELY(x) __builtin_expect(!!(x), 1)
68+
#define UNLIKELY(x) __builtin_expect(!!(x), 0)
69+
70+
#endif /* !_WIN32 */
6571

6672
// get the address of the given string in the environment variable (or NULL)
6773
char *utils_env_var_get_str(const char *envvar, const char *str);

‎test/props/provider_props.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ void createFixedProvider(umf_memory_provider_handle_t *out_provider,
1616

1717
umf_fixed_memory_provider_params_handle_t params = nullptr;
1818
umf_result_t res =
19-
umfFixedMemoryProviderParamsCreate(&params, memory_buffer, buffer_size);
19+
umfFixedMemoryProviderParamsCreate(memory_buffer, buffer_size, &params);
2020
ASSERT_EQ(res, UMF_RESULT_SUCCESS);
2121
ASSERT_NE(params, nullptr);
2222

0 commit comments

Comments
 (0)