Replace Objective-C bridge with Swift ONNX runtime backend

2026-06-09 09:38:43 -04:00
parent b7fb7f7dea
commit 2af97cc6bf
8 changed files with 1648 additions and 124 deletions
@@ -0,0 +1,166 @@
+@preconcurrency import COnnxRuntime
+import Foundation
+
+internal enum ORTNative {
+    nonisolated(unsafe) static let apiBase = OrtGetApiBase()
+    nonisolated(unsafe) static let api = apiBase!.pointee.GetApi(UInt32(ORT_API_VERSION))!
+    nonisolated(unsafe) static let trainingApi = api.pointee.GetTrainingApi?(UInt32(ORT_API_VERSION))
+
+    static func check(_ status: OpaquePointer?, action: String) throws {
+        guard let status else {
+            return
+        }
+
+        let message = api.pointee.GetErrorMessage?(status).map { String(cString: $0) } ?? "Unknown ONNX Runtime error"
+        api.pointee.ReleaseStatus?(status)
+        throw ORTError.runtimeFailure(action: action, message: message)
+    }
+
+    static func withAllocator<Result>(_ body: (UnsafeMutablePointer<OrtAllocator>) throws -> Result) throws -> Result {
+        var allocator: UnsafeMutablePointer<OrtAllocator>?
+        try check(api.pointee.GetAllocatorWithDefaultOptions?(&allocator), action: "GetAllocatorWithDefaultOptions")
+        guard let allocator else {
+            throw ORTError.runtimeFailure(action: "GetAllocatorWithDefaultOptions", message: "Allocator was not returned")
+        }
+        return try body(allocator)
+    }
+
+    static func copyAllocatedString(
+        _ stringPointer: UnsafeMutablePointer<CChar>?,
+        allocator: UnsafeMutablePointer<OrtAllocator>
+    ) throws -> String {
+        guard let stringPointer else {
+            throw ORTError.runtimeFailure(action: "copyAllocatedString", message: "Expected allocated string")
+        }
+
+        defer {
+            _ = api.pointee.AllocatorFree?(allocator, stringPointer)
+        }
+
+        return String(cString: stringPointer)
+    }
+
+    static func tensorElementCount(shape: [Int]) throws -> Int {
+        guard shape.allSatisfy({ $0 >= 0 }) else {
+            throw ORTError.invalidArgument("Tensor shape contains a negative dimension: \(shape)")
+        }
+
+        return shape.reduce(1, *)
+    }
+
+    static func elementByteCount(for type: ORTTensorElementDataType) throws -> Int {
+        switch type {
+        case .float:
+            return MemoryLayout<Float>.stride
+        case .int8:
+            return MemoryLayout<Int8>.stride
+        case .uInt8:
+            return MemoryLayout<UInt8>.stride
+        case .int32:
+            return MemoryLayout<Int32>.stride
+        case .uInt32:
+            return MemoryLayout<UInt32>.stride
+        case .int64:
+            return MemoryLayout<Int64>.stride
+        case .uInt64:
+            return MemoryLayout<UInt64>.stride
+        case .string:
+            throw ORTError.invalidArgument("String tensors do not have a fixed byte width")
+        case .undefined:
+            throw ORTError.invalidArgument("Undefined tensor element type is not supported")
+        }
+    }
+
+    static func makeShapeBuffer(_ shape: [Int]) -> [Int64] {
+        shape.map(Int64.init)
+    }
+
+    static func withCStringArray<Result>(
+        _ strings: [String],
+        _ body: ([UnsafePointer<CChar>?]) throws -> Result
+    ) rethrows -> Result {
+        let storage = strings.map { Array($0.utf8CString) }
+        let pointers = storage.map { $0.withUnsafeBufferPointer { $0.baseAddress } }
+        return try body(pointers)
+    }
+
+    static func fetchTensorInfo(from value: OpaquePointer) throws -> ORTTensorTypeAndShapeInfo {
+        var info: OpaquePointer?
+        try check(api.pointee.GetTensorTypeAndShape?(value, &info), action: "GetTensorTypeAndShape")
+        guard let info else {
+            throw ORTError.runtimeFailure(action: "GetTensorTypeAndShape", message: "Tensor info was not returned")
+        }
+
+        defer {
+            api.pointee.ReleaseTensorTypeAndShapeInfo?(info)
+        }
+
+        var elementType = ONNXTensorElementDataType(rawValue: 0)
+        try check(api.pointee.GetTensorElementType?(info, &elementType), action: "GetTensorElementType")
+
+        var dimensionsCount: Int = 0
+        try check(api.pointee.GetDimensionsCount?(info, &dimensionsCount), action: "GetDimensionsCount")
+
+        var dimensions = Array(repeating: Int64(0), count: dimensionsCount)
+        if dimensionsCount > 0 {
+            try check(
+                api.pointee.GetDimensions?(info, &dimensions, dimensionsCount),
+                action: "GetDimensions"
+            )
+        }
+
+        return ORTTensorTypeAndShapeInfo(
+            elementType: try ORTTensorElementDataType(native: elementType),
+            shape: dimensions.map(Int.init)
+        )
+    }
+
+    static func fetchValueTypeInfo(from value: OpaquePointer) throws -> ORTValueTypeInfo {
+        var typeInfo: OpaquePointer?
+        try check(api.pointee.GetTypeInfo?(value, &typeInfo), action: "GetTypeInfo")
+        guard let typeInfo else {
+            throw ORTError.runtimeFailure(action: "GetTypeInfo", message: "Type info was not returned")
+        }
+
+        defer {
+            api.pointee.ReleaseTypeInfo?(typeInfo)
+        }
+
+        var onnxType = ONNXType(rawValue: 0)
+        try check(api.pointee.GetOnnxTypeFromTypeInfo?(typeInfo, &onnxType), action: "GetOnnxTypeFromTypeInfo")
+
+        var tensorInfoPointer: OpaquePointer?
+        try check(api.pointee.CastTypeInfoToTensorInfo?(typeInfo, &tensorInfoPointer), action: "CastTypeInfoToTensorInfo")
+
+        let tensorInfo = try tensorInfoPointer.map { pointer in
+            try fetchTensorInfoView(from: pointer)
+        }
+
+        return ORTValueTypeInfo(
+            type: try ORTValueType(native: onnxType),
+            tensorTypeAndShapeInfo: tensorInfo
+        )
+    }
+
+    static func fetchTensorInfoView(from info: OpaquePointer) throws -> ORTTensorTypeAndShapeInfo {
+        var elementType = ONNXTensorElementDataType(rawValue: 0)
+        try check(api.pointee.GetTensorElementType?(info, &elementType), action: "GetTensorElementType")
+
+        var dimensionsCount: Int = 0
+        try check(api.pointee.GetDimensionsCount?(info, &dimensionsCount), action: "GetDimensionsCount")
+
+        var dimensions = Array(repeating: Int64(0), count: dimensionsCount)
+        if dimensionsCount > 0 {
+            try check(api.pointee.GetDimensions?(info, &dimensions, dimensionsCount), action: "GetDimensions")
+        }
+
+        return ORTTensorTypeAndShapeInfo(
+            elementType: try ORTTensorElementDataType(native: elementType),
+            shape: dimensions.map(Int.init)
+        )
+    }
+
+    static func boolString(_ value: Bool) -> String {
+        value ? "1" : "0"
+    }
+}