diff --git a/mlx/backend/cuda/eval.cpp b/mlx/backend/cuda/eval.cpp index 661f6222..ef9ee20c 100644 --- a/mlx/backend/cuda/eval.cpp +++ b/mlx/backend/cuda/eval.cpp @@ -80,4 +80,8 @@ void synchronize(Stream s) { cu::get_command_encoder(s).synchronize(); } +void clear_streams() { + cu::get_command_encoders().clear(); +} + } // namespace mlx::core::gpu diff --git a/mlx/backend/gpu/eval.h b/mlx/backend/gpu/eval.h index bdd3f81f..ca86687d 100644 --- a/mlx/backend/gpu/eval.h +++ b/mlx/backend/gpu/eval.h @@ -15,5 +15,6 @@ void new_stream(Stream stream); void eval(array& arr); void finalize(Stream s); void synchronize(Stream s); +void clear_streams(); } // namespace mlx::core::gpu diff --git a/mlx/backend/metal/eval.cpp b/mlx/backend/metal/eval.cpp index 4fa08fc4..6f55976e 100644 --- a/mlx/backend/metal/eval.cpp +++ b/mlx/backend/metal/eval.cpp @@ -87,4 +87,8 @@ void synchronize(Stream s) { metal::get_command_encoder(s).synchronize(); } +void clear_streams() { + metal::get_command_encoders().clear(); +} + } // namespace mlx::core::gpu diff --git a/mlx/backend/no_gpu/eval.cpp b/mlx/backend/no_gpu/eval.cpp index 831cfb39..6d6c32fc 100644 --- a/mlx/backend/no_gpu/eval.cpp +++ b/mlx/backend/no_gpu/eval.cpp @@ -23,4 +23,6 @@ void synchronize(Stream) { throw std::runtime_error("[gpu::synchronize] GPU backend is not available"); } +void clear_streams() {} + } // namespace mlx::core::gpu diff --git a/mlx/scheduler.cpp b/mlx/scheduler.cpp index cabe9919..4937504b 100644 --- a/mlx/scheduler.cpp +++ b/mlx/scheduler.cpp @@ -20,6 +20,10 @@ void synchronize() { synchronize(default_stream(default_device())); } +void clear_streams() { + gpu::clear_streams(); +} + namespace scheduler { Scheduler::Scheduler() { diff --git a/mlx/stream.h b/mlx/stream.h index 54f7d820..b3383bc4 100644 --- a/mlx/stream.h +++ b/mlx/stream.h @@ -39,4 +39,7 @@ MLX_API void synchronize(); /* Synchronize with the provided stream. */ MLX_API void synchronize(Stream); +/* Destroy all streams created in current thread. */ +MLX_API void clear_streams(); + } // namespace mlx::core diff --git a/python/src/stream.cpp b/python/src/stream.cpp index 17e16a41..9b0f61c1 100644 --- a/python/src/stream.cpp +++ b/python/src/stream.cpp @@ -133,6 +133,10 @@ void init_stream(nb::module_& m) { &mx::new_stream, "device"_a, R"pbdoc(Make a new stream on the given device.)pbdoc"); + m.def( + "clear_streams", + &mx::clear_streams, + R"pbdoc(Destroy all streams created in current thread.)pbdoc"); nb::class_(m, "StreamContext", R"pbdoc( A context manager for setting the current device and stream. diff --git a/python/tests/mlx_tests.py b/python/tests/mlx_tests.py index c344e7c8..aa42e6e7 100644 --- a/python/tests/mlx_tests.py +++ b/python/tests/mlx_tests.py @@ -9,6 +9,7 @@ os.environ["MLX_ENABLE_TF32"] = "0" os.environ["MLX_ENABLE_CACHE_THRASHING_CHECK"] = "0" import platform +import sys import unittest from typing import Any, Callable, List, Tuple, Union @@ -18,6 +19,8 @@ import numpy as np class MLXTestRunner(unittest.TestProgram): def __init__(self, *args, **kwargs): + # Do not exit in runTests + kwargs["exit"] = False super().__init__(*args, **kwargs) def createTests(self, *args, **kwargs): @@ -51,6 +54,11 @@ class MLXTestRunner(unittest.TestProgram): filter_and_add(self.test) self.test = filtered_suite + def runTests(self): + super().runTests() + mx.clear_streams() + sys.exit(0 if self.result.wasSuccessful() else 1) + class MLXTestCase(unittest.TestCase): @property diff --git a/python/tests/test_threads.py b/python/tests/test_threads.py index 5e125469..aaf47807 100644 --- a/python/tests/test_threads.py +++ b/python/tests/test_threads.py @@ -7,7 +7,7 @@ import mlx.core as mx import mlx_tests -class TestReduce(mlx_tests.MLXTestCase): +class TestThreads(mlx_tests.MLXTestCase): def test_threadlocal_stream(self): test_stream = mx.new_stream(mx.default_device()) @@ -16,6 +16,7 @@ class TestReduce(mlx_tests.MLXTestCase): with mx.stream(test_stream): x = mx.arange(10) mx.eval(2 * x) + mx.clear_streams() t1 = threading.Thread(target=test_failure) t2 = threading.Thread(target=test_failure) @@ -31,6 +32,7 @@ class TestReduce(mlx_tests.MLXTestCase): x = mx.arange(10) mx.eval(2 * x) self.assertEqual(x.tolist(), list(range(10))) + mx.clear_streams() t1 = threading.Thread(target=test_success) t2 = threading.Thread(target=test_success) diff --git a/tests/scheduler_tests.cpp b/tests/scheduler_tests.cpp index e622ddee..64e060d0 100644 --- a/tests/scheduler_tests.cpp +++ b/tests/scheduler_tests.cpp @@ -48,8 +48,11 @@ TEST_CASE("test default stream in threads") { for (size_t i = 0; i < num_threads; ++i) { threads.emplace_back([&thread_streams, &mtx]() { auto s = default_stream(gpu::is_available() ? Device::gpu : Device::cpu); - std::lock_guard lock(mtx); - thread_streams.insert(s); + { + std::lock_guard lock(mtx); + thread_streams.insert(s); + } + clear_streams(); }); } @@ -80,12 +83,27 @@ TEST_CASE("test access stream in other thread") { } catch (const std::runtime_error&) { error_caught = true; } + clear_streams(); }); t.join(); CHECK(error_caught); } +TEST_CASE("test new stream in threads") { + std::vector threads; + for (int i = 0; i < 1; ++i) { + threads.emplace_back([]() { + auto s = new_stream(default_device()); + eval(arange(10, s)); + clear_streams(); + }); + } + for (auto& t : threads) { + t.join(); + } +} + TEST_CASE("test get streams") { // Initialize default CPU stream before querying default_stream(Device::cpu);