mirror of
https://github.com/ml-explore/mlx.git
synced 2025-06-26 02:33:21 +08:00
4ba0c24a8f
* export and import functions * refactor + works for few primitives * nit * allow primitives with state * nit * nit * simplify serialize / deserialize * fix for constants * python bindings * maybe fix serialize failure case * add example * more primitives, training kind of works * same result for python and c++ * some fixes * fix export * template it up * some simplificatoin * rebase * allow kwargs and multiple functions * exporter * more primitives for exporting * deal with endianness * handle invalid stream * add docstring
42 lines
896 B
C++
42 lines
896 B
C++
// 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) {}
|
|
};
|
|
|
|
/** Get the default stream for the given device. */
|
|
Stream default_stream(Device d);
|
|
|
|
/** Make the stream the default for its device. */
|
|
void set_default_stream(Stream s);
|
|
|
|
/** Make a new stream on the given device. */
|
|
Stream new_stream(Device d);
|
|
|
|
/** Get the stream with the given index. */
|
|
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
|