Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

UCT/IB: Create multithread KSM with mkey_index for symmetric key #4

Closed
wants to merge 1 commit into from
Closed

UCT/IB: Create multithread KSM with mkey_index for symmetric key #4

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
149 changes: 87 additions & 62 deletions src/uct/ib/mlx5/dv/ib_mlx5dv_md.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -462,14 +462,69 @@ UCS_PROFILE_FUNC_ALWAYS(ucs_status_t, uct_ib_mlx5_devx_reg_atomic_key,
&memh->atomic_rkey);
}

static ucs_status_t
uct_ib_mlx5_devx_reg_symmetric(uct_ib_mlx5_md_t *md,
uct_ib_mlx5_devx_mem_t *memh, void *address,
uct_ib_mr_type_t mr_type, unsigned flags)
{
uint32_t start = md->smkey_index;
struct mlx5dv_devx_obj *smkey_mr;
uint32_t symmetric_rkey;
ucs_status_t status;

if ((mr_type != UCT_IB_MR_DEFAULT) ||
!(flags & UCT_MD_MEM_SYMMETRIC_RKEY) ||
!(md->flags & UCT_IB_MLX5_MD_FLAG_MKEY_BY_NAME_RESERVE)) {
return UCS_ERR_UNSUPPORTED;
}

if (memh->smkey_mr != NULL) {
/* Already created by MT memory registration */
return UCS_ERR_ALREADY_EXISTS;
}

/* Best effort, only allocate in the range below the atomic keys. */
while (md->smkey_index < md->super.mkey_by_name_reserve.size) {
status = uct_ib_mlx5_devx_reg_ksm_data(
md, memh, UCT_IB_MR_DEFAULT, 0,
(memh->super.flags & UCT_IB_MEM_ACCESS_REMOTE_ATOMIC),
md->super.mkey_by_name_reserve.base + md->smkey_index, NULL,
&smkey_mr, &symmetric_rkey);
if (status == UCS_OK) {
ucs_trace("%s: symmetric rkey created for addr=%p "
"mkey_index=0x%x smkey_mr=%p rkey=0x%x",
uct_ib_device_name(&md->super.dev), address,
md->super.mkey_by_name_reserve.base + start, smkey_mr,
symmetric_rkey);

memh->smkey_mr = smkey_mr;
memh->super.rkey = symmetric_rkey;
md->smkey_index++;
return UCS_OK;
}

/* Use blocks of 8 mkeys, first mkey creation gives block ownership.
* Try from the start of the next block if any failure.
*/
md->smkey_index = ucs_align_up_pow2(md->smkey_index + 1,
md->super.config.smkey_block_size);
}

ucs_debug("%s: symmetric rkey create failed for addr=%p mkey_index=0x%x",
uct_ib_device_name(&md->super.dev), address,
md->super.mkey_by_name_reserve.base + start);
return UCS_ERR_NO_ELEM;
}

static ucs_status_t
uct_ib_mlx5_devx_reg_mt(uct_ib_mlx5_md_t *md, void *address, size_t length,
int is_atomic, const uct_md_mem_reg_params_t *params,
uint64_t access_flags, uint32_t *mkey_p,
uct_ib_mlx5_devx_ksm_data_t **ksm_data_p)
uct_ib_mlx5_devx_mem_t *memh, uct_ib_mr_type_t mr_type)
{
size_t chunk = md->super.config.mt_reg_chunk;
int mr_num = ucs_div_round_up(length, chunk);
size_t chunk = md->super.config.mt_reg_chunk;
int mr_num = ucs_div_round_up(length, chunk);
unsigned flags = UCT_MD_MEM_REG_FIELD_VALUE(params, flags, FIELD_FLAGS, 0);
uct_ib_mlx5_devx_ksm_data_t *ksm_data;
ucs_status_t status;
int dmabuf_fd;
Expand Down Expand Up @@ -505,15 +560,26 @@ uct_ib_mlx5_devx_reg_mt(uct_ib_mlx5_md_t *md, void *address, size_t length,
goto err_free;
}

status = uct_ib_mlx5_devx_reg_ksm_data_mt(md, is_atomic, address, ksm_data,
length, (uint64_t)address, 0,
"multi-thread key",
&ksm_data->dvmr, mkey_p);
if (status != UCS_OK) {
goto err_dereg;
memh->super.flags |= UCT_IB_MEM_MULTITHREADED;
memh->mrs[mr_type].ksm_data = ksm_data;

status = uct_ib_mlx5_devx_reg_symmetric(md, memh, address, mr_type, flags);
if (status == UCS_OK) {
ksm_data->dvmr = NULL;
*mkey_p = memh->super.rkey;
} else {
status = uct_ib_mlx5_devx_reg_ksm_data_mt(md, is_atomic, address,
ksm_data, length,
(uint64_t)address, 0,
"multi-thread key",
&ksm_data->dvmr, mkey_p);
if (status != UCS_OK) {
memh->super.flags &= ~UCT_IB_MEM_MULTITHREADED;
memh->mrs[mr_type].ksm_data = NULL;
goto err_dereg;
}
}

*ksm_data_p = ksm_data;
return UCS_OK;

err_dereg:
Expand All @@ -535,9 +601,12 @@ uct_ib_mlx5_devx_dereg_mt(uct_ib_mlx5_md_t *md,
ucs_trace("%s: destroy KSM %p", uct_ib_device_name(&md->super.dev),
ksm_data->dvmr);

status = uct_ib_mlx5_devx_obj_destroy(ksm_data->dvmr, "MKEY, KSM");
if (status != UCS_OK) {
return status;
/* Not created if symmetric key was used instead */
if (ksm_data->dvmr != NULL) {
status = uct_ib_mlx5_devx_obj_destroy(ksm_data->dvmr, "MKEY, KSM");
if (status != UCS_OK) {
return status;
}
}

status = uct_ib_md_handle_mr_list_mt(&md->super, 0, ksm_data->length, NULL,
Expand All @@ -558,49 +627,6 @@ uct_ib_mlx5_devx_dereg_mt(uct_ib_mlx5_md_t *md,
return status;
}

static void uct_ib_mlx5_devx_reg_symmetric(uct_ib_mlx5_md_t *md,
uct_ib_mlx5_devx_mem_t *memh,
void *address)
{
uint32_t start = md->smkey_index;
struct mlx5dv_devx_obj *smkey_mr;
uint32_t symmetric_rkey;
ucs_status_t status;

ucs_assert(!(memh->super.flags & UCT_IB_MEM_MULTITHREADED));
/* Best effort, only allocate in the range below the atomic keys. */
while (md->smkey_index < md->super.mkey_by_name_reserve.size) {
status = uct_ib_mlx5_devx_reg_ksm_data_contig(
md, &memh->mrs[UCT_IB_MR_DEFAULT], address, (uint64_t)address,
(memh->super.flags & UCT_IB_MEM_ACCESS_REMOTE_ATOMIC),
md->super.mkey_by_name_reserve.base + md->smkey_index, NULL,
&smkey_mr, &symmetric_rkey);
if (status == UCS_OK) {
ucs_trace("%s: symmetric rkey created for addr=%p "
"mkey_index=0x%x smkey_mr=%p rkey=0x%x",
uct_ib_device_name(&md->super.dev), address,
md->super.mkey_by_name_reserve.base + start, smkey_mr,
symmetric_rkey);

memh->smkey_mr = smkey_mr;
memh->super.rkey = symmetric_rkey;
md->smkey_index++;
return;
}

/* Use blocks of 8 mkeys, first mkey creation gives block ownership.
* Try from the start of the next block if any failure.
*/
md->smkey_index = ucs_align_up_pow2(md->smkey_index + 1,
md->super.config.smkey_block_size);
}

ucs_debug("%s: symmetric rkey create failed for addr=%p mkey_index=0x%x",
uct_ib_device_name(&md->super.dev), address,
md->super.mkey_by_name_reserve.base + start);
}


static ucs_status_t
uct_ib_mlx5_devx_reg_mr(uct_ib_mlx5_md_t *md, uct_ib_mlx5_devx_mem_t *memh,
void *address, size_t length,
Expand All @@ -619,11 +645,10 @@ uct_ib_mlx5_devx_reg_mr(uct_ib_mlx5_md_t *md, uct_ib_mlx5_devx_mem_t *memh,
status = uct_ib_mlx5_devx_reg_mt(md, address, length,
(memh->super.flags &
UCT_IB_MEM_ACCESS_REMOTE_ATOMIC),
params, access_flags, &mkey,
&memh->mrs[mr_type].ksm_data);
params, access_flags, &mkey, memh,
mr_type);
if (status == UCS_OK) {
*rkey_p = *lkey_p = mkey;
memh->super.flags |= UCT_IB_MEM_MULTITHREADED;
return UCS_OK;
} else if (status != UCS_ERR_UNSUPPORTED) {
return status;
Expand Down Expand Up @@ -684,9 +709,9 @@ uct_ib_mlx5_devx_mem_reg(uct_md_h uct_md, void *address, size_t length,
goto err_memh_free;
}

if ((flags & UCT_MD_MEM_SYMMETRIC_RKEY) &&
(md->flags & UCT_IB_MLX5_MD_FLAG_MKEY_BY_NAME_RESERVE)) {
uct_ib_mlx5_devx_reg_symmetric(md, memh, address);
if (!(memh->super.flags & UCT_IB_MEM_MULTITHREADED)) {
uct_ib_mlx5_devx_reg_symmetric(md, memh, address, UCT_IB_MR_DEFAULT,
flags);
}

if (md->super.relaxed_order) {
Expand Down
14 changes: 13 additions & 1 deletion test/gtest/uct/ib/test_ib_md.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ class test_ib_md : public test_md

void test_mkey_pack_mt(bool invalidate);
void test_mkey_pack_mt_internal(unsigned access_mask, bool invalidate);
void test_smkey_reg_atomic(void);

private:
#ifdef HAVE_MLX5_DV
Expand Down Expand Up @@ -200,7 +201,7 @@ void test_ib_md::check_smkeys(uct_rkey_t rkey1, uct_rkey_t rkey2)
uct_ib_md_atomic_rkey(rkey2) - uct_ib_md_direct_rkey(rkey2));
}

UCS_TEST_P(test_ib_md, smkey_reg_atomic)
void test_ib_md::test_smkey_reg_atomic(void)
{
static const size_t size = 8192;
void *buffer;
Expand All @@ -225,6 +226,17 @@ UCS_TEST_P(test_ib_md, smkey_reg_atomic)
ucs_mmap_free(buffer, size);
}

UCS_TEST_P(test_ib_md, smkey_reg_atomic)
{
test_smkey_reg_atomic();
}

UCS_TEST_P(test_ib_md, smkey_reg_atomic_mt, "REG_MT_THRESH=1k",
"REG_MT_CHUNK=1k")
{
test_smkey_reg_atomic();
}

void
test_ib_md::test_mkey_pack_mt_internal(unsigned access_mask, bool invalidate)
{
Expand Down