prepare for new mlx / mlx-c release by davidkoski · Pull Request #227 · ml-explore/mlx-swift

davidkoski · 2025-04-24T22:13:15Z

adopt mlx v0.25.1 and mlx-c v0.2.0
fixes metalKernel() should recreate the kernel object each time #211 (moves metalKernel API back to what it was)
use an evalLock for eval, asyncEval and Stream creation
add withError and withErrorHandler error handling
add function import/export: https://ml-explore.github.io/mlx/build/html/usage/export.html

davidkoski · 2025-04-24T22:14:25Z

Source/MLX/MLXArray+Bytes.swift

+            let data = asDataCopy()
+            return data.data.withUnsafeBytes { ptr in
+                device.makeBuffer(bytes: ptr.baseAddress!, length: ptr.count)
+            }


Unrelated but per question the non-contiguous case wasn't actually making contiguous storage.

davidkoski · 2025-04-24T22:15:01Z

Source/MLX/MLXArray.swift

    /// specialized conversion between integer types -- see ``item(_:)``
    private func itemInt() -> Int {
+        precondition(self.size == 1)
+        eval()


Now includes an explicit eval in order to use the evalLock

davidkoski · 2025-04-24T22:15:54Z

Source/MLX/MLXArray.swift

-        mlx_array_eval(ctx)
+        _ = evalLock.withLock {
+            mlx_array_eval(ctx)
+        }


All eval, asyncEval, Stream need to take a lock -- new version of mlx-c puts more work in the calling thread and it needs to be guarded. To be revisited in the future.

davidkoski · 2025-04-24T22:17:10Z

Source/MLX/MLXFastKernel.swift

-            outputShapes: [[Int]],
-            outputDTypes: [DType],
-            initValue: Float? = nil,
-            verbose: Bool = false


This reverts the API change from #150 -- the new mlx-c API allows us to match the old calls (and python). See #211

davidkoski · 2025-04-24T22:17:53Z

Source/MLX/Ops.swift

 ///     - stride: kernel stride
 ///     - padding: input padding
 ///     - dilation: kernel dilation
+///     - outputPadding: output padding


The convTranspose functions pick up an outputPadding (optional)

davidkoski · 2025-04-24T22:18:27Z

Source/MLX/Ops.swift

 public func gatherMatmul(
    _ a: MLXArray, _ b: MLXArray, lhsIndices: MLXArray? = nil, rhsIndices: MLXArray? = nil,
-    stream: StreamOrDevice = .default
+    sortedIndices: Bool = false, stream: StreamOrDevice = .default


gather_mm takes a hint for sorted indices

davidkoski · 2025-04-24T22:19:13Z

Source/MLX/Stream.swift


    deinit {
-        mlx_stream_free(ctx)
+        _ = evalLock.withLock {


Streams need to use the same lock as they manipulate the same shared structures.

davidkoski · 2025-04-28T20:37:04Z

Source/MLX/ErrorHandler.swift

        fatalError(message.map { String(cString: $0) } ?? "")
    }
+
+/// Evaluate the block with a scoped MLX error handler.


New error handling capabilities -- the import/export code requires this to correctly catch errors from mlx::core and this functionality is generally useful for swift code that wants to catch errors from core rather than fatally assert.

davidkoski · 2025-04-28T20:40:47Z

Source/MLX/Export.swift

+///
+/// ### See Also
+/// - ``exportFunctions(to:shapeless:_:export:)``
+/// - ``importFunction(from:)``


New export/import capability -- the example code is from the tests.

davidkoski · 2025-04-28T20:42:12Z

Source/MLX/Ops.swift

-        mlx_load(&result, path.cString(using: .utf8), stream.ctx)
+        _ = try withError {
+            mlx_load(&result, path.cString(using: .utf8), stream.ctx)
+        }


These can properly catch an error for a missing file.

Note the _ = is because mlx_load returns an int

davidkoski · 2025-04-28T20:44:13Z

Source/MLX/Transforms+Eval.swift

+///
+/// ### See Also
+/// - <doc:lazy-evaluation>
+public func checkedEval(_ values: Any...) throws {


For use when you want eval to throw. Note: this won't catch broadcasting errors as those occur at graph build time, e.g. a + b.

awni · 2025-04-30T20:09:04Z

Source/MLX/DType.swift

        public var eps: Double {
            switch dtype {
-            case .float16: Double(Float16.ulpOfOne)
+            #if !arch(x86_64)


Why just the fp16 guard here?

I ran into this when building apps that depended on it -- Float16 isn't a thing on x86. We don't see it in framework builds, but by default App builds do all architectures when building Release so we encounter this (even though that slice won't run)

Right but presumably bfloat16 isn't a thing on x86 Macs as well. I guess the difference is we hardcode the bf16 values here.

but by default App builds do all architectures when building Release

I wonder how the rest of MLX builds in that case? I think MLX fp16 runs on x86 using the software implementation -- which means we should probably hardcode the values here too?

BFloat isn't a type in swift so this one doesn't come up.

Sure, we could hard code the values here for x86 -- is that a supported configuration? Easy enough to do.

is that a supported configuration?

I mean it's not common.. so this is kind of minor.

Linux x86 is supported and we run those tests regularly. macOS x86 is very uncommon.

Hence I also think it's fine to leave it like this and if anyone asks for that we can add it back.

awni · 2025-04-30T20:13:20Z

Source/MLX/Export.swift

+///   - url: file url to write the `.mlxfn` file
+///   - shapeless: if `true` the function allows inputs with variable shapes
+///   - f: the function to capture
+/// - Returns: a help that records the call


The return comment there is unclear. What is a "help"?

should be "helper" -- I will clarify

Source/MLX/Stream.swift

awni · 2025-04-30T21:24:19Z

Source/MLX/Export.swift

+///   - url: file url to write the `.mlxfn` file
+///   - shapeless: if `true` the function allows inputs with variable shapes
+///   - f: the function to capture
+///   - export: closure for recording the calls


Looks like that parameter doesn't exist?

Or maybe it's supposed to be a return?

I guess I'm wrong about that.. but I'm trying to decipher how the export parameter works.

try exportFunctions(to: url, shapeless: true, f) { export in try export(x) try export(x, x) try export(x, x, x) }

export is the trailing closure -- the { export in ... }

and the export parameter in the closure is a FunctionExporterMultiple

That's makes sense now.. but maybe is slightly confusing. Because the export parameter is both the closure and the name of the parameter to the closure. But if it's a common idiom then let's go for it.

You don't normally see the name of the closure parameter (but you can). Let me think if there is a more expressive name. If not, I think nobody will notice as the trailing closure syntax is idiomatic.

awni · 2025-04-30T21:29:18Z

Source/MLX/Export.swift

+/// // f2 is a callable that represents the loaded function
+/// let f2 = try importFunction(from: url)
+///
+/// let a = MLXArray(10)
+/// let b = MLXArray([5, 10, 20])
+///
+/// // call it -- the shapes and labels have to match
+/// let r = try f2(a, y: b)[0]


Nit: name the function f rather than f2 as there is no f1

awni

Looks awesome. Left a few nits and some questions, please take a look. Otherwise I think this is good to go!

davidkoski · 2025-04-30T21:42:36Z

Looks awesome. Left a few nits and some questions, please take a look. Otherwise I think this is good to go!

Thank you @awni -- appreciated as always!

- fixes #211 (moves metalKernel API back to what it was) - add improved swift error handling - add function import/export - adopt mlx v0.25.1 and mlx-c v0.2.0 - fixes #211 (moves metalKernel API back to what it was) - use an evalLock for eval, asyncEval and Stream creation - add withError and withErrorHandler error handling - add function import/export: https://ml-explore.github.io/mlx/build/html/usage/export.html

davidkoski assigned awni Apr 24, 2025

davidkoski commented Apr 24, 2025

View reviewed changes

davidkoski commented Apr 28, 2025

View reviewed changes

davidkoski marked this pull request as ready for review April 28, 2025 20:44

awni reviewed Apr 30, 2025

View reviewed changes

Source/MLX/Stream.swift Outdated Show resolved Hide resolved

awni reviewed Apr 30, 2025

View reviewed changes

awni approved these changes Apr 30, 2025

View reviewed changes

davidkoski force-pushed the import_functions branch from 3cd0150 to b057e39 Compare May 1, 2025 15:47

davidkoski force-pushed the import_functions branch from b057e39 to 7874c5a Compare May 1, 2025 15:49

davidkoski merged commit bc923f3 into mlx0_24_2 May 1, 2025

davidkoski deleted the import_functions branch May 1, 2025 15:50

Conversation

davidkoski commented Apr 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

davidkoski Apr 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

awni Apr 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

awni left a comment

Choose a reason for hiding this comment

Uh oh!

davidkoski commented Apr 30, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

davidkoski commented Apr 24, 2025 •

edited

Loading

davidkoski Apr 28, 2025 •

edited

Loading

awni Apr 30, 2025 •

edited

Loading