diff --git a/config/algorithms/MPI/Cray-MPICH/alltoallv.json b/config/algorithms/MPI/Cray-MPICH/alltoallv.json new file mode 100644 index 00000000..949513dd --- /dev/null +++ b/config/algorithms/MPI/Cray-MPICH/alltoallv.json @@ -0,0 +1,21 @@ +{ + "default_mpich": { + "desc": "MPICH default algorithm selection", + "version": "8.0.0", + "constraints": [ + { + "key": "count", + "conditions": [ + { + "operator": ">=", + "value": "comm_sz" + } + ] + } + ], + "selection": "auto", + "tags": [ + "default" + ] + } +} \ No newline at end of file diff --git a/config/algorithms/MPI/LibPico/alltoall.json b/config/algorithms/MPI/LibPico/alltoall.json index 3a2ed4be..12366a70 100644 --- a/config/algorithms/MPI/LibPico/alltoall.json +++ b/config/algorithms/MPI/LibPico/alltoall.json @@ -28,6 +28,36 @@ "external" ] }, + "bine_dh_over": { + "desc": "LibPico external Alltoall algorithm . This algorithm follows the distance-halving Bine butterfly communication pattern described in the Bine paper.", + "version": "1.0.0", + "constraints": [ + { + "key": "count", + "conditions": [ + { + "operator": ">=", + "value": "comm_sz" + } + ] + }, + { + "key": "comm_sz", + "conditions": [ + { + "operator": "is_power_of_two", + "value": true + } + ] + } + ], + "selection": "pico", + "tags": [ + "bine", + "external", + "distance-halving" + ] + }, "pairwise_ompi_over": { "desc": "libpico pairwise algorithm, copied from Open MPI.", "version": "1.0.0", @@ -57,4 +87,4 @@ "external" ] } -} +} \ No newline at end of file diff --git a/config/algorithms/MPI/LibPico/alltoallv.json b/config/algorithms/MPI/LibPico/alltoallv.json new file mode 100644 index 00000000..fac609ad --- /dev/null +++ b/config/algorithms/MPI/LibPico/alltoallv.json @@ -0,0 +1,32 @@ +{ + "bine_dh_over": { + "desc": "LibPico external Alltoall algorithm . This algorithm follows the distance-halving Bine butterfly communication pattern described in the Bine paper.", + "version": "1.0.0", + "constraints": [ + { + "key": "count", + "conditions": [ + { + "operator": ">=", + "value": "comm_sz" + } + ] + }, + { + "key": "comm_sz", + "conditions": [ + { + "operator": "is_power_of_two", + "value": true + } + ] + } + ], + "selection": "pico", + "tags": [ + "bine", + "external", + "distance-halving" + ] + } +} \ No newline at end of file diff --git a/config/algorithms/MPI/MPICH/alltoallv.json b/config/algorithms/MPI/MPICH/alltoallv.json new file mode 100644 index 00000000..418c0305 --- /dev/null +++ b/config/algorithms/MPI/MPICH/alltoallv.json @@ -0,0 +1,21 @@ +{ + "default_mpich": { + "desc": "MPICH default Alltoallv algorithm selection", + "version": "4.3.0", + "constraints": [ + { + "key": "count", + "conditions": [ + { + "operator": ">=", + "value": "comm_sz" + } + ] + } + ], + "selection": "auto", + "tags": [ + "default" + ] + } +} \ No newline at end of file diff --git a/config/algorithms/MPI/Open-MPI/alltoallv.json b/config/algorithms/MPI/Open-MPI/alltoallv.json new file mode 100644 index 00000000..cec83bf3 --- /dev/null +++ b/config/algorithms/MPI/Open-MPI/alltoallv.json @@ -0,0 +1,66 @@ +{ + "default_ompi": { + "desc": "Open MPI default Alltoallv algorithm selection.", + "version": "4.1.5", + "constraints": [ + { + "key": "count", + "conditions": [ + { + "operator": ">=", + "value": "comm_sz" + } + ] + } + ], + "selection": 0, + "cuda_support": "yes", + "rocm_support": "ucx", + "tags": [ + "ignore" + ] + }, + "basic_linear_ompi": { + "desc": "Open MPI basic linear Alltoallv algorithm.", + "version": "4.1.5", + "constraints": [ + { + "key": "count", + "conditions": [ + { + "operator": ">=", + "value": "comm_sz" + } + ] + } + ], + "selection": 1, + "cuda_support": "yes", + "rocm_support": "ucx", + "tags": [ + "basic_linear", + "linear" + ] + }, + "pairwise_ompi": { + "desc": "Open MPI pairwise Alltoallv algorithm.", + "version": "4.1.5", + "constraints": [ + { + "key": "count", + "conditions": [ + { + "operator": ">=", + "value": "comm_sz" + } + ] + } + ], + "selection": 2, + "cuda_support": "yes", + "rocm_support": "ucx", + "tags": [ + "pairwise" + ] + } +} \ No newline at end of file diff --git a/config/algorithms/NCCL/LibPico/alltoallv.json b/config/algorithms/NCCL/LibPico/alltoallv.json new file mode 100644 index 00000000..ecb282a2 --- /dev/null +++ b/config/algorithms/NCCL/LibPico/alltoallv.json @@ -0,0 +1,53 @@ +{ + "nccl_spreadout_alltoallv_over": { + "desc": "LibPico external NCCL spread-out Alltoallv algorithm. This implementation uses NCCL point-to-point ncclSend/ncclRecv operations arranged in size - 1 communication rounds.", + "version": "1.0.0", + "constraints": [ + { + "key": "count", + "conditions": [ + { + "operator": ">=", + "value": "comm_sz" + } + ] + } + ], + "selection": "pico", + "tags": [ + "alltoallv", + "nccl", + "cuda", + "gpu", + "spread-out", + "sendrecv", + "external" + ] + }, + + "nccl_fanout_alltoallv_over": { + "desc": "LibPico external NCCL fan-out Alltoallv algorithm. This implementation posts all ncclSend/ncclRecv operations inside a single ncclGroupStart/ncclGroupEnd region.", + "version": "1.0.0", + "constraints": [ + { + "key": "count", + "conditions": [ + { + "operator": ">=", + "value": "comm_sz" + } + ] + } + ], + "selection": "pico", + "tags": [ + "alltoallv", + "nccl", + "cuda", + "gpu", + "fan-out", + "sendrecv", + "external" + ] + } +} \ No newline at end of file diff --git a/config/parse_test.py b/config/parse_test.py index b6f1aeb6..265a5886 100644 --- a/config/parse_test.py +++ b/config/parse_test.py @@ -15,7 +15,7 @@ "type": "object", "properties": { "libpico_version": {"type": "string", "pattern": "^\\d+\\.\\d+\\.\\d+$"}, - "collective": {"type": "string", "enum": ["ALLREDUCE", "ALLTOALL", "ALLGATHER", "BCAST", "GATHER", "REDUCE", "REDUCE_SCATTER", "SCATTER"]}, + "collective": {"type": "string", "enum": ["ALLREDUCE", "ALLTOALL","ALLTOALLV", "ALLGATHER", "BCAST", "GATHER", "REDUCE", "REDUCE_SCATTER", "SCATTER"]}, "MPI_Op": {"type": "string"}, "tags": { "type": "object", diff --git a/config/test/alltoallv.json b/config/test/alltoallv.json new file mode 100644 index 00000000..9ff0fa87 --- /dev/null +++ b/config/test/alltoallv.json @@ -0,0 +1,16 @@ +{ + "libpico_version": "1.0.0", + "collective": "ALLTOALLV", + "MPI_Op": "MPI_SUM", + "tags": { + "include": [ + "internal", + "external" + ], + "exclude": [] + }, + "specific": { + "include": [], + "exclude": [] + } +} \ No newline at end of file diff --git a/include/libpico.h b/include/libpico.h index b2bd445c..637c8fa4 100644 --- a/include/libpico.h +++ b/include/libpico.h @@ -19,6 +19,8 @@ void* rbuf, size_t rcount, MPI_Datatype rdtype, MPI_Comm comm #define ALLTOALL_MPI_ARGS const void *sbuf, size_t scount, MPI_Datatype sdtype, \ void *rbuf, size_t rcount, MPI_Datatype rdtype, MPI_Comm comm +#define ALLTOALLV_MPI_ARGS const void *sbuf, const int scounts[], const int sdispls[], MPI_Datatype sdtype, \ + void *rbuf, const int rcounts[], const int rdispls[], MPI_Datatype rdtype, MPI_Comm comm #define BCAST_MPI_ARGS void *buf, size_t count, MPI_Datatype dtype, int root, MPI_Comm comm #define GATHER_MPI_ARGS const void *sbuf, size_t scount, MPI_Datatype sdtype, \ void *rbuf, size_t rcount, MPI_Datatype rdtype, int root, MPI_Comm comm @@ -38,6 +40,9 @@ ncclDataType_t dtype, ncclComm_t nccl_comm, cudaStream_t stream #define ALLTOALL_NCCL_ARGS const void *sbuf, void* rbuf, size_t count,\ ncclDataType_t dtype, ncclComm_t nccl_comm, cudaStream_t stream +#define ALLTOALLV_NCCL_ARGS const void *sbuf, const int scounts[], const int sdispls[], \ + void *rbuf, const int rcounts[], const int rdispls[], \ + ncclDataType_t dtype, ncclComm_t nccl_comm, cudaStream_t stream #define BCAST_NCCL_ARGS void *buf, size_t count, ncclDataType_t dtype, \ int root, ncclComm_t nccl_comm, cudaStream_t stream #define GATHER_NCCL_ARGS const void *sbuf, void *rbuf, size_t count, ncclDataType_t dtype, \ @@ -87,11 +92,19 @@ int allgather_bine_block_by_block_nccl(ALLGATHER_NCCL_ARGS); int alltoall_pairwise_ompi(ALLTOALL_MPI_ARGS); int alltoall_bine(ALLTOALL_MPI_ARGS); +int alltoall_bine_DH(ALLTOALL_MPI_ARGS); // write here the prototype of NCCL alltoall functions #ifdef PICO_NCCL // int alltoall_nccl_custom(ALLTOALL_NCCL_ARGS) #endif +int alltoallv_bine_DH(ALLTOALLV_MPI_ARGS); + +#ifdef PICO_NCCL +int alltoallv_nccl_spreadout(ALLTOALLV_NCCL_ARGS); +int alltoallv_nccl_fanout(ALLTOALLV_NCCL_ARGS); +#endif + int bcast_linear(BCAST_MPI_ARGS); int bcast_binomial_halving(BCAST_MPI_ARGS); int bcast_binomial_doubling(BCAST_MPI_ARGS); diff --git a/libpico/libpico_alltoall.c b/libpico/libpico_alltoall.c index f0bb8c4f..6eb84869 100644 --- a/libpico/libpico_alltoall.c +++ b/libpico/libpico_alltoall.c @@ -7,11 +7,12 @@ #include #include #include - +#include +#include +#include #include "libpico.h" #include "libpico_utils.h" - /* Alltoall pairwise implementation from Open MPI 5.0.1 base module. * Original file: ompi/mca/coll/base/coll_base_alltoall.c * Original function: ompi_coll_base_alltoall_intra_pairwise @@ -199,7 +200,116 @@ int alltoall_bine(const void *sendbuf, size_t s_count, MPI_Datatype s_dtype, return err; } -// write NCCL implementations here -#ifdef PICO_NCCL +int alltoall_bine_DH(const void *sendbuf, size_t s_count, MPI_Datatype s_dtype, + void *recvbuf, size_t r_count, MPI_Datatype r_dtype, MPI_Comm comm) +{ + assert(s_count == r_count); + assert(s_dtype == r_dtype); + + int r, size, dtype, s, err = MPI_SUCCESS; + char *work_buffer = NULL; + char *send_buffer = NULL; + char *keep_buffer = NULL; + size_t block_size; + size_t header_size; + size_t packet_size; + + err = MPI_Comm_rank(comm, &r); + if (err != MPI_SUCCESS) goto err_hndl; + + err = MPI_Comm_size(comm, &size); + if (err != MPI_SUCCESS) goto err_hndl; + + err = MPI_Type_size(s_dtype, &dtype); + if (err != MPI_SUCCESS) goto err_hndl; + + s = (int)log2((double)size); + assert((1 << s) == size); + + block_size = s_count * (size_t)dtype; + header_size = 2 * sizeof(int); + packet_size = header_size + block_size; + + work_buffer = malloc((size_t)size * packet_size); + send_buffer = malloc((size_t)size * packet_size); + keep_buffer = malloc((size_t)size * packet_size); + + if (work_buffer == NULL || send_buffer == NULL || keep_buffer == NULL) { + err = MPI_ERR_NO_MEM; + goto err_hndl; + } + + for (int i = 0; i < size; i++) { + char *rec = work_buffer + (size_t)i * packet_size; + int src = r; + int dst = i; + + memcpy(rec, &src, sizeof(int)); + memcpy(rec + sizeof(int), &dst, sizeof(int)); + memcpy(rec + header_size, + (const char *)sendbuf + (size_t)i * block_size, + block_size); + } + + for (int step = 0; step < s; step++) { + int partner; + int send_count = 0; + int keep_count = 0; + + if ((r % 2) == 0) + partner = mod(r + (1 - (int)pow(-2, s - step)) / 3, size); + else + partner = mod(r - (1 - (int)pow(-2, s - step)) / 3, size); + + for (int j = 0; j < size; j++) { + char *rec = work_buffer + (size_t)j * packet_size; + int dst; + + memcpy(&dst, rec + sizeof(int), sizeof(int)); + + if (same_prefix_negabinary(dst, partner, s, step + 1)) { + memcpy(send_buffer + (size_t)send_count * packet_size, + rec, + packet_size); + send_count++; + } else { + memcpy(keep_buffer + (size_t)keep_count * packet_size, + rec, + packet_size); + keep_count++; + } + } + + err = MPI_Sendrecv(send_buffer, + (int)((size_t)send_count * packet_size), MPI_BYTE, partner, 0, + work_buffer + (size_t)keep_count * packet_size, + (int)((size_t)send_count * packet_size), MPI_BYTE, partner, 0, + comm, MPI_STATUS_IGNORE); + if (err != MPI_SUCCESS) goto err_hndl; + + memcpy(work_buffer, + keep_buffer, + (size_t)keep_count * packet_size); + } + + for (int i = 0; i < size; i++) { + char *rec = work_buffer + (size_t)i * packet_size; + int src, dst; + + memcpy(&src, rec, sizeof(int)); + memcpy(&dst, rec + sizeof(int), sizeof(int)); + + assert(dst == r); + + memcpy((char *)recvbuf + (size_t)src * block_size, + rec + header_size, + block_size); + } + +err_hndl: + free(keep_buffer); + free(send_buffer); + free(work_buffer); + return err; +} -#endif \ No newline at end of file diff --git a/libpico/libpico_alltoallv.c b/libpico/libpico_alltoallv.c new file mode 100644 index 00000000..d026ada7 --- /dev/null +++ b/libpico/libpico_alltoallv.c @@ -0,0 +1,774 @@ +#include +#include +#include +#include +#include +#include +#include +#include "libpico.h" +#include "libpico_utils.h" +#ifdef PICO_MPI_CUDA_AWARE +#include "support_kernel.h" +#endif +#include + +// This implementation follows the distance-halving Bine butterfly pattern described in the paper +int alltoallv_bine_DH(const void *sendbuf, const int sendcounts[], const int sdispls[], MPI_Datatype sendtype, + void *recvbuf, const int recvcounts[], const int rdispls[], MPI_Datatype recvtype, + MPI_Comm comm) +{ +#ifdef PICO_MPI_CUDA_AWARE + assert(sendtype == recvtype); + void *host_sendbuf = NULL, *host_recvbuf = NULL; + int r, size, dtype, s, err = MPI_SUCCESS; + char *work_buffer = NULL, *send_buffer = NULL, *keep_buffer = NULL; + size_t header_size = 3 * sizeof(int), max_dim_buffer, dim_work = 0; + + err = MPI_Comm_rank(comm, &r); + if (err != MPI_SUCCESS) + goto err_hndl; + + err = MPI_Comm_size(comm, &size); + if (err != MPI_SUCCESS) + goto err_hndl; + + err = MPI_Type_size(sendtype, &dtype); + if (err != MPI_SUCCESS) + goto err_hndl; + + s = (int)log2((double)size); + assert((1 << s) == size); + size_t num_bytes_send = num_bytes_fun(sendcounts,sdispls, dtype, size); + size_t num_bytes_recv = num_bytes_fun(recvcounts,rdispls ,dtype, size); + PICO_TAG_BEGIN("malloc cudaMemcpy"); + if (num_bytes_send > 0){ + host_sendbuf=malloc( num_bytes_send); + if(host_sendbuf==NULL){ + err=MPI_ERR_NO_MEM; + goto err_hndl; + } + BINE_CUDA_CHECK(cudaMemcpy(host_sendbuf,sendbuf, num_bytes_send, cudaMemcpyDeviceToHost)); + } + if (num_bytes_recv > 0){ + host_recvbuf=malloc(num_bytes_recv); + if(host_recvbuf==NULL){ + err=MPI_ERR_NO_MEM; + goto err_hndl; + } + } + PICO_TAG_END("malloc cudaMemcpy"); + size_t local_total_bytes = 0; + size_t local_blocks = 0; + + for (int i = 0; i < size; i++) + { + size_t block_size = (size_t)sendcounts[i] * (size_t)dtype; + + if (block_size > 0) + { + local_total_bytes += block_size; + local_blocks++; + } + } + + max_dim_buffer = local_total_bytes + local_blocks * header_size; + + if (max_dim_buffer == 0) + max_dim_buffer = 1; + + work_buffer = malloc(max_dim_buffer * 3); + if (work_buffer == NULL) + { + err = MPI_ERR_NO_MEM; + goto err_hndl; + } + + send_buffer = work_buffer + max_dim_buffer; + keep_buffer = work_buffer + 2 * max_dim_buffer; + PICO_TAG_BEGIN("init:for"); + for (int i = 0; i < size; i++) + { + int block_size = (int)((size_t)sendcounts[i] * (size_t)dtype); + size_t offset = (size_t)sdispls[i] * (size_t)dtype; + size_t packet_size = header_size + (size_t)block_size; + if (block_size > 0) + { + char *rec = work_buffer + dim_work; + int src = r; + int dst = i; + memcpy(rec, &src, sizeof(int)); + memcpy(rec + sizeof(int), &dst, sizeof(int)); + memcpy(rec + 2 * sizeof(int), &block_size, sizeof(int)); + + memcpy(rec + header_size, + (const char *)host_sendbuf + offset, + (size_t)block_size); + + dim_work += packet_size; + } + } + PICO_TAG_END("init:for"); + for (int step = 0; step < s; step++) + { + int partner, dim_send = 0, dim_keep = 0, dim_recv = 0; + + if ((r % 2) == 0) + partner = mod(r + (1 - (int)pow(-2, s - step)) / 3, size); + else + partner = mod(r - (1 - (int)pow(-2, s - step)) / 3, size); + + size_t skip = 0; + PICO_TAG_BEGIN("step:while"); + while (skip < dim_work) + { + char *rec = work_buffer + skip; + + int dst, block_size, src; + memcpy(&src, rec, sizeof(int)); + memcpy(&dst, rec + sizeof(int), sizeof(int)); + memcpy(&block_size, rec + 2 * sizeof(int), sizeof(int)); + + size_t packet_size = header_size + (size_t)block_size; + int logical_dst = logical_rank_for_bine_dh_root(dst, src, size); + int logical_partner = logical_rank_for_bine_dh_root(partner, src, size); + if (same_prefix_negabinary(logical_dst, logical_partner, s, step + 1)) + { + memcpy(send_buffer + dim_send, rec, packet_size); + dim_send += (int)packet_size; + } + else + { + if (dim_keep != (int)skip) + memmove(work_buffer + dim_keep, rec, packet_size); + + dim_keep += (int)packet_size; + } + + skip += packet_size; + } + PICO_TAG_END("step:while"); + MPI_Request req; + MPI_Status status; + + PICO_TAG_BEGIN("step:isend_data"); + err = MPI_Isend(send_buffer, dim_send, MPI_BYTE, + partner, 1, comm, &req); + PICO_TAG_END("step:isend_data"); + if (err != MPI_SUCCESS) + goto err_hndl; + + PICO_TAG_BEGIN("step:probe"); + err = MPI_Probe(partner, 1, comm, &status); + PICO_TAG_END("step:probe"); + if (err != MPI_SUCCESS) + goto err_hndl; + + err = MPI_Get_count(&status, MPI_BYTE, &dim_recv); + if (err != MPI_SUCCESS) + goto err_hndl; + + size_t needed = (size_t)dim_keep + (size_t)dim_recv; + if (needed == 0) + needed = 1; + + if (needed > max_dim_buffer) + { + PICO_TAG_BEGIN("step:realloc"); + + char *old_buffer = work_buffer; + + char *new_buffer = NULL; + + new_buffer = malloc(needed * 3); + if (new_buffer == NULL) + { + err = MPI_ERR_NO_MEM; + goto err_hndl; + } + + char *new_work_buffer = new_buffer; + char *new_send_buffer = new_buffer + needed; + char *new_keep_buffer = new_buffer + 2 * needed; + PICO_TAG_END("step:realloc"); + PICO_TAG_BEGIN("step:memcpy"); + memcpy(new_work_buffer, work_buffer, (size_t)dim_keep); + PICO_TAG_END("step:memcpy"); + work_buffer = new_work_buffer; + send_buffer = new_send_buffer; + keep_buffer = new_keep_buffer; + + max_dim_buffer = needed; + + PICO_TAG_BEGIN("step:recv_data"); + err = MPI_Recv(work_buffer + dim_keep, dim_recv, MPI_BYTE, + partner, 1, comm, MPI_STATUS_IGNORE); + PICO_TAG_END("step:recv_data"); + if (err != MPI_SUCCESS) + goto err_hndl; + + err = MPI_Wait(&req, MPI_STATUS_IGNORE); + if (err != MPI_SUCCESS) + goto err_hndl; + + free(old_buffer); + } + else + { + + PICO_TAG_BEGIN("step:recv_data"); + err = MPI_Recv(work_buffer + dim_keep, dim_recv, MPI_BYTE, + partner, 1, comm, MPI_STATUS_IGNORE); + PICO_TAG_END("step:recv_data"); + + if (err != MPI_SUCCESS) + goto err_hndl; + + err = MPI_Wait(&req, MPI_STATUS_IGNORE); + if (err != MPI_SUCCESS) + goto err_hndl; + } + + dim_work = (size_t)dim_keep + (size_t)dim_recv; + } + PICO_TAG_BEGIN("permutation"); + size_t skip = 0; + while (skip < dim_work) + { + char *rec = work_buffer + skip; + + int src, dst, block_size; + memcpy(&src, rec, sizeof(int)); + memcpy(&dst, rec + sizeof(int), sizeof(int)); + memcpy(&block_size, rec + 2 * sizeof(int), sizeof(int)); + + size_t packet_size = header_size + (size_t)block_size; + + assert(dst == r); + + size_t offset = (size_t)rdispls[src] * (size_t)dtype; + + if (block_size > 0) + { + memcpy((char *)host_recvbuf + offset, + rec + header_size, + (size_t)block_size); + } + + skip += packet_size; + } + PICO_TAG_END("permutation"); + PICO_TAG_BEGIN("final cudaMemcpy"); + if (num_bytes_recv > 0){ + BINE_CUDA_CHECK(cudaMemcpy(recvbuf,host_recvbuf, num_bytes_recv, cudaMemcpyHostToDevice)); + } + PICO_TAG_END("final cudaMemcpy"); +err_hndl: + free(work_buffer); + free(host_sendbuf); + free(host_recvbuf); + return err; +#else + assert(sendtype == recvtype); + + int r, size, dtype, s, err = MPI_SUCCESS; + char *work_buffer = NULL, *send_buffer = NULL, *keep_buffer = NULL; + size_t header_size = 3 * sizeof(int), max_dim_buffer, dim_work = 0; + + err = MPI_Comm_rank(comm, &r); + if (err != MPI_SUCCESS) + goto err_hndl; + + err = MPI_Comm_size(comm, &size); + if (err != MPI_SUCCESS) + goto err_hndl; + + err = MPI_Type_size(sendtype, &dtype); + if (err != MPI_SUCCESS) + goto err_hndl; + + s = (int)log2((double)size); + assert((1 << s) == size); + + size_t local_total_bytes = 0; + size_t local_blocks = 0; + + for (int i = 0; i < size; i++) + { + size_t block_size = (size_t)sendcounts[i] * (size_t)dtype; + + if (block_size > 0) + { + local_total_bytes += block_size; + local_blocks++; + } + } + + max_dim_buffer = local_total_bytes + local_blocks * header_size; + + if (max_dim_buffer == 0) + max_dim_buffer = 1; + + work_buffer = malloc(max_dim_buffer * 3); + if (work_buffer == NULL) + { + err = MPI_ERR_NO_MEM; + goto err_hndl; + } + + send_buffer = work_buffer + max_dim_buffer; + keep_buffer = work_buffer + 2 * max_dim_buffer; + PICO_TAG_BEGIN("init:for"); + for (int i = 0; i < size; i++) + { + int block_size = (int)((size_t)sendcounts[i] * (size_t)dtype); + size_t offset = (size_t)sdispls[i] * (size_t)dtype; + size_t packet_size = header_size + (size_t)block_size; + if (block_size > 0) + { + char *rec = work_buffer + dim_work; + int src = r; + int dst = i; + + memcpy(rec, &src, sizeof(int)); + memcpy(rec + sizeof(int), &dst, sizeof(int)); + memcpy(rec + 2 * sizeof(int), &block_size, sizeof(int)); + + memcpy(rec + header_size, + (const char *)sendbuf + offset, + (size_t)block_size); + + dim_work += packet_size; + } + } + PICO_TAG_END("init:for"); + for (int step = 0; step < s; step++) + { + int partner, dim_send = 0, dim_keep = 0, dim_recv = 0; + + if ((r % 2) == 0) + partner = mod(r + (1 - (int)pow(-2, s - step)) / 3, size); + else + partner = mod(r - (1 - (int)pow(-2, s - step)) / 3, size); + + size_t skip = 0; + PICO_TAG_BEGIN("step:while"); + while (skip < dim_work) + { + char *rec = work_buffer + skip; + + int dst, block_size, src; + memcpy(&src, rec, sizeof(int)); + memcpy(&dst, rec + sizeof(int), sizeof(int)); + memcpy(&block_size, rec + 2 * sizeof(int), sizeof(int)); + + size_t packet_size = header_size + (size_t)block_size; + int logical_dst = logical_rank_for_bine_dh_root(dst, src, size); + int logical_partner = logical_rank_for_bine_dh_root(partner, src, size); + if (same_prefix_negabinary(logical_dst, logical_partner, s, step + 1)) + { + memcpy(send_buffer + dim_send, rec, packet_size); + dim_send += (int)packet_size; + } + else + { + if (dim_keep != (int)skip) + memmove(work_buffer + dim_keep, rec, packet_size); + + dim_keep += (int)packet_size; + } + + skip += packet_size; + } + PICO_TAG_END("step:while"); + MPI_Request req; + MPI_Status status; + + PICO_TAG_BEGIN("step:isend_data"); + err = MPI_Isend(send_buffer, dim_send, MPI_BYTE, + partner, 1, comm, &req); + PICO_TAG_END("step:isend_data"); + if (err != MPI_SUCCESS) + goto err_hndl; + + PICO_TAG_BEGIN("step:probe"); + err = MPI_Probe(partner, 1, comm, &status); + PICO_TAG_END("step:probe"); + if (err != MPI_SUCCESS) + goto err_hndl; + + err = MPI_Get_count(&status, MPI_BYTE, &dim_recv); + if (err != MPI_SUCCESS) + goto err_hndl; + + size_t needed = (size_t)dim_keep + (size_t)dim_recv; + if (needed == 0) + needed = 1; + + if (needed > max_dim_buffer) + { + PICO_TAG_BEGIN("step:realloc"); + + char *old_buffer = work_buffer; + + char *new_buffer = NULL; + + new_buffer = malloc(needed * 3); + if (new_buffer == NULL) + { + err = MPI_ERR_NO_MEM; + goto err_hndl; + } + + char *new_work_buffer = new_buffer; + char *new_send_buffer = new_buffer + needed; + char *new_keep_buffer = new_buffer + 2 * needed; + PICO_TAG_END("step:realloc"); + PICO_TAG_BEGIN("step:memcpy"); + memcpy(new_work_buffer, work_buffer, (size_t)dim_keep); + PICO_TAG_END("step:memcpy"); + work_buffer = new_work_buffer; + send_buffer = new_send_buffer; + keep_buffer = new_keep_buffer; + + max_dim_buffer = needed; + + PICO_TAG_BEGIN("step:recv_data"); + err = MPI_Recv(work_buffer + dim_keep, dim_recv, MPI_BYTE, + partner, 1, comm, MPI_STATUS_IGNORE); + PICO_TAG_END("step:recv_data"); + if (err != MPI_SUCCESS) + goto err_hndl; + + err = MPI_Wait(&req, MPI_STATUS_IGNORE); + if (err != MPI_SUCCESS) + goto err_hndl; + + free(old_buffer); + } + else + { + + PICO_TAG_BEGIN("step:recv_data"); + err = MPI_Recv(work_buffer + dim_keep, dim_recv, MPI_BYTE, + partner, 1, comm, MPI_STATUS_IGNORE); + PICO_TAG_END("step:recv_data"); + + if (err != MPI_SUCCESS) + goto err_hndl; + + err = MPI_Wait(&req, MPI_STATUS_IGNORE); + if (err != MPI_SUCCESS) + goto err_hndl; + } + + dim_work = (size_t)dim_keep + (size_t)dim_recv; + } + PICO_TAG_BEGIN("permutation"); + size_t skip = 0; + while (skip < dim_work) + { + char *rec = work_buffer + skip; + + int src, dst, block_size; + memcpy(&src, rec, sizeof(int)); + memcpy(&dst, rec + sizeof(int), sizeof(int)); + memcpy(&block_size, rec + 2 * sizeof(int), sizeof(int)); + + size_t packet_size = header_size + (size_t)block_size; + + assert(dst == r); + + size_t offset = (size_t)rdispls[src] * (size_t)dtype; + + if (block_size > 0) + { + memcpy((char *)recvbuf + offset, + rec + header_size, + (size_t)block_size); + } + + skip += packet_size; + } + PICO_TAG_END("permutation"); +err_hndl: + free(work_buffer); + return err; +#endif +} + +// write NCCL implementations here +#ifdef PICO_NCCL + +static int ncclSpreadOutAllToAllvGPU(void *output, + const void *input, + const int *send_bytes, + const int *recv_bytes, + const int *send_byte_displs, + const int *recv_byte_displs, + int rank, + int size, + ncclComm_t comm, + cudaStream_t stream) +{ + /* + * Self-copy: local data rank -> rank. + */ + if (send_bytes[rank] > 0) + { + BINE_CUDA_CHECK(cudaMemcpyAsync( + (char *)output + recv_byte_displs[rank], + (const char *)input + send_byte_displs[rank], + send_bytes[rank], + cudaMemcpyDeviceToDevice, + stream)); + } + + /* + * Spread-out Alltoallv: + * each rank communicates with one send partner and one receive partner per step. + */ + for (int step = 1; step < size; step++) + { + int send_partner = (rank + step) % size; + int recv_partner = (rank - step + size) % size; + + BINE_NCCL_CHECK(ncclGroupStart()); + + if (send_bytes[send_partner] > 0) + { + const char *send_ptr = + (const char *)input + send_byte_displs[send_partner]; + + BINE_NCCL_CHECK(ncclSend(send_ptr, + send_bytes[send_partner], + ncclInt8, + send_partner, + comm, + stream)); + } + + if (recv_bytes[recv_partner] > 0) + { + char *recv_ptr = + (char *)output + recv_byte_displs[recv_partner]; + + BINE_NCCL_CHECK(ncclRecv(recv_ptr, + recv_bytes[recv_partner], + ncclInt8, + recv_partner, + comm, + stream)); + } + + BINE_NCCL_CHECK(ncclGroupEnd()); + } + + return MPI_SUCCESS; +} + +static int ncclAllToAllvGPU(void *output, + const void *input, + const int *send_bytes, + const int *recv_bytes, + const int *send_byte_displs, + const int *recv_byte_displs, + int rank, + int size, + ncclComm_t comm, + cudaStream_t stream) +{ + /* + * Self-copy: local data rank -> rank. + * We handle self locally and skip peer == rank in NCCL P2P. + */ + if (send_bytes[rank] > 0) + { + BINE_CUDA_CHECK(cudaMemcpyAsync( + (char *)output + recv_byte_displs[rank], + (const char *)input + send_byte_displs[rank], + send_bytes[rank], + cudaMemcpyDeviceToDevice, + stream)); + } + + /* + * Fan-out Alltoallv: + * all ncclSend/ncclRecv operations are posted inside one NCCL group. + */ + BINE_NCCL_CHECK(ncclGroupStart()); + + for (int peer = 0; peer < size; peer++) + { + if (peer == rank) + continue; + + if (send_bytes[peer] > 0) + { + const char *send_ptr = + (const char *)input + send_byte_displs[peer]; + + BINE_NCCL_CHECK(ncclSend(send_ptr, + send_bytes[peer], + ncclInt8, + peer, + comm, + stream)); + } + + if (recv_bytes[peer] > 0) + { + char *recv_ptr = + (char *)output + recv_byte_displs[peer]; + + BINE_NCCL_CHECK(ncclRecv(recv_ptr, + recv_bytes[peer], + ncclInt8, + peer, + comm, + stream)); + } + } + + BINE_NCCL_CHECK(ncclGroupEnd()); + + return MPI_SUCCESS; +} + +static int nccl_dtype_size_bytes(ncclDataType_t dtype, int *dtype_size) +{ + switch (dtype) + { + case ncclInt8: + case ncclUint8: + *dtype_size = 1; + return MPI_SUCCESS; + + case ncclFloat16: + *dtype_size = 2; + return MPI_SUCCESS; + + case ncclInt32: + case ncclUint32: + case ncclFloat32: + *dtype_size = 4; + return MPI_SUCCESS; + + case ncclInt64: + case ncclUint64: + case ncclFloat64: + *dtype_size = 8; + return MPI_SUCCESS; + + default: + fprintf(stderr, "Unsupported NCCL datatype in Alltoallv\n"); + return MPI_ERR_TYPE; + } +} + +int alltoallv_nccl_spreadout(ALLTOALLV_NCCL_ARGS) +{ + int rank, size, dtype_size; + int err = MPI_SUCCESS; + + BINE_NCCL_CHECK(ncclCommUserRank(nccl_comm, &rank)); + BINE_NCCL_CHECK(ncclCommCount(nccl_comm, &size)); + + err = nccl_dtype_size_bytes(dtype, &dtype_size); + if (err != MPI_SUCCESS) + return err; + + int *send_bytes = malloc(size * sizeof(int)); + int *recv_bytes = malloc(size * sizeof(int)); + int *send_byte_displs = malloc(size * sizeof(int)); + int *recv_byte_displs = malloc(size * sizeof(int)); + + if (!send_bytes || !recv_bytes || !send_byte_displs || !recv_byte_displs) + { + free(send_bytes); + free(recv_bytes); + free(send_byte_displs); + free(recv_byte_displs); + return MPI_ERR_NO_MEM; + } + + for (int i = 0; i < size; i++) + { + send_bytes[i] = scounts[i] * dtype_size; + recv_bytes[i] = rcounts[i] * dtype_size; + send_byte_displs[i] = sdispls[i] * dtype_size; + recv_byte_displs[i] = rdispls[i] * dtype_size; + } + + err = ncclSpreadOutAllToAllvGPU(rbuf, + sbuf, + send_bytes, + recv_bytes, + send_byte_displs, + recv_byte_displs, + rank, + size, + nccl_comm, + stream); + + free(send_bytes); + free(recv_bytes); + free(send_byte_displs); + free(recv_byte_displs); + + return err; +} + +int alltoallv_nccl_fanout(ALLTOALLV_NCCL_ARGS) +{ + int rank, size, dtype_size; + int err = MPI_SUCCESS; + + BINE_NCCL_CHECK(ncclCommUserRank(nccl_comm, &rank)); + BINE_NCCL_CHECK(ncclCommCount(nccl_comm, &size)); + + err = nccl_dtype_size_bytes(dtype, &dtype_size); + if (err != MPI_SUCCESS) + return err; + + int *send_bytes = malloc(size * sizeof(int)); + int *recv_bytes = malloc(size * sizeof(int)); + int *send_byte_displs = malloc(size * sizeof(int)); + int *recv_byte_displs = malloc(size * sizeof(int)); + + if (!send_bytes || !recv_bytes || !send_byte_displs || !recv_byte_displs) + { + free(send_bytes); + free(recv_bytes); + free(send_byte_displs); + free(recv_byte_displs); + return MPI_ERR_NO_MEM; + } + + for (int i = 0; i < size; i++) + { + send_bytes[i] = scounts[i] * dtype_size; + recv_bytes[i] = rcounts[i] * dtype_size; + send_byte_displs[i] = sdispls[i] * dtype_size; + recv_byte_displs[i] = rdispls[i] * dtype_size; + } + + err = ncclAllToAllvGPU(rbuf, + sbuf, + send_bytes, + recv_bytes, + send_byte_displs, + recv_byte_displs, + rank, + size, + nccl_comm, + stream); + + free(send_bytes); + free(recv_bytes); + free(send_byte_displs); + free(recv_byte_displs); + + return err; +} + +#endif \ No newline at end of file diff --git a/libpico/libpico_utils.h b/libpico/libpico_utils.h index a9b54a03..0a7f3952 100644 --- a/libpico/libpico_utils.h +++ b/libpico/libpico_utils.h @@ -139,7 +139,7 @@ static inline void pico_get_group_config(int *node_size, int *node_rank, int *no } while (0) #ifdef PICO_NCCL -size_t pico_get_nccl_type_size(ncclDataType_t type) +static inline size_t pico_get_nccl_type_size(ncclDataType_t type) { switch (type) { @@ -879,4 +879,40 @@ static inline unsigned int floor_power_of_two(unsigned int n) { return n - (n >> 1); } +static inline int same_prefix_negabinary(int a, int b, int total_bits, int prefix_len) +{ + a = binary_to_negabinary(a); + b = binary_to_negabinary(b); + + int shift = total_bits - prefix_len; + return (a >> shift) == (b >> shift); +} + +static inline int logical_rank_for_bine_dh_root(int x, int root, int size) +{ + if ((root % 2) == 0) + // use the normal rotation + return mod(x - root, size); + else + // use the mirrored rotation + return mod(root - x, size); +} + +static inline size_t num_bytes_fun(const int counts[],const int displs[],int dtype,int size) +{ + size_t num_bytes = 0; + + for (int i = 0; i < size; i++) + { + if (counts[i] <= 0) + continue; + size_t end =((size_t)displs[i] + (size_t)counts[i]) *(size_t)dtype; + if (end > num_bytes) + num_bytes = end; + } + + return num_bytes; +} + + #endif // LIBPICO_UTILS_H diff --git a/pico_core/pico_core_alltoallv.c b/pico_core/pico_core_alltoallv.c new file mode 100644 index 00000000..544dbb24 --- /dev/null +++ b/pico_core/pico_core_alltoallv.c @@ -0,0 +1,39 @@ +#include +#include "pico_core_utils.h" + +int alltoallv_allocator(ALLOCATOR_ARGS) { + *sbuf = malloc(count * type_size); + *rbuf = malloc(count * type_size); + *rbuf_gt = malloc(count * type_size); + + if (*sbuf == NULL || *rbuf == NULL || *rbuf_gt == NULL) { + fprintf(stderr, "Error: failed to allocate host buffers for alltoallv\n"); + return -1; + } + + return 0; +} + +#if defined PICO_MPI_CUDA_AWARE || defined PICO_NCCL +int alltoallv_allocator_cuda(ALLOCATOR_ARGS) { + cudaError_t err; + + PICO_CORE_CUDA_CHECK(cudaMalloc(sbuf, count * type_size), err); + if (err != cudaSuccess) return -1; + + PICO_CORE_CUDA_CHECK(cudaMalloc(rbuf, count * type_size), err); + if (err != cudaSuccess) { + cudaFree(*sbuf); + return -1; + } + + PICO_CORE_CUDA_CHECK(cudaMalloc(rbuf_gt, count * type_size), err); + if (err != cudaSuccess) { + cudaFree(*sbuf); + cudaFree(*rbuf); + return -1; + } + + return 0; +} +#endif \ No newline at end of file diff --git a/pico_core/pico_core_utils.c b/pico_core/pico_core_utils.c index f5dbdbd9..c29e5fcc 100644 --- a/pico_core/pico_core_utils.c +++ b/pico_core/pico_core_utils.c @@ -27,6 +27,7 @@ static inline coll_t get_collective_from_string(const char *coll_str) { CHECK_STR(coll_str, "ALLREDUCE", ALLREDUCE); CHECK_STR(coll_str, "ALLGATHER", ALLGATHER); CHECK_STR(coll_str, "ALLTOALL", ALLTOALL); + CHECK_STR(coll_str, "ALLTOALLV", ALLTOALLV); CHECK_STR(coll_str, "BCAST", BCAST); CHECK_STR(coll_str, "GATHER", GATHER); CHECK_STR(coll_str, "REDUCE", REDUCE); @@ -53,6 +54,8 @@ static inline allocator_func_ptr get_allocator(coll_t collective) { return allgather_allocator; case ALLTOALL: return alltoall_allocator; + case ALLTOALLV: + return alltoallv_allocator; case BCAST: return bcast_allocator; case GATHER: @@ -77,6 +80,8 @@ static inline allocator_func_ptr get_allocator_cuda(coll_t collective) { return allgather_allocator_cuda; case ALLTOALL: return alltoall_allocator_cuda; + case ALLTOALLV: + return alltoallv_allocator_cuda; case BCAST: return bcast_allocator_cuda; case GATHER: @@ -165,7 +170,7 @@ static inline alltoall_func_ptr get_alltoall_function(const char *algorithm) { #ifndef PICO_NCCL CHECK_STR(algorithm, "bine_over", alltoall_bine); CHECK_STR(algorithm, "pairwise_ompi_over", alltoall_pairwise_ompi); - + CHECK_STR(algorithm, "bine_dh_over", alltoall_bine_DH); PICO_CORE_DEBUG_PRINT_STR("MPI_Alltoall"); return alltoall_wrapper; #else @@ -175,6 +180,21 @@ static inline alltoall_func_ptr get_alltoall_function(const char *algorithm) { #endif } +static inline alltoallv_func_ptr get_alltoallv_function(const char *algorithm){ +#ifndef PICO_NCCL + CHECK_STR(algorithm, "bine_dh_over", alltoallv_bine_DH); + + PICO_CORE_DEBUG_PRINT_STR("MPI_Alltoallv"); + return alltoallv_wrapper; +#else + CHECK_STR(algorithm, "nccl_spreadout_alltoallv_over", alltoallv_nccl_spreadout); + CHECK_STR(algorithm, "nccl_fanout_alltoallv_over", alltoallv_nccl_fanout); + + fprintf(stderr, "Failed: unsupported NCCL Alltoallv algorithm: %s\n", algorithm); + exit(EXIT_FAILURE); +#endif +} + /** * @brief Select and returns the appropriate bcast function based * on the algorithm. It returns the default bcast function if the @@ -338,6 +358,9 @@ int get_routine(test_routine_t *test_routine, const char *algorithm) { case ALLTOALL: test_routine->function.alltoall = get_alltoall_function(algorithm); break; + case ALLTOALLV: + test_routine->function.alltoallv = get_alltoallv_function(algorithm); + break; case BCAST: test_routine->function.bcast = get_bcast_function(algorithm); break; @@ -485,6 +508,9 @@ int coll_memcpy_host_to_device(void** d_buf, void** buf, size_t count, size_t ty case ALLTOALL: PICO_CORE_CUDA_CHECK(cudaMemcpy(*d_buf, *buf, count * type_size, cudaMemcpyHostToDevice), err); break; + case ALLTOALLV: + PICO_CORE_CUDA_CHECK(cudaMemcpy(*d_buf, *buf, count * type_size, cudaMemcpyHostToDevice), err); + break; case BCAST: if (rank == 0) { PICO_CORE_CUDA_CHECK(cudaMemcpy(*d_buf, *buf, count * type_size, cudaMemcpyHostToDevice), err); @@ -529,6 +555,9 @@ int coll_memcpy_device_to_host(void** d_buf, void** buf, size_t count, size_t ty case ALLTOALL: PICO_CORE_CUDA_CHECK(cudaMemcpy(*buf, *d_buf, count * type_size, cudaMemcpyDeviceToHost), err); break; + case ALLTOALLV: + PICO_CORE_CUDA_CHECK(cudaMemcpy(*buf, *d_buf, count * type_size, cudaMemcpyDeviceToHost), err); + break; case BCAST: if (rank != 0) { PICO_CORE_CUDA_CHECK(cudaMemcpy(*buf, *d_buf, count * type_size, cudaMemcpyDeviceToHost), err); @@ -582,6 +611,28 @@ int run_coll_once(test_routine_t test_routine, void *sbuf, void *rbuf, rbuf, local_count, dtype, comm); break; + case ALLTOALLV: { + int *scounts = malloc(comm_sz * sizeof(int)); + int *sdispls = malloc(comm_sz * sizeof(int)); + int *rcounts_v = malloc(comm_sz * sizeof(int)); + int *rdispls = malloc(comm_sz * sizeof(int)); + + for (int i = 0; i < comm_sz; i++) { + scounts[i] = (int)local_count; + rcounts_v[i] = (int)local_count; + sdispls[i] = i * (int)local_count; + rdispls[i] = i * (int)local_count; + } + + ret = test_routine.function.alltoallv(sbuf, scounts, sdispls, dtype, + rbuf, rcounts_v, rdispls, dtype, comm); + + free(scounts); + free(sdispls); + free(rcounts_v); + free(rdispls); + break; + } case BCAST: ret = test_routine.function.bcast(sbuf, count, dtype, 0, comm); break; @@ -638,6 +689,29 @@ int test_loop(test_routine_t test_routine, void *sbuf, void *rbuf, size_t count, rbuf, local_count, dtype, comm, iter, times, test_routine); break; + case ALLTOALLV: { + int *scounts = malloc(comm_sz * sizeof(int)); + int *sdispls = malloc(comm_sz * sizeof(int)); + int *rcounts_v = malloc(comm_sz * sizeof(int)); + int *rdispls = malloc(comm_sz * sizeof(int)); + + for (int i = 0; i < comm_sz; i++) { + scounts[i] = (int)local_count; + rcounts_v[i] = (int)local_count; + sdispls[i] = i * (int)local_count; + rdispls[i] = i * (int)local_count; + } + + ret = alltoallv_test_loop(sbuf, scounts, sdispls, dtype, + rbuf, rcounts_v, rdispls, dtype, + comm, iter, times, test_routine); + + free(scounts); + free(sdispls); + free(rcounts_v); + free(rdispls); + break; + } case BCAST: ret = bcast_test_loop(sbuf, count, dtype, 0, comm, iter, times, test_routine); @@ -691,6 +765,39 @@ int test_loop(test_routine_t test_routine, void *sbuf, void *rbuf, size_t count, case ALLTOALL: ret = alltoall_test_loop(sbuf, rbuf, local_count, dtype, nccl_comm, stream, iter, times, test_routine); + case ALLTOALLV: { + int *scounts = malloc(comm_sz * sizeof(int)); + int *sdispls = malloc(comm_sz * sizeof(int)); + int *rcounts_v = malloc(comm_sz * sizeof(int)); + int *rdispls = malloc(comm_sz * sizeof(int)); + + if (scounts == NULL || sdispls == NULL || rcounts_v == NULL || rdispls == NULL) { + free(scounts); + free(sdispls); + free(rcounts_v); + free(rdispls); + fprintf(stderr, "Error: malloc failed in NCCL Alltoallv test_loop\n"); + return -1; + } + + for (int i = 0; i < comm_sz; i++) { + scounts[i] = (int)local_count; + rcounts_v[i] = (int)local_count; + sdispls[i] = i * (int)local_count; + rdispls[i] = i * (int)local_count; + } + + ret = alltoallv_test_loop(sbuf, scounts, sdispls, + rbuf, rcounts_v, rdispls, + dtype, nccl_comm, stream, + iter, times, test_routine); + + free(scounts); + free(sdispls); + free(rcounts_v); + free(rdispls); + break; + } break; case BCAST: ret = bcast_test_loop(sbuf, count, dtype, 0, nccl_comm, stream, @@ -742,6 +849,27 @@ int ground_truth_check(test_routine_t test_routine, void *sbuf, void *rbuf, rbuf_gt, count / (size_t) comm_sz, dtype, comm); GT_CHECK_BUFFER(rbuf, rbuf_gt, count / (size_t) comm_sz, dtype, comm); break; + case ALLTOALLV: { + int *scounts = malloc(comm_sz * sizeof(int)); + int *sdispls = malloc(comm_sz * sizeof(int)); + int *rcounts_v = malloc(comm_sz * sizeof(int)); + int *rdispls = malloc(comm_sz * sizeof(int)); + size_t local_count = count / (size_t)comm_sz; + for (int i = 0; i < comm_sz; i++) { + scounts[i] = (int)local_count; + rcounts_v[i] = (int)local_count; + sdispls[i] = i * (int)local_count; + rdispls[i] = i * (int)local_count; + } + PMPI_Alltoallv(sbuf, scounts, sdispls, dtype, + rbuf_gt, rcounts_v, rdispls, dtype, comm); + GT_CHECK_BUFFER(rbuf, rbuf_gt, count, dtype, comm); + free(scounts); + free(sdispls); + free(rcounts_v); + free(rdispls); + break; + } case BCAST: if(rank == 0) { memcpy(rbuf_gt, sbuf, count * type_size); @@ -1057,7 +1185,8 @@ int rand_sbuf_generator(void *sbuf, MPI_Datatype dtype, size_t count, size_t real_sbuf_count = (test_routine.collective == ALLGATHER || test_routine.collective == GATHER || - test_routine.collective == ALLTOALL) ? + test_routine.collective == ALLTOALL || + test_routine.collective == ALLTOALLV) ? count / (size_t) comm_sz : count; for(size_t i = 0; i < real_sbuf_count; i++) { @@ -1264,10 +1393,11 @@ int debug_sbuf_generator(void *sbuf, MPI_Datatype dtype, size_t count, } size_t real_sbuf_count = - (test_routine.collective == ALLGATHER || - test_routine.collective == GATHER || - test_routine.collective == ALLTOALL) ? - count / (size_t) comm_sz : count; + (test_routine.collective == ALLGATHER || + test_routine.collective == GATHER || + test_routine.collective == ALLTOALL || + test_routine.collective == ALLTOALLV) ? + count / (size_t) comm_sz : count; for(int i=0; i< real_sbuf_count; i++){ if(dtype == MPI_INT64_T) { diff --git a/pico_core/pico_core_utils.h b/pico_core/pico_core_utils.h index 3011ce13..29235f5e 100644 --- a/pico_core/pico_core_utils.h +++ b/pico_core/pico_core_utils.h @@ -72,6 +72,7 @@ typedef enum{ ALLREDUCE = 0, ALLGATHER, ALLTOALL, + ALLTOALLV, BCAST, GATHER, REDUCE, @@ -108,6 +109,7 @@ typedef int (*allocator_func_ptr)(ALLOCATOR_ARGS); int allreduce_allocator(ALLOCATOR_ARGS); int allgather_allocator(ALLOCATOR_ARGS); int alltoall_allocator(ALLOCATOR_ARGS); +int alltoallv_allocator(ALLOCATOR_ARGS); int bcast_allocator(ALLOCATOR_ARGS); int gather_allocator(ALLOCATOR_ARGS); int reduce_allocator(ALLOCATOR_ARGS); @@ -122,6 +124,7 @@ int scatter_allocator(ALLOCATOR_ARGS); typedef int (*allreduce_func_ptr)(ALLREDUCE_MPI_ARGS); typedef int (*allgather_func_ptr)(ALLGATHER_MPI_ARGS); typedef int (*alltoall_func_ptr)(ALLTOALL_MPI_ARGS); +typedef int (*alltoallv_func_ptr)(ALLTOALLV_MPI_ARGS); typedef int (*bcast_func_ptr)(BCAST_MPI_ARGS); typedef int (*gather_func_ptr)(GATHER_MPI_ARGS); typedef int (*reduce_func_ptr)(REDUCE_MPI_ARGS); @@ -137,6 +140,9 @@ static inline int allgather_wrapper(ALLGATHER_MPI_ARGS){ static inline int alltoall_wrapper(ALLTOALL_MPI_ARGS){ return MPI_Alltoall(sbuf, (int)scount, sdtype, rbuf, (int)rcount, rdtype, comm); } +static inline int alltoallv_wrapper(ALLTOALLV_MPI_ARGS){ + return MPI_Alltoallv(sbuf, scounts, sdispls, sdtype,rbuf, rcounts, rdispls, rdtype, comm); +} static inline int bcast_wrapper(BCAST_MPI_ARGS){ return MPI_Bcast(buf, (int)count, dtype, root, comm); } @@ -153,6 +159,7 @@ static inline int scatter_wrapper(SCATTER_MPI_ARGS){ typedef ncclResult_t (*allreduce_func_ptr)(ALLREDUCE_NCCL_ARGS); typedef ncclResult_t (*allgather_func_ptr)(ALLGATHER_NCCL_ARGS); typedef ncclResult_t (*alltoall_func_ptr)(ALLTOALL_NCCL_ARGS); +typedef int (*alltoallv_func_ptr)(ALLTOALLV_NCCL_ARGS); typedef ncclResult_t (*bcast_func_ptr)(BCAST_NCCL_ARGS); typedef ncclResult_t (*gather_func_ptr)(GATHER_NCCL_ARGS); typedef ncclResult_t (*reduce_func_ptr)(REDUCE_NCCL_ARGS); @@ -189,6 +196,7 @@ typedef struct { allreduce_func_ptr allreduce; allgather_func_ptr allgather; alltoall_func_ptr alltoall; + alltoallv_func_ptr alltoallv; bcast_func_ptr bcast; gather_func_ptr gather; reduce_func_ptr reduce; @@ -211,6 +219,7 @@ typedef struct { int allreduce_allocator_cuda(ALLOCATOR_ARGS); int allgather_allocator_cuda(ALLOCATOR_ARGS); int alltoall_allocator_cuda(ALLOCATOR_ARGS); +int alltoallv_allocator_cuda(ALLOCATOR_ARGS); int bcast_allocator_cuda(ALLOCATOR_ARGS); int gather_allocator_cuda(ALLOCATOR_ARGS); int reduce_allocator_cuda(ALLOCATOR_ARGS); @@ -298,6 +307,7 @@ static inline int OP_NAME##_test_loop(ARGS, int iter, double *times, \ DEFINE_TEST_LOOP(allreduce, ALLREDUCE_MPI_ARGS, allreduce(sbuf, rbuf, count, dtype, MPI_SUM, comm)) DEFINE_TEST_LOOP(allgather, ALLGATHER_MPI_ARGS, allgather(sbuf, scount, sdtype, rbuf, rcount, rdtype, comm)) DEFINE_TEST_LOOP(alltoall, ALLTOALL_MPI_ARGS, alltoall(sbuf, scount, sdtype, rbuf, rcount, rdtype, comm)) +DEFINE_TEST_LOOP(alltoallv, ALLTOALLV_MPI_ARGS, alltoallv(sbuf, scounts, sdispls, sdtype, rbuf, rcounts, rdispls, rdtype, comm)) DEFINE_TEST_LOOP(bcast, BCAST_MPI_ARGS, bcast(buf, count, dtype, 0, comm)) DEFINE_TEST_LOOP(gather, GATHER_MPI_ARGS, gather(sbuf, scount, sdtype, rbuf, rcount, rdtype, 0, comm)) DEFINE_TEST_LOOP(reduce, REDUCE_MPI_ARGS, reduce(sbuf, rbuf, count, dtype, MPI_SUM, 0, comm)) @@ -359,6 +369,7 @@ static inline int OP_NAME##_test_loop(ARGS, int iter, double *times, \ DEFINE_TEST_LOOP(allreduce, ALLREDUCE_NCCL_ARGS, allreduce(sbuf, rbuf, count, dtype, ncclSum, nccl_comm, stream)) DEFINE_TEST_LOOP(allgather, ALLGATHER_NCCL_ARGS, allgather(sbuf, rbuf, count, dtype, nccl_comm, stream)) DEFINE_TEST_LOOP(alltoall, ALLTOALL_NCCL_ARGS, alltoall(sbuf, rbuf, count, dtype, nccl_comm, stream)) +DEFINE_TEST_LOOP(alltoallv, ALLTOALLV_NCCL_ARGS, alltoallv(sbuf, scounts, sdispls, rbuf, rcounts, rdispls, dtype, nccl_comm, stream)) DEFINE_TEST_LOOP(bcast, BCAST_NCCL_ARGS, bcast(buf, count, dtype, 0, nccl_comm, stream)) DEFINE_TEST_LOOP(gather, GATHER_NCCL_ARGS, gather(sbuf, rbuf, count, dtype, 0, nccl_comm, stream)) DEFINE_TEST_LOOP(reduce, REDUCE_NCCL_ARGS, reduce(sbuf, rbuf, count, dtype, ncclSum, 0, nccl_comm, stream)) diff --git a/selector/change_dynamic_rules.py b/selector/change_dynamic_rules.py index d58edb11..b1f0b088 100644 --- a/selector/change_dynamic_rules.py +++ b/selector/change_dynamic_rules.py @@ -81,7 +81,7 @@ def modify_dynamic_rule(rule_file: str | os.PathLike, collective_type: str, new_ for i, line in enumerate(lines): if re.search(pattern, line): if i + 4 < len(lines): # Ensure we don't go out of bounds - lines[i+4] = f"0 {new_rule} 0 0 # Algorithm\n" + lines[i+4] = f"0 {int(new_rule)} 0 0 # Algorithm\n" with open(rule_file, 'w') as txt_file: txt_file.writelines(lines) return @@ -110,9 +110,15 @@ def main(): print(f"{__file__}: Environment variables not set.", file=sys.stderr) print(f"DYNAMIC_RULE_FILE={dynamic_rule_file}\nCOLLECTIVE_TYPE={collective_type}\nALGORITHM_DECLARATIONS_DIR={os.getenv('ALGORITHM_DECLARATIONS_DIR')}", file=sys.stderr) sys.exit(1) - new_rule = find_dynamic_rule(algorithm_decls_dirs, collective_type, algorithm) - modify_dynamic_rule(dynamic_rule_file, collective_type, new_rule) - + if not str(new_rule).isdigit(): + print( + f"{__file__}: algorithm {algorithm} for {collective_type} maps to non-Open-MPI rule '{new_rule}'. " + "Resetting Open MPI dynamic rule to default.", + file=sys.stderr) + modify_dynamic_rule(dynamic_rule_file, collective_type, 0) + sys.exit(0) + + modify_dynamic_rule(dynamic_rule_file, collective_type, int(new_rule)) if __name__ == "__main__": main() diff --git a/selector/ompi_dynamic_rules.txt b/selector/ompi_dynamic_rules.txt index 4042af06..fcc915bb 100644 --- a/selector/ompi_dynamic_rules.txt +++ b/selector/ompi_dynamic_rules.txt @@ -1,4 +1,4 @@ -9 # Number of collectives +10 # Number of collectives 0 # Collective ID = ALLGATHER 1 # Number of comm sizes 2 # Comm size @@ -17,6 +17,12 @@ 1 # Number of message sizes 0 0 0 0 # Algorithm +4 # Collective ID = ALLTOALLV +1 # Number of comm sizes +2 # Comm size +1 # Number of message sizes +0 0 0 0 # Algorithm + 6 # Collective ID = BARRIER 1 # Number of comm sizes 2 # Comm size diff --git a/tui/models.py b/tui/models.py index 7ac7f3f2..203c2159 100644 --- a/tui/models.py +++ b/tui/models.py @@ -569,6 +569,7 @@ def from_dict(cls, gpu_json: Dict[str, Any]) -> "GPUSupport": class CollectiveType(Enum): UNKNOWN = 'unknown' ALLTOALL = 'alltoall' + ALLTOALLV = 'alltoallv' ALLREDUCE = 'allreduce' ALLGATHER = 'allgather' BARRIER = 'barrier' @@ -586,6 +587,8 @@ def from_str(cls, value: str): value = value.lower() if value == 'alltoall': return cls.ALLTOALL + elif value == 'alltoallv': + return cls.ALLTOALLV elif value == 'allreduce': return cls.ALLREDUCE elif value == 'allgather': diff --git a/tui/tui/steps/libraries.py b/tui/tui/steps/libraries.py index b3b41237..b0e23a9b 100644 --- a/tui/tui/steps/libraries.py +++ b/tui/tui/steps/libraries.py @@ -77,6 +77,7 @@ def compose(self) -> ComposeResult: "Allgather", "Allreduce", "Alltoall", + "Alltoallv", "Broadcast", "Gather", "Reduce",