partial fix for #237 -- switching Device exhausts many resources (#242)

* partial fix for #237 -- switching Device exhausts many resources

* remove debug code

* remove debug code

* comments

Author: davidkoski

Source/MLX/Device.swift

@@ -28,9 +28,19 @@ public enum DeviceType: String, Hashable, Sendable {
 public final class Device: @unchecked Sendable, Equatable {
 
     let ctx: mlx_device
+    let defaultStream: Stream
 
     init(_ ctx: mlx_device) {
         self.ctx = ctx
+
+        var deviceType = MLX_GPU
+        mlx_device_get_type(&deviceType, ctx)
+        self.defaultStream =
+            switch deviceType {
+            case MLX_CPU: .cpu
+            case MLX_GPU: .gpu
+            default: .gpu
+            }
     }
 
     public init(_ deviceType: DeviceType, index: Int32 = 0) {
@@ -42,19 +52,32 @@ public final class Device: @unchecked Sendable, Equatable {
             cDeviceType = MLX_GPU
         }
         self.ctx = mlx_device_new_type(cDeviceType, index)
+        self.defaultStream =
+            switch deviceType {
+            case .cpu: .cpu
+            case .gpu: .gpu
+            }
     }
 
-    public init() {
+    @available(*, deprecated, message: "please use defaultDevice()")
+    public convenience init() {
         var ctx = mlx_device_new()
         mlx_get_default_device(&ctx)
-        self.ctx = ctx
+        self.init(ctx)
     }
 
     deinit {
         mlx_device_free(ctx)
     }
 
+    /// static CPU device
+    ///
+    /// See ``withDefaultDevice(_:_:)``
     static public let cpu: Device = Device(.cpu)
+
+    /// static GPU device
+    ///
+    /// See ``withDefaultDevice(_:_:)``
     static public let gpu: Device = Device(.gpu)
 
     public var deviceType: DeviceType? {
@@ -67,28 +90,50 @@ public final class Device: @unchecked Sendable, Equatable {
         }
     }
 
+    // support for global default device
     static let _lock = NSLock()
     #if swift(>=5.10)
-        nonisolated(unsafe) static var _defaultDevice = gpu
-        nonisolated(unsafe) static var _defaultStream = Stream(gpu)
+        nonisolated(unsafe) static var _defaultDevice: Device?
     #else
-        static var _defaultDevice = gpu
-        static var _defaultStream = Stream(gpu)
+        static var _defaultDevice: Device?
     #endif
 
-    static public func defaultDevice() -> Device {
+    @TaskLocal static var _tlDefaultDevice = _resolveGlobalDefaultDevice()
+
+    private static func _resolveGlobalDefaultDevice() -> Device {
         _lock.withLock {
-            _defaultDevice
+            _defaultDevice ?? .gpu
         }
     }
 
+    /// Return the current default device.
+    ///
+    /// This is used by ``StreamOrDevice/default`` -- the default stream parameter
+    /// to most functions.
+    static public func defaultDevice() -> Device {
+        _tlDefaultDevice
+    }
+
+    /// Use a device scoped to a task.
+    static public func withDefaultDevice<R>(
+        _ device: Device, _ body: () throws -> R
+    ) rethrows -> R {
+        try $_tlDefaultDevice.withValue(device, operation: body)
+    }
+
+    /// Use a device scoped to a task.
+    static public func withDefaultDevice<R>(
+        _ device: Device, _ body: () async throws -> R
+    ) async rethrows -> R {
+        try await $_tlDefaultDevice.withValue(device, operation: body)
+    }
+
+    /// Return the current default stream.
     static func defaultStream() -> Stream {
-        _lock.withLock {
-            _defaultStream
-        }
+        _tlDefaultDevice.defaultStream
     }
 
-    /// Set the default device.
+    /// Set the default device globally.  Prefer the scoped version, ``withDefaultDevice(_:_:)``.
     ///
     /// For example:
     ///
@@ -99,12 +144,19 @@ public final class Device: @unchecked Sendable, Equatable {
     /// By default this is ``gpu``.
     ///
     /// ### See Also
+    /// - ``withDefaultDevice(_:_:)``
     /// - ``StreamOrDevice/default``
-    static public func setDefault(device: Device) {
+    @available(*, deprecated, message: "please use withDefaultDevice()")
+    static public func setDefault(device: Device?) {
         _lock.withLock {
-            mlx_set_default_device(device.ctx)
+            if let device {
+                // sets the mlx core default device -- only used
+                // by the deprecated init().  this isn't thread
+                // safe or really usable across tasks/threads
+                // but is kept for backward compatibility
+                mlx_set_default_device(device.ctx)
+            }
             _defaultDevice = device
-            _defaultStream = Stream(device)
         }
     }
 
@@ -139,6 +191,7 @@ extension Device: CustomStringConvertible {
 /// - Parameters:
 ///     - device: device to be used
 ///     - fn: function to be executed
+@available(*, deprecated, message: "please use Device.withDefaultDevice()")
 public func using<R>(device: Device, fn: () throws -> R) rethrows -> R {
-    try Stream.withNewDefaultStream(device: device, fn)
+    try Device.withDefaultDevice(device, fn)
 }

Source/MLX/Stream.swift

@@ -35,7 +35,7 @@ public struct StreamOrDevice: Sendable, CustomStringConvertible, Equatable {
     /// This will be ``Device/gpu`` unless ``Device/setDefault(device:)``
     /// sets it otherwise.
     public static var `default`: StreamOrDevice {
-        StreamOrDevice(Stream.defaultStream)
+        StreamOrDevice(Stream.defaultStream ?? Device.defaultStream())
     }
 
     public static func device(_ device: Device) -> StreamOrDevice {
@@ -83,10 +83,10 @@ public final class Stream: @unchecked Sendable, Equatable {
 
     let ctx: mlx_stream
 
-    public static let gpu = Stream(.gpu)
-    public static let cpu = Stream(.cpu)
+    public static let gpu = Stream(mlx_default_gpu_stream_new())
+    public static let cpu = Stream(mlx_default_cpu_stream_new())
 
-    @TaskLocal static var defaultStream = Stream()
+    @TaskLocal static var defaultStream: Stream?
 
     /// Set the ``StreamOrDevice/default`` scoped to a Task.
     public static func withNewDefaultStream<R>(device: Device? = nil, _ body: () throws -> R)

Tests/MLXTests/StreamTests.swift

@@ -31,14 +31,65 @@ class StreamTests: XCTestCase {
     func testUsingDevice() {
         let defaultDevice = Device.defaultDevice()
 
-        using(device: .cpu) {
+        Device.withDefaultDevice(.cpu) {
+            // these _should_ be the same
+            XCTAssertTrue(Device.defaultDevice().description.contains("cpu"))
             XCTAssertTrue(StreamOrDevice.default.description.contains("cpu"))
         }
         XCTAssertEqual(defaultDevice, Device.defaultDevice())
 
-        using(device: .gpu) {
+        Device.withDefaultDevice(.gpu) {
+            XCTAssertTrue(Device.defaultDevice().description.contains("gpu"))
             XCTAssertTrue(StreamOrDevice.default.description.contains("gpu"))
         }
         XCTAssertTrue(StreamOrDevice.default.description.contains("gpu"))
     }
+
+    func testSetUnsetDefaultDevice() {
+        // Issue #237 -- setting an unsetting the default device in a loop
+        // exhausts many resources
+        for _ in 1 ..< 10000 {
+            let defaultDevice = MLX.Device.defaultDevice()
+            MLX.Device.setDefault(device: .cpu)
+            defer {
+                MLX.Device.setDefault(device: defaultDevice)
+            }
+
+            let x = MLXArray(1)
+            let _ = x * x
+        }
+        print("here")
+    }
+
+    func testWithDefaultDevice() {
+        // Issue #237 -- scoped variant
+        for _ in 1 ..< 10000 {
+            Device.withDefaultDevice(.cpu) {
+                Device.withDefaultDevice(.gpu) {
+                    let x = MLXArray(1)
+                    let _ = x * x
+                }
+            }
+        }
+        print("here")
+    }
+
+    func disabledTestCreateStream() {
+        // see https://github.com/ml-explore/mlx/issues/2118
+        for _ in 1 ..< 10000 {
+            let _ = Stream(.cpu)
+        }
+        print("here")
+    }
+
+    func disabledTestCreateStreamScoped() {
+        // see https://github.com/ml-explore/mlx/issues/2118
+        for _ in 1 ..< 10000 {
+            Stream.withNewDefaultStream(device: .cpu) {
+                let x = MLXArray(1)
+                let _ = x * x
+            }
+        }
+    }
+
 }