Optimize memory usage

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"],

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)
+        }
+    }
+}

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<MLXArray>) -> 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<mlx_array>.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)
     }
 }
 

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<Int>, 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)
 }
 

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<Result>(
+        _ body: (UnsafePointer<Int32>?, 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<Int32>.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

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
                     }
                 }

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)
     }