Skip to content

Commit 773a7d3

Browse files
committed
feat: support CCL Gather
1 parent ef4045a commit 773a7d3

10 files changed

Lines changed: 1036 additions & 133 deletions

File tree

‎examples/ccl/gather.cc‎

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

0 commit comments

Comments
 (0)