zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 09:21:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add rudimentary barrier

Browse Source
This commit is contained in:
Angelos Katharopoulos 2024-11-05 11:34:55 -08:00
parent 44f0de2854
commit 6fc00d2c10
1 changed files with 13 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
mlx/distributed/sockets/sockets.cpp
View File

@ -385,6 +385,19 @@ Group Group::split(int color, int key) {
throw std::runtime_error("Splitting not supported yet");
}
void Group::barrier() {
char buff[128];
std::memset(buff, 1, 128);
auto group = std::static_pointer_cast<GroupImpl>(raw_group());
int size = group->size();
int rank = group->rank();
for (int distance = 1; distance <= size / 2; distance *= 2) {
group->send_recv_sum<char>(buff, 128, rank ^ distance);
}
}
Group init(bool strict /* = false */) {
static std::shared_ptr<GroupImpl> global_group = nullptr;