Skip to content

Commit 5be040c

Browse files
hzhouhzhou
committed
ch4/ofi: use pipeline_pool for pack buffer
When we use malloc for pack buffer and when it is used for copying from and to GPU, we potentially can incur expensive registration cost every time. This is the case with ZE. Use pipeline_pool allows reusage of pack buffers.
1 parent 46a12d1 commit 5be040c

File tree

7 files changed

+35
-30
lines changed

7 files changed

+35
-30
lines changed

‎src/mpid/ch4/netmod/ofi/ofi_events.c‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -568,9 +568,8 @@ int MPIDI_OFI_handle_cq_error(int vci, int nic, ssize_t ret)
568568
MPIR_STATUS_SET_CANCEL_BIT(req->status, TRUE);
569569
MPIR_STATUS_SET_COUNT(req->status, 0);
570570
MPIR_Datatype_release_if_not_builtin(MPIDI_OFI_REQUEST(req, datatype));
571-
if ((event_id == MPIDI_OFI_EVENT_RECV_PACK) &&
572-
MPIDI_OFI_REQUEST(req, noncontig.pack.pack_buffer)) {
573-
MPL_free(MPIDI_OFI_REQUEST(req, noncontig.pack.pack_buffer));
571+
if (event_id == MPIDI_OFI_EVENT_RECV_PACK) {
572+
MPIDI_OFI_free_pack_buffer(req);
574573
} else if (event_id == MPIDI_OFI_EVENT_RECV_NOPACK) {
575574
MPL_free(MPIDI_OFI_REQUEST(req, noncontig.nopack.iovs));
576575
}

‎src/mpid/ch4/netmod/ofi/ofi_events.h‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -82,9 +82,8 @@ MPL_STATIC_INLINE_PREFIX int MPIDI_OFI_send_event(int vci,
8282
MPIR_FUNC_ENTER;
8383

8484
/* free the packing buffers and datatype */
85-
if ((event_id == MPIDI_OFI_EVENT_SEND_PACK) &&
86-
(MPIDI_OFI_REQUEST(sreq, noncontig.pack.pack_buffer))) {
87-
MPL_free(MPIDI_OFI_REQUEST(sreq, noncontig.pack.pack_buffer));
85+
if (event_id == MPIDI_OFI_EVENT_SEND_PACK) {
86+
MPIDI_OFI_free_pack_buffer(sreq);
8887
} else if (MPIDI_OFI_ENABLE_PT2PT_NOPACK && (event_id == MPIDI_OFI_EVENT_SEND_NOPACK)) {
8988
MPL_free(MPIDI_OFI_REQUEST(sreq, noncontig.nopack.iovs));
9089
}
@@ -133,7 +132,7 @@ MPL_STATIC_INLINE_PREFIX int MPIDI_OFI_recv_complete(MPIR_Request * rreq, int ev
133132
if (mpi_errno) {
134133
MPIR_ERR_SET(rreq->status.MPI_ERROR, MPI_ERR_TYPE, "**dtypemismatch");
135134
}
136-
MPL_free(MPIDI_OFI_REQUEST(rreq, noncontig.pack.pack_buffer));
135+
MPIDI_OFI_free_pack_buffer(rreq);
137136
} else if (event_id == MPIDI_OFI_EVENT_RECV_NOPACK) {
138137
#ifdef HAVE_ERROR_CHECKING
139138
MPI_Count elements;

‎src/mpid/ch4/netmod/ofi/ofi_impl.h‎

Lines changed: 26 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -834,27 +834,36 @@ MPL_STATIC_INLINE_PREFIX void MPIDI_OFI_gpu_rma_register(const void *buffer, siz
834834
#undef CQ_D_HEAD
835835
#undef CQ_D_TAIL
836836

837-
MPL_STATIC_INLINE_PREFIX int MPIDI_OFI_gpu_malloc_pack_buffer(void **ptr, size_t pack_size)
838-
{
839-
if (MPIDI_OFI_ENABLE_HMEM) {
840-
return MPL_gpu_malloc_host(ptr, pack_size);
841-
} else {
842-
#ifdef MPL_DEFINE_ALIGNED_ALLOC
843-
*ptr = MPL_aligned_alloc(256, pack_size, MPL_MEM_BUFFER);
844-
#else
845-
*ptr = MPL_malloc(pack_size, MPL_MEM_BUFFER);
846-
#endif
847-
return 0;
837+
MPL_STATIC_INLINE_PREFIX void *MPIDI_OFI_malloc_pack_buffer(MPIR_Request * req, MPI_Aint pack_size)
838+
{
839+
void *pack_buf;
840+
bool is_genq;
841+
if (pack_size <= MPIR_CVAR_CH4_OFI_PIPELINE_CHUNK_SZ) {
842+
int vci = MPIR_REQUEST_POOL_FROM_HANDLE(req->handle);
843+
MPIDU_genq_private_pool_alloc_cell(MPIDI_OFI_global.per_vci[vci].pipeline_pool, &pack_buf);
844+
is_genq = true;
845+
}
846+
if (!pack_buf) {
847+
pack_buf = MPL_aligned_alloc(64, pack_size, MPL_MEM_OTHER);
848+
is_genq = false;
848849
}
850+
if (pack_buf) {
851+
MPIDI_OFI_REQUEST(req, noncontig.pack.pack_buffer) = pack_buf;
852+
MPIDI_OFI_REQUEST(req, noncontig.pack.is_genq) = is_genq;
853+
}
854+
return pack_buf;
849855
}
850856

851-
MPL_STATIC_INLINE_PREFIX int MPIDI_OFI_gpu_free_pack_buffer(void *ptr)
857+
MPL_STATIC_INLINE_PREFIX void MPIDI_OFI_free_pack_buffer(MPIR_Request * req)
852858
{
853-
if (MPIDI_OFI_ENABLE_HMEM) {
854-
return MPL_gpu_free_host(ptr);
855-
} else {
856-
MPL_free(ptr);
857-
return 0;
859+
if (MPIDI_OFI_REQUEST(req, noncontig.pack.pack_buffer)) {
860+
if (MPIDI_OFI_REQUEST(req, noncontig.pack.is_genq)) {
861+
int vci = MPIR_REQUEST_POOL_FROM_HANDLE(req->handle);
862+
MPIDU_genq_private_pool_free_cell(MPIDI_OFI_global.per_vci[vci].pipeline_pool,
863+
MPIDI_OFI_REQUEST(req, noncontig.pack.pack_buffer));
864+
} else {
865+
MPL_free(MPIDI_OFI_REQUEST(req, noncontig.pack.pack_buffer));
866+
}
858867
}
859868
}
860869

‎src/mpid/ch4/netmod/ofi/ofi_pre.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -302,6 +302,7 @@ typedef struct {
302302
enum MPIDI_OFI_req_kind kind;
303303
union {
304304
struct {
305+
bool is_genq;
305306
char *pack_buffer;
306307
} pack;
307308
struct {

‎src/mpid/ch4/netmod/ofi/ofi_recv.h‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -224,11 +224,9 @@ MPL_STATIC_INLINE_PREFIX int MPIDI_OFI_do_irecv(void *buf,
224224

225225
/* Unpack */
226226
MPIDI_OFI_REQUEST(rreq, event_id) = MPIDI_OFI_EVENT_RECV_PACK;
227-
MPIDI_OFI_REQUEST(rreq, noncontig.pack.pack_buffer) =
228-
MPL_aligned_alloc(64, data_sz, MPL_MEM_OTHER);
227+
recv_buf = MPIDI_OFI_malloc_pack_buffer(rreq, data_sz);
229228
MPIR_ERR_CHKANDJUMP1(MPIDI_OFI_REQUEST(rreq, noncontig.pack.pack_buffer) == NULL, mpi_errno,
230229
MPI_ERR_OTHER, "**nomem", "**nomem %s", "Recv Pack Buffer alloc");
231-
recv_buf = MPIDI_OFI_REQUEST(rreq, noncontig.pack.pack_buffer);
232230
} else {
233231
MPIDI_OFI_REQUEST(rreq, noncontig.pack.pack_buffer) = NULL;
234232
}

‎src/mpid/ch4/netmod/ofi/ofi_rndv.c‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ int MPIDI_OFI_recv_rndv_event(int vci, struct fi_cq_tagged_entry *wc, MPIR_Reque
9191
/* if we were expecting an eager send, free the unneeded pack_buffer or iovs array */
9292
switch (MPIDI_OFI_REQUEST(rreq, event_id)) {
9393
case MPIDI_OFI_EVENT_RECV_PACK:
94-
MPL_free(MPIDI_OFI_REQUEST(rreq, noncontig.pack.pack_buffer));
94+
MPIDI_OFI_free_pack_buffer(rreq);
9595
break;
9696
case MPIDI_OFI_EVENT_RECV_NOPACK:
9797
MPL_free(MPIDI_OFI_REQUEST(rreq, noncontig.nopack.iovs));

‎src/mpid/ch4/netmod/ofi/ofi_send.h‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -493,7 +493,7 @@ MPL_STATIC_INLINE_PREFIX int MPIDI_OFI_send(const void *buf, MPI_Aint count, MPI
493493

494494
void *data = NULL;
495495
if (need_pack) {
496-
void *pack_buf = MPL_aligned_alloc(64, data_sz, MPL_MEM_OTHER);
496+
void *pack_buf = MPIDI_OFI_malloc_pack_buffer(sreq, data_sz);
497497
MPIR_ERR_CHKANDJUMP1(pack_buf == NULL, mpi_errno,
498498
MPI_ERR_OTHER, "**nomem", "**nomem %s", "Send Pack buffer alloc");
499499

@@ -504,7 +504,6 @@ MPL_STATIC_INLINE_PREFIX int MPIDI_OFI_send(const void *buf, MPI_Aint count, MPI
504504
MPIR_ERR_CHECK(mpi_errno);
505505

506506
data = pack_buf;
507-
MPIDI_OFI_REQUEST(sreq, noncontig.pack.pack_buffer) = pack_buf;
508507
} else {
509508
data = MPIR_get_contig_ptr(buf, dt_true_lb);
510509
MPIDI_OFI_REQUEST(sreq, noncontig.pack.pack_buffer) = NULL;

0 commit comments

Comments
 (0)