Skip to content

Commit ec6caa8

Browse files
committed
feat: support CCL AllToAll
1 parent ef4045a commit ec6caa8

10 files changed

Lines changed: 950 additions & 138 deletions

File tree

‎examples/ccl/all_to_all.cc‎

Lines changed: 293 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,293 @@
1+
/**
2+
* InfiniCCL Example: Thread-per-GPU Single-Node AllToAll
3+
*
4+
* This example creates one native CCL rank per GPU and validates an
5+
* out-of-place all-to-all exchange composed from grouped point-to-point calls.
6+
*/
7+
8+
#include <unistd.h>
9+
10+
#include <algorithm>
11+
#include <array>
12+
#include <atomic>
13+
#include <charconv>
14+
#include <cstdlib>
15+
#include <cstring>
16+
#include <iostream>
17+
#include <limits>
18+
#include <string>
19+
#include <system_error>
20+
#include <thread>
21+
#include <vector>
22+
23+
// Public API
24+
#include "infiniccl.h"
25+
26+
// Example-Specific Utilities
27+
#include "utils.h"
28+
29+
// Internal Headers (Accessible via example-specific include paths, technically
30+
// not public APIs)
31+
#include "backend_manifest.h"
32+
33+
using namespace infini::ccl;
34+
35+
namespace {
36+
37+
struct ScenarioState {
38+
std::atomic<bool> correct{true};
39+
std::atomic<int> completed{0};
40+
};
41+
42+
struct ThreadArgs {
43+
int rank;
44+
int size;
45+
infinicclUniqueId id;
46+
size_t count_per_peer;
47+
int warmup_iterations;
48+
int profile_iterations;
49+
ScenarioState *state;
50+
};
51+
52+
template <typename T>
53+
bool ParsePositiveNumber(const char *text, T *value) {
54+
if (!text || !value) {
55+
return false;
56+
}
57+
58+
T parsed{};
59+
const char *end = text + std::strlen(text);
60+
const auto result = std::from_chars(text, end, parsed);
61+
if (result.ec != std::errc{} || result.ptr != end || parsed <= 0) {
62+
return false;
63+
}
64+
65+
*value = parsed;
66+
return true;
67+
}
68+
69+
float BlockValue(int source_rank, int destination_rank) {
70+
return static_cast<float>((source_rank + 1) * 1000 + destination_rank + 1);
71+
}
72+
73+
void FillInput(std::vector<float> *input, size_t count_per_peer, int world_size,
74+
int rank) {
75+
for (int destination = 0; destination < world_size; ++destination) {
76+
const size_t offset = static_cast<size_t>(destination) * count_per_peer;
77+
std::fill_n(input->begin() + offset, count_per_peer,
78+
BlockValue(rank, destination));
79+
}
80+
}
81+
82+
bool ValidateAllToAll(const std::vector<float> &result, size_t count_per_peer,
83+
int world_size, int rank) {
84+
bool correct = true;
85+
for (int source = 0; source < world_size; ++source) {
86+
const size_t offset = static_cast<size_t>(source) * count_per_peer;
87+
const bool block_correct = Validator::ValidateResult(
88+
result.data() + offset, count_per_peer, BlockValue(source, rank), rank);
89+
correct = block_correct && correct;
90+
}
91+
return correct;
92+
}
93+
94+
void WaitForAll(ScenarioState *state, int world_size) {
95+
state->completed.fetch_add(1, std::memory_order_acq_rel);
96+
while (state->completed.load(std::memory_order_acquire) < world_size) {
97+
std::this_thread::yield();
98+
}
99+
}
100+
101+
void PrintResult(bool correct, const std::vector<float> &result,
102+
size_t count_per_peer, int world_size, double elapsed_ms) {
103+
constexpr const char *kGreen = "\033[32m";
104+
constexpr const char *kRed = "\033[31m";
105+
constexpr const char *kReset = "\033[0m";
106+
107+
std::cout << "\n=== CCL AllToAll Results ===" << std::endl;
108+
std::cout << "Correct: "
109+
<< (correct ? (kGreen + std::string("YES") + kReset)
110+
: (kRed + std::string("NO") + kReset))
111+
<< std::endl;
112+
std::cout << "Average time: " << elapsed_ms << " ms" << std::endl;
113+
std::cout << "Sample receive blocks: ";
114+
for (int source = 0; source < std::min(world_size, 4); ++source) {
115+
const size_t offset = static_cast<size_t>(source) * count_per_peer;
116+
std::cout << "[src" << source << ": " << result[offset] << "] ";
117+
}
118+
std::cout << std::endl;
119+
}
120+
121+
void WorkerThread(ThreadArgs args) {
122+
constexpr Device::Type kDevType =
123+
ListGetBest<DevicePriority>(EnabledDevices{});
124+
using Rt = Runtime<kDevType>;
125+
126+
CHECK_RT(Rt, Rt::SetDevice(args.rank));
127+
128+
std::array<char, 256> hostname{};
129+
if (gethostname(hostname.data(), hostname.size()) != 0) {
130+
std::cerr << "Failed to query the hostname for the AllToAll worker."
131+
<< std::endl;
132+
std::exit(EXIT_FAILURE);
133+
}
134+
hostname.back() = '\0';
135+
std::cout << "[Rank " << args.rank << "] Host: " << hostname.data()
136+
<< " | GPU: " << Device::StringFromType(kDevType) << " | Device "
137+
<< args.rank << std::endl;
138+
139+
infinicclComm_t comm = nullptr;
140+
CHECK_INFINI(infinicclCommInitRank(&comm, args.size, args.id, args.rank));
141+
142+
const size_t total_elements =
143+
args.count_per_peer * static_cast<size_t>(args.size);
144+
const size_t total_bytes = total_elements * sizeof(float);
145+
std::vector<float> h_send(total_elements);
146+
std::vector<float> h_recv(total_elements, 0.0f);
147+
FillInput(&h_send, args.count_per_peer, args.size, args.rank);
148+
149+
float *d_send = nullptr;
150+
float *d_recv = nullptr;
151+
CHECK_RT(Rt, Rt::Malloc(reinterpret_cast<void **>(&d_send), total_bytes));
152+
CHECK_RT(Rt, Rt::Malloc(reinterpret_cast<void **>(&d_recv), total_bytes));
153+
CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), total_bytes,
154+
Rt::MemcpyHostToDevice));
155+
CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
156+
157+
for (int i = 0; i < args.warmup_iterations; ++i) {
158+
CHECK_INFINI(infinicclAllToAll(d_send, d_recv, args.count_per_peer,
159+
infinicclFloat32, comm, nullptr));
160+
}
161+
CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
162+
163+
Timer timer;
164+
for (int i = 0; i < args.profile_iterations; ++i) {
165+
CHECK_INFINI(infinicclAllToAll(d_send, d_recv, args.count_per_peer,
166+
infinicclFloat32, comm, nullptr));
167+
}
168+
CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
169+
const double elapsed_ms =
170+
timer.ElapsedMs() / static_cast<double>(args.profile_iterations);
171+
172+
CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes,
173+
Rt::MemcpyDeviceToHost));
174+
CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
175+
const bool local_correct =
176+
ValidateAllToAll(h_recv, args.count_per_peer, args.size, args.rank);
177+
if (!local_correct) {
178+
args.state->correct.store(false, std::memory_order_relaxed);
179+
}
180+
WaitForAll(args.state, args.size);
181+
182+
if (args.rank == 0) {
183+
PrintResult(args.state->correct.load(std::memory_order_acquire), h_recv,
184+
args.count_per_peer, args.size, elapsed_ms);
185+
}
186+
187+
CHECK_RT(Rt, Rt::Free(d_send));
188+
CHECK_RT(Rt, Rt::Free(d_recv));
189+
CHECK_INFINI(infinicclCommDestroy(comm));
190+
}
191+
192+
void PrintUsage(const char *program) {
193+
std::cout << "Usage: " << program << " [options]\n"
194+
<< "Options:\n"
195+
<< " -g <num_gpus> Number of GPUs (default: 8)\n"
196+
<< " -w <warmup_iters> Warmup iterations (default: 2)\n"
197+
<< " -p <profile_iters> Profile iterations (default: 20)\n"
198+
<< " -n <count_per_peer> Elements sent to each peer (default: "
199+
"1048576)\n";
200+
}
201+
202+
} // namespace
203+
204+
int main(int argc, char **argv) {
205+
int num_gpus = 8;
206+
int warmup_iterations = 2;
207+
int profile_iterations = 20;
208+
size_t count_per_peer = 1 << 20;
209+
210+
int opt = 0;
211+
while ((opt = getopt(argc, argv, "g:w:p:n:h")) != -1) {
212+
bool parsed = false;
213+
switch (opt) {
214+
case 'g':
215+
parsed = ParsePositiveNumber(optarg, &num_gpus);
216+
break;
217+
case 'w':
218+
parsed = ParsePositiveNumber(optarg, &warmup_iterations);
219+
break;
220+
case 'p':
221+
parsed = ParsePositiveNumber(optarg, &profile_iterations);
222+
break;
223+
case 'n':
224+
parsed = ParsePositiveNumber(optarg, &count_per_peer);
225+
break;
226+
case 'h':
227+
PrintUsage(argv[0]);
228+
return EXIT_SUCCESS;
229+
default:
230+
PrintUsage(argv[0]);
231+
return EXIT_FAILURE;
232+
}
233+
234+
if (!parsed) {
235+
std::cerr << "Invalid positive numeric option for AllToAll." << std::endl;
236+
return EXIT_FAILURE;
237+
}
238+
}
239+
240+
if (optind != argc) {
241+
std::cerr << "Unexpected positional argument for AllToAll." << std::endl;
242+
return EXIT_FAILURE;
243+
}
244+
if (static_cast<size_t>(num_gpus) >
245+
std::numeric_limits<size_t>::max() / count_per_peer ||
246+
count_per_peer * static_cast<size_t>(num_gpus) >
247+
std::numeric_limits<size_t>::max() / sizeof(float)) {
248+
std::cerr << "AllToAll buffer size overflows `size_t`." << std::endl;
249+
return EXIT_FAILURE;
250+
}
251+
252+
std::array<char, 256> hostname{};
253+
if (gethostname(hostname.data(), hostname.size()) != 0) {
254+
std::cerr << "Failed to query the hostname for AllToAll." << std::endl;
255+
return EXIT_FAILURE;
256+
}
257+
hostname.back() = '\0';
258+
std::cout << "[Main Process] Host: " << hostname.data()
259+
<< " | Target GPUs: " << num_gpus << std::endl;
260+
std::cout << "[Main Process] Count per peer: " << count_per_peer
261+
<< " floats | Warmup: " << warmup_iterations
262+
<< " | Profile: " << profile_iterations << std::endl;
263+
264+
infinicclUniqueId shared_id{};
265+
CHECK_INFINI(infinicclGetUniqueId(&shared_id));
266+
267+
ScenarioState state;
268+
std::vector<std::thread> threads;
269+
threads.reserve(num_gpus);
270+
for (int rank = 0; rank < num_gpus; ++rank) {
271+
ThreadArgs args{rank, num_gpus, shared_id,
272+
count_per_peer, warmup_iterations, profile_iterations,
273+
&state};
274+
threads.emplace_back(WorkerThread, args);
275+
}
276+
277+
for (auto &thread : threads) {
278+
if (thread.joinable()) {
279+
thread.join();
280+
}
281+
}
282+
283+
const bool correct = state.correct.load(std::memory_order_acquire);
284+
if (correct) {
285+
std::cout << "[Main Process] CCL AllToAll validation passed." << std::endl;
286+
} else {
287+
std::cerr << "[Main Process] CCL AllToAll validation failed." << std::endl;
288+
}
289+
std::cout
290+
<< "[Main Process] All worker threads joined. InfiniCCL finalized safely."
291+
<< std::endl;
292+
return correct ? EXIT_SUCCESS : EXIT_FAILURE;
293+
}

0 commit comments

Comments
 (0)