diff --git a/oshmem/mca/memheap/buddy/memheap_buddy.c b/oshmem/mca/memheap/buddy/memheap_buddy.c index c7715f4dd1d..c6f0fdc5642 100644 --- a/oshmem/mca/memheap/buddy/memheap_buddy.c +++ b/oshmem/mca/memheap/buddy/memheap_buddy.c @@ -469,7 +469,7 @@ static int _do_alloc(uint32_t order, } *p_buff = (void*) addr; - /* no barrier because it is not required by spec! */ + MCA_SPML_CALL(memuse_hook(*p_buff, 1ULL<num_transports = 1; @@ -371,7 +371,7 @@ void mca_spml_ucx_rmkey_unpack(sshmem_mkey_t *mkey, uint32_t segno, int pe, int mkey->u.data, &ucx_mkey->rkey); if (UCS_OK != err) { - SPML_ERROR("failed to unpack rkey"); + SPML_ERROR("failed to unpack rkey: %s", ucs_status_string(err)); goto error_fatal; } @@ -384,6 +384,35 @@ void mca_spml_ucx_rmkey_unpack(sshmem_mkey_t *mkey, uint32_t segno, int pe, int return; } +void mca_spml_ucx_memuse_hook(void *addr, size_t length) +{ + int my_pe; + spml_ucx_mkey_t *ucx_mkey; + ucp_mem_advise_params_t params; + ucs_status_t status; + + if (!(mca_spml_ucx.heap_reg_nb && memheap_is_va_in_segment(addr, HEAP_SEG_INDEX))) { + return; + } + + my_pe = oshmem_my_proc_id(); + ucx_mkey = &mca_spml_ucx.ucp_peers[my_pe].mkeys[HEAP_SEG_INDEX].key; + + params.field_mask = UCP_MEM_ADVISE_PARAM_FIELD_ADDRESS | + UCP_MEM_ADVISE_PARAM_FIELD_LENGTH | + UCP_MEM_ADVISE_PARAM_FIELD_ADVICE; + + params.address = addr; + params.length = length; + params.advice = UCP_MADV_WILLNEED; + + status = ucp_mem_advise(mca_spml_ucx.ucp_context, ucx_mkey->mem_h, ¶ms); + if (UCS_OK != status) { + SPML_ERROR("ucp_mem_advise failed addr %p len %llu : %s", + addr, (unsigned long long)length, ucs_status_string(status)); + } +} + sshmem_mkey_t *mca_spml_ucx_register(void* addr, size_t size, uint64_t shmid, @@ -539,7 +568,7 @@ int mca_spml_ucx_fence(void) err = ucp_worker_flush(mca_spml_ucx.ucp_worker); if (UCS_OK != err) { - SPML_ERROR("fence failed"); + SPML_ERROR("fence failed: %s", ucs_status_string(err)); oshmem_shmem_abort(-1); return OSHMEM_ERROR; } @@ -552,7 +581,7 @@ int mca_spml_ucx_quiet(void) err = ucp_worker_flush(mca_spml_ucx.ucp_worker); if (UCS_OK != err) { - SPML_ERROR("fence failed"); + SPML_ERROR("fence failed: %s", ucs_status_string(err)); oshmem_shmem_abort(-1); return OSHMEM_ERROR; } diff --git a/oshmem/mca/spml/ucx/spml_ucx.h b/oshmem/mca/spml/ucx/spml_ucx.h index 4b248e026a3..b524031d3f2 100644 --- a/oshmem/mca/spml/ucx/spml_ucx.h +++ b/oshmem/mca/spml/ucx/spml_ucx.h @@ -108,6 +108,8 @@ extern sshmem_mkey_t *mca_spml_ucx_register(void* addr, int *count); extern int mca_spml_ucx_deregister(sshmem_mkey_t *mkeys); +extern void mca_spml_ucx_memuse_hook(void *addr, size_t length); + extern void mca_spml_ucx_rmkey_unpack(sshmem_mkey_t *mkey, uint32_t segno, int pe, int tr_id); extern void mca_spml_ucx_rmkey_free(sshmem_mkey_t *mkey); diff --git a/oshmem/mca/spml/yoda/spml_yoda.c b/oshmem/mca/spml/yoda/spml_yoda.c index 68e8b5acf49..91dbf138ed2 100644 --- a/oshmem/mca/spml/yoda/spml_yoda.c +++ b/oshmem/mca/spml/yoda/spml_yoda.c @@ -65,6 +65,7 @@ mca_spml_yoda_module_t mca_spml_yoda = { mca_spml_yoda_fence, mca_spml_base_rmkey_unpack, mca_spml_base_rmkey_free, + mca_spml_base_memuse_hook, (void *)&mca_spml_yoda }