diff --git a/examples/ccl/all_to_all.cc b/examples/ccl/all_to_all.cc new file mode 100644 index 0000000..643d7c2 --- /dev/null +++ b/examples/ccl/all_to_all.cc @@ -0,0 +1,330 @@ +/** + * InfiniCCL Example: Thread-per-GPU Single-Node AllToAll + * + * This example creates one native CCL rank per GPU and validates an + * out-of-place all-to-all exchange composed from grouped point-to-point calls. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// Public API +#include "infiniccl.h" + +// Example-Specific Utilities +#include "utils.h" + +// Internal Headers (Accessible via example-specific include paths, technically +// not public APIs) +#include "backend_manifest.h" + +using namespace infini::ccl; + +namespace { + +struct ScenarioState { + std::atomic correct{true}; + std::atomic completed{0}; +}; + +struct ThreadArgs { + int rank; + int size; + infinicclUniqueId id; + size_t count_per_peer; + int warmup_iterations; + int profile_iterations; + ScenarioState *state; +}; + +template +bool ParsePositiveNumber(const char *text, T *value) { + if (!text || !value) { + return false; + } + + T parsed{}; + const char *end = text + std::strlen(text); + const auto result = std::from_chars(text, end, parsed); + if (result.ec != std::errc{} || result.ptr != end || parsed <= 0) { + return false; + } + + *value = parsed; + return true; +} + +float BlockValue(int source_rank, int destination_rank) { + return static_cast((source_rank + 1) * 1000 + destination_rank + 1); +} + +void FillInput(std::vector *input, size_t count_per_peer, int world_size, + int rank) { + for (int destination = 0; destination < world_size; ++destination) { + const size_t offset = static_cast(destination) * count_per_peer; + std::fill_n(input->begin() + offset, count_per_peer, + BlockValue(rank, destination)); + } +} + +bool ValidateAllToAll(const std::vector &result, size_t count_per_peer, + int world_size, int rank) { + bool correct = true; + for (int source = 0; source < world_size; ++source) { + const size_t offset = static_cast(source) * count_per_peer; + const bool block_correct = Validator::ValidateResult( + result.data() + offset, count_per_peer, BlockValue(source, rank), rank); + correct = block_correct && correct; + } + return correct; +} + +void PrintAllToAllMetrics(size_t count_per_peer, int world_size, + double elapsed_ms) { + constexpr double kBytesPerMiB = 1024.0 * 1024.0; + constexpr double kBytesPerGB = 1.0e9; + const double peer_bytes = static_cast(count_per_peer) * sizeof(float); + const double total_bytes = peer_bytes * static_cast(world_size); + const auto original_flags = std::cout.flags(); + const auto original_precision = std::cout.precision(); + + std::cout << "Data size per peer: " << count_per_peer << " floats (" + << std::fixed << std::setprecision(2) << peer_bytes / kBytesPerMiB + << " MiB)" << std::endl; + std::cout << "Total data per rank: " + << count_per_peer * static_cast(world_size) << " floats (" + << total_bytes / kBytesPerMiB << " MiB)" << std::endl; + std::cout << "Time: " << std::setprecision(3) << elapsed_ms << " ms" + << std::endl; + if (elapsed_ms > 0.0) { + const double algorithm_bandwidth = + total_bytes / kBytesPerGB / (elapsed_ms / 1000.0); + const double bus_bandwidth = algorithm_bandwidth * + static_cast(world_size - 1) / + static_cast(world_size); + std::cout << "Throughput: " << std::setprecision(2) << bus_bandwidth + << " GB/s (Bus BW)" << std::endl; + std::cout << "Alg Bandwidth: " << algorithm_bandwidth << " GB/s" + << std::endl; + } else { + std::cout << "Throughput: N/A (Bus BW)" << std::endl; + std::cout << "Alg Bandwidth: N/A" << std::endl; + } + + std::cout.flags(original_flags); + std::cout.precision(original_precision); +} + +void WaitForAll(ScenarioState *state, int world_size) { + state->completed.fetch_add(1, std::memory_order_acq_rel); + while (state->completed.load(std::memory_order_acquire) < world_size) { + std::this_thread::yield(); + } +} + +void PrintResult(bool correct, const std::vector &result, + size_t count_per_peer, int world_size, double elapsed_ms) { + constexpr const char *kGreen = "\033[32m"; + constexpr const char *kRed = "\033[31m"; + constexpr const char *kReset = "\033[0m"; + + std::cout << "\n=== CCL AllToAll Results ===" << std::endl; + std::cout << "Correct: " + << (correct ? (kGreen + std::string("YES") + kReset) + : (kRed + std::string("NO") + kReset)) + << std::endl; + PrintAllToAllMetrics(count_per_peer, world_size, elapsed_ms); + std::cout << "Sample receive blocks: "; + for (int source = 0; source < std::min(world_size, 4); ++source) { + const size_t offset = static_cast(source) * count_per_peer; + std::cout << "[src" << source << ": " << result[offset] << "] "; + } + std::cout << std::endl; +} + +void WorkerThread(ThreadArgs args) { + constexpr Device::Type kDevType = + ListGetBest(EnabledDevices{}); + using Rt = Runtime; + + CHECK_RT(Rt, Rt::SetDevice(args.rank)); + + std::array hostname{}; + if (gethostname(hostname.data(), hostname.size()) != 0) { + std::cerr << "Failed to query the hostname for the AllToAll worker." + << std::endl; + std::exit(EXIT_FAILURE); + } + hostname.back() = '\0'; + std::cout << "[Rank " << args.rank << "] Host: " << hostname.data() + << " | GPU: " << Device::StringFromType(kDevType) << " | Device " + << args.rank << std::endl; + + infinicclComm_t comm = nullptr; + CHECK_INFINI(infinicclCommInitRank(&comm, args.size, args.id, args.rank)); + + const size_t total_elements = + args.count_per_peer * static_cast(args.size); + const size_t total_bytes = total_elements * sizeof(float); + std::vector h_send(total_elements); + std::vector h_recv(total_elements, 0.0f); + FillInput(&h_send, args.count_per_peer, args.size, args.rank); + + float *d_send = nullptr; + float *d_recv = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_send), total_bytes)); + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_recv), total_bytes)); + CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), total_bytes, + Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + for (int i = 0; i < args.warmup_iterations; ++i) { + CHECK_INFINI(infinicclAllToAll(d_send, d_recv, args.count_per_peer, + infinicclFloat32, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + Timer timer; + for (int i = 0; i < args.profile_iterations; ++i) { + CHECK_INFINI(infinicclAllToAll(d_send, d_recv, args.count_per_peer, + infinicclFloat32, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const double elapsed_ms = + timer.ElapsedMs() / static_cast(args.profile_iterations); + + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes, + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const bool local_correct = + ValidateAllToAll(h_recv, args.count_per_peer, args.size, args.rank); + if (!local_correct) { + args.state->correct.store(false, std::memory_order_relaxed); + } + WaitForAll(args.state, args.size); + + if (args.rank == 0) { + PrintResult(args.state->correct.load(std::memory_order_acquire), h_recv, + args.count_per_peer, args.size, elapsed_ms); + } + + CHECK_RT(Rt, Rt::Free(d_send)); + CHECK_RT(Rt, Rt::Free(d_recv)); + CHECK_INFINI(infinicclCommDestroy(comm)); +} + +void PrintUsage(const char *program) { + std::cout << "Usage: " << program << " [options]\n" + << "Options:\n" + << " -g Number of GPUs (default: 8)\n" + << " -w Warmup iterations (default: 2)\n" + << " -p Profile iterations (default: 20)\n" + << " -n Elements sent to each peer (default: " + "1048576)\n"; +} + +} // namespace + +int main(int argc, char **argv) { + int num_gpus = 8; + int warmup_iterations = 2; + int profile_iterations = 20; + size_t count_per_peer = 1 << 20; + + int opt = 0; + while ((opt = getopt(argc, argv, "g:w:p:n:h")) != -1) { + bool parsed = false; + switch (opt) { + case 'g': + parsed = ParsePositiveNumber(optarg, &num_gpus); + break; + case 'w': + parsed = ParsePositiveNumber(optarg, &warmup_iterations); + break; + case 'p': + parsed = ParsePositiveNumber(optarg, &profile_iterations); + break; + case 'n': + parsed = ParsePositiveNumber(optarg, &count_per_peer); + break; + case 'h': + PrintUsage(argv[0]); + return EXIT_SUCCESS; + default: + PrintUsage(argv[0]); + return EXIT_FAILURE; + } + + if (!parsed) { + std::cerr << "Invalid positive numeric option for AllToAll." << std::endl; + return EXIT_FAILURE; + } + } + + if (optind != argc) { + std::cerr << "Unexpected positional argument for AllToAll." << std::endl; + return EXIT_FAILURE; + } + if (static_cast(num_gpus) > + std::numeric_limits::max() / count_per_peer || + count_per_peer * static_cast(num_gpus) > + std::numeric_limits::max() / sizeof(float)) { + std::cerr << "AllToAll buffer size overflows `size_t`." << std::endl; + return EXIT_FAILURE; + } + + std::array hostname{}; + if (gethostname(hostname.data(), hostname.size()) != 0) { + std::cerr << "Failed to query the hostname for AllToAll." << std::endl; + return EXIT_FAILURE; + } + hostname.back() = '\0'; + std::cout << "[Main Process] Host: " << hostname.data() + << " | Target GPUs: " << num_gpus << std::endl; + std::cout << "[Main Process] Count per peer: " << count_per_peer + << " floats | Warmup: " << warmup_iterations + << " | Profile: " << profile_iterations << std::endl; + + infinicclUniqueId shared_id{}; + CHECK_INFINI(infinicclGetUniqueId(&shared_id)); + + ScenarioState state; + std::vector threads; + threads.reserve(num_gpus); + for (int rank = 0; rank < num_gpus; ++rank) { + ThreadArgs args{rank, num_gpus, shared_id, + count_per_peer, warmup_iterations, profile_iterations, + &state}; + threads.emplace_back(WorkerThread, args); + } + + for (auto &thread : threads) { + if (thread.joinable()) { + thread.join(); + } + } + + const bool correct = state.correct.load(std::memory_order_acquire); + if (correct) { + std::cout << "[Main Process] CCL AllToAll validation passed." << std::endl; + } else { + std::cerr << "[Main Process] CCL AllToAll validation failed." << std::endl; + } + std::cout + << "[Main Process] All worker threads joined. InfiniCCL finalized safely." + << std::endl; + return correct ? EXIT_SUCCESS : EXIT_FAILURE; +} diff --git a/examples/ccl_mpi_hybrid/all_to_all.cc b/examples/ccl_mpi_hybrid/all_to_all.cc new file mode 100644 index 0000000..0c877cc --- /dev/null +++ b/examples/ccl_mpi_hybrid/all_to_all.cc @@ -0,0 +1,313 @@ +/** + * InfiniCCL Example: AllToAll (OpenMPI + CCL Hybrid) + * + * This example first uses AllToAll through an OpenMPI inter communicator to + * distribute a native CCL unique ID. It then initializes a native CCL + * communicator and validates an out-of-place GPU all-to-all exchange. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// Public API +#include "infiniccl.h" + +// Example-Specific Utilities +#include "utils.h" + +// Internal Headers (Accessible via example-specific include paths, technically +// not public APIs) +#include "backend_manifest.h" +#include "device.h" +#include "runtime.h" +#include "traits.h" + +using namespace infini::ccl; + +namespace { + +bool ParseLocalRank(const char *text, int *local_rank) { + if (!text || !local_rank) { + return false; + } + + int parsed = -1; + const char *end = text + std::strlen(text); + const auto result = std::from_chars(text, end, parsed); + if (result.ec != std::errc{} || result.ptr != end || parsed < 0) { + return false; + } + + *local_rank = parsed; + return true; +} + +float BlockValue(int source_rank, int destination_rank) { + return static_cast((source_rank + 1) * 1000 + destination_rank + 1); +} + +void FillInput(std::vector *input, size_t count_per_peer, int world_size, + int rank) { + for (int destination = 0; destination < world_size; ++destination) { + const size_t offset = static_cast(destination) * count_per_peer; + std::fill_n(input->begin() + offset, count_per_peer, + BlockValue(rank, destination)); + } +} + +bool ValidateAllToAll(const std::vector &result, size_t count_per_peer, + int world_size, int rank) { + bool correct = true; + for (int source = 0; source < world_size; ++source) { + const size_t offset = static_cast(source) * count_per_peer; + const bool block_correct = Validator::ValidateResult( + result.data() + offset, count_per_peer, BlockValue(source, rank), rank); + correct = block_correct && correct; + } + return correct; +} + +void PrintAllToAllMetrics(size_t count_per_peer, int world_size, + double elapsed_ms) { + constexpr double kBytesPerMiB = 1024.0 * 1024.0; + constexpr double kBytesPerGB = 1.0e9; + const double peer_bytes = static_cast(count_per_peer) * sizeof(float); + const double total_bytes = peer_bytes * static_cast(world_size); + const auto original_flags = std::cout.flags(); + const auto original_precision = std::cout.precision(); + + std::cout << "Data size per peer: " << count_per_peer << " floats (" + << std::fixed << std::setprecision(2) << peer_bytes / kBytesPerMiB + << " MiB)" << std::endl; + std::cout << "Total data per rank: " + << count_per_peer * static_cast(world_size) << " floats (" + << total_bytes / kBytesPerMiB << " MiB)" << std::endl; + std::cout << "Time: " << std::setprecision(3) << elapsed_ms << " ms" + << std::endl; + if (elapsed_ms > 0.0) { + const double algorithm_bandwidth = + total_bytes / kBytesPerGB / (elapsed_ms / 1000.0); + const double bus_bandwidth = algorithm_bandwidth * + static_cast(world_size - 1) / + static_cast(world_size); + std::cout << "Throughput: " << std::setprecision(2) << bus_bandwidth + << " GB/s (Bus BW)" << std::endl; + std::cout << "Alg Bandwidth: " << algorithm_bandwidth << " GB/s" + << std::endl; + } else { + std::cout << "Throughput: N/A (Bus BW)" << std::endl; + std::cout << "Alg Bandwidth: N/A" << std::endl; + } + + std::cout.flags(original_flags); + std::cout.precision(original_precision); +} + +void PrintResult(bool correct, const std::vector &result, + size_t count_per_peer, int world_size, double elapsed_ms) { + constexpr const char *kGreen = "\033[32m"; + constexpr const char *kRed = "\033[31m"; + constexpr const char *kReset = "\033[0m"; + + std::cout << "\n=== Hybrid CCL AllToAll Results ===" << std::endl; + std::cout << "Correct: " + << (correct ? (kGreen + std::string("YES") + kReset) + : (kRed + std::string("NO") + kReset)) + << std::endl; + PrintAllToAllMetrics(count_per_peer, world_size, elapsed_ms); + std::cout << "Sample receive blocks: "; + for (int source = 0; source < std::min(world_size, 4); ++source) { + const size_t offset = static_cast(source) * count_per_peer; + std::cout << "[src" << source << ": " << result[offset] << "] "; + } + std::cout << std::endl; +} + +bool RunAllToAllExample(int argc, char **argv) { + constexpr Device::Type kDevType = + ListGetBest(EnabledDevices{}); + using Rt = Runtime; + + constexpr int kWarmupIterations = 2; + constexpr int kProfileIterations = 20; + constexpr size_t kCountPerPeer = 1 << 20; + + CHECK_INFINI(infinicclInit(&argc, &argv)); + + int rank = -1; + int size = 0; + CHECK_INFINI(infinicclGetRank(&rank)); + CHECK_INFINI(infinicclGetSize(&size)); + if (size <= 0) { + std::cerr << "Invalid world size for hybrid AllToAll." << std::endl; + std::exit(EXIT_FAILURE); + } + + int local_rank = -1; + if (!ParseLocalRank(std::getenv("OMPI_COMM_WORLD_LOCAL_RANK"), &local_rank)) { + std::cerr << "Missing or invalid `OMPI_COMM_WORLD_LOCAL_RANK`." + << std::endl; + std::exit(EXIT_FAILURE); + } + CHECK_RT(Rt, Rt::SetDevice(local_rank)); + + std::array hostname{}; + if (gethostname(hostname.data(), hostname.size()) != 0) { + std::cerr << "Failed to query the hostname for hybrid AllToAll." + << std::endl; + std::exit(EXIT_FAILURE); + } + hostname.back() = '\0'; + std::cout << "[Rank " << rank << "] Host: " << hostname.data() + << " | GPU: " << Device::StringFromType(kDevType) << " | Device " + << local_rank << std::endl; + + infinicclComm_t comm = nullptr; + CHECK_INFINI(infinicclCommInitAll(&comm, size, nullptr)); + + const size_t world_size = static_cast(size); + const size_t id_bytes = sizeof(infinicclUniqueId); + if (id_bytes > std::numeric_limits::max() / world_size) { + std::cerr << "AllToAll bootstrap buffer size overflows `size_t`." + << std::endl; + std::exit(EXIT_FAILURE); + } + const size_t all_id_bytes = id_bytes * world_size; + + // Bootstrap the native communicator with AllToAll itself. Rank 0 repeats + // its unique ID in every destination block; each rank then takes the block + // received from source rank 0. With only the OpenMPI inter communicator + // initialized, the CCL provider delegates this call to `MPI_Alltoall`. + infinicclUniqueId id{}; + if (rank == 0) { + CHECK_INFINI(infinicclGetUniqueId(&id)); + } + std::vector h_id_send(all_id_bytes, 0); + if (rank == 0) { + for (int destination = 0; destination < size; ++destination) { + std::memcpy( + h_id_send.data() + static_cast(destination) * id_bytes, &id, + id_bytes); + } + } + + uint8_t *d_id_send = nullptr; + uint8_t *d_id_recv = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_id_send), all_id_bytes)); + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_id_recv), all_id_bytes)); + CHECK_RT(Rt, Rt::Memcpy(d_id_send, h_id_send.data(), all_id_bytes, + Rt::MemcpyHostToDevice)); + CHECK_INFINI(infinicclAllToAll(d_id_send, d_id_recv, id_bytes, infinicclUInt8, + comm, nullptr)); + CHECK_RT(Rt, Rt::Memcpy(&id, d_id_recv, id_bytes, Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + CHECK_RT(Rt, Rt::Free(d_id_send)); + CHECK_RT(Rt, Rt::Free(d_id_recv)); + + CHECK_INFINI(infinicclCommInitRank(&comm, size, id, rank)); + + if (kCountPerPeer > std::numeric_limits::max() / world_size || + kCountPerPeer * world_size > + std::numeric_limits::max() / sizeof(float)) { + std::cerr << "Hybrid AllToAll buffer size overflows `size_t`." << std::endl; + std::exit(EXIT_FAILURE); + } + const size_t total_elements = kCountPerPeer * world_size; + const size_t total_bytes = total_elements * sizeof(float); + std::vector h_send(total_elements); + std::vector h_recv(total_elements, 0.0f); + FillInput(&h_send, kCountPerPeer, size, rank); + + float *d_send = nullptr; + float *d_recv = nullptr; + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_send), total_bytes)); + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_recv), total_bytes)); + CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), total_bytes, + Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + for (int i = 0; i < kWarmupIterations; ++i) { + CHECK_INFINI(infinicclAllToAll(d_send, d_recv, kCountPerPeer, + infinicclFloat32, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + Timer timer; + for (int i = 0; i < kProfileIterations; ++i) { + CHECK_INFINI(infinicclAllToAll(d_send, d_recv, kCountPerPeer, + infinicclFloat32, comm, nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const double elapsed_ms = + timer.ElapsedMs() / static_cast(kProfileIterations); + + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes, + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const bool local_correct = + ValidateAllToAll(h_recv, kCountPerPeer, size, rank); + + // Use native AllToAll once more to exchange every rank's validation flag. + // This makes all ranks derive the same final process status without direct + // MPI or vendor-library calls in the example. + std::vector h_status_send(world_size, local_correct ? 1 : 0); + std::vector h_status_recv(world_size, 0); + int32_t *d_status_send = nullptr; + int32_t *d_status_recv = nullptr; + const size_t status_bytes = world_size * sizeof(int32_t); + CHECK_RT(Rt, + Rt::Malloc(reinterpret_cast(&d_status_send), status_bytes)); + CHECK_RT(Rt, + Rt::Malloc(reinterpret_cast(&d_status_recv), status_bytes)); + CHECK_RT(Rt, Rt::Memcpy(d_status_send, h_status_send.data(), status_bytes, + Rt::MemcpyHostToDevice)); + CHECK_INFINI(infinicclAllToAll(d_status_send, d_status_recv, 1, + infinicclInt32, comm, nullptr)); + CHECK_RT(Rt, Rt::Memcpy(h_status_recv.data(), d_status_recv, status_bytes, + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const bool correct = + local_correct && std::all_of(h_status_recv.begin(), h_status_recv.end(), + [](int32_t value) { return value == 1; }); + + if (rank == 0) { + PrintResult(correct, h_recv, kCountPerPeer, size, elapsed_ms); + } + + CHECK_RT(Rt, Rt::Free(d_status_send)); + CHECK_RT(Rt, Rt::Free(d_status_recv)); + CHECK_RT(Rt, Rt::Free(d_send)); + CHECK_RT(Rt, Rt::Free(d_recv)); + CHECK_INFINI(infinicclCommDestroy(comm)); + CHECK_INFINI(infinicclFinalize()); + + if (rank == 0) { + if (correct) { + std::cout << "[Main Process] Hybrid CCL AllToAll validation passed." + << std::endl; + } else { + std::cerr << "[Main Process] Hybrid CCL AllToAll validation failed." + << std::endl; + } + std::cout << "InfiniCCL finalized." << std::endl; + } + return correct; +} + +} // namespace + +int main(int argc, char **argv) { + return RunAllToAllExample(argc, argv) ? EXIT_SUCCESS : EXIT_FAILURE; +} diff --git a/examples/mpi/all_to_all.cc b/examples/mpi/all_to_all.cc index a1347ca..95dbdf7 100644 --- a/examples/mpi/all_to_all.cc +++ b/examples/mpi/all_to_all.cc @@ -1,14 +1,23 @@ /** - * InfiniCCL Example: AllToAll - * * This example demonstrates the planned API for performing a - * collective all-to-all exchange across multiple GPUs and nodes. + * InfiniCCL Example: AllToAll (MPI Backend) + * + * This example exchanges a distinct block with every rank, validates every + * received source block, and propagates validation failures to all ranks. */ #include #include -#include +#include +#include +#include +#include +#include +#include #include +#include +#include +#include #include // Public API @@ -26,163 +35,236 @@ using namespace infini::ccl; -void RunAllToAllExample(int argc, char **argv, int warmup_iter, - int profile_iter, const size_t kCountPerPeer) { +namespace { + +bool ParseLocalRank(const char *text, int *local_rank) { + if (!text || !local_rank) { + return false; + } + + int parsed = -1; + const char *end = text + std::strlen(text); + const auto result = std::from_chars(text, end, parsed); + if (result.ec != std::errc{} || result.ptr != end || parsed < 0) { + return false; + } + + *local_rank = parsed; + return true; +} + +float BlockValue(int source_rank, int destination_rank) { + return static_cast((source_rank + 1) * 1000 + destination_rank + 1); +} + +void FillInput(std::vector *input, size_t count_per_peer, int world_size, + int rank) { + for (int destination = 0; destination < world_size; ++destination) { + const size_t offset = static_cast(destination) * count_per_peer; + std::fill_n(input->begin() + offset, count_per_peer, + BlockValue(rank, destination)); + } +} + +bool ValidateAllToAll(const std::vector &result, size_t count_per_peer, + int world_size, int rank) { + bool correct = true; + for (int source = 0; source < world_size; ++source) { + const size_t offset = static_cast(source) * count_per_peer; + const bool block_correct = Validator::ValidateResult( + result.data() + offset, count_per_peer, BlockValue(source, rank), rank); + correct = block_correct && correct; + } + return correct; +} + +void PrintAllToAllMetrics(size_t count_per_peer, int world_size, + double elapsed_ms) { + constexpr double kBytesPerMiB = 1024.0 * 1024.0; + constexpr double kBytesPerGB = 1.0e9; + const double peer_bytes = static_cast(count_per_peer) * sizeof(float); + const double total_bytes = peer_bytes * static_cast(world_size); + const auto original_flags = std::cout.flags(); + const auto original_precision = std::cout.precision(); + + std::cout << "Data size per peer: " << count_per_peer << " floats (" + << std::fixed << std::setprecision(2) << peer_bytes / kBytesPerMiB + << " MiB)" << std::endl; + std::cout << "Total data per rank: " + << count_per_peer * static_cast(world_size) << " floats (" + << total_bytes / kBytesPerMiB << " MiB)" << std::endl; + std::cout << "Time: " << std::setprecision(3) << elapsed_ms << " ms" + << std::endl; + if (elapsed_ms > 0.0) { + const double algorithm_bandwidth = + total_bytes / kBytesPerGB / (elapsed_ms / 1000.0); + const double bus_bandwidth = algorithm_bandwidth * + static_cast(world_size - 1) / + static_cast(world_size); + std::cout << "Throughput: " << std::setprecision(2) << bus_bandwidth + << " GB/s (Bus BW)" << std::endl; + std::cout << "Alg Bandwidth: " << algorithm_bandwidth << " GB/s" + << std::endl; + } else { + std::cout << "Throughput: N/A (Bus BW)" << std::endl; + std::cout << "Alg Bandwidth: N/A" << std::endl; + } + + std::cout.flags(original_flags); + std::cout.precision(original_precision); +} + +bool RunAllToAllExample(int argc, char **argv, int warmup_iterations, + int profile_iterations, size_t count_per_peer) { constexpr Device::Type kDevType = ListGetBest(EnabledDevices{}); using Rt = Runtime; CHECK_INFINI(infinicclInit(&argc, &argv)); - int rank, size; + int rank = -1; + int size = 0; CHECK_INFINI(infinicclGetRank(&rank)); CHECK_INFINI(infinicclGetSize(&size)); + if (size <= 0) { + std::cerr << "Invalid world size for MPI AllToAll." << std::endl; + std::exit(EXIT_FAILURE); + } - char hostname[256]; - gethostname(hostname, sizeof(hostname)); - - // Map local rank to GPU device. - // Note: this is just for info printing. In practice, this part is not needed. - const char *local_rank_str = std::getenv("OMPI_COMM_WORLD_LOCAL_RANK"); - int local_rank = 0; - if (local_rank_str != nullptr) { - local_rank = std::atoi(local_rank_str); + int local_rank = -1; + if (!ParseLocalRank(std::getenv("OMPI_COMM_WORLD_LOCAL_RANK"), &local_rank)) { + std::cerr << "Missing or invalid `OMPI_COMM_WORLD_LOCAL_RANK`." + << std::endl; + std::exit(EXIT_FAILURE); } + CHECK_RT(Rt, Rt::SetDevice(local_rank)); - std::cout << "[Rank " << rank << "] Host: " << hostname - << " | GPU: " << Device::StringFromType(kDevType) << " " - << " | Device " << local_rank << std::endl; + std::array hostname{}; + if (gethostname(hostname.data(), hostname.size()) != 0) { + std::cerr << "Failed to query the hostname for MPI AllToAll." << std::endl; + std::exit(EXIT_FAILURE); + } + hostname.back() = '\0'; + std::cout << "[Rank " << rank << "] Host: " << hostname.data() + << " | GPU: " << Device::StringFromType(kDevType) << " | Device " + << local_rank << std::endl; - // Setup communicator infinicclComm_t comm = nullptr; CHECK_INFINI(infinicclCommInitAll(&comm, size, nullptr)); - // Prepare Data - // For AllToAll, `total_elements_per_rank = count_per_peer * world_size`. - const size_t kTotalCount = kCountPerPeer * static_cast(size); - - std::vector h_send(kTotalCount); - std::vector h_recv(kTotalCount, 0.0f); - - // Layout: block per destination rank. - // `block[dst]` is what this rank sends to dst. - // Fill with value encoding `src`/`dst` for easy validation. - for (int dst = 0; dst < size; ++dst) { - float v = static_cast(rank * 1000 + dst); - size_t off = static_cast(dst) * kCountPerPeer; - for (size_t i = 0; i < kCountPerPeer; ++i) { - h_send[off + i] = v; - } + const size_t world_size = static_cast(size); + if (count_per_peer > std::numeric_limits::max() / world_size || + count_per_peer * world_size > + std::numeric_limits::max() / sizeof(float)) { + std::cerr << "MPI AllToAll buffer size overflows `size_t`." << std::endl; + std::exit(EXIT_FAILURE); } + const size_t total_elements = count_per_peer * world_size; + const size_t total_bytes = total_elements * sizeof(float); + + std::vector h_send(total_elements); + std::vector h_recv(total_elements, 0.0f); + FillInput(&h_send, count_per_peer, size, rank); float *d_send = nullptr; float *d_recv = nullptr; - size_t total_bytes = kTotalCount * sizeof(*d_send); - - CHECK_RT(Rt, Rt::Malloc((void **)&d_send, total_bytes)); - CHECK_RT(Rt, Rt::Malloc((void **)&d_recv, total_bytes)); + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_send), total_bytes)); + CHECK_RT(Rt, Rt::Malloc(reinterpret_cast(&d_recv), total_bytes)); CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), total_bytes, Rt::MemcpyHostToDevice)); - CHECK_RT(Rt, Rt::Memcpy(d_recv, h_recv.data(), total_bytes, - Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); if (rank == 0) { - std::cout << "\n=== Performing AllToAll on GPU Memory ===" << std::endl; - std::cout << "Count per peer: " << kCountPerPeer << " floats" << std::endl; - std::cout << "Total per rank: " << kTotalCount << " floats (" + std::cout << "\n=== Performing MPI AllToAll on GPU Memory ===" << std::endl; + std::cout << "Count per peer: " << count_per_peer << " floats" << std::endl; + std::cout << "Total per rank: " << total_elements << " floats (" << total_bytes / 1024 / 1024 << " MB)" << std::endl; - std::cout << "Warm-up iterations: " << warmup_iter << std::endl; - std::cout << "Profile iterations: " << profile_iter << std::endl; + std::cout << "Warm-up iterations: " << warmup_iterations << std::endl; + std::cout << "Profile iterations: " << profile_iterations << std::endl; } - CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); - - // Warm-up and D2H transfer the answer. - CHECK_INFINI(infinicclAllToAll(d_send, d_recv, kCountPerPeer, - infinicclFloat32, comm, nullptr)); - CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes, - Rt::MemcpyDeviceToHost)); - - for (int i = 1; i < warmup_iter; ++i) { - CHECK_INFINI(infinicclAllToAll(d_send, d_recv, kCountPerPeer, + for (int i = 0; i < warmup_iterations; ++i) { + CHECK_INFINI(infinicclAllToAll(d_send, d_recv, count_per_peer, infinicclFloat32, comm, nullptr)); } CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); - // Profiling Timer timer; - - for (int i = 0; i < profile_iter; ++i) { - CHECK_INFINI(infinicclAllToAll(d_send, d_recv, kCountPerPeer, + for (int i = 0; i < profile_iterations; ++i) { + CHECK_INFINI(infinicclAllToAll(d_send, d_recv, count_per_peer, infinicclFloat32, comm, nullptr)); } - CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const double elapsed_ms = + timer.ElapsedMs() / static_cast(profile_iterations); + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes, Rt::MemcpyDeviceToHost)); - double elapsed = timer.ElapsedMs() / static_cast(profile_iter); - - // Result Validation - // recv `block[src]` on rank r should come from `src`'s send `block[dst=r]`. - bool correct = true; - int error_count = 0; - - for (int src = 0; src < size; ++src) { - float expected = static_cast(src * 1000 + rank); - size_t off = static_cast(src) * kCountPerPeer; - - bool block_ok = Validator::ValidateResult(h_recv.data() + off, - kCountPerPeer, expected, rank); - - correct = correct && block_ok; - } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const bool local_correct = + ValidateAllToAll(h_recv, count_per_peer, size, rank); + + // Exchange each rank's validation flag with every destination so all ranks + // derive the same final process status. + std::vector h_status_send(world_size, local_correct ? 1 : 0); + std::vector h_status_recv(world_size, 0); + int32_t *d_status_send = nullptr; + int32_t *d_status_recv = nullptr; + const size_t status_bytes = world_size * sizeof(int32_t); + CHECK_RT(Rt, + Rt::Malloc(reinterpret_cast(&d_status_send), status_bytes)); + CHECK_RT(Rt, + Rt::Malloc(reinterpret_cast(&d_status_recv), status_bytes)); + CHECK_RT(Rt, Rt::Memcpy(d_status_send, h_status_send.data(), status_bytes, + Rt::MemcpyHostToDevice)); + CHECK_INFINI(infinicclAllToAll(d_status_send, d_status_recv, 1, + infinicclInt32, comm, nullptr)); + CHECK_RT(Rt, Rt::Memcpy(h_status_recv.data(), d_status_recv, status_bytes, + Rt::MemcpyDeviceToHost)); + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + const bool correct = + local_correct && std::all_of(h_status_recv.begin(), h_status_recv.end(), + [](int32_t value) { return value == 1; }); if (rank == 0) { - const char *GREEN = "\033[32m"; - const char *RED = "\033[31m"; - const char *RESET = "\033[0m"; - - std::cout << "\n=== AllToAll Results ===" << std::endl; + constexpr const char *kGreen = "\033[32m"; + constexpr const char *kRed = "\033[31m"; + constexpr const char *kReset = "\033[0m"; + std::cout << "\n=== MPI AllToAll Results ===" << std::endl; std::cout << "Correct: " - << (correct ? (GREEN + std::string("YES") + RESET) - : (RED + std::string("NO") + RESET)); - if (!correct) { - std::cout << " (" << error_count << " errors)"; - } - std::cout << std::endl; - - std::cout << "Sample recv blocks: "; - for (int src = 0; src < std::min(size, 4); ++src) { - size_t off = static_cast(src) * kCountPerPeer; - std::cout << "[src" << src << ": " << h_recv[off] << "] "; + << (correct ? (kGreen + std::string("YES") + kReset) + : (kRed + std::string("NO") + kReset)) + << std::endl; + PrintAllToAllMetrics(count_per_peer, size, elapsed_ms); + std::cout << "Sample receive blocks: "; + for (int source = 0; source < std::min(size, 4); ++source) { + const size_t offset = static_cast(source) * count_per_peer; + std::cout << "[src" << source << ": " << h_recv[offset] << "] "; } std::cout << std::endl; } - // Metrics Reporting (Only from rank 0 for cleaner output) - if (rank == 0) { - Metrics metrics{elapsed, total_bytes, size}; - metrics.Print(); - } - - // Cleanup + CHECK_RT(Rt, Rt::Free(d_status_send)); + CHECK_RT(Rt, Rt::Free(d_status_recv)); CHECK_RT(Rt, Rt::Free(d_send)); CHECK_RT(Rt, Rt::Free(d_recv)); - CHECK_INFINI(infinicclCommDestroy(comm)); CHECK_INFINI(infinicclFinalize()); if (rank == 0) { std::cout << "InfiniCCL finalized." << std::endl; } + return correct; } -int main(int argc, char **argv) { - int warmup_iters = 2; - int profile_iters = 20; - size_t count_per_peer = 1 << 18; - - RunAllToAllExample(argc, argv, warmup_iters, profile_iters, count_per_peer); +} // namespace - return EXIT_SUCCESS; +int main(int argc, char **argv) { + constexpr int kWarmupIterations = 2; + constexpr int kProfileIterations = 20; + constexpr size_t kCountPerPeer = 1 << 18; + return RunAllToAllExample(argc, argv, kWarmupIterations, kProfileIterations, + kCountPerPeer) + ? EXIT_SUCCESS + : EXIT_FAILURE; } diff --git a/src/backends/ccl/common/impl/all_to_all.h b/src/backends/ccl/common/impl/all_to_all.h new file mode 100644 index 0000000..287d58d --- /dev/null +++ b/src/backends/ccl/common/impl/all_to_all.h @@ -0,0 +1,118 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_ALL_TO_ALL_H_ +#define INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_ALL_TO_ALL_H_ + +#include +#include + +#include "backends/ccl/common/api.h" +#include "backends/ccl/common/comm_instance.h" +#include "base/all_to_all.h" +#include "communicator.h" +#include "data_type_impl.h" +#include "logging.h" + +namespace infini::ccl { + +template +struct DeferredAllToAll { + using type = AllToAll; +}; + +template +class CclAllToAllImpl { + public: + static ReturnStatus Apply(const void *send_buff, void *recv_buff, + size_t count, DataType data_type, + Communicator *comm, void *stream) { + using Api = CclApi; + using TypeMap = CclTypeMap; + using CommInstance = CclCommInstance; + + const bool has_native_comm = comm && comm->intra_comm() && + comm->intra_comm_backend() == backend && + comm->device_type() == device; + if (!has_native_comm) { + if (!comm || !comm->inter_comm() || + comm->inter_comm_backend() != BackendType::kOmpi) { + return ReturnStatus::kInternalError; + } + + using FallbackOperation = typename DeferredAllToAll::type; + if constexpr (BackendEnabled::value) { + return AllToAllImpl::Apply( + send_buff, recv_buff, count, data_type, comm, stream); + } + + return ReturnStatus::kInternalError; + } + + if (send_buff == recv_buff) { + LOG("In-place buffers are not supported by the native CCL `AllToAll` " + "implementation."); + return ReturnStatus::kNotSupported; + } + if (comm->size() <= 0 || comm->rank() < 0 || comm->rank() >= comm->size()) { + LOG("Invalid rank or world size for native CCL `AllToAll`."); + return ReturnStatus::kInternalError; + } + + auto *instance = static_cast(comm->intra_comm()); + if (!instance->handle) { + return ReturnStatus::kInternalError; + } + + typename Api::DataType native_type{}; + if (!TypeMap::ToBackendDataType(data_type, &native_type)) { + return ReturnStatus::kNotSupported; + } + + const size_t type_size = kDataTypeToSize.at(data_type); + if (count > std::numeric_limits::max() / type_size) { + LOG("Per-peer byte size overflows `size_t` for native CCL " + "`AllToAll`."); + return ReturnStatus::kInvalidArgument; + } + const size_t peer_bytes = count * type_size; + const size_t world_size = static_cast(comm->size()); + if (peer_bytes > std::numeric_limits::max() / world_size) { + LOG("Total byte size overflows `size_t` for native CCL `AllToAll`."); + return ReturnStatus::kInvalidArgument; + } + + const auto *send_bytes = static_cast(send_buff); + auto *recv_bytes = static_cast(recv_buff); + auto native_stream = reinterpret_cast(stream); + + ReturnStatus status = Api::Check(Api::GroupStart()); + if (status != ReturnStatus::kSuccess) { + return status; + } + + ReturnStatus first_error = ReturnStatus::kSuccess; + for (int peer = 0; peer < comm->size(); ++peer) { + const size_t offset = static_cast(peer) * peer_bytes; + status = Api::Check(Api::Send(send_bytes + offset, count, native_type, + peer, instance->handle, native_stream)); + if (first_error == ReturnStatus::kSuccess && + status != ReturnStatus::kSuccess) { + first_error = status; + } + + status = Api::Check(Api::Recv(recv_bytes + offset, count, native_type, + peer, instance->handle, native_stream)); + if (first_error == ReturnStatus::kSuccess && + status != ReturnStatus::kSuccess) { + first_error = status; + } + } + + const ReturnStatus group_end_status = Api::Check(Api::GroupEnd()); + return first_error != ReturnStatus::kSuccess ? first_error + : group_end_status; + } +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_ALL_TO_ALL_H_ diff --git a/src/backends/ccl/mccl/api.h b/src/backends/ccl/mccl/api.h index a5d2bcd..2e97324 100644 --- a/src/backends/ccl/mccl/api.h +++ b/src/backends/ccl/mccl/api.h @@ -49,6 +49,20 @@ struct McclApi { return mcclAllReduce(send_buff, recv_buff, count, data_type, op, comm, stream); } + + static Result GroupStart() { return mcclGroupStart(); } + + static Result GroupEnd() { return mcclGroupEnd(); } + + static Result Send(const void *send_buff, size_t count, DataType data_type, + int peer, Comm comm, Stream stream) { + return mcclSend(send_buff, count, data_type, peer, comm, stream); + } + + static Result Recv(void *recv_buff, size_t count, DataType data_type, + int peer, Comm comm, Stream stream) { + return mcclRecv(recv_buff, count, data_type, peer, comm, stream); + } }; } // namespace infini::ccl diff --git a/src/backends/ccl/mccl/impl/all_to_all.h b/src/backends/ccl/mccl/impl/all_to_all.h new file mode 100644 index 0000000..3f124ef --- /dev/null +++ b/src/backends/ccl/mccl/impl/all_to_all.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_ALL_TO_ALL_H_ +#define INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_ALL_TO_ALL_H_ + +#include "backends/ccl/common/impl/all_to_all.h" + +namespace infini::ccl { + +template +class AllToAllImpl + : public CclAllToAllImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_ALL_TO_ALL_H_ diff --git a/src/backends/ccl/nccl/api.h b/src/backends/ccl/nccl/api.h index e7b6119..8dd49a2 100644 --- a/src/backends/ccl/nccl/api.h +++ b/src/backends/ccl/nccl/api.h @@ -46,6 +46,20 @@ struct NcclApi { return ncclAllReduce(send_buff, recv_buff, count, data_type, op, comm, stream); } + + static Result GroupStart() { return ncclGroupStart(); } + + static Result GroupEnd() { return ncclGroupEnd(); } + + static Result Send(const void *send_buff, size_t count, DataType data_type, + int peer, Comm comm, Stream stream) { + return ncclSend(send_buff, count, data_type, peer, comm, stream); + } + + static Result Recv(void *recv_buff, size_t count, DataType data_type, + int peer, Comm comm, Stream stream) { + return ncclRecv(recv_buff, count, data_type, peer, comm, stream); + } }; } // namespace infini::ccl diff --git a/src/backends/ccl/nccl/impl/all_to_all.h b/src/backends/ccl/nccl/impl/all_to_all.h new file mode 100644 index 0000000..f135eb0 --- /dev/null +++ b/src/backends/ccl/nccl/impl/all_to_all.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_ALL_TO_ALL_H_ +#define INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_ALL_TO_ALL_H_ + +#include "backends/ccl/common/impl/all_to_all.h" + +namespace infini::ccl { + +template +class AllToAllImpl + : public CclAllToAllImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_ALL_TO_ALL_H_ diff --git a/src/backends/mpi/ompi/impl/all_to_all.h b/src/backends/mpi/ompi/impl/all_to_all.h index f3f386d..1e4a1c2 100644 --- a/src/backends/mpi/ompi/impl/all_to_all.h +++ b/src/backends/mpi/ompi/impl/all_to_all.h @@ -1,15 +1,18 @@ #ifndef INFINI_CCL_BACKENDS_MPI_OMPI_IMPL_ALL_TO_ALL_H_ #define INFINI_CCL_BACKENDS_MPI_OMPI_IMPL_ALL_TO_ALL_H_ +#include #include +#include #include "backends/mpi/ompi/checks.h" #include "backends/mpi/ompi/comm_instance.h" -#include "backends/mpi/ompi/type_map.h" #include "base/all_to_all.h" #include "communicator.h" +#include "data_type_impl.h" #include "dispatcher.h" #include "logging.h" +#include "runtime.h" namespace infini::ccl { @@ -23,50 +26,63 @@ class AllToAllImpl { ListGetBest(ActiveDevices{}); using Rt = Runtime; - auto *inst = static_cast(comm->inter_comm()); - - if (!inst || inst->handle == MPI_COMM_NULL) { + if (!comm || !comm->inter_comm() || + comm->inter_comm_backend() != BackendType::kOmpi) { LOG("Invalid OpenMPI communicator instance for `AllToAll`."); return ReturnStatus::kInternalError; } + auto *inst = static_cast(comm->inter_comm()); + if (inst->handle == MPI_COMM_NULL) { + LOG("Invalid OpenMPI communicator handle for `AllToAll`."); + return ReturnStatus::kInternalError; + } + if (comm->size() <= 0) { + LOG("Invalid world size for `AllToAll`."); + return ReturnStatus::kInternalError; + } - if (count > static_cast(std::numeric_limits::max())) { - LOG("count exceeds MPI `int` range for `AllToAll`."); + size_t type_size = kDataTypeToSize.at(data_type); + if (count > std::numeric_limits::max() / type_size) { + LOG("Per-peer byte size overflows `size_t` for `AllToAll`."); + return ReturnStatus::kInvalidArgument; + } + size_t peer_bytes = count * type_size; + if (peer_bytes > static_cast(std::numeric_limits::max())) { + LOG("Per-peer byte count exceeds MPI `int` range for `AllToAll`."); return ReturnStatus::kInvalidArgument; } - int mpi_count = static_cast(count); size_t world_size = static_cast(comm->size()); - MPI_Datatype mpi_type = DataTypeToOmpiType(data_type); - size_t type_size = kDataTypeToSize.at(data_type); - size_t total_count = count * world_size; - size_t total_bytes = total_count * type_size; + if (peer_bytes > std::numeric_limits::max() / world_size) { + LOG("Total byte size overflows `size_t` for `AllToAll`."); + return ReturnStatus::kInvalidArgument; + } + size_t total_bytes = peer_bytes * world_size; + int mpi_peer_bytes = static_cast(peer_bytes); // Handle GPU Memory (Staging Pattern) // Note: we simply use host-staging for now. - void *host_sendbuf = malloc(total_bytes); - void *host_recvbuf = malloc(total_bytes); + std::unique_ptr host_sendbuf( + std::malloc(total_bytes), &std::free); + std::unique_ptr host_recvbuf( + std::malloc(total_bytes), &std::free); if (!host_sendbuf || !host_recvbuf) { - free(host_sendbuf); - free(host_recvbuf); LOG("Failed to allocate host buffers for `AllToAll` staging."); return ReturnStatus::kSystemError; } - CHECK_STATUS(Rt, Rt::Memcpy(host_sendbuf, send_buff, total_bytes, + CHECK_STATUS(Rt, Rt::Memcpy(host_sendbuf.get(), send_buff, total_bytes, Rt::MemcpyDeviceToHost)); CHECK_STATUS(Rt, Rt::StreamSynchronize(static_cast(stream))); - INFINI_CHECK_MPI(MPI_Alltoall(host_sendbuf, mpi_count, mpi_type, - host_recvbuf, mpi_count, mpi_type, + INFINI_CHECK_MPI(MPI_Alltoall(host_sendbuf.get(), mpi_peer_bytes, MPI_BYTE, + host_recvbuf.get(), mpi_peer_bytes, MPI_BYTE, inst->handle)); - CHECK_STATUS(Rt, Rt::Memcpy(recv_buff, host_recvbuf, total_bytes, + CHECK_STATUS(Rt, Rt::Memcpy(recv_buff, host_recvbuf.get(), total_bytes, Rt::MemcpyHostToDevice)); - free(host_sendbuf); - free(host_recvbuf); return ReturnStatus::kSuccess; } }; diff --git a/src/base/all_to_all.h b/src/base/all_to_all.h index c549a55..85331bc 100644 --- a/src/base/all_to_all.h +++ b/src/base/all_to_all.h @@ -1,8 +1,8 @@ #ifndef INFINI_CCL_BASE_ALL_TO_ALL_H_ #define INFINI_CCL_BASE_ALL_TO_ALL_H_ -#include "comm_impl.h" #include "communicator.h" +#include "data_type_impl.h" #include "logging.h" #include "operation.h" #include "return_status_impl.h" @@ -19,7 +19,14 @@ class AllToAll : public Operation { static ReturnStatus Execute(const void *send_buff, void *recv_buff, size_t count, DataType datatype, void *comm_handle, void *stream) { - if (HasInvalidArgs(send_buff, recv_buff, datatype, comm_handle)) { + if (HasInvalidRequiredArgs(datatype, comm_handle)) { + return ReturnStatus::kInvalidArgument; + } + if (count == 0) { + return ReturnStatus::kSuccess; + } + if (!send_buff || !recv_buff) { + LOG("Invalid buffer pointer for `AllToAll`."); return ReturnStatus::kInvalidArgument; } auto *comm = static_cast(comm_handle); @@ -28,17 +35,12 @@ class AllToAll : public Operation { } private: - static bool HasInvalidArgs(const void *send_buff, void *recv_buff, - DataType datatype, void *comm_handle) { + static bool HasInvalidRequiredArgs(DataType datatype, void *comm_handle) { if (!comm_handle) { // TODO(lzm): change to use `glog`. LOG("Invalid communicator handle for `AllToAll`."); return true; } - if (!send_buff || !recv_buff) { - LOG("Invalid buffer pointer for `AllToAll`."); - return true; - } if (datatype < DataType::kChar || datatype >= DataType::kNumTypes) { LOG("Invalid data type for `AllToAll`."); return true;