distributed_impl.h

Go to the documentation of this file.

// Copyright © 2024 Apple Inc.
 
#pragma once
 
#include "mlx/distributed/distributed.h"
 
namespace mlx::core::distributed::detail {

class GroupImpl {
 public:
  virtual ~GroupImpl() {}
 
  virtual int rank() = 0;
  virtual int size() = 0;
  virtual std::shared_ptr<GroupImpl> split(int color, int key = -1) = 0;
 
  virtual void all_sum(const array& input, array& output, Stream stream) = 0;
  virtual void all_gather(const array& input, array& output, Stream stream) = 0;
  virtual void send(const array& input, int dst, Stream stream) = 0;
  virtual void recv(array& out, int src, Stream stream) = 0;
};
 
/* Perform an all reduce sum operation */
void all_sum(Group group, const array& input, array& output, Stream stream);
 
/* Perform an all gather operation */
void all_gather(Group group, const array& input, array& output, Stream stream);

void send(Group group, const array& input, int dst, Stream stream);

void recv(Group group, array& out, int src, Stream stream);
 
} // namespace mlx::core::distributed::detail