Skip to content

Commit a42977f

Browse files
hppritchahppritcha
authored
Merge pull request #6707 from hoopoepg/topic/alloc-with-hint-realloc-inplace-v4.0
ALLOC_WITH_HINT: added inplace realloc - v4.0
2 parents bd602cc + 456c5b9 commit a42977f

File tree

4 files changed

+91
-13
lines changed

4 files changed

+91
-13
lines changed

‎oshmem/mca/spml/ucx/spml_ucx.c

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -596,7 +596,8 @@ static int mca_spml_ucx_ctx_create_common(long options, mca_spml_ucx_ctx_t **ucx
596596
{
597597
ucp_worker_params_t params;
598598
ucp_ep_params_t ep_params;
599-
size_t i, j, nprocs = oshmem_num_procs();
599+
size_t i, nprocs = oshmem_num_procs();
600+
int j;
600601
ucs_status_t err;
601602
spml_ucx_mkey_t *ucx_mkey;
602603
sshmem_mkey_t *mkey;

‎oshmem/mca/sshmem/ucx/sshmem_ucx.h

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,10 +49,19 @@ sshmem_ucx_shadow_allocator_t *sshmem_ucx_shadow_create(unsigned count);
4949
void sshmem_ucx_shadow_destroy(sshmem_ucx_shadow_allocator_t *allocator);
5050
int sshmem_ucx_shadow_alloc(sshmem_ucx_shadow_allocator_t *allocator,
5151
unsigned count, unsigned *index);
52+
53+
/* Reallocate existing allocated buffer. If possible - used inplace
54+
* reallocation.
55+
* Parameter 'inplace' - out, in case if zero - new buffer was allocated
56+
* (inplace is not possible), user should remove original buffer after data
57+
* is copied, else (if inplace == 0) - no additional action required */
58+
int sshmem_ucx_shadow_realloc(sshmem_ucx_shadow_allocator_t *allocator,
59+
unsigned count, unsigned old_index, unsigned *index,
60+
int *inplace);
5261
int sshmem_ucx_shadow_free(sshmem_ucx_shadow_allocator_t *allocator,
5362
unsigned index);
54-
size_t sshmem_ucx_shadow_size(sshmem_ucx_shadow_allocator_t *allocator,
55-
unsigned index);
63+
unsigned sshmem_ucx_shadow_size(sshmem_ucx_shadow_allocator_t *allocator,
64+
unsigned index);
5665

5766
END_C_DECLS
5867

‎oshmem/mca/sshmem/ucx/sshmem_ucx_module.c

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -189,7 +189,6 @@ static uct_ib_device_mem_h alloc_device_mem(mca_spml_ucx_t *spml, size_t size,
189189
uct_md_h uct_md;
190190
void *address;
191191
size_t length;
192-
int ret;
193192

194193
uct_md = ucp_context_find_tl_md(spml->ucp_context, "mlx5");
195194
if (uct_md == NULL) {
@@ -336,7 +335,7 @@ static unsigned sshmem_ucx_memheap_ptr2index(map_segment_t *s, void *ptr)
336335
return ((char*)ptr - (char*)s->super.va_base) / ALLOC_ELEM_SIZE;
337336
}
338337

339-
void sshmem_ucx_memheap_wordcopy(void *dst, void *src, size_t size)
338+
static void sshmem_ucx_memheap_wordcopy(void *dst, void *src, size_t size)
340339
{
341340
const size_t count = (size + sizeof(uint64_t) - 1) / sizeof(uint64_t);
342341
uint64_t *dst64 = (uint64_t*)dst;
@@ -353,8 +352,9 @@ static int sshmem_ucx_memheap_realloc(map_segment_t *s, size_t size,
353352
void* old_ptr, void** new_ptr)
354353
{
355354
mca_sshmem_ucx_segment_context_t *ctx = s->context;
356-
unsigned alloc_count, index;
355+
unsigned alloc_count, index, old_index, old_alloc_count;
357356
int res;
357+
int inplace;
358358

359359
if (size > s->seg_size) {
360360
return OSHMEM_ERR_OUT_OF_RESOURCE;
@@ -371,18 +371,24 @@ static int sshmem_ucx_memheap_realloc(map_segment_t *s, size_t size,
371371
/* Allocate new element. Zero-size allocation should still return a unique
372372
* pointer, so allocate 1 byte */
373373
alloc_count = max((size + ALLOC_ELEM_SIZE - 1) / ALLOC_ELEM_SIZE, 1);
374-
res = sshmem_ucx_shadow_alloc(ctx->shadow_allocator, alloc_count, &index);
374+
375+
if (!old_ptr) {
376+
res = sshmem_ucx_shadow_alloc(ctx->shadow_allocator, alloc_count, &index);
377+
} else {
378+
old_index = sshmem_ucx_memheap_ptr2index(s, old_ptr);
379+
res = sshmem_ucx_shadow_realloc(ctx->shadow_allocator, alloc_count,
380+
old_index, &index, &inplace);
381+
}
382+
375383
if (res != OSHMEM_SUCCESS) {
376384
return res;
377385
}
378386

379387
*new_ptr = sshmem_ucx_memheap_index2ptr(s, index);
380388

381389
/* Copy to new segment and release old*/
382-
if (old_ptr) {
383-
unsigned old_index = sshmem_ucx_memheap_ptr2index(s, old_ptr);
384-
unsigned old_alloc_count = sshmem_ucx_shadow_size(ctx->shadow_allocator,
385-
old_index);
390+
if (old_ptr && !inplace) {
391+
old_alloc_count = sshmem_ucx_shadow_size(ctx->shadow_allocator, old_index);
386392
sshmem_ucx_memheap_wordcopy(*new_ptr, old_ptr,
387393
min(size, old_alloc_count * ALLOC_ELEM_SIZE));
388394
sshmem_ucx_shadow_free(ctx->shadow_allocator, old_index);

‎oshmem/mca/sshmem/ucx/sshmem_ucx_shadow.c

Lines changed: 64 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,68 @@ static void sshmem_ucx_shadow_merge_blocks(sshmem_ucx_shadow_allocator_t *alloca
107107
}
108108
}
109109

110+
111+
112+
int sshmem_ucx_shadow_realloc(sshmem_ucx_shadow_allocator_t *allocator,
113+
unsigned count, unsigned old_index, unsigned *index,
114+
int *inplace)
115+
{
116+
sshmem_ucx_shadow_alloc_elem_t *elem = &allocator->elems[old_index];
117+
unsigned old_count = elem->block_size;
118+
sshmem_ucx_shadow_alloc_elem_t *end;
119+
sshmem_ucx_shadow_alloc_elem_t *next;
120+
121+
assert(count > 0);
122+
assert(!sshmem_ucx_shadow_is_free(elem));
123+
124+
*inplace = 1;
125+
126+
if (count == old_count) {
127+
*index = old_index;
128+
return OSHMEM_SUCCESS;
129+
}
130+
131+
if (count < old_count) {
132+
/* requested block is shorter than allocated block
133+
* then just cut current buffer */
134+
sshmem_ucx_shadow_set_elem(elem + count,
135+
SSHMEM_UCX_SHADOW_ELEM_FLAG_FREE,
136+
elem->block_size - count);
137+
elem->block_size = count;
138+
*index = old_index;
139+
sshmem_ucx_shadow_merge_blocks(allocator);
140+
return OSHMEM_SUCCESS;
141+
}
142+
143+
assert(count > old_count);
144+
145+
end = &allocator->elems[allocator->num_elems];
146+
next = &elem[old_count];
147+
/* try to check if next element is free & has enough length */
148+
if ((next < end) && /* non-last element? */
149+
sshmem_ucx_shadow_is_free(next) && /* next is free */
150+
(old_count + next->block_size >= count))
151+
{
152+
assert(elem < next);
153+
assert(elem + count > next);
154+
assert(elem + count <= end);
155+
assert(next + next->block_size <= end);
156+
157+
if (old_count + next->block_size > count) {
158+
sshmem_ucx_shadow_set_elem(elem + count, SSHMEM_UCX_SHADOW_ELEM_FLAG_FREE,
159+
old_count + next->block_size - count);
160+
}
161+
162+
sshmem_ucx_shadow_set_elem(next, 0, 0);
163+
elem->block_size = count;
164+
*index = old_index;
165+
return OSHMEM_SUCCESS;
166+
}
167+
168+
*inplace = 0;
169+
return sshmem_ucx_shadow_alloc(allocator, count, index);
170+
}
171+
110172
int sshmem_ucx_shadow_free(sshmem_ucx_shadow_allocator_t *allocator,
111173
unsigned index)
112174
{
@@ -117,8 +179,8 @@ int sshmem_ucx_shadow_free(sshmem_ucx_shadow_allocator_t *allocator,
117179
return OSHMEM_SUCCESS;
118180
}
119181

120-
size_t sshmem_ucx_shadow_size(sshmem_ucx_shadow_allocator_t *allocator,
121-
unsigned index)
182+
unsigned sshmem_ucx_shadow_size(sshmem_ucx_shadow_allocator_t *allocator,
183+
unsigned index)
122184
{
123185
sshmem_ucx_shadow_alloc_elem_t *elem = &allocator->elems[index];
124186

0 commit comments

Comments
 (0)