Skip to content

Commit 752bc4e

Browse files
committed
feat: support CCL ReduceScatter
1 parent ef4045a commit 752bc4e

10 files changed

Lines changed: 1012 additions & 67 deletions

File tree

‎examples/ccl/reduce_scatter.cc‎

Lines changed: 360 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,360 @@
1+
/**
2+
* InfiniCCL Example: Thread-per-GPU Single-Node ReduceScatter
3+
*
4+
* This example creates one native CCL rank per GPU, then validates
5+
* out-of-place and canonical in-place ReduceScatter without an MPI launcher.
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+
struct ScenarioState {
40+
std::atomic<bool> correct{true};
41+
std::atomic<int> completed{0};
42+
};
43+
44+
struct ThreadArgs {
45+
int rank;
46+
int size;
47+
infinicclUniqueId id;
48+
size_t recv_count;
49+
int warmup_iter;
50+
int profile_iter;
51+
ScenarioState *out_of_place;
52+
ScenarioState *in_place;
53+
};
54+
55+
template <typename T>
56+
bool ParsePositiveNumber(const char *text, T *value) {
57+
if (!text || !value) {
58+
return false;
59+
}
60+
61+
T parsed{};
62+
const char *end = text + std::strlen(text);
63+
const auto result = std::from_chars(text, end, parsed);
64+
if (result.ec != std::errc{} || result.ptr != end || parsed <= 0) {
65+
return false;
66+
}
67+
68+
*value = parsed;
69+
return true;
70+
}
71+
72+
void FillInput(std::vector<float> *input, size_t recv_count, int world_size,
73+
int rank) {
74+
for (int dest_rank = 0; dest_rank < world_size; ++dest_rank) {
75+
const float value =
76+
static_cast<float>(rank + 1) * static_cast<float>(dest_rank + 1);
77+
const size_t offset = static_cast<size_t>(dest_rank) * recv_count;
78+
std::fill_n(input->begin() + offset, recv_count, value);
79+
}
80+
}
81+
82+
float ExpectedValue(int rank, int world_size) {
83+
const float rank_sum = static_cast<float>(world_size) *
84+
(static_cast<float>(world_size) + 1.0f) / 2.0f;
85+
return rank_sum * static_cast<float>(rank + 1);
86+
}
87+
88+
void WaitForScenario(ScenarioState *state, int world_size) {
89+
state->completed.fetch_add(1, std::memory_order_release);
90+
while (state->completed.load(std::memory_order_acquire) < world_size) {
91+
std::this_thread::yield();
92+
}
93+
}
94+
95+
void PrintReduceScatterMetrics(size_t recv_count, int world_size,
96+
double elapsed_ms) {
97+
constexpr double kBytesPerMiB = 1024.0 * 1024.0;
98+
constexpr double kBytesPerGB = 1.0e9;
99+
const size_t total_count = recv_count * static_cast<size_t>(world_size);
100+
const double recv_bytes = static_cast<double>(recv_count) * sizeof(float);
101+
const double total_bytes = static_cast<double>(total_count) * sizeof(float);
102+
const auto original_flags = std::cout.flags();
103+
const auto original_precision = std::cout.precision();
104+
105+
std::cout << "Received data per rank: " << recv_count << " floats ("
106+
<< std::fixed << std::setprecision(2) << recv_bytes / kBytesPerMiB
107+
<< " MiB)" << std::endl;
108+
std::cout << "Total reduced data: " << total_count << " floats ("
109+
<< total_bytes / kBytesPerMiB << " MiB)" << std::endl;
110+
std::cout << "Time: " << std::setprecision(3) << elapsed_ms
111+
<< " ms" << std::endl;
112+
if (world_size > 0 && elapsed_ms > 0.0 && std::isfinite(elapsed_ms)) {
113+
const double algorithm_bandwidth =
114+
total_bytes / kBytesPerGB / (elapsed_ms / 1000.0);
115+
const double bus_bandwidth = algorithm_bandwidth *
116+
static_cast<double>(world_size - 1) /
117+
static_cast<double>(world_size);
118+
std::cout << "Throughput: " << std::setprecision(2)
119+
<< bus_bandwidth << " GB/s (Bus BW)" << std::endl;
120+
std::cout << "Alg Bandwidth: " << algorithm_bandwidth << " GB/s"
121+
<< std::endl;
122+
} else {
123+
std::cout << "Throughput: N/A (Bus BW)" << std::endl;
124+
std::cout << "Alg Bandwidth: N/A" << std::endl;
125+
}
126+
127+
std::cout.flags(original_flags);
128+
std::cout.precision(original_precision);
129+
}
130+
131+
void PrintResult(const char *scenario, bool correct,
132+
const std::vector<float> &result, float expected,
133+
size_t recv_count, int world_size, double elapsed_ms) {
134+
constexpr const char *kGreen = "\033[32m";
135+
constexpr const char *kRed = "\033[31m";
136+
constexpr const char *kReset = "\033[0m";
137+
138+
std::cout << "\n=== " << scenario
139+
<< " CCL ReduceScatter Results ===" << std::endl;
140+
std::cout << "Correct: "
141+
<< (correct ? (kGreen + std::string("YES") + kReset)
142+
: (kRed + std::string("NO") + kReset))
143+
<< std::endl;
144+
std::cout << "Expect: " << expected << std::endl;
145+
std::cout << "Actual: " << result.front() << std::endl;
146+
PrintReduceScatterMetrics(recv_count, world_size, elapsed_ms);
147+
}
148+
149+
void WorkerThread(ThreadArgs args) {
150+
constexpr Device::Type kDevType =
151+
ListGetBest<DevicePriority>(EnabledDevices{});
152+
using Rt = Runtime<kDevType>;
153+
154+
CHECK_RT(Rt, Rt::SetDevice(args.rank));
155+
156+
std::array<char, 256> hostname{};
157+
if (gethostname(hostname.data(), hostname.size()) != 0) {
158+
std::cerr << "Failed to query the hostname for the ReduceScatter worker."
159+
<< std::endl;
160+
std::exit(EXIT_FAILURE);
161+
}
162+
hostname.back() = '\0';
163+
std::cout << "[Rank " << args.rank << "] Host: " << hostname.data()
164+
<< " | GPU: " << Device::StringFromType(kDevType) << " | Device "
165+
<< args.rank << std::endl;
166+
167+
infinicclComm_t comm = nullptr;
168+
CHECK_INFINI(infinicclCommInitRank(&comm, args.size, args.id, args.rank));
169+
170+
const size_t send_count = args.recv_count * static_cast<size_t>(args.size);
171+
const size_t send_bytes = send_count * sizeof(float);
172+
const size_t recv_bytes = args.recv_count * sizeof(float);
173+
const size_t local_offset = static_cast<size_t>(args.rank) * args.recv_count;
174+
175+
std::vector<float> h_send(send_count);
176+
std::vector<float> h_recv(args.recv_count, 0.0f);
177+
FillInput(&h_send, args.recv_count, args.size, args.rank);
178+
179+
float *d_send = nullptr;
180+
float *d_recv = nullptr;
181+
CHECK_RT(Rt, Rt::Malloc(reinterpret_cast<void **>(&d_send), send_bytes));
182+
CHECK_RT(Rt, Rt::Malloc(reinterpret_cast<void **>(&d_recv), recv_bytes));
183+
184+
auto run_scenario = [&](bool in_place) {
185+
CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), send_bytes,
186+
Rt::MemcpyHostToDevice));
187+
CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
188+
189+
float *active_recv = in_place ? d_send + local_offset : d_recv;
190+
auto restore_local_block = [&]() {
191+
if (in_place) {
192+
CHECK_RT(Rt, Rt::Memcpy(active_recv, h_send.data() + local_offset,
193+
recv_bytes, Rt::MemcpyHostToDevice));
194+
CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
195+
}
196+
};
197+
198+
for (int i = 0; i < args.warmup_iter; ++i) {
199+
restore_local_block();
200+
CHECK_INFINI(infinicclReduceScatter(d_send, active_recv, args.recv_count,
201+
infinicclFloat32, infinicclSum, comm,
202+
nullptr));
203+
}
204+
CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
205+
206+
double elapsed_ms = 0.0;
207+
if (in_place) {
208+
for (int i = 0; i < args.profile_iter; ++i) {
209+
restore_local_block();
210+
Timer timer;
211+
CHECK_INFINI(infinicclReduceScatter(d_send, active_recv,
212+
args.recv_count, infinicclFloat32,
213+
infinicclSum, comm, nullptr));
214+
CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
215+
elapsed_ms += timer.ElapsedMs();
216+
}
217+
elapsed_ms /= static_cast<double>(args.profile_iter);
218+
} else {
219+
Timer timer;
220+
for (int i = 0; i < args.profile_iter; ++i) {
221+
CHECK_INFINI(infinicclReduceScatter(d_send, active_recv,
222+
args.recv_count, infinicclFloat32,
223+
infinicclSum, comm, nullptr));
224+
}
225+
CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
226+
elapsed_ms = timer.ElapsedMs() / static_cast<double>(args.profile_iter);
227+
}
228+
229+
CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), active_recv, recv_bytes,
230+
Rt::MemcpyDeviceToHost));
231+
CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
232+
233+
const float expected = ExpectedValue(args.rank, args.size);
234+
const bool correct = Validator::ValidateResult(
235+
h_recv.data(), args.recv_count, expected, args.rank);
236+
ScenarioState *state = in_place ? args.in_place : args.out_of_place;
237+
if (!correct) {
238+
state->correct.store(false, std::memory_order_relaxed);
239+
}
240+
WaitForScenario(state, args.size);
241+
242+
if (args.rank == 0) {
243+
PrintResult(in_place ? "In-Place" : "Out-of-Place",
244+
state->correct.load(std::memory_order_acquire), h_recv,
245+
expected, args.recv_count, args.size, elapsed_ms);
246+
}
247+
};
248+
249+
run_scenario(false);
250+
run_scenario(true);
251+
252+
CHECK_RT(Rt, Rt::Free(d_send));
253+
CHECK_RT(Rt, Rt::Free(d_recv));
254+
CHECK_INFINI(infinicclCommDestroy(comm));
255+
}
256+
257+
void PrintUsage(const char *program) {
258+
std::cout << "Usage: " << program << " [options]\n"
259+
<< "Options:\n"
260+
<< " -g <num_gpus> Number of GPUs (default: 8)\n"
261+
<< " -w <warmup_iters> Warmup iterations (default: 2)\n"
262+
<< " -p <profile_iters> Profile iterations (default: 20)\n"
263+
<< " -n <recv_count> Elements received per rank (default: "
264+
"1048576)\n";
265+
}
266+
267+
} // namespace
268+
269+
int main(int argc, char **argv) {
270+
int num_gpus = 8;
271+
int warmup_iters = 2;
272+
int profile_iters = 20;
273+
size_t recv_count = 1 << 20;
274+
275+
int opt = 0;
276+
while ((opt = getopt(argc, argv, "g:w:p:n:h")) != -1) {
277+
bool parsed = false;
278+
switch (opt) {
279+
case 'g':
280+
parsed = ParsePositiveNumber(optarg, &num_gpus);
281+
break;
282+
case 'w':
283+
parsed = ParsePositiveNumber(optarg, &warmup_iters);
284+
break;
285+
case 'p':
286+
parsed = ParsePositiveNumber(optarg, &profile_iters);
287+
break;
288+
case 'n':
289+
parsed = ParsePositiveNumber(optarg, &recv_count);
290+
break;
291+
case 'h':
292+
PrintUsage(argv[0]);
293+
return EXIT_SUCCESS;
294+
default:
295+
PrintUsage(argv[0]);
296+
return EXIT_FAILURE;
297+
}
298+
299+
if (!parsed) {
300+
std::cerr << "Invalid positive numeric option for ReduceScatter."
301+
<< std::endl;
302+
return EXIT_FAILURE;
303+
}
304+
}
305+
306+
if (optind != argc) {
307+
std::cerr << "Unexpected positional argument for ReduceScatter."
308+
<< std::endl;
309+
return EXIT_FAILURE;
310+
}
311+
if (recv_count > std::numeric_limits<size_t>::max() / sizeof(float) ||
312+
static_cast<size_t>(num_gpus) >
313+
std::numeric_limits<size_t>::max() / recv_count ||
314+
recv_count * static_cast<size_t>(num_gpus) >
315+
std::numeric_limits<size_t>::max() / sizeof(float)) {
316+
std::cerr << "ReduceScatter buffer size overflows `size_t`." << std::endl;
317+
return EXIT_FAILURE;
318+
}
319+
320+
std::array<char, 256> hostname{};
321+
if (gethostname(hostname.data(), hostname.size()) != 0) {
322+
std::cerr << "Failed to query the hostname for ReduceScatter." << std::endl;
323+
return EXIT_FAILURE;
324+
}
325+
hostname.back() = '\0';
326+
std::cout << "[Main Process] Host: " << hostname.data()
327+
<< " | Target GPUs: " << num_gpus << std::endl;
328+
329+
infinicclUniqueId shared_id{};
330+
CHECK_INFINI(infinicclGetUniqueId(&shared_id));
331+
332+
ScenarioState out_of_place;
333+
ScenarioState in_place;
334+
std::vector<std::thread> threads;
335+
threads.reserve(num_gpus);
336+
for (int rank = 0; rank < num_gpus; ++rank) {
337+
ThreadArgs args{rank, num_gpus, shared_id, recv_count,
338+
warmup_iters, profile_iters, &out_of_place, &in_place};
339+
threads.emplace_back(WorkerThread, args);
340+
}
341+
342+
for (auto &thread : threads) {
343+
if (thread.joinable()) {
344+
thread.join();
345+
}
346+
}
347+
348+
const bool correct = out_of_place.correct.load(std::memory_order_acquire) &&
349+
in_place.correct.load(std::memory_order_acquire);
350+
if (correct) {
351+
std::cout << "[Main Process] All ReduceScatter scenarios passed."
352+
<< std::endl;
353+
} else {
354+
std::cerr << "[Main Process] ReduceScatter validation failed." << std::endl;
355+
}
356+
std::cout
357+
<< "[Main Process] All worker threads joined. InfiniCCL finalized safely."
358+
<< std::endl;
359+
return correct ? EXIT_SUCCESS : EXIT_FAILURE;
360+
}

0 commit comments

Comments
 (0)