Apply the same thread-local approach to the CPU (#3537)

Author: angeloskath

mlx/backend/cpu/encoder.cpp

@@ -2,13 +2,22 @@
 
 #include "mlx/backend/cpu/encoder.h"
 
+#include <fmt/format.h>
+
 namespace mlx::core::cpu {
 
+std::unordered_map<int, CommandEncoder>& get_command_encoders() {
+  static thread_local std::unordered_map<int, CommandEncoder> encoders;
+  return encoders;
+}
+
 CommandEncoder& get_command_encoder(Stream stream) {
-  static std::unordered_map<int, CommandEncoder> encoder_map;
-  auto it = encoder_map.find(stream.index);
-  if (it == encoder_map.end()) {
-    it = encoder_map.emplace(stream.index, stream).first;
+  auto& encoders = get_command_encoders();
+  auto it = encoders.find(stream.index);
+  if (it == encoders.end()) {
+    throw std::runtime_error(
+        fmt::format(
+            "There is no Stream(cpu, {}) in current thread.", stream.index));
   }
   return it->second;
 }

mlx/backend/cpu/encoder.h

@@ -63,5 +63,6 @@ struct MLX_API CommandEncoder {
 };
 
 MLX_API CommandEncoder& get_command_encoder(Stream stream);
+std::unordered_map<int, CommandEncoder>& get_command_encoders();
 
 } // namespace mlx::core::cpu

mlx/backend/cpu/eval.cpp

@@ -7,6 +7,15 @@
 
 namespace mlx::core::cpu {
 
+void new_stream(Stream s) {
+  auto& encoders = get_command_encoders();
+  encoders.try_emplace(s.index, s);
+}
+
+void clear_streams() {
+  get_command_encoders().clear();
+}
+
 void eval(array& arr) {
   auto s = arr.primitive().stream();
 

mlx/backend/cpu/eval.h

@@ -7,6 +7,8 @@
 
 namespace mlx::core::cpu {
 
+void new_stream(Stream s);
 void eval(array& arr);
+void clear_streams();
 
 } // namespace mlx::core::cpu

mlx/scheduler.cpp

@@ -1,6 +1,7 @@
 // Copyright © 2023 Apple Inc.
 
 #include "mlx/scheduler.h"
+#include "mlx/backend/cpu/eval.h"
 #include "mlx/backend/gpu/eval.h"
 
 namespace mlx::core {
@@ -25,6 +26,7 @@ void synchronize() {
 }
 
 void clear_streams() {
+  cpu::clear_streams();
   gpu::clear_streams();
 }
 

mlx/stream.cpp

@@ -2,6 +2,7 @@
 
 #include "mlx/stream.h"
 #include "mlx/backend/cpu/device_info.h"
+#include "mlx/backend/cpu/eval.h"
 #include "mlx/backend/gpu/device_info.h"
 #include "mlx/backend/gpu/eval.h"
 
@@ -70,6 +71,8 @@ Stream new_stream(Device d) {
   auto& s = streams.emplace_back(index, d);
   if (d == Device::gpu) {
     gpu::new_stream(s);
+  } else {
+    cpu::new_stream(s);
   }
   return s;
 }

python/tests/test_threads.py

@@ -9,37 +9,43 @@
 
 class TestThreads(mlx_tests.MLXTestCase):
     def test_threadlocal_stream(self):
+        raised = [False, False]
         test_stream = mx.new_stream(mx.default_device())
 
-        def test_failure():
+        def test_failure(i):
             with self.assertRaises(RuntimeError):
                 with mx.stream(test_stream):
                     x = mx.arange(10)
                     mx.eval(2 * x)
             mx.clear_streams()
+            raised[i] = True
 
-        t1 = threading.Thread(target=test_failure)
-        t2 = threading.Thread(target=test_failure)
+        t1 = threading.Thread(target=test_failure, args=(0,))
+        t2 = threading.Thread(target=test_failure, args=(1,))
         t1.start()
         t2.start()
         t1.join()
         t2.join()
+        self.assertTrue(all(raised))
 
+        raised = [True, True]
         test_stream = mx.new_thread_local_stream(mx.default_device())
 
-        def test_success():
+        def test_success(i):
             with mx.stream(test_stream):
                 x = mx.arange(10)
                 mx.eval(2 * x)
                 self.assertEqual(x.tolist(), list(range(10)))
             mx.clear_streams()
+            raised[i] = False
 
-        t1 = threading.Thread(target=test_success)
-        t2 = threading.Thread(target=test_success)
+        t1 = threading.Thread(target=test_success, args=(0,))
+        t2 = threading.Thread(target=test_success, args=(1,))
         t1.start()
         t2.start()
         t1.join()
         t2.join()
+        self.assertFalse(any(raised))
 
 
 if __name__ == "__main__":