Added AMX acceleration.

ZachNagengast · May 31, 2024 · bf50730 · bf50730
1 parent 07a610f
commit bf50730
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 16 deletions.
diff --git a/Sources/SimilaritySearchKit/Core/Embeddings/Metrics/DistanceMetrics.swift b/Sources/SimilaritySearchKit/Core/Embeddings/Metrics/DistanceMetrics.swift
@@ -6,6 +6,7 @@
 //
 
 import Foundation
+import Accelerate
 
 /// A struct implementing the `DistanceMetricProtocol` using the dot product.
 ///
@@ -20,10 +21,22 @@ public struct DotProduct: DistanceMetricProtocol {
         return sortedScores(scores: scores, topK: resultsCount)
     }
 
-    public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float {
-        let dotProduct = zip(firstEmbedding, secondEmbedding).map(*).reduce(0, +)
-        return dotProduct
-    }
+//    public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float {
+//        let dotProduct = zip(firstEmbedding, secondEmbedding).map(*).reduce(0, +)
+//        return dotProduct
+//    }
+
+	public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float {
+		// Ensure the embeddings have the same length
+		precondition(firstEmbedding.count == secondEmbedding.count, "Embeddings must have the same length")
+
+		var dotProduct: Float = 0
+
+		// Calculate dot product using Accelerate
+		vDSP_dotpr(firstEmbedding, 1, secondEmbedding, 1, &dotProduct, vDSP_Length(firstEmbedding.count))
+
+		return dotProduct
+	}
 
 }
 
@@ -40,14 +53,36 @@ public struct CosineSimilarity: DistanceMetricProtocol {
         return sortedScores(scores: scores, topK: resultsCount)
     }
 
-    public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float {
-        // Calculate cosine distance
-        let dotProduct = zip(firstEmbedding, secondEmbedding).map(*).reduce(0, +)
-        let firstMagnitude = sqrt(firstEmbedding.map { $0 * $0 }.reduce(0, +))
-        let secondMagnitude = sqrt(secondEmbedding.map { $0 * $0 }.reduce(0, +))
-
-        return dotProduct / (firstMagnitude * secondMagnitude)
-    }
+//    public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float {
+//        // Calculate cosine distance
+//        let dotProduct = zip(firstEmbedding, secondEmbedding).map(*).reduce(0, +)
+//        let firstMagnitude = sqrt(firstEmbedding.map { $0 * $0 }.reduce(0, +))
+//        let secondMagnitude = sqrt(secondEmbedding.map { $0 * $0 }.reduce(0, +))
+//
+//        return dotProduct / (firstMagnitude * secondMagnitude)
+//    }
+
+	public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float {
+		// Ensure the embeddings have the same length
+		precondition(firstEmbedding.count == secondEmbedding.count, "Embeddings must have the same length")
+
+		var dotProduct: Float = 0
+		var firstMagnitude: Float = 0
+		var secondMagnitude: Float = 0
+
+		// Calculate dot product and magnitudes using Accelerate
+		vDSP_dotpr(firstEmbedding, 1, secondEmbedding, 1, &dotProduct, vDSP_Length(firstEmbedding.count))
+		vDSP_svesq(firstEmbedding, 1, &firstMagnitude, vDSP_Length(firstEmbedding.count))
+		vDSP_svesq(secondEmbedding, 1, &secondMagnitude, vDSP_Length(secondEmbedding.count))
+
+		// Take square root of magnitudes
+		firstMagnitude = sqrt(firstMagnitude)
+		secondMagnitude = sqrt(secondMagnitude)
+
+		// Return cosine similarity
+		return dotProduct / (firstMagnitude * secondMagnitude)
+	}
+
 }
 
 /// A struct implementing the `DistanceMetricProtocol` using Euclidean distance.
@@ -64,10 +99,23 @@ public struct EuclideanDistance: DistanceMetricProtocol {
         return sortedDistances(distances: distances, topK: resultsCount)
     }
 
-    public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float {
-        let squaredDifferences = zip(firstEmbedding, secondEmbedding).map { ($0 - $1) * ($0 - $1) }
-        return sqrt(squaredDifferences.reduce(0, +))
-    }
+//    public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float {
+//        let squaredDifferences = zip(firstEmbedding, secondEmbedding).map { ($0 - $1) * ($0 - $1) }
+//        return sqrt(squaredDifferences.reduce(0, +))
+//    }
+
+	public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float {
+		// Ensure the embeddings have the same length
+		precondition(firstEmbedding.count == secondEmbedding.count, "Embeddings must have the same length")
+
+		var distance: Float = 0
+
+		// Calculate squared differences and sum them using Accelerate
+		vDSP_distancesq(firstEmbedding, 1, secondEmbedding, 1, &distance, vDSP_Length(firstEmbedding.count))
+
+		// Return the square root of the summed squared differences
+		return sqrt(distance)
+	}
 }
 
 // MARK: - Helpers

diff --git a/Sources/SimilaritySearchKit/Core/Index/SimilarityIndex.swift b/Sources/SimilaritySearchKit/Core/Index/SimilarityIndex.swift
@@ -319,6 +319,7 @@ extension SimilarityIndex {
 
 @available(macOS 13.0, iOS 16.0, *)
 extension SimilarityIndex {
+
     public func saveIndex(toDirectory path: URL? = nil, name: String? = nil) throws -> URL {
         let indexName = name ?? self.indexName
         let basePath: URL