distributed.h

Go to the documentation of this file.

// Copyright © 2024 Apple Inc.
 
#pragma once
 
#include <memory>
 
#include "mlx/array.h"
 
namespace mlx::core::distributed {
 
/* Check if a communication backend is available */
bool is_available();
 
struct Group {
  Group(std::shared_ptr<void> group) : group_(group) {}
 
  int rank();
  int size();
 
  Group split(int color, int key = -1);
 
  const std::shared_ptr<void>& raw_group() {
    return group_;
  }
 
 private:
  std::shared_ptr<void> group_{nullptr};
};
 
Group init(bool strict = false);
 
namespace detail {
 
/* Return the communication stream. */
Stream communication_stream();
 
/* Perform an all reduce sum operation */
void all_sum(Group group, const array& input, array& output);
 
/* Perform an all reduce sum operation */
void all_gather(Group group, const array& input, array& output);
 
} // namespace detail
 
} // namespace mlx::core::distributed