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 device context manager (#679)

Browse Source
This commit is contained in:
Diogo
2024-02-14 17:14:58 -05:00
committed by GitHub
parent ccf1645995
commit 35431a4ac8
15 changed files with 230 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
mlx/utils.h
View File

@@ -2,6 +2,8 @@
#pragma once
#include <variant>
#include "array.h"
#include "device.h"
#include "dtype.h"
@@ -9,6 +11,30 @@
namespace mlx::core {
using StreamOrDevice = std::variant<std::monostate, Stream, Device>;
Stream to_stream(StreamOrDevice s);
struct StreamContext {
public:
StreamContext(StreamOrDevice s) : _stream(default_stream(default_device())) {
if (std::holds_alternative<std::monostate>(s)) {
throw std::runtime_error(
"[StreamContext] Invalid argument, please specify a stream or device.");
}
auto _s = to_stream(s);
set_default_device(_s.device);
set_default_stream(_s);
}
~StreamContext() {
set_default_device(_stream.device);
set_default_stream(_stream);
}
private:
Stream _stream;
};
struct PrintFormatter {
inline void print(std::ostream& os, bool val);
inline void print(std::ostream& os, int16_t val);