Add a convenience for making local streams in python (#3355)

python/src/random.cpp

@@ -18,10 +18,15 @@ using namespace nb::literals;
 
 class PyKeySequence {
  public:
-  PyKeySequence() {
-    // Destroy state before the python interpreter exits.
-    auto atexit = nb::module_::import_("atexit");
-    atexit.attr("register")(nb::cpp_function([this]() { state_.reset(); }));
+  ~PyKeySequence() {
+    if (state_.has_value()) {
+      nb::gil_scoped_acquire gil;
+      state_.reset();
+    }
+  }
+
+  void reset() {
+    state_.reset();
   }
 
   void seed(uint64_t seed) {
@@ -521,4 +526,10 @@ void init_random(nb::module_& parent_module) {
             array:
               The generated random permutation or randomly permuted input array.
       )pbdoc");
+
+  // Ensure the main thread cleanup will happen before the interpreter goes
+  // away. As a result if the other threads join the main thread we should have
+  // a clean tear-down.
+  auto atexit = nb::module_::import_("atexit");
+  atexit.attr("register")(nb::cpp_function([]() { default_key().reset(); }));
 }

python/src/stream.cpp

@@ -41,6 +41,26 @@ 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,
@@ -49,6 +69,11 @@ 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) {
@@ -61,7 +86,29 @@ void init_stream(nb::module_& m) {
             s == nb::cast<mx::Stream>(other);
       });
 
+  nb::class_<PyThreadLocalStream>(
+      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(
+          "__repr__",
+          [](const PyThreadLocalStream& s) {
+            std::ostringstream os;
+            os << "ThreadLocalStream(" << s.device << ")";
+            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();
+      });
+
   nb::implicitly_convertible<mx::Device::DeviceType, mx::Device>();
+  nb::implicitly_convertible<PyThreadLocalStream, mx::Stream>();
 
   m.def(
       "default_stream",

python/src/transforms.cpp

@@ -1463,12 +1463,22 @@ void init_transforms(nb::module_& m) {
          bool shapeless) {
         // Make sure each thread using mx.compile would clear its compile cache
         // before python interpreter exits.
-        static thread_local auto clear_cache = []() {
-          auto atexit = nb::module_::import_("atexit");
-          atexit.attr("register")(
-              nb::cpp_function(&mx::detail::compile_clear_cache));
-          return true;
+        struct ThreadCleanup {
+          ~ThreadCleanup() {
+            if (!mx::detail::compile_cache_empty()) {
+              nb::gil_scoped_acquire gil;
+              mx::detail::compile_clear_cache();
+            }
+          }
         };
+        static thread_local auto clear_cache = []() {
+          // Ensure it is created
+          mx::detail::compile_clear_cache();
+
+          // Ensure it will be cleaned up
+          return ThreadCleanup{};
+        }();
+
         return mlx_func(
             nb::cpp_function(PyCompiledFun{fun, inputs, outputs, shapeless}),
             fun,
@@ -1542,4 +1552,10 @@ void init_transforms(nb::module_& m) {
           A callable that recomputes intermediate states during gradient
           computation.
       )pbdoc");
+
+  // Ensure the main thread cleanup will happen before the interpreter goes
+  // away. As a result if the other threads join the main thread we should have
+  // a clean tear-down.
+  auto atexit = nb::module_::import_("atexit");
+  atexit.attr("register")(nb::cpp_function(&mx::detail::compile_clear_cache));
 }