diff --git a/docs/src/python/devices_and_streams.rst b/docs/src/python/devices_and_streams.rst index 78dd3e56..9d0d15f2 100644 --- a/docs/src/python/devices_and_streams.rst +++ b/docs/src/python/devices_and_streams.rst @@ -14,8 +14,10 @@ Devices and Streams set_default_device default_stream new_stream + new_thread_local_stream set_default_stream stream synchronize + clear_streams device_count device_info diff --git a/mlx/backend/no_gpu/eval.cpp b/mlx/backend/no_gpu/eval.cpp index 6d6c32fc..fa93a233 100644 --- a/mlx/backend/no_gpu/eval.cpp +++ b/mlx/backend/no_gpu/eval.cpp @@ -9,7 +9,10 @@ namespace mlx::core::gpu { void init() {} -void new_stream(Stream) {} +void new_stream(Stream) { + throw std::invalid_argument( + "[new_stream] Cannot make gpu stream without gpu backend."); +} void eval(array&) { throw std::runtime_error("[gpu::eval] GPU backend is not available"); diff --git a/mlx/stream.cpp b/mlx/stream.cpp index ee1db016..8d248d56 100644 --- a/mlx/stream.cpp +++ b/mlx/stream.cpp @@ -6,6 +6,7 @@ #include "mlx/scheduler.h" #include +#include #include #include @@ -29,6 +30,11 @@ auto& all_streams() { return streams_and_mtx; } +auto& thread_local_streams() { + static std::tuple, std::mutex> streams_and_mtx; + return streams_and_mtx; +} + } // namespace Stream default_stream(Device d) { @@ -58,10 +64,6 @@ std::vector get_streams() { } Stream new_stream(Device d) { - if (!gpu::is_available() && d == Device::gpu) { - throw std::invalid_argument( - "[new_stream] Cannot make gpu stream without gpu backend."); - } auto& [streams, mtx] = all_streams(); std::unique_lock lock(mtx); int index = streams.size(); @@ -73,4 +75,20 @@ Stream new_stream(Device d) { return s; } +ThreadLocalStream new_thread_local_stream(Device d) { + auto& [streams, mtx] = thread_local_streams(); + std::lock_guard lock(mtx); + int index = streams.size(); + return streams.emplace_back(index, d); +} + +Stream stream_from_thread_local_stream(ThreadLocalStream tls) { + static thread_local std::map streams; + auto it = streams.find(tls); + if (it == streams.end()) { + it = streams.emplace(tls, new_stream(tls.device)).first; + } + return it->second; +} + } // namespace mlx::core diff --git a/mlx/stream.h b/mlx/stream.h index b3383bc4..ae333e03 100644 --- a/mlx/stream.h +++ b/mlx/stream.h @@ -21,6 +21,10 @@ struct MLX_API Stream { } }; +struct MLX_API ThreadLocalStream : public Stream { + using Stream::Stream; +}; + /** Get the default stream of current thread for the given device. */ MLX_API Stream default_stream(Device d); @@ -30,6 +34,12 @@ MLX_API void set_default_stream(Stream s); /** Make a new stream on the given device. */ MLX_API Stream new_stream(Device d); +/** Make a new stream that will be unique per thread. */ +MLX_API ThreadLocalStream new_thread_local_stream(Device d); + +/** Get the stream for current thread from ThreadLocalStream. */ +MLX_API Stream stream_from_thread_local_stream(ThreadLocalStream tls); + /** Get all available streams. */ MLX_API std::vector get_streams(); diff --git a/mlx/utils.cpp b/mlx/utils.cpp index c257961d..239e6603 100644 --- a/mlx/utils.cpp +++ b/mlx/utils.cpp @@ -17,6 +17,8 @@ Stream to_stream(StreamOrDevice s) { return default_stream(default_device()); } else if (std::holds_alternative(s)) { return default_stream(std::get(s)); + } else if (std::holds_alternative(s)) { + return stream_from_thread_local_stream(std::get(s)); } else { return std::get(s); } diff --git a/mlx/utils.h b/mlx/utils.h index 134f2efc..7835a970 100644 --- a/mlx/utils.h +++ b/mlx/utils.h @@ -13,7 +13,8 @@ namespace mlx::core { -using StreamOrDevice = std::variant; +using StreamOrDevice = + std::variant; MLX_API Stream to_stream(StreamOrDevice s); MLX_API Stream to_stream(StreamOrDevice s, Device default_); diff --git a/python/src/stream.cpp b/python/src/stream.cpp index 9b0f61c1..8764eaf7 100644 --- a/python/src/stream.cpp +++ b/python/src/stream.cpp @@ -41,26 +41,6 @@ class PyStreamContext { mx::StreamContext* _inner; }; -class PyThreadLocalStream { - public: - PyThreadLocalStream(mx::Device d) : device(d) {} - - mx::Stream stream() const { - thread_local std::unordered_map - streams; - - auto it = streams.find(this); - if (it == streams.end()) { - auto result = streams.emplace(this, mx::new_stream(device)); - it = result.first; - } - - return it->second; - } - - mx::Device device; -}; - void init_stream(nb::module_& m) { nb::class_( m, @@ -69,11 +49,6 @@ void init_stream(nb::module_& m) { A stream for running operations on a given device. )pbdoc") .def_ro("device", &mx::Stream::device) - .def( - "__init__", - [](mx::Stream* s, const PyThreadLocalStream& tls) { - return new (s) mx::Stream(tls.stream()); - }) .def( "__repr__", [](const mx::Stream& s) { @@ -86,29 +61,28 @@ void init_stream(nb::module_& m) { s == nb::cast(other); }); - nb::class_( + nb::class_( m, "ThreadLocalStream", R"pbdoc( A stream that will be unique per thread and can be used to run operations on a given device. )pbdoc") - .def_ro("device", &PyThreadLocalStream::device) - .def(nb::init()) + .def_ro("device", &mx::ThreadLocalStream::device) .def( "__repr__", - [](const PyThreadLocalStream& s) { + [](const mx::ThreadLocalStream& s) { std::ostringstream os; - os << "ThreadLocalStream(" << s.device << ")"; + os << "ThreadLocalStream(" << s.device << ", " << s.index << ")"; return os.str(); }) - .def("__eq__", [](const PyThreadLocalStream& s, const nb::object& other) { - auto s_other = mx::default_stream(mx::default_device()); - return nb::try_cast(other, s_other) && - s_other == s.stream(); - }); + .def( + "__eq__", + [](const mx::ThreadLocalStream& s, const nb::object& other) { + return nb::isinstance(other) && + s == nb::cast(other); + }); nb::implicitly_convertible(); - nb::implicitly_convertible(); m.def( "default_stream", @@ -133,6 +107,11 @@ void init_stream(nb::module_& m) { &mx::new_stream, "device"_a, R"pbdoc(Make a new stream on the given device.)pbdoc"); + m.def( + "new_thread_local_stream", + &mx::new_thread_local_stream, + "device"_a, + R"pbdoc(Make a new stream that will be unique per thread.)pbdoc"); m.def( "clear_streams", &mx::clear_streams, diff --git a/python/tests/test_threads.py b/python/tests/test_threads.py index aaf47807..b142e901 100644 --- a/python/tests/test_threads.py +++ b/python/tests/test_threads.py @@ -25,7 +25,7 @@ class TestThreads(mlx_tests.MLXTestCase): t1.join() t2.join() - test_stream = mx.ThreadLocalStream(mx.default_device()) + test_stream = mx.new_thread_local_stream(mx.default_device()) def test_success(): with mx.stream(test_stream): diff --git a/tests/scheduler_tests.cpp b/tests/scheduler_tests.cpp index 64e060d0..93b6818e 100644 --- a/tests/scheduler_tests.cpp +++ b/tests/scheduler_tests.cpp @@ -104,6 +104,28 @@ TEST_CASE("test new stream in threads") { } } +TEST_CASE("test thread local stream") { + auto s = new_thread_local_stream(default_device()); + int result = sum(arange(10, s)).item(); + + std::atomic finished = 0; + std::vector threads; + int num_threads = 4; + for (int i = 0; i < 4; ++i) { + threads.emplace_back([&]() { + int r = sum(arange(10, s)).item(); + CHECK_EQ(result, r); + finished += 1; + clear_streams(); + }); + } + + for (auto& t : threads) { + t.join(); + } + CHECK_EQ(finished, num_threads); +} + TEST_CASE("test get streams") { // Initialize default CPU stream before querying default_stream(Device::cpu);