Adds bufferSize API to be able to derive memory occupied by a single model

Source/Cmlx/include-framework/mlx-c-memory.h

@@ -30,6 +30,7 @@ extern "C" {
 
 int mlx_clear_cache(void);
 int mlx_get_active_memory(size_t* res);
+int mlx_get_array_buffer_size(size_t* res, const mlx_vector_array arrays);
 int mlx_get_cache_memory(size_t* res);
 int mlx_get_memory_limit(size_t* res);
 int mlx_get_peak_memory(size_t* res);

Source/Cmlx/include-framework/mlx-memory.h

@@ -4,8 +4,10 @@
 #pragma once
 
 #include <cstdlib>
+#include <vector>
 
 #include <Cmlx/mlx-api.h>
+#include <Cmlx/mlx-array.h>
 
 namespace mlx::core {
 
@@ -16,6 +18,14 @@ namespace mlx::core {
  * */
 MLX_API size_t get_active_memory();
 
+/* Get the size of the buffers backing the given arrays in bytes.
+ *
+ * Each unique buffer is counted once. The full allocator size of each buffer
+ * is used, which can exceed the logical size of its arrays. The arrays must be
+ * evaluated before calling this function.
+ * */
+MLX_API size_t get_array_buffer_size(const std::vector<array>& arrays);
+
 /* Get the peak amount of used memory in bytes.
  *
  * The maximum memory used recorded from the beginning of the program

Source/Cmlx/include/mlx/c/memory.h

@@ -30,6 +30,7 @@ extern "C" {
 
 int mlx_clear_cache(void);
 int mlx_get_active_memory(size_t* res);
+int mlx_get_array_buffer_size(size_t* res, const mlx_vector_array arrays);
 int mlx_get_cache_memory(size_t* res);
 int mlx_get_memory_limit(size_t* res);
 int mlx_get_peak_memory(size_t* res);

Source/MLX/Documentation.docc/Articles/running-on-ios.md

@@ -64,6 +64,20 @@ print(startMemory.delta(endMemory).description)
 It may be interesting to print the current memory statistics during evaluation if
 you want to see performance over time.
 
+These statistics are process wide. If an application holds several models
+e.g. a language model, an embedding model and an image model -- use
+``Memory/bufferSize(of:)-(some_Collection<MLXArray>)`` to attribute memory to
+one of them:
+
+```swift
+let weights = model.parameters().flattened().map { $0.1 }
+eval(weights)
+print("model: \(Memory.bufferSize(of: weights) / 1024)K")
+```
+
+Each unique buffer is counted once, so tied weights or arrays that view into a
+larger buffer are not double counted.
+
 Decreasing the cache limit to 0 will result in decreased performance due to the
 lack of buffer reuse, but it will also result in smaller memory use.
 Tune this value for your needs.

Source/MLX/Memory.swift

@@ -133,6 +133,14 @@ public enum Memory {
         /// See ``Memory/peakMemory``.
         public var peakMemory: Int
 
+        /// Create a snapshot -- typically produced by ``Memory/snapshot()`` but
+        /// this allows recorded or synthesized values to be used as well.
+        public init(activeMemory: Int, cacheMemory: Int, peakMemory: Int) {
+            self.activeMemory = activeMemory
+            self.cacheMemory = cacheMemory
+            self.peakMemory = peakMemory
+        }
+
         /// Compute the difference between two snapshots:
         ///
         /// ```swift
@@ -196,6 +204,53 @@ public enum Memory {
         return result
     }
 
+    /// Get the size in bytes of the buffers backing the given arrays.
+    ///
+    /// Unlike ``activeMemory``, a process wide counter, this measures a
+    /// specific set of arrays and can be used to attribute memory to e.g. the
+    /// weights of one model:
+    ///
+    /// ```swift
+    /// let weights = model.parameters().flattened().map { $0.1 }
+    /// eval(weights)
+    /// let bytes = Memory.bufferSize(of: weights)
+    /// ```
+    ///
+    /// Each unique buffer is counted once, so arrays that share storage (views,
+    /// tied weights, slices) do not inflate the total. The full allocator size
+    /// of each buffer is reported, which can be slightly larger than the sum of
+    /// `MLXArray/nbytes` of its arrays.
+    ///
+    /// - Important: The arrays must be evaluated; see `eval(_:)`. Unevaluated
+    /// arrays have no buffer to measure and produce an MLX error, see
+    /// ``withError(_:)-6g4wn``.
+    ///
+    /// - Parameter arrays: evaluated arrays to measure
+    /// - Returns: the size in bytes of the unique buffers backing `arrays`
+    ///
+    /// ### See Also
+    /// - ``activeMemory``
+    /// - ``snapshot()``
+    public static func bufferSize(of arrays: some Collection<MLXArray>) -> Int {
+        let vector_array = new_mlx_vector_array(arrays)
+        defer { mlx_vector_array_free(vector_array) }
+
+        var result: size_t = 0
+        mlx_get_array_buffer_size(&result, vector_array)
+        return result
+    }
+
+    /// Get the size in bytes of the buffers backing the given arrays.
+    ///
+    /// A variadic convenience for the `Collection` variant:
+    ///
+    /// ```swift
+    /// let bytes = Memory.bufferSize(of: keys, values)
+    /// ```
+    public static func bufferSize(of arrays: MLXArray...) -> Int {
+        bufferSize(of: arrays)
+    }
+
     /// Get the peak amount of active memory in bytes.
     ///
     /// The maximum memory used is recorded from the beginning of the program

Tests/MLXTests/MemoryTests.swift

@@ -12,4 +12,36 @@ class MemoryTests: XCTestCase {
             print(x * x)
         }
     }
+
+    func testBufferSize() {
+        let a = MLXArray([1, 2, 3, 4] as [Float])
+        let view = a[..<1]
+        eval(a, view)
+
+        let size = Memory.bufferSize(of: a)
+
+        // the allocator size is at least the logical size
+        XCTAssertGreaterThanOrEqual(size, a.nbytes)
+
+        // no arrays, no memory
+        XCTAssertEqual(Memory.bufferSize(of: [MLXArray]()), 0)
+
+        // each unique buffer is counted once -- the view shares a's buffer
+        XCTAssertEqual(Memory.bufferSize(of: view), size)
+        XCTAssertEqual(Memory.bufferSize(of: a, a), size)
+        XCTAssertEqual(Memory.bufferSize(of: a, view), size)
+
+        let b = MLXArray([5, 6] as [Float])
+        eval(b)
+        XCTAssertEqual(Memory.bufferSize(of: a, b), size + Memory.bufferSize(of: b))
+    }
+
+    func testBufferSizeUnevaluated() throws {
+        let a = MLXArray([1, 2, 3, 4] as [Float])
+        eval(a)
+
+        // an unevaluated array has no buffer to measure
+        let lazy = a + 1
+        XCTAssertThrowsError(try withError { _ in Memory.bufferSize(of: lazy) })
+    }
 }