partial fix for #237 -- switching Device exhausts many resources #242
Conversation
| @@ -31,14 +31,65 @@ class StreamTests: XCTestCase { | |||
| func testUsingDevice() { | |||
| let defaultDevice = Device.defaultDevice() | |||
|
|
|||
| using(device: .cpu) { | |||
| Device.withDefaultDevice(.cpu) { | |||
There was a problem hiding this comment.
Switch to the newer API (older is deprecated and is just a cover on this).
| XCTAssertTrue(StreamOrDevice.default.description.contains("gpu")) | ||
| } | ||
| XCTAssertTrue(StreamOrDevice.default.description.contains("gpu")) | ||
| } | ||
|
|
||
| func testSetUnsetDefaultDevice() { |
There was a problem hiding this comment.
This replicates the failure from #237 -- it works with this change.
| } | ||
|
|
||
| func disabledTestCreateStream() { | ||
| // see https://github.com/ml-explore/mlx/issues/2118 |
There was a problem hiding this comment.
Future tests once ml-explore/mlx#2118 is fixed and in a swift build.
| @@ -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()) | |||
There was a problem hiding this comment.
Use a task local stream or fall back to the current default stream attached to the device.
| 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()) |
There was a problem hiding this comment.
To prevent circularity in init between Device and Stream
| @@ -28,9 +28,19 @@ public enum DeviceType: String, Hashable, Sendable { | |||
| public final class Device: @unchecked Sendable, Equatable { | |||
|
|
|||
| let ctx: mlx_device | |||
| let defaultStream: Stream | |||
There was a problem hiding this comment.
Now the defaultStream is scoped to the Device rather than being a global.
| } | ||
|
|
||
| public init() { | ||
| @available(*, deprecated, message: "please use defaultDevice()") | ||
| public convenience init() { |
There was a problem hiding this comment.
This still works but why create a new one when we already have one laying around
| #endif | ||
|
|
||
| static public func defaultDevice() -> Device { | ||
| @TaskLocal static var _tlDefaultDevice = _resolveGlobalDefaultDevice() |
There was a problem hiding this comment.
A thread local Device -- this lets us manipulate the "global" in a task-safe manner, just like we have done elsewhere.
| /// - ``StreamOrDevice/default`` | ||
| static public func setDefault(device: Device) { | ||
| @available(*, deprecated, message: "please use withDefaultDevice()") | ||
| static public func setDefault(device: Device?) { |
There was a problem hiding this comment.
This still works but callers should use the task-scoped variant. Kept for backward compatibility.
This is a workaround for issues seen in the latest tag, I think ultimately caused by:
This uses statically defined streams when switching device between
.cpuand.gpu. Additionally it adds a task-scoped override of the defaultDevicealong the lines of what we have done recently forStreamand error handling.