usearch 2.21.3 → 2.25.0

package/numkong/swift/MaxSim.swift

@@ -0,0 +1,151 @@
+//  MaxSim.swift
+//  NumKong
+//
+//  Created by Ash Vardanian on March 14, 2026.
+//
+
+import CNumKong
+
+// MARK: - MaxSim Protocol
+
+/// Element type that supports MaxSim (late-interaction) scoring with packed representations.
+public protocol NumKongMaxSimElement {
+    associatedtype MaxSimOutput
+    static func _nk_maxsim_packed_size(_ vectorCount: Int, _ depth: Int) -> Int
+    static func _nk_maxsim_pack(
+        _ vectors: UnsafePointer<Self>, _ vectorCount: Int, _ depth: Int, _ stride: Int,
+        _ packed: UnsafeMutableRawPointer)
+    static func _nk_maxsim_packed(
+        _ queryPacked: UnsafeRawPointer, _ docPacked: UnsafeRawPointer, _ queryCount: Int, _ docCount: Int,
+        _ depth: Int, _ result: UnsafeMutablePointer<MaxSimOutput>)
+}
+
+// MARK: - MaxSimPackedMatrix
+
+/// Owns a packed representation of multi-vector embeddings for MaxSim scoring.
+public final class MaxSimPackedMatrix<Element: NumKongMaxSimElement>: @unchecked Sendable {
+    public let vectorCount: Int
+    public let depth: Int
+    public let byteCount: Int
+
+    @usableFromInline
+    let rawPointer: UnsafeMutableRawPointer
+
+    @usableFromInline
+    init(vectorCount: Int, depth: Int, byteCount: Int, rawPointer: UnsafeMutableRawPointer) {
+        self.vectorCount = vectorCount
+        self.depth = depth
+        self.byteCount = byteCount
+        self.rawPointer = rawPointer
+    }
+
+    deinit {
+        rawPointer.deallocate()
+    }
+
+    /// Packs a matrix view into the MaxSim-optimized layout.
+    public convenience init(packing matrix: MatrixView<Element>) throws {
+        guard matrix.rows > 0 && matrix.cols > 0 else {
+            throw NumKongMatrixError.invalidDimensions
+        }
+        let bytes = Element._nk_maxsim_packed_size(matrix.rows, matrix.cols)
+        guard bytes > 0 else { throw NumKongMatrixError.packedBufferTooSmall }
+        let ptr = UnsafeMutableRawPointer.allocate(byteCount: bytes, alignment: 64)
+        Element._nk_maxsim_pack(matrix.baseAddress, matrix.rows, matrix.cols, matrix.rowStrideBytes, ptr)
+        self.init(vectorCount: matrix.rows, depth: matrix.cols, byteCount: bytes, rawPointer: ptr)
+    }
+
+    /// Computes the MaxSim score between this (query) and a document's packed matrix.
+    public func score(_ document: MaxSimPackedMatrix<Element>) -> Element.MaxSimOutput {
+        let ptr = UnsafeMutablePointer<Element.MaxSimOutput>.allocate(capacity: 1)
+        let raw = UnsafeMutableRawPointer(ptr)
+        raw.initializeMemory(as: UInt8.self, repeating: 0, count: MemoryLayout<Element.MaxSimOutput>.size)
+        defer { ptr.deallocate() }
+        Element._nk_maxsim_packed(
+            UnsafeRawPointer(rawPointer),
+            UnsafeRawPointer(document.rawPointer),
+            vectorCount,
+            document.vectorCount,
+            depth,
+            ptr
+        )
+        return ptr.pointee
+    }
+}
+
+// MARK: - Float32 MaxSim Conformance
+
+extension Float32: NumKongMaxSimElement {
+    public typealias MaxSimOutput = Float64
+
+    public static func _nk_maxsim_packed_size(_ vectorCount: Int, _ depth: Int) -> Int {
+        Int(nk_maxsim_packed_size_f32(UInt64(vectorCount), UInt64(depth)))
+    }
+
+    public static func _nk_maxsim_pack(
+        _ vectors: UnsafePointer<Float32>, _ vectorCount: Int, _ depth: Int, _ stride: Int,
+        _ packed: UnsafeMutableRawPointer
+    ) {
+        nk_maxsim_pack_f32(vectors, UInt64(vectorCount), UInt64(depth), UInt64(stride), packed)
+    }
+
+    public static func _nk_maxsim_packed(
+        _ queryPacked: UnsafeRawPointer, _ docPacked: UnsafeRawPointer, _ queryCount: Int, _ docCount: Int,
+        _ depth: Int, _ result: UnsafeMutablePointer<Float64>
+    ) {
+        nk_maxsim_packed_f32(queryPacked, docPacked, UInt64(queryCount), UInt64(docCount), UInt64(depth), result)
+    }
+}
+
+// MARK: - BFloat16 MaxSim Conformance
+
+extension BFloat16: NumKongMaxSimElement {
+    public typealias MaxSimOutput = Float32
+
+    public static func _nk_maxsim_packed_size(_ vectorCount: Int, _ depth: Int) -> Int {
+        Int(nk_maxsim_packed_size_bf16(UInt64(vectorCount), UInt64(depth)))
+    }
+
+    public static func _nk_maxsim_pack(
+        _ vectors: UnsafePointer<BFloat16>, _ vectorCount: Int, _ depth: Int, _ stride: Int,
+        _ packed: UnsafeMutableRawPointer
+    ) {
+        let cPtr = UnsafeRawPointer(vectors).assumingMemoryBound(to: nk_bf16_t.self)
+        nk_maxsim_pack_bf16(cPtr, UInt64(vectorCount), UInt64(depth), UInt64(stride), packed)
+    }
+
+    public static func _nk_maxsim_packed(
+        _ queryPacked: UnsafeRawPointer, _ docPacked: UnsafeRawPointer, _ queryCount: Int, _ docCount: Int,
+        _ depth: Int, _ result: UnsafeMutablePointer<Float32>
+    ) {
+        nk_maxsim_packed_bf16(queryPacked, docPacked, UInt64(queryCount), UInt64(docCount), UInt64(depth), result)
+    }
+}
+
+// MARK: - Float16 MaxSim Conformance
+
+#if !arch(x86_64)
+@available(macOS 11.0, iOS 14.0, tvOS 14.0, watchOS 7.0, *)
+extension Float16: NumKongMaxSimElement {
+    public typealias MaxSimOutput = Float32
+
+    public static func _nk_maxsim_packed_size(_ vectorCount: Int, _ depth: Int) -> Int {
+        Int(nk_maxsim_packed_size_f16(UInt64(vectorCount), UInt64(depth)))
+    }
+
+    public static func _nk_maxsim_pack(
+        _ vectors: UnsafePointer<Float16>, _ vectorCount: Int, _ depth: Int, _ stride: Int,
+        _ packed: UnsafeMutableRawPointer
+    ) {
+        let cPtr = UnsafeRawPointer(vectors).assumingMemoryBound(to: nk_f16_t.self)
+        nk_maxsim_pack_f16(cPtr, UInt64(vectorCount), UInt64(depth), UInt64(stride), packed)
+    }
+
+    public static func _nk_maxsim_packed(
+        _ queryPacked: UnsafeRawPointer, _ docPacked: UnsafeRawPointer, _ queryCount: Int, _ docCount: Int,
+        _ depth: Int, _ result: UnsafeMutablePointer<Float32>
+    ) {
+        nk_maxsim_packed_f16(queryPacked, docPacked, UInt64(queryCount), UInt64(docCount), UInt64(depth), result)
+    }
+}
+#endif

package/numkong/swift/NumKong.swift

@@ -0,0 +1,283 @@
+//  NumKong.swift
+//  NumKong
+//
+//  Created by Ash Vardanian on March 14, 2026.
+//
+
+import CNumKong
+
+// MARK: - Geospatial Protocols
+
+/// A type that can compute SIMD-accelerated Haversine (great-circle) distances.
+public protocol NumKongHaversine: BinaryFloatingPoint {
+    static func haversine(
+        aLat: UnsafeBufferPointer<Self>,
+        aLon: UnsafeBufferPointer<Self>,
+        bLat: UnsafeBufferPointer<Self>,
+        bLon: UnsafeBufferPointer<Self>,
+        result: UnsafeMutableBufferPointer<Self>
+    ) -> Bool
+
+    static func haversine<A: Sequence, B: Sequence, C: Sequence, D: Sequence>(
+        aLat: A, aLon: B, bLat: C, bLon: D
+    ) -> [Self]?
+    where A.Element == Self, B.Element == Self, C.Element == Self, D.Element == Self
+}
+
+/// A type that can compute SIMD-accelerated Vincenty (ellipsoidal) geodesic distances.
+public protocol NumKongVincenty: BinaryFloatingPoint {
+    static func vincenty(
+        aLat: UnsafeBufferPointer<Self>,
+        aLon: UnsafeBufferPointer<Self>,
+        bLat: UnsafeBufferPointer<Self>,
+        bLon: UnsafeBufferPointer<Self>,
+        result: UnsafeMutableBufferPointer<Self>
+    ) -> Bool
+
+    static func vincenty<A: Sequence, B: Sequence, C: Sequence, D: Sequence>(
+        aLat: A, aLon: B, bLat: C, bLon: D
+    ) -> [Self]?
+    where A.Element == Self, B.Element == Self, C.Element == Self, D.Element == Self
+}
+
+/// Convenience alias for types supporting both Haversine and Vincenty geospatial distances.
+public typealias NumKongGeospatial = NumKongHaversine & NumKongVincenty
+
+extension Float64: NumKongHaversine {
+    @inlinable @inline(__always)
+    public static func haversine(
+        aLat: UnsafeBufferPointer<Float64>,
+        aLon: UnsafeBufferPointer<Float64>,
+        bLat: UnsafeBufferPointer<Float64>,
+        bLon: UnsafeBufferPointer<Float64>,
+        result: UnsafeMutableBufferPointer<Float64>
+    ) -> Bool {
+        let n = aLat.count
+        guard
+            n > 0 && n == aLon.count && n == bLat.count && n == bLon.count
+                && n == result.count
+        else {
+            return false
+        }
+        nk_haversine_f64(
+            aLat.baseAddress!,
+            aLon.baseAddress!,
+            bLat.baseAddress!,
+            bLon.baseAddress!,
+            UInt64(n),
+            result.baseAddress!
+        )
+        return true
+    }
+}
+
+extension Float64: NumKongVincenty {
+    @inlinable @inline(__always)
+    public static func vincenty(
+        aLat: UnsafeBufferPointer<Float64>,
+        aLon: UnsafeBufferPointer<Float64>,
+        bLat: UnsafeBufferPointer<Float64>,
+        bLon: UnsafeBufferPointer<Float64>,
+        result: UnsafeMutableBufferPointer<Float64>
+    ) -> Bool {
+        let n = aLat.count
+        guard
+            n > 0 && n == aLon.count && n == bLat.count && n == bLon.count
+                && n == result.count
+        else {
+            return false
+        }
+        nk_vincenty_f64(
+            aLat.baseAddress!,
+            aLon.baseAddress!,
+            bLat.baseAddress!,
+            bLon.baseAddress!,
+            UInt64(n),
+            result.baseAddress!
+        )
+        return true
+    }
+}
+
+extension Float32: NumKongHaversine {
+    @inlinable @inline(__always)
+    public static func haversine(
+        aLat: UnsafeBufferPointer<Float32>,
+        aLon: UnsafeBufferPointer<Float32>,
+        bLat: UnsafeBufferPointer<Float32>,
+        bLon: UnsafeBufferPointer<Float32>,
+        result: UnsafeMutableBufferPointer<Float32>
+    ) -> Bool {
+        let n = aLat.count
+        guard
+            n > 0 && n == aLon.count && n == bLat.count && n == bLon.count
+                && n == result.count
+        else {
+            return false
+        }
+        nk_haversine_f32(
+            aLat.baseAddress!,
+            aLon.baseAddress!,
+            bLat.baseAddress!,
+            bLon.baseAddress!,
+            UInt64(n),
+            result.baseAddress!
+        )
+        return true
+    }
+}
+
+extension Float32: NumKongVincenty {
+    @inlinable @inline(__always)
+    public static func vincenty(
+        aLat: UnsafeBufferPointer<Float32>,
+        aLon: UnsafeBufferPointer<Float32>,
+        bLat: UnsafeBufferPointer<Float32>,
+        bLon: UnsafeBufferPointer<Float32>,
+        result: UnsafeMutableBufferPointer<Float32>
+    ) -> Bool {
+        let n = aLat.count
+        guard
+            n > 0 && n == aLon.count && n == bLat.count && n == bLon.count
+                && n == result.count
+        else {
+            return false
+        }
+        nk_vincenty_f32(
+            aLat.baseAddress!,
+            aLon.baseAddress!,
+            bLat.baseAddress!,
+            bLon.baseAddress!,
+            UInt64(n),
+            result.baseAddress!
+        )
+        return true
+    }
+}
+
+// MARK: - Geospatial Sequence Extensions
+
+extension Float64 {
+    @inlinable @inline(__always)
+    public static func haversine<A: Sequence, B: Sequence, C: Sequence, D: Sequence>(
+        aLat: A, aLon: B, bLat: C, bLon: D
+    ) -> [Float64]?
+    where A.Element == Float64, B.Element == Float64, C.Element == Float64, D.Element == Float64 {
+        _nkWithGeoQuad(aLat, aLon, bLat, bLon) { a, b, c, d, r, n in
+            nk_haversine_f64(a, b, c, d, UInt64(n), r)
+        }
+    }
+
+    @inlinable @inline(__always)
+    public static func vincenty<A: Sequence, B: Sequence, C: Sequence, D: Sequence>(
+        aLat: A, aLon: B, bLat: C, bLon: D
+    ) -> [Float64]?
+    where A.Element == Float64, B.Element == Float64, C.Element == Float64, D.Element == Float64 {
+        _nkWithGeoQuad(aLat, aLon, bLat, bLon) { a, b, c, d, r, n in
+            nk_vincenty_f64(a, b, c, d, UInt64(n), r)
+        }
+    }
+}
+
+extension Float32 {
+    @inlinable @inline(__always)
+    public static func haversine<A: Sequence, B: Sequence, C: Sequence, D: Sequence>(
+        aLat: A, aLon: B, bLat: C, bLon: D
+    ) -> [Float32]?
+    where A.Element == Float32, B.Element == Float32, C.Element == Float32, D.Element == Float32 {
+        _nkWithGeoQuad(aLat, aLon, bLat, bLon) { a, b, c, d, r, n in
+            nk_haversine_f32(a, b, c, d, UInt64(n), r)
+        }
+    }
+
+    @inlinable @inline(__always)
+    public static func vincenty<A: Sequence, B: Sequence, C: Sequence, D: Sequence>(
+        aLat: A, aLon: B, bLat: C, bLon: D
+    ) -> [Float32]?
+    where A.Element == Float32, B.Element == Float32, C.Element == Float32, D.Element == Float32 {
+        _nkWithGeoQuad(aLat, aLon, bLat, bLon) { a, b, c, d, r, n in
+            nk_vincenty_f32(a, b, c, d, UInt64(n), r)
+        }
+    }
+}
+
+// MARK: - Geospatial Functions
+
+@inlinable @inline(__always)
+public func haversine<A: Sequence>(
+    aLat: A, aLon: A, bLat: A, bLon: A
+) -> [Float64]?
+where A.Element == Float64 {
+    Float64.haversine(aLat: aLat, aLon: aLon, bLat: bLat, bLon: bLon)
+}
+
+@inlinable @inline(__always)
+public func haversine<A: Sequence>(
+    aLat: A, aLon: A, bLat: A, bLon: A
+) -> [Float32]?
+where A.Element == Float32 {
+    Float32.haversine(aLat: aLat, aLon: aLon, bLat: bLat, bLon: bLon)
+}
+
+@inlinable @inline(__always)
+public func vincenty<A: Sequence>(
+    aLat: A, aLon: A, bLat: A, bLon: A
+) -> [Float64]?
+where A.Element == Float64 {
+    Float64.vincenty(aLat: aLat, aLon: aLon, bLat: bLat, bLon: bLon)
+}
+
+@inlinable @inline(__always)
+public func vincenty<A: Sequence>(
+    aLat: A, aLon: A, bLat: A, bLon: A
+) -> [Float32]?
+where A.Element == Float32 {
+    Float32.vincenty(aLat: aLat, aLon: aLon, bLat: bLat, bLon: bLon)
+}
+
+// MARK: - Capabilities
+
+/// Bitmask of SIMD instruction sets detected at runtime.
+public enum Capabilities {
+    public static var available: UInt64 { nk_capabilities() }
+
+    public static let neon: UInt64 = 1 << 1
+    public static let haswell: UInt64 = 1 << 2
+    public static let skylake: UInt64 = 1 << 3
+    public static let neonHalf: UInt64 = 1 << 4
+    public static let neonSDot: UInt64 = 1 << 5
+    public static let neonFhm: UInt64 = 1 << 6
+    public static let icelake: UInt64 = 1 << 7
+    public static let genoa: UInt64 = 1 << 8
+    public static let neonBfDot: UInt64 = 1 << 9
+    public static let sve: UInt64 = 1 << 10
+    public static let sveHalf: UInt64 = 1 << 11
+    public static let sveSDot: UInt64 = 1 << 12
+    public static let alder: UInt64 = 1 << 13
+    public static let sveBfDot: UInt64 = 1 << 14
+    public static let sve2: UInt64 = 1 << 15
+    public static let v128Relaxed: UInt64 = 1 << 16
+    public static let sapphire: UInt64 = 1 << 17
+    public static let sapphireAmx: UInt64 = 1 << 18
+    public static let rvv: UInt64 = 1 << 19
+    public static let rvvHalf: UInt64 = 1 << 20
+    public static let rvvBf16: UInt64 = 1 << 21
+    public static let graniteAmx: UInt64 = 1 << 22
+    public static let turin: UInt64 = 1 << 23
+    public static let sme: UInt64 = 1 << 24
+    public static let sme2: UInt64 = 1 << 25
+    public static let smeF64: UInt64 = 1 << 26
+    public static let smeFa64: UInt64 = 1 << 27
+    public static let sve2p1: UInt64 = 1 << 28
+    public static let sme2p1: UInt64 = 1 << 29
+    public static let smeHalf: UInt64 = 1 << 30
+    public static let smeBf16: UInt64 = 1 << 31
+    public static let smeLut2: UInt64 = 1 << 32
+    public static let rvvBB: UInt64 = 1 << 33
+    public static let sierra: UInt64 = 1 << 34
+    public static let smeBi32: UInt64 = 1 << 35
+    public static let loongsonAsx: UInt64 = 1 << 36
+    public static let powerVsx: UInt64 = 1 << 37
+    public static let diamond: UInt64 = 1 << 38
+    public static let neonFp8: UInt64 = 1 << 39
+}