stream.h

Go to the documentation of this file.

// Copyright © 2023 Apple Inc.
 
#pragma once
 
#include "mlx/device.h"
 
namespace mlx::core {
 
struct Stream {
  int index;
  Device device;
  explicit Stream(int index, Device device) : index(index), device(device) {}
};

Stream default_stream(Device d);

void set_default_stream(Stream s);

Stream new_stream(Device d);

Stream get_stream(int index);
 
inline bool operator==(const Stream& lhs, const Stream& rhs) {
  return lhs.index == rhs.index;
}
 
inline bool operator!=(const Stream& lhs, const Stream& rhs) {
  return !(lhs == rhs);
}
 
/* Synchronize with the default stream. */
void synchronize();
 
/* Synchronize with the provided stream. */
void synchronize(Stream);
 
} // namespace mlx::core