zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Adds send/recv ops in distributed (#1366)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2024-08-26 23:01:37 -07:00
committed by GitHub
parent 1d94ac3f90
commit cdb59faea6
13 changed files with 345 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
mlx/distributed/primitives.h
View File

@@ -97,4 +97,39 @@ class AllGather : public DistPrimitive {
DEFINE_PRINT(AllGather);
};
class Send : public DistPrimitive {
public:
Send(Stream stream, Group group, int dst)
: DistPrimitive(stream, group), dst_(dst) {}
void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
override;
void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
override;
std::pair<std::vector<array>, std::vector<int>> vmap(
const std::vector<array>& inputs,
const std::vector<int>& axes) override;
DEFINE_PRINT(Send);
private:
int dst_;
};
class Recv : public DistPrimitive {
public:
Recv(Stream stream, Group group, int src)
: DistPrimitive(stream, group), src_(src) {}
void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
override;
void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
override;
DEFINE_PRINT(Recv);
private:
int src_;
};
} // namespace mlx::core::distributed