diff --git a/Package.swift b/Package.swift index 941a108a5..fbd165bfa 100644 --- a/Package.swift +++ b/Package.swift @@ -361,6 +361,8 @@ let package = Package( exclude: mlxSwiftExcludes, swiftSettings: [ .enableExperimentalFeature("StrictConcurrency") + .swiftLanguageMode(.v6), + .enableExperimentalFeature("Lifetimes"), ] ), .target( @@ -434,6 +436,12 @@ let package = Package( path: "Source/Examples", sources: ["CustomFunctionExample.swift"] ), + .executableTarget( + name: "Benchmarks", + dependencies: ["MLX", "MLXNN"], + path: "Source/Benchmarks", + sources: ["Benchmarks.swift"] + ), .executableTarget( name: "CustomFunctionExampleSimple", dependencies: ["MLX"], diff --git a/Source/Benchmarks/Benchmarks.swift b/Source/Benchmarks/Benchmarks.swift new file mode 100644 index 000000000..959b90c7b --- /dev/null +++ b/Source/Benchmarks/Benchmarks.swift @@ -0,0 +1,130 @@ +// Copyright © 2026 Apple Inc. + +import Dispatch +import Foundation +import MLX +import MLXNN + +private struct BenchmarkResult { + let name: String + let iterations: Int + let minMs: Double + let medianMs: Double + let meanMs: Double +} + +private func measure( + name: String, warmup: Int = 5, iterations: Int = 100, _ body: () -> Void +) -> BenchmarkResult { + for _ in 0 ..< warmup { body() } + + var samples = [Double]() + samples.reserveCapacity(iterations) + for _ in 0 ..< iterations { + let start = DispatchTime.now().uptimeNanoseconds + body() + let end = DispatchTime.now().uptimeNanoseconds + samples.append(Double(end - start) / 1_000_000) + } + samples.sort() + + return BenchmarkResult( + name: name, + iterations: iterations, + minMs: samples.first!, + medianMs: samples[samples.count / 2], + meanMs: samples.reduce(0, +) / Double(iterations)) +} + +private func printResults(_ results: [BenchmarkResult]) { + print("Benchmark Iters Min ms Median ms Mean ms") + for result in results { + let name = result.name.padding(toLength: 44, withPad: " ", startingAt: 0) + let values = String( + format: "%6d %10.4f %10.4f %10.4f", result.iterations, result.minMs, + result.medianMs, result.meanMs) + print("\(name) \(values)") + } +} + +private func graphConstruction() -> BenchmarkResult { + measure(name: "Graph construction (200 elementwise ops)") { + var x = MLXArray.zeros([4, 4]) + for _ in 0 ..< 200 { x = x + 1 } + precondition(x.shape == [4, 4]) + } +} + +private func realizedEval() -> BenchmarkResult { + let x = MLXArray.ones([4, 4]) + eval(x) + return measure(name: "eval() of an already-realized array") { eval(x) } +} + +private func mlpForward() -> [BenchmarkResult] { + let l1 = Linear(128, 256) + let l2 = Linear(256, 256) + let l3 = Linear(256, 10) + let input = MLXRandom.normal([32, 128]) + + func forward(_ x: MLXArray) -> MLXArray { + l3(relu(l2(relu(l1(x))))) + } + + let compiled = compile(inputs: [l1, l2, l3], forward) + let expected = forward(input) + let actual = compiled(input) + eval(expected, actual) + precondition(allClose(expected, actual).item(Bool.self)) + + return [ + measure(name: "MLP forward + eval (uncompiled)") { eval(forward(input)) }, + measure(name: "MLP forward + eval (compiled)") { eval(compiled(input)) }, + ] +} + +private func evalContention(device: Device) -> [BenchmarkResult] { + let workerCount = 4 + let workPerWorker = 8 + let streams = (0 ..< workerCount).map { _ in Stream(device) } + + let runWorker: @Sendable (Int) -> Void = { worker in + let stream = StreamOrDevice.stream(streams[worker]) + var x = full([64, 64], values: Float(worker), stream: stream) + for _ in 0 ..< workPerWorker { + x = add(x, 1, stream: stream) + x = multiply(x, 1.0001, stream: stream) + } + eval(x) + } + + let sequential = measure(name: "4 graph + eval workers (sequential)", iterations: 25) { + for worker in 0 ..< workerCount { runWorker(worker) } + } + let concurrent = measure(name: "4 graph + eval workers (concurrent)", iterations: 25) { + DispatchQueue.concurrentPerform(iterations: workerCount, execute: runWorker) + } + return [sequential, concurrent] +} + +@main +private struct Benchmarks { + static func main() { + let requested = CommandLine.arguments.dropFirst().first + let device: Device + switch requested { + case nil: device = Device.defaultDevice() + case "gpu": device = .gpu + case "cpu": device = .cpu + default: + fatalError("usage: Benchmarks [cpu|gpu]") + } + + Device.withDefaultDevice(device) { + var results = [graphConstruction(), realizedEval()] + results.append(contentsOf: mlpForward()) + results.append(contentsOf: evalContention(device: device)) + printResults(results) + } + } +} diff --git a/Source/MLX/Cmlx+Util.swift b/Source/MLX/Cmlx+Util.swift index 61854afa0..55a2b9adc 100644 --- a/Source/MLX/Cmlx+Util.swift +++ b/Source/MLX/Cmlx+Util.swift @@ -5,8 +5,34 @@ import Foundation // return a +1 mlx_vector_array containing the given arrays func new_mlx_vector_array(_ arrays: some Collection) -> mlx_vector_array { - withExtendedLifetime(arrays) { - mlx_vector_array_new_data(arrays.map { $0.ctx }, arrays.count) + guard !arrays.isEmpty else { + return mlx_vector_array_new_data(nil, 0) + } + + return withExtendedLifetime(arrays) { + if arrays.count <= mlxInteropStackBufferCapacity { + return withUnsafeTemporaryAllocation(of: mlx_array.self, capacity: arrays.count) { + buffer in + var index = buffer.startIndex + for array in arrays { + buffer.initializeElement(at: index, to: array.ctx) + buffer.formIndex(after: &index) + } + return mlx_vector_array_new_data(buffer.baseAddress, buffer.count) + } + } + + let buffer = UnsafeMutableBufferPointer.allocate(capacity: arrays.count) + var initializedCount = 0 + defer { + buffer.baseAddress?.deinitialize(count: initializedCount) + buffer.deallocate() + } + for array in arrays { + buffer.initializeElement(at: initializedCount, to: array.ctx) + initializedCount += 1 + } + return mlx_vector_array_new_data(buffer.baseAddress, buffer.count) } } diff --git a/Source/MLX/Factory.swift b/Source/MLX/Factory.swift index 913540511..d4f56cf7e 100644 --- a/Source/MLX/Factory.swift +++ b/Source/MLX/Factory.swift @@ -875,7 +875,9 @@ public func full( ) -> MLXArray { var result = mlx_array_new() let values = values.asMLXArray(dtype: nil) - mlx_full(&result, shape.asInt32, shape.count, values.ctx, type.dtype.cmlxDtype, stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_full(&result, shape, count, values.ctx, type.dtype.cmlxDtype, stream.ctx) + } return MLXArray(result) } @@ -905,7 +907,9 @@ public func full( _ shape: some Collection, values: MLXArray, dtype: DType, stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_full(&result, shape.asInt32, shape.count, values.ctx, dtype.cmlxDtype, stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_full(&result, shape, count, values.ctx, dtype.cmlxDtype, stream.ctx) + } return MLXArray(result) } @@ -935,7 +939,9 @@ public func full( ) -> MLXArray { var result = mlx_array_new() let values = values.asMLXArray(dtype: nil) - mlx_full(&result, shape.asInt32, shape.count, values.ctx, values.dtype.cmlxDtype, stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_full(&result, shape, count, values.ctx, values.dtype.cmlxDtype, stream.ctx) + } return MLXArray(result) } diff --git a/Source/MLX/Foundation+Util.swift b/Source/MLX/Foundation+Util.swift index 5748dc9c9..5ab888770 100644 --- a/Source/MLX/Foundation+Util.swift +++ b/Source/MLX/Foundation+Util.swift @@ -2,6 +2,48 @@ import Foundation +/// Keep the common shape/axis case on the stack without allowing an untrusted collection size to +/// produce an arbitrarily large stack allocation. +@usableFromInline +let mlxInteropStackBufferCapacity = 64 + +extension Collection where Element == Int { + + /// Calls `body` with a scoped `Int32` representation suitable for C APIs. + /// + /// Unlike ``asInt32``, this does not allocate a temporary `Array`. The pointer must not be + /// stored or otherwise escape `body`. + @inlinable + func withInt32Buffer( + _ body: (UnsafePointer?, Int) throws -> Result + ) rethrows -> Result { + guard !isEmpty else { return try body(nil, 0) } + + if count <= mlxInteropStackBufferCapacity { + return try withUnsafeTemporaryAllocation(of: Int32.self, capacity: count) { buffer in + var destination = buffer.startIndex + for value in self { + buffer.initializeElement(at: destination, to: Int32(value)) + buffer.formIndex(after: &destination) + } + return try body(buffer.baseAddress, buffer.count) + } + } + + let buffer = UnsafeMutableBufferPointer.allocate(capacity: count) + var initializedCount = 0 + defer { + buffer.baseAddress?.deinitialize(count: initializedCount) + buffer.deallocate() + } + for value in self { + buffer.initializeElement(at: initializedCount, to: Int32(value)) + initializedCount += 1 + } + return try body(buffer.baseAddress, buffer.count) + } +} + extension [Int] { /// Convenience to coerce array of `Int` to `Int32` -- Cmlx uses `Int32` for many things but it is diff --git a/Source/MLX/MLXArray+Bytes.swift b/Source/MLX/MLXArray+Bytes.swift index b4d1d62dc..9970a3663 100644 --- a/Source/MLX/MLXArray+Bytes.swift +++ b/Source/MLX/MLXArray+Bytes.swift @@ -38,7 +38,16 @@ extension MLXArray { ?? self.size } - func copy(from: UnsafeRawBufferPointer, toContiguous output: UnsafeMutableRawBufferPointer) { + /// Copy `self`'s (possibly non-contiguous) backing, starting at `base`, into a + /// contiguous `output` buffer. + /// + /// `base` is treated as a bare starting point, not a pre-sized region: for the + /// non-contiguous case the actual reachable byte range (which can extend *before* + /// `base` when strides are negative) is computed internally and bounds-checked via + /// `RawSpan`, rather than trusting unchecked pointer arithmetic to stay in bounds. + func copy(from base: UnsafeRawPointer, toContiguous output: UnsafeMutableRawBufferPointer) { + guard !output.isEmpty else { return } + let contiguousDimension = self.contiguousToDimension() let shape = self.shape let strides = self.internalStrides @@ -69,22 +78,18 @@ extension MLXArray { // the index of the current source item var index = Array.init(repeating: 0, count: ndim) - // output pointer - var dest = output.baseAddress! - - while true { - // compute the source index by multiplying the index by the - // stride for each dimension - - // note: in the case where the array has negative strides / offset - // the base pointer we have will have the offset already applied, - // e.g. asStrided(a, [3, 3], strides: [-3, -1], offset: 8) + // Keep this in step with the odometer below. Recomputing it from every index and + // stride made traversal O(numberOfChunks * rank). + var sourceIndex = 0 - let sourceIndex = zip(index, strides).reduce(0) { $0 + ($1.0 * $1.1) } + // output byte offset + var destOffset = 0 - // convert to byte pointer - let src = from.baseAddress! + sourceIndex * itemSize - dest.copyMemory(from: src, byteCount: destItemSize) + while true { + // offset relative to the span's own start -- always >= 0 by construction, + // since minSourceIndex is the true minimum reachable sourceIndex + let spanOffset = (sourceIndex - minSourceIndex) * itemSize + let chunk = source.extracting(spanOffset ..< (spanOffset + destItemSize)) // next output address dest += destItemSize @@ -99,9 +104,11 @@ extension MLXArray { } index[dimension] = 0 + sourceIndex -= (shape[dimension] - 1) * strides[dimension] } else { // just increment the dimension and we are done index[dimension] += 1 + sourceIndex += strides[dimension] break } } diff --git a/Source/MLX/MLXArray+Ops.swift b/Source/MLX/MLXArray+Ops.swift index f940d2970..279ca23ef 100644 --- a/Source/MLX/MLXArray+Ops.swift +++ b/Source/MLX/MLXArray+Ops.swift @@ -2562,7 +2562,9 @@ extension MLXArray { -> MLXArray { var result = mlx_array_new() - mlx_reshape(&result, ctx, newShape.asInt32, newShape.count, stream.ctx) + newShape.withInt32Buffer { shape, count in + mlx_reshape(&result, ctx, shape, count, stream.ctx) + } return MLXArray(result) } @@ -2584,7 +2586,9 @@ extension MLXArray { /// - ``reshaped(_:stream:)`` public func reshaped(_ newShape: Int..., stream: StreamOrDevice = .default) -> MLXArray { var result = mlx_array_new() - mlx_reshape(&result, ctx, newShape.asInt32, newShape.count, stream.ctx) + newShape.withInt32Buffer { shape, count in + mlx_reshape(&result, ctx, shape, count, stream.ctx) + } return MLXArray(result) } diff --git a/Source/MLX/MLXArray+View.swift b/Source/MLX/MLXArray+View.swift new file mode 100644 index 000000000..64c21861e --- /dev/null +++ b/Source/MLX/MLXArray+View.swift @@ -0,0 +1,94 @@ +// Copyright © 2024 Apple Inc. + +import Cmlx + +extension MLXArray { + + /// A move-only, flat view of this array's contents as `Scalar`, converting the ``DType`` + /// if needed. + /// + /// This is shorthand for the common ``asArray(_:)`` / ``asType(_:stream:)`` + ``asArray(_:)`` + /// patterns, but lets you pull out single values or ranges (and iterate a zero-copy `Span`) + /// without eagerly copying the whole array into a `[Scalar]`. + /// + /// ```swift + /// let flat = prediction.view(Float.self) + /// let x = flat[3] // one value + /// let head = flat[0 ..< 4] // a range, copied out + /// flat.withSpan { span in ... } // zero-copy bulk read + /// ``` + /// + /// ### See Also + /// - ``MLXArrayOf`` + /// - ``asArray(_:)`` + public func view(_ type: Scalar.Type = Scalar.self) -> MLXArrayOf { + MLXArrayOf(self) + } +} + +/// A move-only, flat (1d) view over an ``MLXArray`` as a specific `Scalar` type. +/// +/// Created via ``MLXArray/view(_:)``. The array is converted to `Scalar`'s ``DType`` and made +/// contiguous up front; reads then copy out of the backing directly with no further MLX calls. +/// +/// ### See Also +/// - ``MLXArray/view(_:)`` +public struct MLXArrayOf: ~Copyable { + + /// The backing: matches `Scalar`'s ``DType``, contiguous, and evaluated. + public let values: MLXArray + + /// Number of elements in the flattened view. + public let count: Int + + public init(_ array: MLXArray) { + var values = array.asType(Scalar.self) + values.eval() + // reads index the backing directly, so it must be contiguous + if values.contiguousToDimension() != 0 { + values = values.contiguous() + values.eval() + } + self.values = values + self.count = values.size + } + + /// A zero-copy `Span` over the contents. Prefer this over per-element subscripting in + /// hot loops. + public var span: Span { + @_lifetime(borrow self) + borrowing get { + let base = unsafe UnsafeRawPointer(mlx_array_data_uint8(values.ctx)!) + .assumingMemoryBound(to: Scalar.self) + let buffer = unsafe UnsafeBufferPointer(start: base, count: count) + let span = unsafe Span(_unsafeElements: buffer) + // the backing lives as long as `self` holds `values`, not as long as the local buffer + return unsafe _overrideLifetime(span, borrowing: self) + } + } + + /// Copy out a single value at `index` in the flattened contents. + public subscript(_ index: Int) -> Scalar { + precondition(index >= 0 && index < count, "index \(index) out of bounds 0..<\(count)") + return span[index] + } + + /// Copy out a contiguous `range` of the flattened contents. + public subscript(_ range: Range) -> [Scalar] { + precondition( + range.lowerBound >= 0 && range.upperBound <= count, + "range \(range) out of bounds 0..<\(count)") + guard !range.isEmpty else { return [] } + + return withExtendedLifetime(values) { + let base = unsafe UnsafeRawPointer(mlx_array_data_uint8(values.ctx)!) + .assumingMemoryBound(to: Scalar.self) + .advanced(by: range.lowerBound) + let source = unsafe UnsafeBufferPointer(start: base, count: range.count) + return Array(source) + } + } + + /// Copy out the entire flattened contents as a `[Scalar]`. + public func asArray() -> [Scalar] { values.asArray(Scalar.self) } +} diff --git a/Source/MLX/Nested.swift b/Source/MLX/Nested.swift index 59c11457b..634355684 100644 --- a/Source/MLX/Nested.swift +++ b/Source/MLX/Nested.swift @@ -581,7 +581,7 @@ public indirect enum NestedItem: IndentedDescription { var index = index let sorted = dictionary.sorted { lhs, rhs in - String(describing: lhs) < String(describing: rhs) + String(describing: lhs.0) < String(describing: rhs.0) } for (key, element) in sorted { diff --git a/Source/MLX/Ops+Array.swift b/Source/MLX/Ops+Array.swift index 4b90ffbf6..a10eefe42 100644 --- a/Source/MLX/Ops+Array.swift +++ b/Source/MLX/Ops+Array.swift @@ -1470,7 +1470,9 @@ public func reshaped( -> MLXArray { var result = mlx_array_new() - mlx_reshape(&result, array.ctx, newShape.asInt32, newShape.count, stream.ctx) + newShape.withInt32Buffer { shape, count in + mlx_reshape(&result, array.ctx, shape, count, stream.ctx) + } return MLXArray(result) } @@ -1490,7 +1492,9 @@ public func reshaped(_ array: MLXArray, _ newShape: Int..., stream: StreamOrDevi -> MLXArray { var result = mlx_array_new() - mlx_reshape(&result, array.ctx, newShape.asInt32, newShape.count, stream.ctx) + newShape.withInt32Buffer { shape, count in + mlx_reshape(&result, array.ctx, shape, count, stream.ctx) + } return MLXArray(result) } @@ -1735,7 +1739,9 @@ public func sum( stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_sum_axes(&result, array.ctx, axes.asInt32, axes.count, keepDims, stream.ctx) + axes.withInt32Buffer { axes, count in + mlx_sum_axes(&result, array.ctx, axes, count, keepDims, stream.ctx) + } return MLXArray(result) } diff --git a/Source/MLX/Ops.swift b/Source/MLX/Ops.swift index b7e95a330..e23bf6daa 100644 --- a/Source/MLX/Ops.swift +++ b/Source/MLX/Ops.swift @@ -385,11 +385,13 @@ public func asStrided( } var result = mlx_array_new() - mlx_as_strided( - &result, - array.ctx, shape.asInt32, shape.count, resolvedStrides, resolvedStrides.count, - offset, - stream.ctx) + shape.withInt32Buffer { shape, shapeCount in + mlx_as_strided( + &result, + array.ctx, shape, shapeCount, resolvedStrides, resolvedStrides.count, + offset, + stream.ctx) + } return MLXArray(result) } @@ -486,7 +488,9 @@ public func broadcast( -> MLXArray { var result = mlx_array_new() - mlx_broadcast_to(&result, array.ctx, shape.asInt32, shape.count, stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_broadcast_to(&result, array.ctx, shape, count, stream.ctx) + } return MLXArray(result) } @@ -1335,7 +1339,9 @@ public func expandedDimensions( -> MLXArray { var result = mlx_array_new() - mlx_expand_dims_axes(&result, array.ctx, axes.asInt32, axes.count, stream.ctx) + axes.withInt32Buffer { axes, count in + mlx_expand_dims_axes(&result, array.ctx, axes, count, stream.ctx) + } return MLXArray(result) } diff --git a/Source/MLX/Random.swift b/Source/MLX/Random.swift index a89c5d2cd..b899e0ebf 100644 --- a/Source/MLX/Random.swift +++ b/Source/MLX/Random.swift @@ -157,9 +157,11 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_uniform( - &result, lb.ctx, ub.ctx, shape.asInt32, shape.count, type.dtype.cmlxDtype, key.ctx, - stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_uniform( + &result, lb.ctx, ub.ctx, shape, count, type.dtype.cmlxDtype, key.ctx, + stream.ctx) + } return MLXArray(result) } @@ -182,9 +184,11 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_uniform( - &result, lb.ctx, ub.ctx, shape.asInt32, shape.count, type.dtype.cmlxDtype, key.ctx, - stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_uniform( + &result, lb.ctx, ub.ctx, shape, count, type.dtype.cmlxDtype, key.ctx, + stream.ctx) + } return MLXArray(result) } @@ -214,9 +218,11 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_uniform( - &result, low.ctx, high.ctx, shape.asInt32, shape.count, type.dtype.cmlxDtype, key.ctx, - stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_uniform( + &result, low.ctx, high.ctx, shape, count, type.dtype.cmlxDtype, key.ctx, + stream.ctx) + } return MLXArray(result) } @@ -245,10 +251,12 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_uniform( - &result, low.ctx, high.ctx, shape.asInt32, shape.count, dtype.cmlxDtype, key.ctx, - stream.ctx - ) + shape.withInt32Buffer { shape, count in + mlx_random_uniform( + &result, low.ctx, high.ctx, shape, count, dtype.cmlxDtype, key.ctx, + stream.ctx + ) + } return MLXArray(result) } @@ -284,9 +292,11 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_normal( - &result, shape.asInt32, shape.count, type.dtype.cmlxDtype, loc, scale, key.ctx, - stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_normal( + &result, shape, count, type.dtype.cmlxDtype, loc, scale, key.ctx, + stream.ctx) + } return MLXArray(result) } @@ -320,8 +330,10 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_normal( - &result, shape.asInt32, shape.count, dtype.cmlxDtype, loc, scale, key.ctx, stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_normal( + &result, shape, count, dtype.cmlxDtype, loc, scale, key.ctx, stream.ctx) + } return MLXArray(result) } @@ -350,9 +362,11 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_multivariate_normal( - &result, mean.ctx, covariance.ctx, shape.asInt32, shape.count, - dtype.cmlxDtype, key.ctx, stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_multivariate_normal( + &result, mean.ctx, covariance.ctx, shape, count, + dtype.cmlxDtype, key.ctx, stream.ctx) + } return MLXArray(result) } @@ -419,9 +433,11 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_randint( - &result, lb.ctx, ub.ctx, shape.asInt32, shape.count, T.dtype.cmlxDtype, key.ctx, - stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_randint( + &result, lb.ctx, ub.ctx, shape, count, T.dtype.cmlxDtype, key.ctx, + stream.ctx) + } return MLXArray(result) } @@ -449,10 +465,12 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_randint( - &result, low.ctx, high.ctx, shape.asInt32, shape.count, low.dtype.cmlxDtype, key.ctx, - stream.ctx - ) + shape.withInt32Buffer { shape, count in + mlx_random_randint( + &result, low.ctx, high.ctx, shape, count, low.dtype.cmlxDtype, key.ctx, + stream.ctx + ) + } return MLXArray(result) } @@ -480,9 +498,11 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_randint( - &result, low.ctx, high.ctx, shape.asInt32, shape.count, type.dtype.cmlxDtype, key.ctx, - stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_randint( + &result, low.ctx, high.ctx, shape, count, type.dtype.cmlxDtype, key.ctx, + stream.ctx) + } return MLXArray(result) } @@ -510,7 +530,9 @@ public enum MLXRandom { let p = MLXArray(0.5) let key = resolve(key: key) var result = mlx_array_new() - mlx_random_bernoulli(&result, p.ctx, shape.asInt32, shape.count, key.ctx, stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_bernoulli(&result, p.ctx, shape, count, key.ctx, stream.ctx) + } return MLXArray(result) } @@ -541,7 +563,9 @@ public enum MLXRandom { let shape = shape.map { Array($0) } ?? p.shape let key = resolve(key: key) var result = mlx_array_new() - mlx_random_bernoulli(&result, p.ctx, shape.asInt32, shape.count, key.ctx, stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_bernoulli(&result, p.ctx, shape, count, key.ctx, stream.ctx) + } return MLXArray(result) } @@ -574,9 +598,11 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_truncated_normal( - &result, lb.ctx, ub.ctx, shape.asInt32, shape.count, type.dtype.cmlxDtype, key.ctx, - stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_truncated_normal( + &result, lb.ctx, ub.ctx, shape, count, type.dtype.cmlxDtype, key.ctx, + stream.ctx) + } return MLXArray(result) } @@ -599,9 +625,11 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_truncated_normal( - &result, lb.ctx, ub.ctx, shape.asInt32, shape.count, type.dtype.cmlxDtype, key.ctx, - stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_truncated_normal( + &result, lb.ctx, ub.ctx, shape, count, type.dtype.cmlxDtype, key.ctx, + stream.ctx) + } return MLXArray(result) } @@ -630,9 +658,11 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_truncated_normal( - &result, low.ctx, high.ctx, shape.asInt32, shape.count, type.dtype.cmlxDtype, key.ctx, - stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_truncated_normal( + &result, low.ctx, high.ctx, shape, count, type.dtype.cmlxDtype, key.ctx, + stream.ctx) + } return MLXArray(result) } @@ -660,10 +690,12 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_truncated_normal( - &result, low.ctx, high.ctx, shape.asInt32, shape.count, dtype.cmlxDtype, key.ctx, - stream.ctx - ) + shape.withInt32Buffer { shape, count in + mlx_random_truncated_normal( + &result, low.ctx, high.ctx, shape, count, dtype.cmlxDtype, key.ctx, + stream.ctx + ) + } return MLXArray(result) } @@ -690,8 +722,10 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_gumbel( - &result, shape.asInt32, shape.count, type.dtype.cmlxDtype, key.ctx, stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_gumbel( + &result, shape, count, type.dtype.cmlxDtype, key.ctx, stream.ctx) + } return MLXArray(result) } @@ -717,7 +751,9 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_gumbel(&result, shape.asInt32, shape.count, dtype.cmlxDtype, key.ctx, stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_gumbel(&result, shape, count, dtype.cmlxDtype, key.ctx, stream.ctx) + } return MLXArray(result) } @@ -752,8 +788,10 @@ public enum MLXRandom { if let shape { var result = mlx_array_new() - mlx_random_categorical_shape( - &result, logits.ctx, axis.int32, shape.asInt32, shape.count, key.ctx, stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_categorical_shape( + &result, logits.ctx, axis.int32, shape, count, key.ctx, stream.ctx) + } return MLXArray(result) } else { @@ -813,8 +851,10 @@ public enum MLXRandom { let key = resolve(key: key) var result = mlx_array_new() - mlx_random_laplace( - &result, shape.asInt32, shape.count, dtype.cmlxDtype, loc, scale, key.ctx, stream.ctx) + shape.withInt32Buffer { shape, count in + mlx_random_laplace( + &result, shape, count, dtype.cmlxDtype, loc, scale, key.ctx, stream.ctx) + } return MLXArray(result) } diff --git a/Source/MLX/Stream.swift b/Source/MLX/Stream.swift index 6f2d2aa40..d32e1107f 100644 --- a/Source/MLX/Stream.swift +++ b/Source/MLX/Stream.swift @@ -52,7 +52,7 @@ public struct StreamOrDevice: Sendable, CustomStringConvertible, Equatable { public static let gpu = device(.gpu) public static func stream(_ stream: Stream) -> StreamOrDevice { - StreamOrDevice(Device.defaultStream()) + StreamOrDevice(stream) } /// Internal context -- used with Cmlx calls. diff --git a/Source/MLX/Transforms+Compile.swift b/Source/MLX/Transforms+Compile.swift index e52835380..078a574a1 100644 --- a/Source/MLX/Transforms+Compile.swift +++ b/Source/MLX/Transforms+Compile.swift @@ -3,6 +3,52 @@ import Cmlx import Foundation +private let compileConfigurationGeneration = Mutex(UInt64(0)) + +/// Mutable values read by the persistent tracer closure. Keeping these in a separate object +/// avoids a retain cycle between `CompiledFunction` and the C closure that it owns. +private final class CompileTraceState: @unchecked Sendable { + let f: ([MLXArray]) -> [MLXArray] + let outputs: [any Updatable] + + var argumentsCount = 0 + var stateInputs: [MLXArray] = [] + + init(outputs: [any Updatable], _ f: @escaping ([MLXArray]) -> [MLXArray]) { + self.f = f + self.outputs = outputs + } + + /// Called synchronously by MLX while `CompiledFunction.lock` is held. + func trace(_ tracers: [MLXArray]) -> [MLXArray] { + let tracerArguments = Array(tracers.prefix(argumentsCount)) + let savedStateInputs = stateInputs.map { $0.copyContext() } + + for (state, tracer) in zip(stateInputs, tracers.dropFirst(argumentsCount)) { + state._updateInternal(tracer) + } + + // A trace temporarily installs tracer arrays in caller-owned state. Always restore the + // original arrays before returning, including when this function gains throwing work in + // the future. + defer { + for (state, saved) in zip(stateInputs, savedStateInputs) { + state._updateInternal(saved) + } + } + + // The function may return one of the mutable state wrappers directly. Snapshot its MLX + // context before the defer below restores caller-owned state, otherwise the returned + // wrapper would be rewired to an uncaptured original compile input. + let result = f(tracerArguments).map { $0.copyContext() } + let stateOutputTracers = outputs.flatMap { $0.innerState() }.map { $0.copyContext() } + return result + stateOutputTracers + } +} + +// `@unchecked Sendable`: `f`, `inputs`, and `outputs` are plain (non-`@Sendable`) stored +// values used directly outside of `lock` (during `init`), so the compiler can't verify +// this structurally even though `call(_:)` fully serializes access via `lock`. // Note: this is all immutable state -- the `id` property is only set at init time final class CompiledFunction: @unchecked (Sendable) { @@ -20,6 +66,14 @@ final class CompiledFunction: @unchecked (Sendable) { let shapeless: Bool + private let traceState: CompileTraceState + + /// Persistent wrapper returned by `mlx_detail_compile`. The actual compiled graphs remain + /// keyed by `id` in MLX's compiler cache; retaining this wrapper avoids rebuilding the Swift + /// trampoline and C++ `std::function` wrappers on every cache hit. + private var compiled: mlx_closure? + private var compiledGeneration: UInt64? + init( inputs: [any Updatable], outputs: [any Updatable], shapeless: Bool, _ f: @escaping ([MLXArray]) -> [MLXArray] @@ -28,12 +82,19 @@ final class CompiledFunction: @unchecked (Sendable) { self.inputs = inputs self.outputs = outputs self.shapeless = shapeless + self.traceState = CompileTraceState(outputs: outputs, f) self.id = UInt(bitPattern: Unmanaged.passUnretained(self).toOpaque()) } deinit { - // remove the compiled structure from the back end - mlx_detail_compile_erase(id) + // Serialize destruction with application of other MLX transform closures. The tracer + // closure only retains `traceState` (not `self`), so freeing it here cannot form a cycle. + evalLock.withLock { + if let compiled { + mlx_closure_free(compiled) + } + mlx_detail_compile_erase(id) + } } func call(_ arguments: [MLXArray]) -> [MLXArray] { @@ -42,92 +103,88 @@ final class CompiledFunction: @unchecked (Sendable) { } } - func innerCall(_ arguments: [MLXArray]) -> [MLXArray] { - let stateInputs = inputs.flatMap { $0.innerState() } - let argumentsCount = arguments.count - - // inner function to hande the compilation. this is called - // once per compile (typically once overall, but can be called - // again if the conditions for recompile change) - func inner(tracers: [MLXArray]) -> [MLXArray] { - - // put the tracers in their appropriate places: - // - arguments to the function - // - inner state - - let tracerArguments = Array(tracers.prefix(argumentsCount)) - - // save a snapshot of the inner state - let savedStateInputs = stateInputs.map { $0.copyContext() } - - // replace the inner state with the tracers - for (s, tracer) in zip(stateInputs, tracers[argumentsCount...]) { - s._updateInternal(tracer) - } - - // call the function with the tracer arguments - // and the state holding tracers - let result = f(tracerArguments) - - // recapture the state as it may have changed - let stateOutputTracers = outputs.flatMap { $0.innerState() }.map { $0.copyContext() } - - // put the original values back in the state - for (s, saved) in zip(stateInputs, savedStateInputs) { - s._updateInternal(saved) - } - - // return the result of the function and the state - return result + stateOutputTracers + private func buildCompiledClosure() -> mlx_closure? { + let traceState = traceState + let innerClosure = new_mlx_closure { tracers in + traceState.trace(tracers) } - - let innerClosure = new_mlx_closure(inner(tracers:)) defer { mlx_closure_free(innerClosure) } - // note: this will use the cached compile (via the id) - // but will be able to re-evaluate with fresh state if needed - evalLock.lock() var compiled = mlx_closure_new() let compileStatus = mlx_detail_compile(&compiled, innerClosure, id, shapeless, [], 0) - defer { - mlx_closure_free(compiled) - evalLock.unlock() - } // mlx_error was already dispatched on failure: // • outside withError — fatalError was called; we won't reach here - // • inside withError — error is stored in the ErrorBox; return [] so + // • inside withError — error is stored in the ErrorBox; return nil so // withError can throw instead of crashing downstream guard compileStatus == 0 else { - return [] + mlx_closure_free(compiled) + return nil } - let innerInputs = arguments + stateInputs - let innerInputsVector = new_mlx_vector_array(innerInputs) - defer { mlx_vector_array_free(innerInputsVector) } + return compiled + } + + func innerCall(_ arguments: [MLXArray]) -> [MLXArray] { + traceState.stateInputs = inputs.flatMap { $0.innerState() } + traceState.argumentsCount = arguments.count + + return evalLock.withLock { + let generation = compileConfigurationGeneration.withLock { $0 } + + // `mlx_detail_compile` observes whether compilation is enabled when it creates the + // wrapper. Rebuild only after the public mode setter is used so a cached wrapper does + // not permanently preserve an earlier enabled/disabled mode. + if compiledGeneration != generation { + if let compiled { + mlx_closure_free(compiled) + self.compiled = nil + } + compiledGeneration = generation + } - // will compile the function (if needed) and evaluate the - // compiled graph - var resultVector = mlx_vector_array_new() - let applyStatus = mlx_closure_apply(&resultVector, compiled, innerInputsVector) - defer { mlx_vector_array_free(resultVector) } + if compiled == nil { + guard let built = buildCompiledClosure() else { + compiledGeneration = nil + return [] + } + compiled = built + } - guard applyStatus == 0 else { - return [] - } + guard let compiled else { + return [] + } - let resultsPlusStateOutput = mlx_vector_array_values(resultVector) + let innerInputs = arguments + traceState.stateInputs + let innerInputsVector = new_mlx_vector_array(innerInputs) + defer { mlx_vector_array_free(innerInputsVector) } - // push the stateOutput into the state - let stateOutput = outputs.flatMap { $0.innerState() } + // This compiles on a cache miss (including a new shape/dtype) and evaluates the graph. + var resultVector = mlx_vector_array_new() + let applyStatus = mlx_closure_apply(&resultVector, compiled, innerInputsVector) + defer { mlx_vector_array_free(resultVector) } - for (s, newValues) in zip(stateOutput, resultsPlusStateOutput.suffix(stateOutput.count)) { - s._updateInternal(newValues) - } + guard applyStatus == 0 else { + // MLX marks a cache entry non-empty before tracing it. If tracing fails, remove + // that potentially incomplete entry so a later call can retry cleanly. + mlx_detail_compile_erase(id) + return [] + } - let resultLength = resultsPlusStateOutput.count - stateOutput.count - let results = Array(resultsPlusStateOutput.prefix(resultLength)) - return results + let resultsPlusStateOutput = mlx_vector_array_values(resultVector) + + // push the stateOutput into the state + let stateOutput = outputs.flatMap { $0.innerState() } + + for (state, newValues) in zip( + stateOutput, resultsPlusStateOutput.suffix(stateOutput.count)) + { + state._updateInternal(newValues) + } + + let resultLength = resultsPlusStateOutput.count - stateOutput.count + return Array(resultsPlusStateOutput.prefix(resultLength)) + } } } @@ -232,9 +289,15 @@ public func compile( /// /// Default is enabled. public func compile(enable: Bool = true) { - if enable { - mlx_enable_compile() - } else { - mlx_disable_compile() + evalLock.withLock { + let status: Int32 + if enable { + status = mlx_enable_compile() + } else { + status = mlx_disable_compile() + } + if status == 0 { + compileConfigurationGeneration.withLock { $0 &+= 1 } + } } } diff --git a/Source/MLX/Transforms+Eval.swift b/Source/MLX/Transforms+Eval.swift index 46fe9c593..06e752009 100644 --- a/Source/MLX/Transforms+Eval.swift +++ b/Source/MLX/Transforms+Eval.swift @@ -13,6 +13,7 @@ let evalLock = NSRecursiveLock() /// ### See Also /// - public func eval(_ arrays: MLXArray...) { + guard !arrays.isEmpty else { return } let vector_array = new_mlx_vector_array(arrays) _ = evalLock.withLock { mlx_eval(vector_array) @@ -25,6 +26,7 @@ public func eval(_ arrays: MLXArray...) { /// ### See Also /// - public func eval(_ arrays: some Collection) { + guard !arrays.isEmpty else { return } let vector_array = new_mlx_vector_array(arrays) _ = evalLock.withLock { mlx_eval(vector_array) @@ -38,6 +40,7 @@ public func eval(_ arrays: some Collection) { /// - /// - ``asyncEval(_:)-(Collection)`` public func asyncEval(_ arrays: some Collection) { + guard !arrays.isEmpty else { return } let vector_array = new_mlx_vector_array(arrays) _ = evalLock.withLock { mlx_async_eval(vector_array) diff --git a/Source/MLX/Transforms+Internal.swift b/Source/MLX/Transforms+Internal.swift index 394aa4e98..297cb1ab6 100644 --- a/Source/MLX/Transforms+Internal.swift +++ b/Source/MLX/Transforms+Internal.swift @@ -2,28 +2,67 @@ import Cmlx import Foundation +import Synchronization // see Transforms+Variants for generated grad() functions -private func valueAndGradient( - apply valueAndGrad: mlx_closure_value_and_grad, arrays: some Collection -) - -> ([MLXArray], [MLXArray]) -{ - let input_vector = new_mlx_vector_array(arrays) - defer { mlx_vector_array_free(input_vector) } +/// Copy the backing MLX contexts in a parameter tree so later `_updateInternal` calls on the +/// original wrappers cannot mutate the snapshot. Used by MLXNN to restore a model after tracing. +@_documentation(visibility: internal) +public func _snapshotArrayContexts( + _ parameters: NestedDictionary +) -> NestedDictionary { + parameters.mapValues { $0.copyContext() } +} + +/// Owns an MLX value-and-gradient transform for the lifetime of the Swift function returned to +/// the caller. The transform is immutable after construction and applications are serialized by +/// `evalLock`, so it is safe to reuse across calls and threads. +private final class ValueAndGradientTransform { + private let transform: mlx_closure_value_and_grad + + init?(_ f: @escaping ([MLXArray]) -> [MLXArray], argumentNumbers: some Collection) { + var transform = mlx_closure_value_and_grad_new() + let closure = new_mlx_closure(f) + let argumentNumbers = argumentNumbers.asInt32 + let status = evalLock.withLock { + let status = mlx_value_and_grad( + &transform, closure, argumentNumbers, argumentNumbers.count) + mlx_closure_free(closure) + return status + } - var r0 = mlx_vector_array_new() - var r1 = mlx_vector_array_new() + guard status == 0 else { + mlx_closure_value_and_grad_free(transform) + return nil + } + self.transform = transform + } - _ = evalLock.withLock { - mlx_closure_value_and_grad_apply(&r0, &r1, valueAndGrad, input_vector) + deinit { + _ = evalLock.withLock { + mlx_closure_value_and_grad_free(transform) + } } - defer { mlx_vector_array_free(r0) } - defer { mlx_vector_array_free(r1) } + func call(_ arrays: some Collection) -> ([MLXArray], [MLXArray]) { + let inputVector = new_mlx_vector_array(arrays) + defer { mlx_vector_array_free(inputVector) } - return (mlx_vector_array_values(r0), mlx_vector_array_values(r1)) + var values = mlx_vector_array_new() + var gradients = mlx_vector_array_new() + defer { mlx_vector_array_free(values) } + defer { mlx_vector_array_free(gradients) } + + let status = evalLock.withLock { + mlx_closure_value_and_grad_apply(&values, &gradients, transform, inputVector) + } + guard status == 0 else { + return ([], []) + } + + return (mlx_vector_array_values(values), mlx_vector_array_values(gradients)) + } } func buildGradient(_ f: @escaping ([MLXArray]) -> [MLXArray], argumentNumbers: some Collection) @@ -31,16 +70,12 @@ func buildGradient(_ f: @escaping ([MLXArray]) -> [MLXArray], argumentNumbers: s [MLXArray] ) -> [MLXArray] { - { (arrays: [MLXArray]) in - var vag = mlx_closure_value_and_grad_new() - - let closure = new_mlx_closure(f) - mlx_value_and_grad(&vag, closure, argumentNumbers.asInt32, argumentNumbers.count) - mlx_closure_free(closure) - - defer { mlx_closure_value_and_grad_free(vag) } + guard let transform = ValueAndGradientTransform(f, argumentNumbers: argumentNumbers) else { + return { _ in [] } + } - return valueAndGradient(apply: vag, arrays: arrays).1 + return { (arrays: [MLXArray]) in + transform.call(arrays).1 } } @@ -49,16 +84,79 @@ func buildValueAndGradient( ) -> ( [MLXArray] ) -> ([MLXArray], [MLXArray]) { - { (arrays: [MLXArray]) in - var vag = mlx_closure_value_and_grad_new() + guard let transform = ValueAndGradientTransform(f, argumentNumbers: argumentNumbers) else { + return { _ in ([], []) } + } - let closure = new_mlx_closure(f) - mlx_value_and_grad(&vag, closure, argumentNumbers.asInt32, argumentNumbers.count) - mlx_closure_free(closure) + return { (arrays: [MLXArray]) in + transform.call(arrays) + } +} + +/// Caches the transform used by the nested-parameter plus array-input overload. The parameter +/// topology is part of the transform because its tracer closure reconstructs that topology; a +/// topology change therefore creates a fresh transform while ordinary value/shape changes reuse +/// the existing one. +private final class NestedArrayValueAndGradientTransform { + typealias Parameters = NestedDictionary - defer { mlx_closure_value_and_grad_free(vag) } + let f: (Parameters, [MLXArray]) -> [MLXArray] + let lock = Mutex(()) - return valueAndGradient(apply: vag, arrays: arrays) + private var topology: NestedDictionary? + private var transform: ValueAndGradientTransform? + + init(_ f: @escaping (Parameters, [MLXArray]) -> [MLXArray]) { + self.f = f + } + + func call(_ parameters: Parameters, _ extraArrays: [MLXArray]) -> ([MLXArray], Parameters) { + var result: ([MLXArray], Parameters) = ([], parameters) + lock.withLock { _ in + result = innerCall(parameters, extraArrays) + } + return result + } + + private func innerCall( + _ parameters: Parameters, _ extraArrays: [MLXArray] + ) -> ([MLXArray], Parameters) { + let currentTopology = parameters.mapValues { _ in false } + let flattenedParameters = parameters.flattenedValues() + + if topology != currentTopology || transform == nil { + let parameterCount = flattenedParameters.count + let parameterTemplate = parameters + let f = f + + transform = ValueAndGradientTransform( + { inputs in + let flatParameters = Array(inputs.prefix(parameterCount)) + let parameters = parameterTemplate.replacingValues(with: flatParameters) + let extras = Array(inputs.dropFirst(parameterCount)) + return f(parameters, extras) + }, argumentNumbers: 0 ..< parameterCount) + topology = transform == nil ? nil : currentTopology + } + + guard let transform else { + return ([], parameters) + } + + let (values, flatGradients) = transform.call(flattenedParameters + extraArrays) + let gradients = parameters.replacingValues(with: flatGradients) + return (values, gradients) + } +} + +func buildValueAndGradient( + _ f: @escaping (NestedDictionary, [MLXArray]) -> [MLXArray] +) -> (NestedDictionary, [MLXArray]) -> ( + [MLXArray], NestedDictionary +) { + let transform = NestedArrayValueAndGradientTransform(f) + return { parameters, arrays in + transform.call(parameters, arrays) } } @@ -72,16 +170,7 @@ func buildValueAndGradient( [MLXArray], NestedDictionary ) in - // capture the state so that we can unflatten - let flattenedParameters = parameters.flattened() - let flattenedKeys = flattenedParameters.map { $0.0 } - let flattenedArrays = flattenedParameters.map { $0.1 } - - // function to unflatten back into the NestedDictionary - func unflattened(_ arrays: [MLXArray]) -> NestedDictionary { - let tuples = zip(flattenedKeys, arrays).map { ($0.0, $0.1) } - return NestedDictionary.unflattened(tuples) - } + let flattenedArrays = parameters.flattenedValues() // this goes in the closure and is wrapped by mlx_value_and_grad // @@ -93,21 +182,18 @@ func buildValueAndGradient( // arg indexes to indicate which ones to grad (it should work // as is) func inner(flattenedArrays: [MLXArray]) -> [MLXArray] { - let parameters = unflattened(flattenedArrays) + let parameters = parameters.replacingValues(with: flattenedArrays) return f(parameters, arrays) } - var vag = mlx_closure_value_and_grad_new() - - let closure = new_mlx_closure(inner) - mlx_value_and_grad( - &vag, closure, Array(Int32(0) ..< Int32(flattenedArrays.count)), flattenedArrays.count) - mlx_closure_free(closure) - - defer { mlx_closure_value_and_grad_free(vag) } + guard let transform = ValueAndGradientTransform( + inner, argumentNumbers: 0 ..< flattenedArrays.count) + else { + return ([], parameters) + } - let (values, flatGradients) = valueAndGradient(apply: vag, arrays: flattenedArrays) - let gradients = unflattened(flatGradients) + let (values, flatGradients) = transform.call(flattenedArrays) + let gradients = parameters.replacingValues(with: flatGradients) return (values, gradients) } diff --git a/Source/MLX/Transforms.swift b/Source/MLX/Transforms.swift index 3e73a8f8b..f2966d106 100644 --- a/Source/MLX/Transforms.swift +++ b/Source/MLX/Transforms.swift @@ -87,3 +87,17 @@ public func valueAndGrad( ) { buildValueAndGradient(f) } + +/// Returns a reusable value-and-gradient function for nested parameters and additional array +/// inputs. Only the nested parameters are differentiated; `arrays` are ordinary transform inputs +/// so their values may change without rebuilding the transform. +/// +/// If the parameter topology changes, the returned function rebuilds its transform for the new +/// topology. Changes to parameter values, shapes, or dtypes preserve the transform object. +public func valueAndGrad( + _ f: @escaping (NestedDictionary, [MLXArray]) -> [MLXArray] +) -> (NestedDictionary, [MLXArray]) -> ( + [MLXArray], NestedDictionary +) { + buildValueAndGradient(f) +} diff --git a/Source/MLXNN/ValueAndGrad.swift b/Source/MLXNN/ValueAndGrad.swift index c5c3525cf..2caabd820 100644 --- a/Source/MLXNN/ValueAndGrad.swift +++ b/Source/MLXNN/ValueAndGrad.swift @@ -59,11 +59,16 @@ public func valueAndGrad( // arrays we can capture the result of the valueAndGrad and use it // over and over func inner(parameters: ModuleParameters, arrays: [MLXArray]) -> [MLXArray] { + let savedParameters = _snapshotArrayContexts(model.trainableParameters()) model.update(parameters: parameters) + defer { model.update(parameters: savedParameters) } return [f(model, arrays[0], arrays[1])] } - let vg = valueAndGrad(inner) + // Select MLX's array-specialized overload explicitly. It passes `a1`/`a2` as transform + // inputs and reuses the value-and-gradient transform while the parameter topology is stable. + let vg: (ModuleParameters, [MLXArray]) -> ([MLXArray], ModuleParameters) = + valueAndGrad(inner) // outer function func wrapped(model: Model, a1: MLXArray, a2: MLXArray) -> (MLXArray, ModuleParameters) { @@ -102,11 +107,16 @@ public func valueAndGrad( ) -> (Model, [MLXArray]) -> ([MLXArray], ModuleParameters) { func inner(parameters: ModuleParameters, arrays: [MLXArray]) -> [MLXArray] { + let savedParameters = _snapshotArrayContexts(model.trainableParameters()) model.update(parameters: parameters) + defer { model.update(parameters: savedParameters) } return f(model, arrays) } - let vg = valueAndGrad(inner) + // Extra arrays are transform inputs rather than captured values, so changing a batch does + // not rebuild the value-and-gradient transform. + let vg: (ModuleParameters, [MLXArray]) -> ([MLXArray], ModuleParameters) = + valueAndGrad(inner) func wrapped(model: Model, arrays: [MLXArray]) -> ([MLXArray], ModuleParameters) { vg(model.trainableParameters(), arrays) @@ -148,7 +158,9 @@ public func valueAndGrad( func wrapped(model: Model, arguments: Arguments) -> ([MLXArray], ModuleParameters) { func inner(parameters: ModuleParameters, _ extra: ()) -> [MLXArray] { + let savedParameters = _snapshotArrayContexts(model.trainableParameters()) model.update(parameters: parameters) + defer { model.update(parameters: savedParameters) } return f(model, arguments) } diff --git a/Source/MLXOptimizers/Optimizers.swift b/Source/MLXOptimizers/Optimizers.swift index a01083b85..38cfb5b91 100644 --- a/Source/MLXOptimizers/Optimizers.swift +++ b/Source/MLXOptimizers/Optimizers.swift @@ -122,6 +122,9 @@ open class OptimizerBase: Optimizer { final func apply(gradients: ModuleParameters, modelParameters: ModuleParameters) -> ModuleParameters { + prepareForUpdate() + defer { finishUpdate() } + let (p, s) = gradients.mapValues(modelParameters, stateStorage) { gradient, parameter, state in // handle optionality of the visitor params @@ -133,6 +136,15 @@ open class OptimizerBase: Optimizer { return p } + /// Prepare optimizer-wide state used by all parameters in an update. + /// + /// Most optimizers do not need optimizer-wide state. Subclasses such as Adam use this hook + /// to construct shared scalar expressions once instead of once per parameter. + func prepareForUpdate() {} + + /// Release any transient optimizer-wide state prepared by ``prepareForUpdate()``. + func finishUpdate() {} + open func applySingle(gradient: MLXArray, parameter: MLXArray, state: State) -> ( MLXArray, State ) { @@ -171,29 +183,27 @@ public struct TupleState: Updatable { } } -/// State container for Adam-style optimizers that need first and second moments -/// plus an update step for optional bias correction. +/// State container for Adam-style optimizers that need first and second moments. +/// +/// The update step is optimizer-wide rather than per-parameter, so models with many parameters do +/// not create and retain an identical scalar graph for every parameter. public struct AdamState: Updatable { let values: (MLXArray, MLXArray) - var step: MLXArray - init(_ values: (MLXArray, MLXArray), step: MLXArray) { + init(_ values: (MLXArray, MLXArray)) { self.values = values - self.step = step } - init(_ a: MLXArray, _ b: MLXArray, step: MLXArray) { + init(_ a: MLXArray, _ b: MLXArray) { self.values = (a, b) - self.step = step } init(zeros array: MLXArray) { self.values = (MLXArray.zeros(like: array), MLXArray.zeros(like: array)) - self.step = MLXArray(0) } public func innerState() -> [MLXArray] { - [values.0, values.1, step] + [values.0, values.1] } } @@ -408,7 +418,22 @@ open class Adam: OptimizerBase { /// The epsilon added to the denominator to improve numerical stability public var eps: Float = 1e-8 /// If `true`, apply bias correction to the first and second moments - public var biasCorrection = false + public var biasCorrection = false { + didSet { + // A disabled optimizer deliberately does not build a step graph. If correction is + // enabled later, begin its correction schedule with the next update. + if biasCorrection && !oldValue { + step = MLXArray(0) + } + } + } + + /// A single step shared by every parameter. It is included in `innerState()` only when bias + /// correction needs it, avoiding unused scalar state and graph work in the default mode. + var step = MLXArray(0) + + /// Bias-correction factors prepared once for all parameters in the current update. + private var preparedCorrection: (first: MLXArray, second: MLXArray)? /// Initialize the optimizer. /// - Parameters: @@ -430,29 +455,57 @@ open class Adam: OptimizerBase { AdamState(zeros: parameter) } + override open func innerState() -> [MLXArray] { + let parameterState = super.innerState() + return biasCorrection ? parameterState + [step] : parameterState + } + + override func prepareForUpdate() { + guard biasCorrection else { return } + step = step + 1 + preparedCorrection = correctionFactors(step: step) + } + + override func finishUpdate() { + preparedCorrection = nil + } + + private func correctionFactors(step: MLXArray) -> (first: MLXArray, second: MLXArray) { + let (b1, b2) = betas + return ( + learningRate / (1 - pow(b1, step)), + rsqrt(1 - pow(b2, step)) + ) + } + override open func applySingle(gradient: MLXArray, parameter: MLXArray, state: AdamState) -> ( MLXArray, AdamState ) { let (b1, b2) = betas - var state = state var (m, v) = state.values - state.step = state.step + 1 m = b1 * m + (1 - b1) * gradient v = b2 * v + (1 - b2) * square(gradient) let update: MLXArray if biasCorrection { - let step = state.step - let c1 = learningRate / (1 - pow(b1, step)) - let c2 = rsqrt(1 - pow(b2, step)) + // `apply(gradients:modelParameters:)` prepares these once per whole update. A direct + // `applySingle` call is itself one optimizer update, so advance the shared step here. + let correction: (first: MLXArray, second: MLXArray) + if let preparedCorrection { + correction = preparedCorrection + } else { + step = step + 1 + correction = correctionFactors(step: step) + } + let (c1, c2) = correction update = (c1 * m) / (sqrt(v) * c2 + eps) } else { update = learningRate * m / (sqrt(v) + eps) } - return (parameter - update, AdamState(m, v, step: state.step)) + return (parameter - update, AdamState(m, v)) } } @@ -874,11 +927,9 @@ open class Muon: OptimizerBaseArrayState { public func clipGradNorm(gradients: some Collection, maxNorm: Float) -> ( [MLXArray], MLXArray ) { - let normSquared = gradients.reduce(MLXArray(0)) { $0 + $1.square().sum() } - let totalNorm = sqrt(normSquared) - let normalizer = maxNorm / (totalNorm + 1e-6) - - let clippedGradients = gradients.map { which(totalNorm .< maxNorm, $0, $0 * normalizer) } + let totalNorm = globalGradientNorm(gradients) + let scale = minimum(Float(1), maxNorm / (totalNorm + 1e-6)) + let clippedGradients = gradients.map { $0 * scale } return (clippedGradients, totalNorm) } @@ -895,11 +946,17 @@ public func clipGradNorm(gradients: some Collection, maxNorm: Float) - public func clipGradNorm(gradients: ModuleParameters, maxNorm: Float) -> ( ModuleParameters, MLXArray ) { - let normSquared = gradients.reduce(MLXArray(0)) { $0 + $1.square().sum() } - let totalNorm = sqrt(normSquared) - let normalizer = maxNorm / (totalNorm + 1e-6) - - let clippedGradients = gradients.mapValues { which(totalNorm .< maxNorm, $0, $0 * normalizer) } + let totalNorm = globalGradientNorm(gradients.flattenedValues()) + let scale = minimum(Float(1), maxNorm / (totalNorm + 1e-6)) + let clippedGradients = gradients.mapValues { $0 * scale } return (clippedGradients, totalNorm) } + +/// Build a shallow reduction graph for the squared norms. Stacking the scalar reductions lets MLX +/// perform one flat sum instead of constructing a left-deep chain of additions in Swift. +private func globalGradientNorm(_ gradients: some Collection) -> MLXArray { + let squaredNorms = gradients.map { square($0).sum() } + guard !squaredNorms.isEmpty else { return MLXArray(Float(0)) } + return sqrt(stacked(squaredNorms).sum()) +} diff --git a/Tests/MLXTests/MLXArrayTests.swift b/Tests/MLXTests/MLXArrayTests.swift index 715a77590..045d7ed66 100644 --- a/Tests/MLXTests/MLXArrayTests.swift +++ b/Tests/MLXTests/MLXArrayTests.swift @@ -3,6 +3,7 @@ import Foundation import XCTest +import Cmlx @testable import MLX class MLXArrayTests: XCTestCase { @@ -94,6 +95,75 @@ class MLXArrayTests: XCTestCase { XCTAssertEqual(s_arr, expected) } + func testAsArrayNegativeStride() { + // negative strides + a nonzero offset: this is the case that was + // previously untested and relied on raw pointer arithmetic outside + // any range the Swift-side code verified + let a = MLXArray(0 ..< 16, [4, 4]) + let s = asStrided(a, [4, 4], strides: [-4, -1], offset: 15) + + let expected: [Int32] = Array((0 ..< 16).reversed()).map { Int32($0) } + assertEqual(s, MLXArray(expected, [4, 4])) + + let s_arr = s.asArray(Int32.self) + XCTAssertEqual(s_arr, expected) + } + + func testAsArrayMixedStrides() { + let a = MLXArray(0 ..< 12) + let s = asStrided(a, [3, 4], strides: [-1, 3], offset: 2) + + let expected: [Int32] = [2, 5, 8, 11, 1, 4, 7, 10, 0, 3, 6, 9] + XCTAssertEqual(s.asArray(Int32.self), expected) + } + + func testViewMakesNonContiguousArrayContiguousWithoutChangingShape() { + let a = MLXArray(0 ..< 12, [4, 3]) + let transposed = asStrided(a, [3, 4], strides: [1, 3]) + let view = transposed.view(Int32.self) + + XCTAssertEqual(view.values.shape, [3, 4]) + XCTAssertEqual(view.values.contiguousToDimension(), 0) + XCTAssertEqual(view.count, 12) + XCTAssertEqual(view[2 ..< 7], [6, 9, 1, 4, 7]) + XCTAssertEqual(view.asArray(), [0, 3, 6, 9, 1, 4, 7, 10, 2, 5, 8, 11]) + + let empty = MLXArray([Int32]()).view(Int32.self) + XCTAssertEqual(empty.count, 0) + XCTAssertEqual(empty[0 ..< 0], []) + } + + func testScopedInteropBuffersHandleEmptyAndPopulatedCollections() { + let empty = [Int]().withInt32Buffer { pointer, count in + XCTAssertNil(pointer) + return count + } + XCTAssertEqual(empty, 0) + + Array(0 ..< 5)[1 ..< 4].withInt32Buffer { pointer, count in + XCTAssertEqual(count, 3) + XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer, count: count)), [1, 2, 3]) + } + Array(0 ... mlxInteropStackBufferCapacity).withInt32Buffer { pointer, count in + XCTAssertEqual(count, mlxInteropStackBufferCapacity + 1) + XCTAssertEqual(pointer?[count - 1], Int32(mlxInteropStackBufferCapacity)) + } + + let emptyVector = new_mlx_vector_array([MLXArray]()) + defer { mlx_vector_array_free(emptyVector) } + XCTAssertEqual(mlx_vector_array_size(emptyVector), 0) + + let arrays = [MLXArray(1), MLXArray(2)] + let vector = new_mlx_vector_array(arrays) + defer { mlx_vector_array_free(vector) } + XCTAssertEqual(mlx_vector_array_values(vector).map { $0.item(Int.self) }, [1, 2]) + + let manyArrays = (0 ... mlxInteropStackBufferCapacity).map { MLXArray($0) } + let largeVector = new_mlx_vector_array(manyArrays) + defer { mlx_vector_array_free(largeVector) } + XCTAssertEqual(mlx_vector_array_size(largeVector), mlxInteropStackBufferCapacity + 1) + } + func testContiguousStrides() { XCTAssertEqual(contiguousStrides(shape: [1, 1, 1]), [1, 1, 1]) XCTAssertEqual(contiguousStrides(shape: [4, 4]), [4, 1]) diff --git a/Tests/MLXTests/NestedTests.swift b/Tests/MLXTests/NestedTests.swift index 3e9d83b55..2eae04599 100644 --- a/Tests/MLXTests/NestedTests.swift +++ b/Tests/MLXTests/NestedTests.swift @@ -81,6 +81,30 @@ class NestedTests: XCTestCase { XCTAssertEqual(n2, expected) } + func testReplacingValuesDoesNotDescribeValues() { + final class DescriptionProbe: CustomStringConvertible { + var wasDescribed = false + + var description: String { + wasDescribed = true + return "probe" + } + } + + let first = DescriptionProbe() + let second = DescriptionProbe() + var nested = NestedDictionary() + nested["b"] = .value(first) + nested["a"] = .value(second) + + let replacements = [DescriptionProbe(), DescriptionProbe()] + _ = nested.replacingValues(with: replacements) + + XCTAssertFalse(first.wasDescribed) + XCTAssertFalse(second.wasDescribed) + XCTAssertTrue(replacements.allSatisfy { !$0.wasDescribed }) + } + func testMap2() { // map 2 parallel structures var d1 = NestedDictionary() diff --git a/Tests/MLXTests/OptimizerTests.swift b/Tests/MLXTests/OptimizerTests.swift index d773a5563..7e87d0891 100644 --- a/Tests/MLXTests/OptimizerTests.swift +++ b/Tests/MLXTests/OptimizerTests.swift @@ -167,6 +167,65 @@ class OptimizerTests: XCTestCase { let v = (1 - Float(0.999)) * square(gradient) let expected = parameter - (c1 * m) / (sqrt(v) * c2 + Float(1e-8)) assertEqual(result.0, expected, atol: 1e-6) + XCTAssertEqual(optimizer.step.item(Int.self), 1) + } + + func testAdamSharedStepState() { + let model = TwoParameterModel() + let gradients = model.parameters().mapValues { key, parameter in + key == "bias" ? MLXArray.ones(like: parameter) : 2 * MLXArray.ones(like: parameter) + } + let optimizer = Adam(learningRate: 0.1, biasCorrection: true) + + optimizer.update(model: model, gradients: gradients) + optimizer.update(model: model, gradients: gradients) + eval(model, optimizer) + + // Two moment arrays per parameter plus one optimizer-wide scalar step. + let state = optimizer.innerState() + XCTAssertEqual(state.count, 5) + XCTAssertEqual(state.filter { $0.ndim == 0 }.count, 1) + XCTAssertEqual(state.last?.item(Int.self), 2) + + // With constant gradients, bias-corrected moments produce a 0.1 update on both steps, + // independent of each gradient's magnitude. + let expected = MLXArray([Float(-0.2), -0.2, -0.2]) + assertEqual(model.bias, expected, atol: 1e-5) + assertEqual(model.weight, expected, atol: 1e-5) + } + + func testCompiledAdamBiasCorrectionAdvancesSharedStep() { + let model = TwoParameterModel() + let optimizer = Adam(learningRate: 0.1, biasCorrection: true) + + func step(_ gradient: MLXArray) -> MLXArray { + let gradients = model.parameters().mapValues { MLXArray.ones(like: $0) * gradient } + optimizer.update(model: model, gradients: gradients) + return model.weight + } + + let compiledStep = MLX.compile( + inputs: [model, optimizer], outputs: [model, optimizer], step) + _ = compiledStep(MLXArray(1 as Float)) + _ = compiledStep(MLXArray(1 as Float)) + eval(optimizer) + XCTAssertEqual(optimizer.step.item(Int.self), 2) + } + + func testAdamWithoutBiasCorrectionHasNoStepState() { + let model = TwoParameterModel() + let gradients = model.parameters().mapValues { MLXArray.ones(like: $0) } + let optimizer = Adam(learningRate: 0.1, biasCorrection: false) + + optimizer.update(model: model, gradients: gradients) + optimizer.update(model: model, gradients: gradients) + eval(model, optimizer) + + // Only first and second moments are state; no unused scalar step graph is retained. + let state = optimizer.innerState() + XCTAssertEqual(state.count, 4) + XCTAssertTrue(state.allSatisfy { $0.shape == [3] }) + XCTAssertEqual(optimizer.step.item(Int.self), 0) } func testAdamW() { @@ -195,6 +254,22 @@ class OptimizerTests: XCTestCase { assertEqual(result.0, expected, atol: 1e-6) } + func testAdamWUsesSharedStepState() { + let model = TwoParameterModel() + let gradients = model.parameters().mapValues { MLXArray.ones(like: $0) } + let optimizer = AdamW( + learningRate: 0.1, weightDecay: 0.01, biasCorrection: true) + + optimizer.update(model: model, gradients: gradients) + optimizer.update(model: model, gradients: gradients) + eval(model, optimizer) + + let state = optimizer.innerState() + XCTAssertEqual(state.count, 5) + XCTAssertEqual(state.filter { $0.ndim == 0 }.count, 1) + XCTAssertEqual(state.last?.item(Int.self), 2) + } + func testAdamax() { checkShape(optimizer: Adamax(learningRate: 0.1)) checkTrain(optimizer: Adamax(learningRate: 0.1)) @@ -268,4 +343,49 @@ class OptimizerTests: XCTestCase { XCTAssertGreaterThan(abs(p2).sum().item(Float.self), 0) } + func testClipGradNormCollectionAboveAndBelowThreshold() { + let gradients = [MLXArray([3.0] as [Float]), MLXArray([4.0] as [Float])] + + let (unchanged, normBelowThreshold) = clipGradNorm( + gradients: gradients, maxNorm: 10) + eval(unchanged, normBelowThreshold) + XCTAssertEqual(normBelowThreshold.item(Float.self), 5, accuracy: 1e-6) + assertEqual(unchanged, gradients) + + let (clipped, normAboveThreshold) = clipGradNorm(gradients: gradients, maxNorm: 2) + eval(clipped, normAboveThreshold) + let scale = Float(2) / (Float(5) + Float(1e-6)) + XCTAssertEqual(normAboveThreshold.item(Float.self), 5, accuracy: 1e-6) + assertEqual(clipped[0], MLXArray([3 * scale] as [Float]), atol: 1e-6) + assertEqual(clipped[1], MLXArray([4 * scale] as [Float]), atol: 1e-6) + } + + func testClipGradNormModuleParametersAndEmptyCollections() { + let gradients = ModuleParameters(values: [ + "first": .value(MLXArray([3.0] as [Float])), + "second": .value(MLXArray([4.0] as [Float])), + ]) + let (clipped, norm) = clipGradNorm(gradients: gradients, maxNorm: 2) + eval(clipped, norm) + + let scale = Float(2) / (Float(5) + Float(1e-6)) + XCTAssertEqual(norm.item(Float.self), 5, accuracy: 1e-6) + assertEqual( + clipped[unwrapping: "first"]!, MLXArray([3 * scale] as [Float]), atol: 1e-6) + assertEqual( + clipped[unwrapping: "second"]!, MLXArray([4 * scale] as [Float]), atol: 1e-6) + + let (emptyArrays, emptyArrayNorm) = clipGradNorm( + gradients: [MLXArray](), maxNorm: 1) + eval(emptyArrayNorm) + XCTAssertTrue(emptyArrays.isEmpty) + XCTAssertEqual(emptyArrayNorm.item(Float.self), 0) + + let (emptyParameters, emptyParameterNorm) = clipGradNorm( + gradients: ModuleParameters(), maxNorm: 1) + eval(emptyParameterNorm) + XCTAssertTrue(emptyParameters.isEmpty) + XCTAssertEqual(emptyParameterNorm.item(Float.self), 0) + } + } diff --git a/Tests/MLXTests/StreamTests.swift b/Tests/MLXTests/StreamTests.swift index ec3a35de7..594b62f75 100644 --- a/Tests/MLXTests/StreamTests.swift +++ b/Tests/MLXTests/StreamTests.swift @@ -28,6 +28,11 @@ class StreamTests: XCTestCase { XCTAssertEqual(s3.deviceType, .cpu) } + func testExplicitStreamUsesProvidedStream() { + let stream = Stream(.cpu) + XCTAssertTrue(StreamOrDevice.stream(stream).stream === stream) + } + func testUsingDevice() { let defaultDevice = Device.defaultDevice() diff --git a/Tests/MLXTests/TransformTests.swift b/Tests/MLXTests/TransformTests.swift index d06999c7d..86efb215f 100644 --- a/Tests/MLXTests/TransformTests.swift +++ b/Tests/MLXTests/TransformTests.swift @@ -67,6 +67,80 @@ class TransformTests: XCTestCase { XCTAssertEqual(grad[0].item(), Float(2 * 1.5)) } + func testGradReusableAcrossValuesAndShapes() { + let gradient = grad { (x: MLXArray) in x.square().sum() } + + XCTAssertEqual(gradient(MLXArray(2)).item(Float.self), 4) + XCTAssertEqual(gradient(MLXArray(-3)).item(Float.self), -6) + XCTAssertEqual( + gradient(MLXArray([Float(1), 2, 3])).asArray(Float.self), [2, 4, 6]) + } + + func testNestedValueAndGradUsesCurrentExtraArrays() { + var parameters = ModuleParameters() + parameters["weight"] = .value(MLXArray(Float(2))) + + let transformed = valueAndGrad { + (parameters: ModuleParameters, arrays: [MLXArray]) -> [MLXArray] in + guard case .value(let weight)? = parameters["weight"] else { + XCTFail("missing weight") + return [] + } + return [weight * arrays[0]] + } + + let (firstValue, firstGradients) = transformed(parameters, [MLXArray(Float(3))]) + XCTAssertEqual(firstValue[0].item(Float.self), 6) + guard case .value(let firstGradient)? = firstGradients["weight"] else { + return XCTFail("missing first gradient") + } + XCTAssertEqual(firstGradient.item(Float.self), 3) + + let (secondValue, secondGradients) = transformed(parameters, [MLXArray(Float(7))]) + XCTAssertEqual(secondValue[0].item(Float.self), 14) + guard case .value(let secondGradient)? = secondGradients["weight"] else { + return XCTFail("missing second gradient") + } + XCTAssertEqual(secondGradient.item(Float.self), 7) + } + + func testNestedValueAndGradRebuildsForChangedParameterTopology() { + let transformed = valueAndGrad { + (parameters: ModuleParameters, arrays: [MLXArray]) -> [MLXArray] in + let weight: MLXArray + if case .value(let value)? = parameters["weight"] { + weight = value + } else if case .dictionary(let layer)? = parameters["layer"], + case .value(let value)? = layer["weight"] + { + weight = value + } else { + XCTFail("missing weight") + return [] + } + return [weight * arrays[0]] + } + + var flat = ModuleParameters() + flat["weight"] = .value(MLXArray(Float(2))) + let (_, flatGradients) = transformed(flat, [MLXArray(Float(3))]) + guard case .value(let flatGradient)? = flatGradients["weight"] else { + return XCTFail("gradient did not preserve flat topology") + } + XCTAssertEqual(flatGradient.item(Float.self), 3) + + var nested = ModuleParameters() + nested["layer"] = .dictionary(["weight": .value(MLXArray(Float(5)))]) + let (nestedValue, nestedGradients) = transformed(nested, [MLXArray(Float(4))]) + XCTAssertEqual(nestedValue[0].item(Float.self), 20) + guard case .dictionary(let layer)? = nestedGradients["layer"], + case .value(let nestedGradient)? = layer["weight"] + else { + return XCTFail("gradient did not preserve nested topology") + } + XCTAssertEqual(nestedGradient.item(Float.self), 4) + } + func testValueAndGradNested() { // valueAndGrad on a nested structure, e.g. parameters. // this isn't a real model but can exercise the @@ -112,6 +186,28 @@ class TransformTests: XCTestCase { } } + func testModelValueAndGradRestoresParametersBeforeEvaluation() { + final class ScaleModel: Module { + let weight = MLXArray(Float(2)) + } + + let model = ScaleModel() + let transformed = valueAndGrad(model: model) { model, arrays in + [(model.weight * arrays[0]).square()] + } + + let (values, gradients) = transformed(model, [MLXArray(Float(3))]) + + // Model state must not retain the temporary VJP primals after the callback returns. + eval(model) + XCTAssertEqual(model.weight.item(Float.self), 2) + XCTAssertEqual(values[0].item(Float.self), 36) + guard case .value(let weightGradient)? = gradients["weight"] else { + return XCTFail("missing weight gradient") + } + XCTAssertEqual(weightGradient.item(Float.self), 36) + } + func testCompile() { func f(inputs: [MLXArray]) -> [MLXArray] { [square(inputs[0] * inputs[1])] @@ -217,6 +313,68 @@ class TransformTests: XCTestCase { XCTAssertEqual(state.o!.item(Float.self), -8) } + func testCompiledStateFreshAcrossRecompile() { + let state = CompileTestState() + + let compiled = compile(inputs: [state]) { inputs in + [inputs[0] + state.y] + } + + let first = compiled([MLXArray(Float(5))]) + XCTAssertEqual(first[0].item(Float.self), 7) + + state.y = MLXArray(Float(10)) + let retraced = compiled([MLXArray([Float(1), 2, 3])]) + XCTAssertEqual(retraced[0].asArray(Float.self), [11, 12, 13]) + + state.y = MLXArray(Float(100)) + let originalShape = compiled([MLXArray(Float(5))]) + XCTAssertEqual(originalShape[0].item(Float.self), 105) + } + + func testCompiledClosureDoesNotRetainItself() { + final class LifetimeToken {} + + weak var weakToken: LifetimeToken? + var compiledFunction: (@Sendable ([MLXArray]) -> [MLXArray])? + + do { + let token = LifetimeToken() + weakToken = token + compiledFunction = compile { inputs in + _ = token + return [inputs[0] + 1] + } + XCTAssertEqual(compiledFunction?([MLXArray(Float(1))])[0].item(Float.self), 2) + } + + XCTAssertNotNil(weakToken) + compiledFunction = nil + XCTAssertNil(weakToken) + } + + func testCompiledClosureTracksCompileModeChanges() { + var traceCount = 0 + let compiled = compile { (x: MLXArray) -> MLXArray in + traceCount += 1 + return x + 1 + } + + _ = compiled(MLXArray(Float(1))) + _ = compiled(MLXArray(Float(2))) + XCTAssertEqual(traceCount, 1) + + compile(enable: false) + defer { compile(enable: true) } + _ = compiled(MLXArray(Float(3))) + _ = compiled(MLXArray(Float(4))) + XCTAssertEqual(traceCount, 3) + + compile(enable: true) + _ = compiled(MLXArray(Float(5))) + XCTAssertEqual(traceCount, 3) + } + func testCompiledRandom() { func f(_ bias: MLXArray) -> MLXArray { MLXRandom.uniform(0 ..< 1, [4]) + bias @@ -371,13 +529,17 @@ class TransformTests: XCTestCase { } } + // Exercise concurrent calls to one cached wrapper, rather than constructing an + // independent compiled function in every task. + let compiledSwiglu = compileSwiglu() + await withTaskGroup(of: Void.self) { group in for _ in 0 ..< 10 { group.addTask { withRandomState(.init()) { let x = MLXRandom.normal([1024, 1024]) let y = MLXRandom.normal(x.shape) - let _ = compileSwiglu()(x, y) + let _ = compiledSwiglu(x, y) } } }