Skip to content

Commit 794fcf3

Browse files
GordonYang1GordonYang1
committed
feat: support Reduce with OpenMPI backend implementation
1 parent 5374d9d commit 794fcf3

4 files changed

Lines changed: 328 additions & 0 deletions

File tree

‎examples/reduce.cc‎

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
/**
2+
* InfiniCCL Example: Reduce
3+
*
4+
* This example demonstrates the API for performing a collective
5+
* reduction across multiple accelerators and nodes, where only
6+
* `root` receives the reduced result.
7+
*/
8+
9+
#include <unistd.h>
10+
11+
#include <iostream>
12+
#include <vector>
13+
14+
// Public API
15+
#include "infiniccl.h"
16+
17+
// Example-Specific Utilities
18+
#include "utils.h"
19+
20+
// Internal Headers (Accessible via example-specific include paths, technically
21+
// not public APIs)
22+
#include "backend_manifest.h"
23+
#include "device.h"
24+
#include "runtime.h"
25+
#include "traits.h"
26+
27+
using namespace infini::ccl;
28+
29+
void RunReduceExample(int argc, char **argv, int warmup_iter, int profile_iter,
30+
const size_t kNumElements) {
31+
constexpr Device::Type kDevType =
32+
ListGetBest<DevicePriority>(EnabledDevices{});
33+
using Rt = Runtime<kDevType>;
34+
35+
CHECK_INFINI(infiniInit(&argc, &argv));
36+
37+
int rank, size;
38+
CHECK_INFINI(infiniGetRank(&rank));
39+
CHECK_INFINI(infiniGetSize(&size));
40+
41+
char hostname[256];
42+
gethostname(hostname, sizeof(hostname));
43+
44+
// Map local rank to GPU device.
45+
// Note: this is just for info printing. In practice, this part is not needed.
46+
const char *local_rank_str = std::getenv("OMPI_COMM_WORLD_LOCAL_RANK");
47+
int local_rank = 0;
48+
if (local_rank_str != nullptr) {
49+
local_rank = std::atoi(local_rank_str);
50+
}
51+
52+
std::cout << "[Rank " << rank << "] Host: " << hostname
53+
<< " | GPU: " << Device::StringFromType(kDevType) << " "
54+
<< " | Device " << local_rank << std::endl;
55+
56+
// Setup Communicator
57+
infiniComm_t comm = nullptr;
58+
CHECK_INFINI(infiniCommInitAll(&comm, size, nullptr));
59+
60+
// Root of the Reduce
61+
constexpr int kRoot = 0;
62+
63+
// Prepare Data
64+
std::vector<float> h_send(kNumElements);
65+
std::vector<float> h_recv(kNumElements, 0.0f);
66+
67+
// Initialize: each rank provides its (rank + 1) as data.
68+
for (size_t i = 0; i < kNumElements; i++) {
69+
h_send[i] = static_cast<float>(rank + 1);
70+
}
71+
72+
float *d_send, *d_recv;
73+
size_t total_bytes = kNumElements * sizeof(*d_send);
74+
CHECK_RT(Rt, Rt::Malloc((void **)&d_send, total_bytes));
75+
CHECK_RT(Rt, Rt::Malloc((void **)&d_recv, total_bytes));
76+
CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), total_bytes,
77+
Rt::MemcpyHostToDevice));
78+
CHECK_RT(Rt, Rt::Memcpy(d_recv, h_recv.data(), total_bytes,
79+
Rt::MemcpyHostToDevice));
80+
81+
if (rank == kRoot) {
82+
std::cout << "\n=== Performing Reduce on GPU Memory ===" << std::endl;
83+
std::cout << "Data size: " << kNumElements << " floats ("
84+
<< total_bytes / 1024 / 1024 << " MB)" << std::endl;
85+
std::cout << "Operation: Sum" << std::endl;
86+
std::cout << "Root Rank: " << kRoot << std::endl;
87+
std::cout << "Warm-up iterations: " << warmup_iter << std::endl;
88+
std::cout << "Profile iterations: " << profile_iter << std::endl;
89+
}
90+
91+
CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
92+
93+
// Warm-up and D2H transfer the answer on `root`.
94+
CHECK_INFINI(infiniReduce(d_send, d_recv, kNumElements, infiniFloat32,
95+
infiniSum, kRoot, comm, nullptr));
96+
if (rank == kRoot) {
97+
CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes,
98+
Rt::MemcpyDeviceToHost));
99+
}
100+
101+
for (int i = 1; i < warmup_iter; ++i) {
102+
CHECK_INFINI(infiniReduce(d_send, d_recv, kNumElements, infiniFloat32,
103+
infiniSum, kRoot, comm, nullptr));
104+
}
105+
CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
106+
107+
// Profiling
108+
Timer timer;
109+
110+
for (int i = 0; i < profile_iter; ++i) {
111+
CHECK_INFINI(infiniReduce(d_send, d_recv, kNumElements, infiniFloat32,
112+
infiniSum, kRoot, comm, nullptr));
113+
}
114+
115+
CHECK_RT(Rt, Rt::StreamSynchronize(nullptr));
116+
double elapsed = timer.elapsed_ms() / static_cast<double>(profile_iter);
117+
118+
// Result Validation (only meaningful on `root`).
119+
if (rank == kRoot) {
120+
float expected = 0.0f;
121+
for (int r = 0; r < size; ++r) {
122+
expected += static_cast<float>(r + 1);
123+
}
124+
125+
Validator::ValidateResult(h_recv.data(), kNumElements, expected, rank, true,
126+
"Reduce");
127+
128+
Metrics metrics{elapsed, total_bytes, size};
129+
metrics.Print();
130+
}
131+
132+
// Cleanup
133+
CHECK_RT(Rt, Rt::Free(d_send));
134+
CHECK_RT(Rt, Rt::Free(d_recv));
135+
136+
CHECK_INFINI(infiniCommDestroy(comm));
137+
CHECK_INFINI(infiniFinalize());
138+
139+
if (rank == kRoot) {
140+
std::cout << "InfiniCCL finalized." << std::endl;
141+
}
142+
}
143+
144+
int main(int argc, char **argv) {
145+
int warmup_iters = 2;
146+
int profile_iters = 20;
147+
size_t num_elements = 1 << 20;
148+
149+
RunReduceExample(argc, argv, warmup_iters, profile_iters, num_elements);
150+
151+
return EXIT_SUCCESS;
152+
}

‎include/comm.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,11 @@ infiniResult_t infiniBroadcast(const void *sendbuff, void *recvbuff,
4848
infiniResult_t infiniBcast(void *buff, size_t count, infiniDataType_t datatype,
4949
int root, infiniComm_t comm, void *stream);
5050

51+
infiniResult_t infiniReduce(const void *sendbuff, void *recvbuff,
52+
size_t count, infiniDataType_t datatype,
53+
infiniRedOp_t op, int root, infiniComm_t comm,
54+
void *stream);
55+
5156
infiniResult_t infiniAllGather(const void *sendbuff, void *recvbuff,
5257
size_t count, infiniDataType_t datatype,
5358
infiniComm_t comm, void *stream);

‎src/base/reduce.h‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
#ifndef INFINI_CCL_BASE_REDUCE_H_
2+
#define INFINI_CCL_BASE_REDUCE_H_
3+
4+
#include "comm_impl.h"
5+
#include "communicator.h"
6+
#include "data_type_impl.h"
7+
#include "logging.h"
8+
#include "operation.h"
9+
#include "return_status_impl.h"
10+
11+
namespace infini::ccl {
12+
13+
template <BackendType backend_type, Device::Type device_type>
14+
struct ReduceImpl;
15+
16+
class Reduce : public Operation<Reduce> {
17+
public:
18+
template <BackendType backend_type, Device::Type device_type>
19+
static ReturnStatus Execute(const void *send_buff, void *recv_buff,
20+
size_t count, DataType datatype,
21+
ReductionOpType op, int root, void *comm_handle,
22+
void *stream) {
23+
if (!comm_handle) {
24+
LOG("Invalid communicator handle for `Reduce`.");
25+
return ReturnStatus::kInvalidArgument;
26+
}
27+
28+
auto *comm = static_cast<Communicator *>(comm_handle);
29+
if (HasInvalidArgs(send_buff, recv_buff, datatype, op, root, comm)) {
30+
return ReturnStatus::kInvalidArgument;
31+
}
32+
if (count == 0) {
33+
return ReturnStatus::kSuccess;
34+
}
35+
36+
return ReduceImpl<backend_type, device_type>::Apply(
37+
send_buff, recv_buff, count, datatype, op, root, comm, stream);
38+
}
39+
40+
private:
41+
static bool HasInvalidArgs(const void *send_buff, void *recv_buff,
42+
DataType datatype, ReductionOpType op, int root,
43+
Communicator *comm) {
44+
if (datatype < DataType::kChar || datatype >= DataType::kNumTypes) {
45+
LOG("Invalid data type for `Reduce`.");
46+
return true;
47+
}
48+
if (op < ReductionOpType::kSum || op >= ReductionOpType::kNumRedOps) {
49+
LOG("Invalid reduction operation for `Reduce`.");
50+
return true;
51+
}
52+
if (root < 0 || root >= comm->size()) {
53+
LOG("Invalid root rank for `Reduce`.");
54+
return true;
55+
}
56+
if (!send_buff) {
57+
LOG("Invalid send buffer pointer for `Reduce`.");
58+
return true;
59+
}
60+
if (comm->rank() == root && !recv_buff) {
61+
LOG("Invalid root receive buffer pointer for `Reduce`.");
62+
return true;
63+
}
64+
return false;
65+
}
66+
};
67+
68+
} // namespace infini::ccl
69+
70+
#endif // INFINI_CCL_BASE_REDUCE_H_

‎src/ompi/impl/reduce.h‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
#ifndef INFINI_CCL_OMPI_IMPL_REDUCE_H_
2+
#define INFINI_CCL_OMPI_IMPL_REDUCE_H_
3+
4+
#include <cstdlib>
5+
#include <limits>
6+
7+
#include "base/reduce.h"
8+
#include "communicator.h"
9+
#include "data_type_impl.h"
10+
#include "dispatcher.h"
11+
#include "logging.h"
12+
#include "ompi/checks.h"
13+
#include "ompi/comm_instance.h"
14+
#include "ompi/type_map.h"
15+
#include "runtime.h"
16+
17+
namespace infini::ccl {
18+
19+
template <Device::Type device_type>
20+
class ReduceImpl<BackendType::kOmpi, device_type> {
21+
public:
22+
static ReturnStatus Apply(const void *send_buff, void *recv_buff,
23+
size_t count, DataType data_type,
24+
ReductionOpType op, int root, Communicator *comm,
25+
void *stream) {
26+
constexpr Device::Type kDev =
27+
ListGetBest<DevicePriority>(ActiveDevices<Reduce>{});
28+
using Rt = Runtime<kDev>;
29+
30+
auto *inst = static_cast<OmpiInstance *>(comm->inter_comm());
31+
if (!inst || inst->handle == MPI_COMM_NULL) {
32+
LOG("Invalid OpenMPI communicator instance for `Reduce`.");
33+
return ReturnStatus::kInternalError;
34+
}
35+
36+
MPI_Datatype mpi_type = DataTypeToOmpiType(data_type);
37+
MPI_Op mpi_op = RedOpToOmpiOp(op);
38+
39+
if (count > static_cast<size_t>(std::numeric_limits<int>::max())) {
40+
LOG("`count` exceeds MPI int range for `Reduce`.");
41+
return ReturnStatus::kInvalidArgument;
42+
}
43+
int mpi_count = static_cast<int>(count);
44+
45+
size_t type_size = kDataTypeToSize.at(data_type);
46+
size_t total_bytes = count * type_size;
47+
const bool is_root = comm->rank() == root;
48+
49+
// Host staging buffers. Only `root` allocates the receive side, since
50+
// `MPI_Reduce` writes the output only on `root`.
51+
void *host_sendbuf = std::malloc(total_bytes);
52+
void *host_recvbuf = is_root ? std::malloc(total_bytes) : nullptr;
53+
if (!host_sendbuf || (is_root && !host_recvbuf)) {
54+
std::free(host_sendbuf);
55+
std::free(host_recvbuf);
56+
LOG("Failed to allocate host buffers for `Reduce` staging.");
57+
return ReturnStatus::kSystemError;
58+
}
59+
60+
CHECK_STATUS(Rt, Rt::Memcpy(host_sendbuf, send_buff, total_bytes,
61+
Rt::MemcpyDeviceToHost));
62+
CHECK_STATUS(Rt, Rt::StreamSynchronize(static_cast<Rt::Stream>(stream)));
63+
64+
INFINI_CHECK_MPI(MPI_Reduce(host_sendbuf, host_recvbuf, mpi_count, mpi_type,
65+
mpi_op, root, inst->handle));
66+
67+
if (is_root) {
68+
if (op == ReductionOpType::kAvg) {
69+
float scale = 1.0f / static_cast<float>(comm->size());
70+
71+
DispatchFunc<kDev, AllTypes>(data_type, [&](auto dtype) {
72+
using T = typename decltype(dtype)::type;
73+
74+
T *typed_buf = static_cast<T *>(host_recvbuf);
75+
76+
// Simply do the averaging on the CPU before the H2D copy.
77+
for (size_t i = 0; i < count; ++i) {
78+
// TODO(lzm): should later use the unified `Cast` function instead
79+
// of `static_cast` to support CPU custom types.
80+
typed_buf[i] *= static_cast<T>(scale);
81+
}
82+
});
83+
}
84+
85+
CHECK_STATUS(Rt, Rt::Memcpy(recv_buff, host_recvbuf, total_bytes,
86+
Rt::MemcpyHostToDevice));
87+
}
88+
89+
std::free(host_sendbuf);
90+
std::free(host_recvbuf);
91+
92+
return ReturnStatus::kSuccess;
93+
}
94+
};
95+
96+
template <>
97+
struct BackendEnabled<Reduce, BackendType::kOmpi> : std::true_type {};
98+
99+
} // namespace infini::ccl
100+
101+
#endif // INFINI_CCL_OMPI_IMPL_REDUCE_H_

0 commit comments

Comments
 (0)