ThreadLocalStream in C++ (#3405)

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

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");

mlx/stream.cpp

@@ -6,6 +6,7 @@
 #include "mlx/scheduler.h"
 
 #include <array>
+#include <map>
 #include <optional>
 #include <shared_mutex>
 
@@ -29,6 +30,11 @@ auto& all_streams() {
   return streams_and_mtx;
 }
 
+auto& thread_local_streams() {
+  static std::tuple<std::vector<ThreadLocalStream>, std::mutex> streams_and_mtx;
+  return streams_and_mtx;
+}
+
 } // namespace
 
 Stream default_stream(Device d) {
@@ -58,10 +64,6 @@ std::vector<Stream> 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<ThreadLocalStream, Stream> 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

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<Stream> get_streams();
 

mlx/utils.cpp

@@ -17,6 +17,8 @@ Stream to_stream(StreamOrDevice s) {
     return default_stream(default_device());
   } else if (std::holds_alternative<Device>(s)) {
     return default_stream(std::get<Device>(s));
+  } else if (std::holds_alternative<ThreadLocalStream>(s)) {
+    return stream_from_thread_local_stream(std::get<ThreadLocalStream>(s));
   } else {
     return std::get<Stream>(s);
   }

mlx/utils.h

@@ -13,7 +13,8 @@
 
 namespace mlx::core {
 
-using StreamOrDevice = std::variant<std::monostate, Stream, Device>;
+using StreamOrDevice =
+    std::variant<std::monostate, Stream, ThreadLocalStream, Device>;
 MLX_API Stream to_stream(StreamOrDevice s);
 MLX_API Stream to_stream(StreamOrDevice s, Device default_);
 

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<const PyThreadLocalStream*, mx::Stream>
-        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_<mx::Stream>(
       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<mx::Stream>(other);
       });
 
-  nb::class_<PyThreadLocalStream>(
+  nb::class_<mx::ThreadLocalStream>(
       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<mx::Device>())
+      .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<mx::Stream>(other, s_other) &&
-            s_other == s.stream();
-      });
+      .def(
+          "__eq__",
+          [](const mx::ThreadLocalStream& s, const nb::object& other) {
+            return nb::isinstance<mx::ThreadLocalStream>(other) &&
+                s == nb::cast<mx::ThreadLocalStream>(other);
+          });
 
   nb::implicitly_convertible<mx::Device::DeviceType, mx::Device>();
-  nb::implicitly_convertible<PyThreadLocalStream, mx::Stream>();
 
   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,

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):

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<int>();
+
+  std::atomic<int> finished = 0;
+  std::vector<std::thread> threads;
+  int num_threads = 4;
+  for (int i = 0; i < 4; ++i) {
+    threads.emplace_back([&]() {
+      int r = sum(arange(10, s)).item<int>();
+      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);