Skip to content

Commit 3816cd5

Browse files
voltjiaZiminli
voltjia
authored and
Ziminli
committed
feat(nccl): add all-gather communication
1 parent 0df90c6 commit 3816cd5

6 files changed

Lines changed: 244 additions & 17 deletions

File tree

‎examples/ccl/all_gather.cc‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
/**
2+
* InfiniCCL Example: Thread-per-GPU Single-Node AllGather
3+
*
4+
* This example validates out-of-place and in-place AllGather across two GPUs
5+
* through InfiniCCL's native CCL backend without an MPI launcher.
6+
*/
7+
8+
#include <algorithm>
9+
#include <atomic>
10+
#include <cstddef>
11+
#include <cstdlib>
12+
#include <iostream>
13+
#include <thread>
14+
#include <vector>
15+
16+
#include "backend_manifest.h"
17+
#include "infiniccl.h"
18+
#include "utils.h"
19+
20+
using namespace infini::ccl;
21+
22+
namespace {
23+
24+
constexpr int kRankCount = 2;
25+
constexpr size_t kNumElements = 1 << 10;
26+
27+
struct ThreadArgs {
28+
int rank;
29+
infinicclUniqueId id;
30+
std::atomic_bool* all_correct;
31+
};
32+
33+
bool Validate(const std::vector<float>& output, int rank, const char* mode) {
34+
for (int source = 0; source < kRankCount; ++source) {
35+
const float expected = static_cast<float>(source + 1);
36+
const size_t offset = static_cast<size_t>(source) * kNumElements;
37+
for (size_t i = 0; i < kNumElements; ++i) {
38+
if (output[offset + i] != expected) {
39+
std::cerr << mode << " validation failed on rank " << rank
40+
<< " at source rank " << source << ", element " << i
41+
<< ": expected " << expected << ", got " << output[offset + i]
42+
<< "." << std::endl;
43+
return false;
44+
}
45+
}
46+
}
47+
return true;
48+
}
49+
50+
void WorkerThread(ThreadArgs args) {
51+
constexpr Device::Type kDevType =
52+
ListGetBest<DevicePriority>(EnabledDevices{});
53+
using Rt = Runtime<kDevType>;
54+
55+
CHECK_RT(Rt, Rt::SetDevice(args.rank));
56+
57+
infinicclComm_t comm = nullptr;
58+
CHECK_INFINI(infinicclCommInitRank(&comm, kRankCount, args.id, args.rank));
59+
60+
std::vector<float> host_send(kNumElements, static_cast<float>(args.rank + 1));
61+
std::vector<float> host_recv(kNumElements * kRankCount, 0.0f);
62+
float* device_send = nullptr;
63+
float* device_recv = nullptr;
64+
const size_t send_bytes = kNumElements * sizeof(float);
65+
const size_t recv_bytes = send_bytes * kRankCount;
66+
67+
CHECK_RT(Rt, Rt::Malloc(reinterpret_cast<void**>(&device_send), send_bytes));
68+
CHECK_RT(Rt, Rt::Malloc(reinterpret_cast<void**>(&device_recv), recv_bytes));
69+
CHECK_RT(Rt, Rt::Memcpy(device_send, host_send.data(), send_bytes,
70+
Rt::MemcpyHostToDevice));
71+
72+
CHECK_INFINI(infinicclAllGather(device_send, device_recv, kNumElements,
73+
infinicclFloat32, comm, nullptr));
74+
CHECK_RT(Rt, Rt::Memcpy(host_recv.data(), device_recv, recv_bytes,
75+
Rt::MemcpyDeviceToHost));
76+
if (!Validate(host_recv, args.rank, "Out-of-place AllGather")) {
77+
args.all_correct->store(false, std::memory_order_relaxed);
78+
}
79+
80+
std::fill(host_recv.begin(), host_recv.end(), 0.0f);
81+
std::fill_n(host_recv.begin() + static_cast<size_t>(args.rank) * kNumElements,
82+
kNumElements, static_cast<float>(args.rank + 1));
83+
CHECK_RT(Rt, Rt::Memcpy(device_recv, host_recv.data(), recv_bytes,
84+
Rt::MemcpyHostToDevice));
85+
86+
float* local_block =
87+
device_recv + static_cast<size_t>(args.rank) * kNumElements;
88+
CHECK_INFINI(infinicclAllGather(local_block, device_recv, kNumElements,
89+
infinicclFloat32, comm, nullptr));
90+
CHECK_RT(Rt, Rt::Memcpy(host_recv.data(), device_recv, recv_bytes,
91+
Rt::MemcpyDeviceToHost));
92+
if (!Validate(host_recv, args.rank, "In-place AllGather")) {
93+
args.all_correct->store(false, std::memory_order_relaxed);
94+
}
95+
96+
CHECK_RT(Rt, Rt::Free(device_send));
97+
CHECK_RT(Rt, Rt::Free(device_recv));
98+
CHECK_INFINI(infinicclCommDestroy(comm));
99+
}
100+
101+
} // namespace
102+
103+
int main() {
104+
infinicclUniqueId shared_id;
105+
CHECK_INFINI(infinicclGetUniqueId(&shared_id));
106+
107+
std::atomic_bool all_correct{true};
108+
std::vector<std::thread> threads;
109+
threads.reserve(kRankCount);
110+
for (int rank = 0; rank < kRankCount; ++rank) {
111+
threads.emplace_back(WorkerThread,
112+
ThreadArgs{rank, shared_id, &all_correct});
113+
}
114+
for (auto& thread : threads) {
115+
thread.join();
116+
}
117+
118+
if (!all_correct.load(std::memory_order_relaxed)) {
119+
return EXIT_FAILURE;
120+
}
121+
std::cout << "AllGather validation passed." << std::endl;
122+
return EXIT_SUCCESS;
123+
}

‎examples/mpi/all_gather.cc‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424

2525
using namespace infini::ccl;
2626

27-
void RunAllGatherExample(int argc, char **argv, int warmup_iter,
27+
bool RunAllGatherExample(int argc, char **argv, int warmup_iter,
2828
int profile_iter, const size_t kNumElements) {
2929
constexpr Device::Type kDevType =
3030
ListGetBest<DevicePriority>(EnabledDevices{});
@@ -113,14 +113,14 @@ void RunAllGatherExample(int argc, char **argv, int warmup_iter,
113113

114114
// Result Validation
115115
bool correct = true;
116-
int error_count = 0;
117116

118117
for (int src_rank = 0; src_rank < size; ++src_rank) {
119118
float expected = static_cast<float>(src_rank + 1);
120119
size_t offset = static_cast<size_t>(src_rank) * kNumElements;
121120

122-
Validator::ValidateResult(h_recv.data() + offset, kNumElements, expected,
123-
rank);
121+
correct = Validator::ValidateResult(h_recv.data() + offset, kNumElements,
122+
expected, rank) &&
123+
correct;
124124
}
125125

126126
if (rank == 0) {
@@ -132,7 +132,6 @@ void RunAllGatherExample(int argc, char **argv, int warmup_iter,
132132
std::cout << "Correct: "
133133
<< (correct ? (GREEN + std::string("YES") + RESET)
134134
: (RED + std::string("NO") + RESET));
135-
if (!correct) std::cout << " (" << error_count << " errors)";
136135
std::cout << std::endl;
137136

138137
std::cout << "Sample blocks: ";
@@ -159,14 +158,16 @@ void RunAllGatherExample(int argc, char **argv, int warmup_iter,
159158
if (rank == 0) {
160159
std::cout << "InfiniCCL finalized." << std::endl;
161160
}
161+
return correct;
162162
}
163163

164164
int main(int argc, char **argv) {
165165
int warmup_iters = 2;
166166
int profile_iters = 20;
167167
size_t num_elements = 1 << 20;
168168

169-
RunAllGatherExample(argc, argv, warmup_iters, profile_iters, num_elements);
170-
171-
return EXIT_SUCCESS;
169+
return RunAllGatherExample(argc, argv, warmup_iters, profile_iters,
170+
num_elements)
171+
? EXIT_SUCCESS
172+
: EXIT_FAILURE;
172173
}

‎src/backends/ccl/common/impl/all_gather.h‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
#ifndef INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_ALL_GATHER_H_
2+
#define INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_ALL_GATHER_H_
3+
4+
#include "backends/ccl/common/api.h"
5+
#include "backends/ccl/common/comm_instance.h"
6+
#include "base/all_gather.h"
7+
#include "communicator.h"
8+
9+
namespace infini::ccl {
10+
11+
template <BackendType backend, Device::Type device>
12+
class CclAllGatherImpl {
13+
public:
14+
static ReturnStatus Apply(const void* send_buff, void* recv_buff,
15+
size_t count, DataType data_type,
16+
Communicator* comm, void* stream) {
17+
using Api = CclApi<backend, device>;
18+
using TypeMap = CclTypeMap<backend, device>;
19+
using CommInstance = CclCommInstance<Api>;
20+
21+
if (!comm || !comm->intra_comm() || comm->intra_comm_backend() != backend ||
22+
comm->device_type() != device) {
23+
return ReturnStatus::kInternalError;
24+
}
25+
26+
auto* intra = static_cast<CommInstance*>(comm->intra_comm());
27+
if (!intra->handle) {
28+
return ReturnStatus::kInternalError;
29+
}
30+
31+
typename Api::DataType ccl_type{};
32+
if (!TypeMap::ToBackendDataType(data_type, &ccl_type)) {
33+
return ReturnStatus::kNotSupported;
34+
}
35+
36+
return Api::Check(
37+
Api::AllGather(send_buff, recv_buff, count, ccl_type, intra->handle,
38+
reinterpret_cast<typename Api::Stream>(stream)));
39+
}
40+
};
41+
42+
} // namespace infini::ccl
43+
44+
#endif // INFINI_CCL_BACKENDS_CCL_COMMON_IMPL_ALL_GATHER_H_

‎src/backends/ccl/nccl/api.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,13 @@ struct NcclApi {
4949
stream);
5050
}
5151

52+
static Result AllGather(const void* send_buff, void* recv_buff,
53+
size_t send_count, DataType data_type, Comm comm,
54+
Stream stream) {
55+
return ncclAllGather(send_buff, recv_buff, send_count, data_type, comm,
56+
stream);
57+
}
58+
5259
static Result Send(const void* send_buff, size_t count, DataType data_type,
5360
int peer, Comm comm, Stream stream) {
5461
return ncclSend(send_buff, count, data_type, peer, comm, stream);

‎src/backends/ccl/nccl/impl/all_gather.h‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
#ifndef INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_ALL_GATHER_H_
2+
#define INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_ALL_GATHER_H_
3+
4+
#include "backends/ccl/common/impl/all_gather.h"
5+
6+
namespace infini::ccl {
7+
8+
template <Device::Type device>
9+
class AllGatherImpl<BackendType::kNccl, device>
10+
: public CclAllGatherImpl<BackendType::kNccl, device> {};
11+
12+
template <>
13+
struct BackendEnabled<AllGather, BackendType::kNccl> : std::true_type {};
14+
15+
} // namespace infini::ccl
16+
17+
#endif // INFINI_CCL_BACKENDS_CCL_NCCL_IMPL_ALL_GATHER_H_

‎src/base/all_gather.h‎

Lines changed: 44 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,30 +19,65 @@ class AllGather : public Operation<AllGather> {
1919
static ReturnStatus Execute(const void *send_buff, void *recv_buff,
2020
size_t count, DataType datatype,
2121
void *comm_handle, void *stream) {
22-
if (HasInvalidArgs(send_buff, recv_buff, datatype, comm_handle)) {
22+
if (!comm_handle) {
23+
LOG("Invalid communicator handle for `AllGather`.");
2324
return ReturnStatus::kInvalidArgument;
2425
}
26+
2527
auto *comm = static_cast<Communicator *>(comm_handle);
28+
if (HasInvalidArgs(send_buff, recv_buff, count, datatype)) {
29+
return ReturnStatus::kInvalidArgument;
30+
}
31+
if (count == 0) {
32+
return ReturnStatus::kSuccess;
33+
}
34+
35+
if (!comm->HasBackend(backend_type) || comm->device_type() != device_type) {
36+
using DispatchKey =
37+
typename BackendDependentType<backend_type, AllGather>::type;
38+
const BackendType comm_backend =
39+
Operation<DispatchKey>::FindSupportedBackend(
40+
comm->device_type(),
41+
{comm->HasBackend(backend_type) ? backend_type
42+
: BackendType::kCount,
43+
comm->intra_comm_backend(), comm->inter_comm_backend()});
44+
if (comm_backend == BackendType::kCount) {
45+
if (comm->intra_comm_backend() == BackendType::kCount &&
46+
comm->inter_comm_backend() == BackendType::kCount) {
47+
LOG("No initialized backend is available for `AllGather`.");
48+
return ReturnStatus::kInternalError;
49+
}
50+
return ReturnStatus::kNotSupported;
51+
}
52+
53+
return Operation<DispatchKey>::Call(comm_backend, comm->device_type(),
54+
send_buff, recv_buff, count, datatype,
55+
comm_handle, stream);
56+
}
57+
2658
return AllGatherImpl<backend_type, device_type>::Apply(
2759
send_buff, recv_buff, count, datatype, comm, stream);
2860
}
2961

3062
private:
63+
template <BackendType, typename T>
64+
struct BackendDependentType {
65+
using type = T;
66+
};
67+
3168
static bool HasInvalidArgs(const void *send_buff, void *recv_buff,
32-
DataType datatype, void *comm_handle) {
33-
if (!comm_handle) {
34-
// TODO(lzm): change to use `glog`.
35-
LOG("Invalid communicator handle for `AllGather`.");
69+
size_t count, DataType datatype) {
70+
if (datatype < DataType::kChar || datatype >= DataType::kNumTypes) {
71+
LOG("Invalid data type for `AllGather`.");
3672
return true;
3773
}
74+
if (count == 0) {
75+
return false;
76+
}
3877
if (!send_buff || !recv_buff) {
3978
LOG("Invalid buffer pointer for `AllGather`.");
4079
return true;
4180
}
42-
if (datatype < DataType::kChar || datatype >= DataType::kNumTypes) {
43-
LOG("Invalid data type for `AllGather`.");
44-
return true;
45-
}
4681
return false;
4782
}
4883
};

0 commit comments

Comments
 (0)