ThreadLocalStream in C++ (#3405)

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,