Pull-based error boundary: typed MLXError, native throws, poison propagation (#270)

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

@@ -12,6 +12,10 @@ extern "C" {
  */
 /**@{*/
 
+/* -------------------------------------------------------------------------
+ * Legacy push-based handler API (unchanged, retained for source compat).
+ * ---------------------------------------------------------------------- */
+
 typedef void (*mlx_error_handler_func)(const char* msg, void* data);
 
 /**
@@ -32,6 +36,63 @@ void _mlx_error(const char* file, const int line, const char* fmt, ...);
  */
 #define mlx_error(...) _mlx_error(__FILE__, __LINE__, __VA_ARGS__)
 
+/* -------------------------------------------------------------------------
+ * New pull-based, structured, thread-local error state.
+ *
+ * Every generated binding stores the error here (in addition to invoking the
+ * legacy handler) before returning a non-zero status. A language binding that
+ * checks the status can then pull a *typed* error out of this thread's slot and
+ * surface it natively (e.g. a Swift `throw`) instead of relying on a global
+ * callback that has no frame to throw from.
+ * ---------------------------------------------------------------------- */
+
+/**
+ * Error classification, derived from the C++ exception type (and, until
+ * mx::core grows typed exceptions, from message inspection for OOM/IO).
+ */
+typedef enum mlx_error_code_ {
+  MLX_ERROR_NONE = 0,
+  MLX_ERROR_INVALID_ARGUMENT, /* std::invalid_argument: shape/dtype/axis */
+  MLX_ERROR_OUT_OF_RANGE,     /* std::out_of_range: indexing */
+  MLX_ERROR_OUT_OF_MEMORY,    /* std::bad_alloc / Metal allocation failure */
+  MLX_ERROR_IO,               /* load/save/format failures */
+  MLX_ERROR_RUNTIME,          /* std::runtime_error and other std::exception */
+  MLX_ERROR_UNKNOWN           /* catch (...) : non-std throw */
+} mlx_error_code;
+
+/**
+ * Code of the most recent error on the *calling thread*, or MLX_ERROR_NONE.
+ * Does not clear the state.
+ */
+mlx_error_code mlx_last_error_code(void);
+
+/**
+ * Message of the most recent error on the calling thread, or "".
+ * The returned pointer is owned by MLX and remains valid until the next failing
+ * MLX call on this thread or a call to mlx_clear_last_error().
+ */
+const char* mlx_last_error_message(void);
+
+/**
+ * Clear the calling thread's error state. Bindings call this after consuming
+ * an error so a subsequent successful call is not misread as a failure.
+ */
+void mlx_clear_last_error(void);
+
+/**
+ * Store a classified error for the calling thread and invoke the legacy
+ * handler. Used by generated bindings; also usable directly.
+ * Macro variant passes __FILE__/__LINE__ like mlx_error().
+ */
+void _mlx_error_with_code(
+    mlx_error_code code,
+    const char* file,
+    const int line,
+    const char* fmt,
+    ...);
+#define mlx_error_with_code(code, ...) \
+  _mlx_error_with_code(code, __FILE__, __LINE__, __VA_ARGS__)
+
 /**@}*/
 
 #ifdef __cplusplus

Source/MLX/Cmlx+Error.swift

@@ -0,0 +1,97 @@
+// Copyright © 2024 Apple Inc.
+
+import Cmlx
+import Foundation
+
+/// A structured error raised by the MLX backend.
+///
+/// Replaces the message-only `MLXError.caught(String)`. The `code` lets callers
+/// react programmatically — e.g. evict the GPU cache and retry on ``Code/outOfMemory``,
+/// or surface a validation message on ``Code/invalidArgument`` — while `message`
+/// preserves the original C++ `what()` text (with the originating `file:line`).
+public struct MLXError: LocalizedError, Sendable, Equatable, CustomStringConvertible {
+
+    /// Classification of the failure, mirroring `mlx_error_code` in mlx-c.
+    public enum Code: Sendable, Equatable {
+        /// Shape / dtype / axis mismatch (`std::invalid_argument`). Usually a
+        /// programmer error, but recoverable when inputs are externally sourced.
+        case invalidArgument
+        /// Out-of-range index (`std::out_of_range`).
+        case outOfRange
+        /// Allocation failure, including Metal `[metal::malloc]` (`std::bad_alloc`).
+        /// Typically recoverable: free buffers / shrink the batch and retry.
+        case outOfMemory
+        /// Load / save / format failure — corrupt safetensors, missing file, bad GGUF.
+        case io
+        /// Any other `std::exception` / `std::runtime_error`.
+        case runtime
+        /// A non-`std::exception` throw crossed the boundary (`catch (...)`).
+        case unknown
+
+        init(_ raw: mlx_error_code) {
+            switch raw {
+            case MLX_ERROR_INVALID_ARGUMENT: self = .invalidArgument
+            case MLX_ERROR_OUT_OF_RANGE: self = .outOfRange
+            case MLX_ERROR_OUT_OF_MEMORY: self = .outOfMemory
+            case MLX_ERROR_IO: self = .io
+            case MLX_ERROR_RUNTIME: self = .runtime
+            default: self = .unknown
+            }
+        }
+    }
+
+    public let code: Code
+    public let message: String
+
+    public var errorDescription: String? { description }
+    public var description: String { "MLX \(code): \(message)" }
+}
+
+/// Install the pull-mode barrier: a no-op *global* handler so the legacy push
+/// path never exits the process. Errors then flow exclusively through status
+/// codes and the thread-local slot into ``checkStatus``. Without this call the
+/// historical behaviour (push handler -> fatalError/exit) is fully preserved,
+/// so the pull model is strictly opt-in.
+public func installPullErrorBarrier() {
+    setErrorHandler({ _, _ in }, data: nil, dtor: nil)
+}
+
+/// Consume the calling thread's mlx-c error slot and throw it natively.
+///
+/// This is the pull side of the exception boundary. mlx-c stores classified
+/// error state in thread-local storage *before* returning a non-zero status;
+/// here — where a real Swift frame exists — we read that state and `throw`,
+/// which a C callback never could. Because the slot is per-thread, an error
+/// raised while evaluating on a Metal completion thread or a `DispatchQueue`
+/// worker is reported on *that* thread and surfaces through the status code to
+/// whoever synchronizes on it, closing the task-local gap in the old design.
+///
+/// - Parameter status: the `Int32` returned by any `mlx_*` C function.
+@inline(__always)
+func checkStatus(_ status: Int32, file: StaticString = #fileID, line: UInt = #line) throws {
+    guard status != 0 else { return }
+
+    let code = MLXError.Code(mlx_last_error_code())
+    let message =
+        mlx_last_error_message().map { String(cString: $0) } ?? "unknown MLX error"
+    mlx_clear_last_error()
+
+    throw MLXError(code: code, message: message)
+}
+
+/// Non-throwing bridge for code paths that are not yet `throws` (e.g. operator
+/// overloads). Records the error into the given ``MLXArray`` as a *poison*
+/// value so the first use — including the next `try eval()` touching it —
+/// rethrows the original error, killing the zombie-value cascade. See
+/// ``MLXArray/poison(_:)``.
+@inline(__always)
+func checkStatus(_ status: Int32, poisoning array: MLXArray) {
+    guard status != 0 else { return }
+
+    let code = MLXError.Code(mlx_last_error_code())
+    let message =
+        mlx_last_error_message().map { String(cString: $0) } ?? "unknown MLX error"
+    mlx_clear_last_error()
+
+    array.poison(MLXError(code: code, message: message))
+}

Source/MLX/ErrorHandler.swift

@@ -215,16 +215,10 @@ public func withError<R>(_ body: () async throws -> R) async throws -> R {
     try await errorHandler.withError({ _ in try await body() })
 }
 
-/// Error type for caught errors during ``withError(_:)-6g4wn``.
-public enum MLXError: LocalizedError, Sendable, Equatable {
-    case caught(String)
-
-    public var errorDescription: String? {
-        switch self {
-        case .caught(let message): "MLX Error: \(message)"
-        }
-    }
-}
+// Note: `MLXError` is now the structured error type declared in
+// Cmlx+Error.swift, carrying a typed `code` classified by the mlx-c
+// exception boundary. The former `enum MLXError { case caught(String) }`
+// is superseded; `withError` below now surfaces the structured type.
 
 /// Boxed error type usable with ``withError(_:)-2wfiu``.
 ///
@@ -369,7 +363,7 @@ private final class ErrorHandler: @unchecked Sendable {
 
         @Sendable
         func errorHandler(_ message: String) {
-            errorBox.firstError = MLXError.caught(message)
+            errorBox.firstError = MLXError(code: .runtime, message: message)
         }
 
         return try withErrorHandler(errorHandler) {
@@ -384,7 +378,7 @@ private final class ErrorHandler: @unchecked Sendable {
 
         @Sendable
         func errorHandler(_ message: String) {
-            errorBox.firstError = MLXError.caught(message)
+            errorBox.firstError = MLXError(code: .runtime, message: message)
         }
 
         return try await withErrorHandler(errorHandler) {

Source/MLX/IO.swift

@@ -133,8 +133,8 @@ public func loadArrays(url: URL, stream: StreamOrDevice = .cpu) throws -> [Strin
         defer { mlx_map_string_to_array_free(r0) }
         defer { mlx_map_string_to_string_free(r1) }
 
-        _ = try withError {
-            mlx_load_safetensors(&r0, &r1, path.cString(using: .utf8), stream.ctx)
+        try withErrorHandler({ _ in /* suppress push path; pull below */ }) {
+            try checkStatus(mlx_load_safetensors(&r0, &r1, path.cString(using: .utf8), stream.ctx))
         }
 
         return mlx_map_array_values(r0)

Source/MLX/MLXArray+Ops.swift

@@ -40,10 +40,19 @@ extension MLXArray {
     /// - <doc:arithmetic>
     /// - ``add(_:_:stream:)``
     public static func + (lhs: MLXArray, rhs: MLXArray) -> MLXArray {
+        // poison propagation: a failed upstream op rides inside the value so
+        // the next throwing sync point rethrows the original, first error
+        if let error = lhs.poisonError ?? rhs.poisonError {
+            let result = MLXArray(mlx_array_new())
+            result.poison(error)
+            return result
+        }
         let s = StreamOrDevice.default
         var result = mlx_array_new()
-        mlx_add(&result, lhs.ctx, rhs.ctx, s.ctx)
-        return MLXArray(result)
+        let status = mlx_add(&result, lhs.ctx, rhs.ctx, s.ctx)
+        let array = MLXArray(result)
+        checkStatus(status, poisoning: array)
+        return array
     }
 
     /// Element-wise addition.

Source/MLX/MLXArray+Poison.swift

@@ -0,0 +1,56 @@
+// Copyright © 2024 Apple Inc.
+
+import Cmlx
+import Foundation
+
+/// Poison propagation for the non-throwing operator API.
+///
+/// Operators like `a + b` route through `operator +` and cannot `throw`. Rather
+/// than return a valid-looking `MLXArray` backed by an empty `mlx_array` (the
+/// current behaviour, which produces cascading secondary errors), we attach the
+/// original ``MLXError`` to the result. The error then rides *inside the value*:
+/// the next `try eval()`, `try item()`, or `try asArray()` that touches it
+/// rethrows the original, first error — preserving attribution and eliminating
+/// the zombie cascade.
+///
+/// Storage is a side-table keyed by object identity so `MLXArray`'s layout and
+/// `Cmlx` ownership are untouched.
+extension MLXArray {
+
+    private static let poisonTable = PoisonTable()
+
+    /// Attach an error to this array. Idempotent — the first error wins,
+    /// matching the "first error" semantics of the old `ErrorBox`.
+    func poison(_ error: MLXError) {
+        Self.poisonTable.set(self, error)
+    }
+
+    /// The attached error, if this array (or the op that produced it) failed.
+    public var poisonError: MLXError? {
+        Self.poisonTable.get(self)
+    }
+
+    /// Rethrow the attached error if present. Called by the throwing sync points
+    /// before handing the array to the backend.
+    func throwIfPoisoned() throws {
+        if let error = poisonError { throw error }
+    }
+}
+
+/// Thread-safe identity-keyed side table. `NSMapTable` with weak keys releases
+/// entries when the `MLXArray` is deallocated, so poison never leaks.
+private final class PoisonTable: @unchecked Sendable {
+    private let lock = NSLock()
+    private let table = NSMapTable<MLXArray, NSError>.weakToStrongObjects()
+
+    func set(_ key: MLXArray, _ error: MLXError) {
+        lock.withLock {
+            guard table.object(forKey: key) == nil else { return } // first wins
+            table.setObject(error as NSError, forKey: key)
+        }
+    }
+
+    func get(_ key: MLXArray) -> MLXError? {
+        lock.withLock { table.object(forKey: key) as? MLXError }
+    }
+}

Source/MLX/Ops.swift

@@ -43,9 +43,18 @@ public func add(
     _ a: some ScalarOrArray, _ b: some ScalarOrArray, stream: StreamOrDevice = .default
 ) -> MLXArray {
     let (a, b) = toArrays(a, b)
+    // poison propagation: a failed upstream op rides inside the value so the
+    // next throwing sync point rethrows the original, first error
+    if let error = a.poisonError ?? b.poisonError {
+        let result = MLXArray(mlx_array_new())
+        result.poison(error)
+        return result
+    }
     var result = mlx_array_new()
-    mlx_add(&result, a.ctx, b.ctx, stream.ctx)
-    return MLXArray(result)
+    let status = mlx_add(&result, a.ctx, b.ctx, stream.ctx)
+    let array = MLXArray(result)
+    checkStatus(status, poisoning: array)
+    return array
 }
 
 @available(*, deprecated, renamed: "addMM(_:_:_:alpha:beta:stream:)")