multi-threading (#300)

Author: RuiyuZhu

Benchmarks/PrivateInformationRetrievalBenchmark/PirBenchmark.swift

@@ -1,4 +1,4 @@
-// Copyright 2024-2025 Apple Inc. and the Swift Homomorphic Encryption project authors
+// Copyright 2024-2026 Apple Inc. and the Swift Homomorphic Encryption project authors
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -24,9 +24,13 @@ nonisolated(unsafe) let benchmarks: () -> Void = {
     pirProcessBenchmark(PirUtil<Bfv<UInt32>>.self)()
     pirProcessBenchmark(PirUtil<Bfv<UInt64>>.self)()
 
-    indexPirBenchmark(PirUtil<Bfv<UInt32>>.self)()
-    indexPirBenchmark(PirUtil<Bfv<UInt64>>.self)()
+    indexPirBenchmark(PirUtil<Bfv<UInt32>>.self, callOptions: .multiThreaded)()
+    indexPirBenchmark(PirUtil<Bfv<UInt32>>.self, callOptions: .singleThreaded)()
+    indexPirBenchmark(PirUtil<Bfv<UInt64>>.self, callOptions: .multiThreaded)()
+    indexPirBenchmark(PirUtil<Bfv<UInt64>>.self, callOptions: .singleThreaded)()
 
-    keywordPirBenchmark(PirUtil<Bfv<UInt32>>.self)()
-    keywordPirBenchmark(PirUtil<Bfv<UInt64>>.self)()
+    keywordPirBenchmark(PirUtil<Bfv<UInt32>>.self, callOptions: .multiThreaded)()
+    keywordPirBenchmark(PirUtil<Bfv<UInt32>>.self, callOptions: .singleThreaded)()
+    keywordPirBenchmark(PirUtil<Bfv<UInt64>>.self, callOptions: .multiThreaded)()
+    keywordPirBenchmark(PirUtil<Bfv<UInt64>>.self, callOptions: .singleThreaded)()
 }

Package.resolved

@@ -81,6 +81,7 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-algorithms", from: "1.2.0"),
         .package(url: "https://github.com/apple/swift-argument-parser.git", from: "1.2.0"),
         .package(url: "https://github.com/apple/swift-async-algorithms.git", from: "1.0.2"),
+        .package(url: "https://github.com/apple/swift-collections.git", from: "1.3.0"),
         .package(url: "https://github.com/apple/swift-crypto.git", from: "3.10.0"),
         .package(url: "https://github.com/apple/swift-log.git", from: "1.0.0"),
         .package(url: "https://github.com/apple/swift-numerics", from: "1.0.0"),
@@ -102,6 +103,7 @@ let package = Package(
         .target(
             name: "HomomorphicEncryption",
             dependencies: [
+                .product(name: "Algorithms", package: "swift-algorithms"),
                 .product(name: "AsyncAlgorithms", package: "swift-async-algorithms"),
                 .product(name: "Crypto", package: "swift-crypto"),
                 .product(name: "_CryptoExtras", package: "swift-crypto"),
@@ -123,6 +125,7 @@ let package = Package(
             name: "PrivateInformationRetrieval",
             dependencies: ["HomomorphicEncryption",
                            .product(name: "AsyncAlgorithms", package: "swift-async-algorithms"),
+                           .product(name: "Collections", package: "swift-collections"),
                            .product(name: "Numerics", package: "swift-numerics")],
             swiftSettings: librarySettings),
         .target(

Package.swift

@@ -12,7 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+public import Algorithms
 public import ModularArithmetic
+import Foundation
 
 /// Brakerski-Fan-Vercauteren cryptosystem.
 public enum Bfv<T: ScalarType>: HeScheme {
@@ -218,6 +220,97 @@ public enum Bfv<T: ScalarType>: HeScheme {
 
     // MARK: Inner product
 
+    @inlinable
+    static func validateInnerProductInput(_ lhs: CanonicalCiphertext,
+                                          _ rhs: CanonicalCiphertext) throws
+    {
+        try validateEquality(of: lhs.context, and: rhs.context)
+        guard lhs.polys.count == freshCiphertextPolyCount, lhs.correctionFactor == 1 else {
+            throw HeError.invalidCiphertext(lhs)
+        }
+        guard rhs.polys.count == freshCiphertextPolyCount, rhs.correctionFactor == 1 else {
+            throw HeError.invalidCiphertext(rhs)
+        }
+    }
+
+    @inlinable
+    public static func innerProductAsync(_ lhs: some Collection<CanonicalCiphertext>,
+                                         _ rhs: some Collection<CanonicalCiphertext>) async throws
+        -> CanonicalCiphertext
+    {
+        // swiftlint:disable:next local_doc_comment
+        /// Computes accumulator += lhs * rhs
+        @Sendable
+        func lazyMultiply(
+            _ lhs: CanonicalCiphertext,
+            _ rhs: CanonicalCiphertext,
+            to accumulator: inout [Array2d<T.DoubleWidth>]) throws
+        {
+            try validateInnerProductInput(lhs, rhs)
+
+            let lhsPolys = try computeBehzPolys(ciphertext: lhs)
+            let rhsPolys = try computeBehzPolys(ciphertext: rhs)
+            PolyRq.addingLazyProduct(lhsPolys[0], rhsPolys[0], to: &accumulator[0])
+            PolyRq.addingLazyProduct(lhsPolys[0], rhsPolys[1], to: &accumulator[1])
+            PolyRq.addingLazyProduct(lhsPolys[1], rhsPolys[0], to: &accumulator[1])
+            PolyRq.addingLazyProduct(lhsPolys[1], rhsPolys[1], to: &accumulator[2])
+        }
+
+        let firstCiphertext = lhs[lhs.startIndex]
+        let rnsTool = firstCiphertext.context.getRnsTool(moduliCount: firstCiphertext.moduli.count)
+        let moduliCount = rnsTool.qBskContext.moduli.count
+        let poly = firstCiphertext.polys[0]
+        let maxProductCount = rnsTool.qBskContext.maxLazyProductAccumulationCount() / 2
+        let polyContext = rnsTool.qBskContext
+
+        let pairs = Array(zip(lhs, rhs))
+        let chunks = pairs.evenlyChunked(in: Util.activeProcessorCount)
+
+        // Each task accumulates its chunk independently then reduces to [0, modulus).
+        let partials = try await chunks.concurrentMap(ordered: false) { chunk in
+            var accumulator = Array(
+                repeating: Array2d(
+                    data: Array(repeating: T.DoubleWidth(0), count: moduliCount * poly.degree),
+                    rowCount: moduliCount, columnCount: poly.degree),
+                count: 3)
+            var reduceCount = 0
+            for (lhsCipher, rhsCipher) in chunk {
+                try lazyMultiply(lhsCipher, rhsCipher, to: &accumulator)
+                reduceCount += 1
+                if reduceCount >= maxProductCount {
+                    reduceCount = 0
+                    reduceInPlace(accumulator: &accumulator, polyContext: polyContext)
+                }
+            }
+            reduceInPlace(accumulator: &accumulator, polyContext: polyContext)
+            return accumulator
+        }
+
+        // Merge: values are in [0, modulus) after reduceInPlace, so summing fits in DoubleWidth.
+        var merged = partials[0]
+        for partial in partials.dropFirst() {
+            for polyIndex in merged.indices {
+                for i in merged[polyIndex].data.indices {
+                    #if DEBUG
+                    let (result, didOverflow) = merged[polyIndex].data[i]
+                        .addingReportingOverflow(partial[polyIndex].data[i])
+                    assert(!didOverflow, "Overflow in polynomial addition at polyIndex \(polyIndex), index \(i)")
+                    merged[polyIndex].data[i] = result
+                    #else
+                    merged[polyIndex].data[i] &+= partial[polyIndex].data[i]
+                    #endif
+                }
+            }
+        }
+
+        var sum = try EvalCiphertext(
+            context: firstCiphertext.context,
+            polys: Array(repeating: .zero(context: rnsTool.qBskContext), count: 3),
+            correctionFactor: 1)
+        reduceToCiphertext(accumulator: merged, result: &sum)
+        return try dropExtendedBase(from: sum)
+    }
+
     @inlinable
     public static func innerProduct(_ lhs: some Collection<CanonicalCiphertext>,
                                     _ rhs: some Collection<CanonicalCiphertext>) throws -> CanonicalCiphertext
@@ -229,13 +322,7 @@ public enum Bfv<T: ScalarType>: HeScheme {
             _ rhs: CanonicalCiphertext,
             to accumulator: inout [Array2d<T.DoubleWidth>]) throws
         {
-            try validateEquality(of: lhs.context, and: rhs.context)
-            guard lhs.polys.count == freshCiphertextPolyCount, lhs.correctionFactor == 1 else {
-                throw HeError.invalidCiphertext(lhs)
-            }
-            guard rhs.polys.count == freshCiphertextPolyCount, rhs.correctionFactor == 1 else {
-                throw HeError.invalidCiphertext(rhs)
-            }
+            try validateInnerProductInput(lhs, rhs)
 
             let lhsPolys = try computeBehzPolys(ciphertext: lhs)
             let rhsPolys = try computeBehzPolys(ciphertext: rhs)
@@ -306,26 +393,89 @@ public enum Bfv<T: ScalarType>: HeScheme {
         }
     }
 
+    /// Computes accumulator += ciphertext * plaintext
     @inlinable
-    public static func innerProduct(ciphertexts: some Collection<EvalCiphertext>,
-                                    plaintexts: some Collection<EvalPlaintext?>) throws -> EvalCiphertext
+    static func lazyMultiply(
+        ciphertext: EvalCiphertext,
+        plaintext: EvalPlaintext,
+        to accumulator: inout [Array2d<T.DoubleWidth>]) throws
     {
-        // swiftlint:disable:next local_doc_comment
-        /// Computes accumulator += ciphertext * plaintext
-        func lazyMultiply(
-            ciphertext: EvalCiphertext,
-            plaintext: EvalPlaintext,
-            to accumulator: inout [Array2d<T.DoubleWidth>]) throws
-        {
-            try validateEquality(of: ciphertext.context, and: plaintext.context)
-            guard ciphertext.moduli.count == plaintext.moduli.count else {
-                throw HeError.incompatibleCiphertextAndPlaintext(ciphertext: ciphertext, plaintext: plaintext)
+        try validateEquality(of: ciphertext.context, and: plaintext.context)
+        guard ciphertext.moduli.count == plaintext.moduli.count else {
+            throw HeError.incompatibleCiphertextAndPlaintext(ciphertext: ciphertext, plaintext: plaintext)
+        }
+        for (polyIndex, ciphertextPoly) in ciphertext.polys.enumerated() {
+            PolyRq.addingLazyProduct(ciphertextPoly, plaintext.poly, to: &accumulator[polyIndex])
+        }
+    }
+
+    @inlinable
+    public static func innerProductAsync(ciphertexts: some Collection<EvalCiphertext>,
+                                         plaintexts: some Collection<EvalPlaintext?>) async throws -> EvalCiphertext
+    {
+        precondition(plaintexts.count == ciphertexts.count)
+        guard var result = ciphertexts.first else {
+            preconditionFailure("Empty ciphertexts")
+        }
+        precondition(ciphertexts.allSatisfy { $0.polys.count == result.polys.count },
+                     "All ciphertexts must have the same polynomial count")
+        let poly = result.polys[0]
+        let maxProductCount = poly.context.maxLazyProductAccumulationCount()
+        let polyContext = result.polyContext()
+        let polyCount = result.polys.count
+        let dataCount = poly.data.count
+        let moduliCount = poly.moduli.count
+        let degree = poly.degree
+
+        let pairs = Array(zip(ciphertexts, plaintexts))
+        let chunks = pairs.evenlyChunked(in: Util.activeProcessorCount)
+
+        // Each task accumulates its chunk independently then reduces to [0, modulus).
+        let partials = try await chunks.concurrentMap(ordered: false) { chunk in
+            var accumulator = Array(
+                repeating: Array2d(
+                    data: Array(repeating: T.DoubleWidth(0), count: dataCount),
+                    rowCount: moduliCount, columnCount: degree),
+                count: polyCount)
+            var reduceCount = 0
+            for (ciphertext, plaintext) in chunk {
+                guard let plaintext else { continue }
+                try lazyMultiply(ciphertext: ciphertext, plaintext: plaintext, to: &accumulator)
+                reduceCount += 1
+                if reduceCount >= maxProductCount {
+                    reduceCount = 0
+                    reduceInPlace(accumulator: &accumulator, polyContext: polyContext)
+                }
             }
-            for (polyIndex, ciphertextPoly) in ciphertext.polys.enumerated() {
-                PolyRq.addingLazyProduct(ciphertextPoly, plaintext.poly, to: &accumulator[polyIndex])
+            reduceInPlace(accumulator: &accumulator, polyContext: polyContext)
+            return accumulator
+        }
+
+        // Merge: values are in [0, modulus) after reduceInPlace, so summing fits in DoubleWidth.
+        var merged = partials[0]
+        for partial in partials.dropFirst() {
+            for polyIndex in merged.indices {
+                for i in merged[polyIndex].data.indices {
+                    #if DEBUG
+                    let (result, didOverflow) = merged[polyIndex].data[i]
+                        .addingReportingOverflow(partial[polyIndex].data[i])
+                    assert(!didOverflow, "Overflow in polynomial addition at polyIndex \(polyIndex), index \(i)")
+                    merged[polyIndex].data[i] = result
+                    #else
+                    merged[polyIndex].data[i] &+= partial[polyIndex].data[i]
+                    #endif
+                }
             }
         }
 
+        reduceToCiphertext(accumulator: merged, result: &result)
+        return result
+    }
+
+    @inlinable
+    public static func innerProduct(ciphertexts: some Collection<EvalCiphertext>,
+                                    plaintexts: some Collection<EvalPlaintext?>) throws -> EvalCiphertext
+    {
         precondition(plaintexts.count == ciphertexts.count)
         guard var result = ciphertexts.first else {
             preconditionFailure("Empty ciphertexts")

Sources/HomomorphicEncryption/Bfv/Bfv.swift

@@ -922,7 +922,7 @@ public protocol HeScheme: Sendable {
     /// - Parameters:
     ///   - ciphertext: Ciphertext to transform.
     ///   - element: Galois element of the transformation. Must be odd in `[1, 2 * N - 1]` where `N` is the RLWE ring
-    /// dimension, given by ``EncryptionParameters/polyDegree``.
+    ///   dimension, given by ``EncryptionParameters/polyDegree``.
     ///   - key: Evaluation key. Must contain Galois element `element`.
     /// - Throws: Error upon failure to apply the Galois transformation.
     /// - seealso: ``applyGaloisAsync(ciphertext:element:using:)``  for an async version of this API

Sources/HomomorphicEncryption/HeScheme.swift

@@ -135,3 +135,82 @@ extension Array where Element: ScalarType {
         return width32.reduce(Width32<Self.Element>(1)) { $0 * $1 }
     }
 }
+
+extension Collection where Self.Element: Sendable {
+    @inlinable
+    package func concurrentMap<T: Sendable>(
+        ordered: Bool = true,
+        _ transform: @Sendable @escaping (Element) async throws -> T) async throws -> [T]
+    {
+        // Fast path for empty collection
+        if isEmpty {
+            return []
+        }
+
+        // We use enumerated indices (Int) for stable ordering.
+        return try await withThrowingTaskGroup(of: (Int, T).self) { group in
+            for (index, element) in self.enumerated() {
+                group.addTask {
+                    try await (index, transform(element))
+                }
+            }
+
+            if ordered {
+                var orderedBuffer = [T?](repeating: nil, count: self.count)
+                for try await (index, value) in group {
+                    orderedBuffer[index] = value
+                }
+                return orderedBuffer.compactMap(\.self)
+            }
+
+            var unorderedResults: [T] = []
+            unorderedResults.reserveCapacity(self.count)
+            for try await (_, value) in group {
+                unorderedResults.append(value)
+            }
+            return unorderedResults
+        }
+    }
+
+    @inlinable
+    package mutating func concurrentConsumingMap<T: Sendable>(
+        ordered: Bool = true,
+        _ transform: @Sendable @escaping (consuming Element) async throws -> T) async throws -> [T]
+    {
+        // Fast path for empty collection
+        if isEmpty {
+            return []
+        }
+
+        // We use enumerated indices (Int) for stable ordering.
+        return try await withThrowingTaskGroup(of: (Int, T).self) { group in
+            for (index, element) in self.enumerated() {
+                group.addTask {
+                    try await (index, transform(element))
+                }
+            }
+
+            if ordered {
+                var orderedBuffer = [T?](repeating: nil, count: self.count)
+                for try await (index, value) in group {
+                    orderedBuffer[index] = value
+                }
+                return orderedBuffer.compactMap(\.self)
+            }
+
+            var unorderedResults: [T] = []
+            unorderedResults.reserveCapacity(self.count)
+            for try await (_, value) in group {
+                unorderedResults.append(value)
+            }
+            return unorderedResults
+        }
+    }
+}
+
+@usableFromInline
+enum Util {
+    @usableFromInline static var activeProcessorCount: Int {
+        ProcessInfo.processInfo.activeProcessorCount
+    }
+}