partial fix for #237 -- switching Device exhausts many resources#242
partial fix for #237 -- switching Device exhausts many resources#242davidkoski merged 4 commits intomainfrom
Conversation
| 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")) | ||
| } | ||
|
|
||
| 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.
| /// 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
| 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.