From 850990b496bf1678ac716ddec4e4f37da174764a Mon Sep 17 00:00:00 2001 From: Fabian Boemer Date: Mon, 26 Aug 2024 11:08:18 -0700 Subject: [PATCH] Adds PNNS Benchmarks. --- ...ivateNearestNeighborsSearchBenchmark.swift | 359 ++++++++++++++++++ Package.swift | 14 + 2 files changed, 373 insertions(+) create mode 100644 Benchmarks/PrivateNearestNeighborsSearchBenchmark/PrivateNearestNeighborsSearchBenchmark.swift diff --git a/Benchmarks/PrivateNearestNeighborsSearchBenchmark/PrivateNearestNeighborsSearchBenchmark.swift b/Benchmarks/PrivateNearestNeighborsSearchBenchmark/PrivateNearestNeighborsSearchBenchmark.swift new file mode 100644 index 00000000..5214934c --- /dev/null +++ b/Benchmarks/PrivateNearestNeighborsSearchBenchmark/PrivateNearestNeighborsSearchBenchmark.swift @@ -0,0 +1,359 @@ +// Copyright 2024 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. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Benchmarks for Pnns functions. +// These benchmarks can be triggered with +// `swift package benchmark --target PNNSBenchmark` +// for more readable results + +@preconcurrency import Benchmark +import Foundation +import HomomorphicEncryption +import HomomorphicEncryptionProtobuf +import PrivateNearestNeighborsSearch +import PrivateNearestNeighborsSearchProtobuf + +@usableFromInline let benchmarkConfiguration = Benchmark.Configuration( + metrics: [ + .wallClock, + .mallocCountTotal, + .peakMemoryResident, + .evaluationKeySize, + .evaluationKeyCount, + .querySize, + .queryCiphertextCount, + .responseSize, + .responseCiphertextCount, + .noiseBudget, + ], + maxDuration: .seconds(5)) + +struct DatabaseConfig { + let rowCount: Int + let vectorDimension: Int + let metadataCount: Int + + init(rowCount: Int, vectorDimension: Int, metadataCount: Int = 0) { + self.rowCount = rowCount + self.vectorDimension = vectorDimension + self.metadataCount = metadataCount + } +} + +func getDatabaseForTesting(config: DatabaseConfig) -> Database { + let rows = (0.. Int { + try Int( + noiseBudget(using: secretKey, variableTime: true) * Double( + noiseBudgetScale)) + } +} + +struct ProcessBenchmarkContext { + let database: Database + let contexts: [Context] + let serverConfig: ServerConfig + + init(databaseConfig: DatabaseConfig, + parameterConfig: EncryptionParametersConfig) throws + { + let plaintextModuli = try Scheme.Scalar.generatePrimes( + significantBitCounts: parameterConfig.plaintextModulusBits, + preferringSmall: true, + nttDegree: parameterConfig.polyDegree) + let coefficientModuli = try Scheme.Scalar.generatePrimes( + significantBitCounts: parameterConfig.coefficientModulusBits, + preferringSmall: false, + nttDegree: parameterConfig.polyDegree) + + let encryptionParams = try EncryptionParameters( + polyDegree: parameterConfig.polyDegree, + plaintextModulus: plaintextModuli[0], + coefficientModuli: coefficientModuli, + errorStdDev: ErrorStdDev.stdDev32, + securityLevel: SecurityLevel.quantum128) + + let evaluationKeyConfig = try MatrixMultiplication.evaluationKeyConfig( + plaintextMatrixDimensions: MatrixDimensions( + rowCount: databaseConfig.rowCount, + columnCount: databaseConfig.vectorDimension), + encryptionParameters: encryptionParams) + let scalingFactor = ClientConfig + .maxScalingFactor( + distanceMetric: .cosineSimilarity, + vectorDimension: databaseConfig.vectorDimension, + plaintextModuli: Array(plaintextModuli[1...])) + let clientConfig = try ClientConfig( + encryptionParams: encryptionParams, + scalingFactor: scalingFactor, + queryPacking: .denseRow, + vectorDimension: databaseConfig.vectorDimension, + evaluationKeyConfig: evaluationKeyConfig, + distanceMetric: .cosineSimilarity) + let babyStepGiantStep = BabyStepGiantStep(vectorDimension: databaseConfig.vectorDimension) + let serverConfig = ServerConfig( + clientConfig: clientConfig, + databasePacking: .diagonal(babyStepGiantStep: babyStepGiantStep)) + self.serverConfig = serverConfig + + self.database = getDatabaseForTesting(config: databaseConfig) + self.contexts = try serverConfig.encryptionParameters.map { encryptionParams in + try Context(encryptionParameters: encryptionParams) + } + } +} + +func processBenchmark(_: Scheme.Type) -> () -> Void { + { + let databaseConfig = DatabaseConfig( + rowCount: 4096, + vectorDimension: 128, + metadataCount: 0) + let encryptionConfig = EncryptionParametersConfig( + polyDegree: 4096, + // use plaintextModulusBits: [16, 17] for plaintext CRT + plaintextModulusBits: [17], + coefficientModulusBits: [27, 28, 28]) + + let benchmarkName = [ + "Process", + String(describing: Scheme.self), + encryptionConfig.description, + "rowCount=\(databaseConfig.rowCount)", + "vectorDimension=\(databaseConfig.vectorDimension)", + "metadataCount=\(databaseConfig.metadataCount)", + ].joined(separator: "/") + // swiftlint:disable closure_parameter_position + Benchmark(benchmarkName, configuration: benchmarkConfiguration) { ( + benchmark, + benchmarkContext: ProcessBenchmarkContext) in + for _ in benchmark.scaledIterations { + try blackHole( + benchmarkContext.database + .process( + config: benchmarkContext.serverConfig, + contexts: benchmarkContext.contexts)) + } + } setup: { + try ProcessBenchmarkContext( + databaseConfig: databaseConfig, + parameterConfig: encryptionConfig) + } + // swiftlint:enable closure_parameter_position + } +} + +struct PnnsBenchmarkContext { + let processedDatabase: ProcessedDatabase + let server: Server + let client: Client + let secretKey: SecretKey + let evaluationKey: Scheme.EvaluationKey + let evaluationKeyCount: Int + let query: Query + let evaluationKeySize: Int + let querySize: Int + let queryCiphertextCount: Int + let responseSize: Int + let responseCiphertextCount: Int + let noiseBudget: Int + + init(databaseConfig: DatabaseConfig, + parameterConfig: EncryptionParametersConfig, + queryCount: Int) throws + { + let plaintextModuli = try Scheme.Scalar.generatePrimes( + significantBitCounts: parameterConfig.plaintextModulusBits, + preferringSmall: true, + nttDegree: parameterConfig.polyDegree) + let coefficientModuli = try Scheme.Scalar.generatePrimes( + significantBitCounts: parameterConfig.coefficientModulusBits, + preferringSmall: false, + nttDegree: parameterConfig.polyDegree) + let encryptionParams = try EncryptionParameters( + polyDegree: parameterConfig.polyDegree, + plaintextModulus: plaintextModuli[0], + coefficientModuli: coefficientModuli, + errorStdDev: ErrorStdDev.stdDev32, + securityLevel: SecurityLevel.quantum128) + + let evaluationKeyConfig = try MatrixMultiplication.evaluationKeyConfig( + plaintextMatrixDimensions: MatrixDimensions( + rowCount: databaseConfig.rowCount, + columnCount: databaseConfig.vectorDimension), + encryptionParameters: encryptionParams) + let scalingFactor = ClientConfig + .maxScalingFactor( + distanceMetric: .cosineSimilarity, + vectorDimension: databaseConfig.vectorDimension, + plaintextModuli: plaintextModuli) + let clientConfig = try ClientConfig( + encryptionParams: encryptionParams, + scalingFactor: scalingFactor, + queryPacking: .denseRow, + vectorDimension: databaseConfig.vectorDimension, + evaluationKeyConfig: evaluationKeyConfig, + distanceMetric: .cosineSimilarity, + extraPlaintextModuli: Array(plaintextModuli[1...])) + + let babyStepGiantStep = BabyStepGiantStep(vectorDimension: databaseConfig.vectorDimension) + let serverConfig = ServerConfig( + clientConfig: clientConfig, + databasePacking: .diagonal(babyStepGiantStep: babyStepGiantStep)) + + let database = getDatabaseForTesting(config: databaseConfig) + let contexts = try clientConfig.encryptionParameters + .map { encryptionParams in try Context(encryptionParameters: encryptionParams) } + self.processedDatabase = try database.process(config: serverConfig, contexts: contexts) + self.client = try Client(config: clientConfig, contexts: contexts) + self.server = try Server(database: processedDatabase, config: serverConfig) + self.secretKey = try client.generateSecretKey() + self.evaluationKey = try client.generateEvaluationKey(using: secretKey) + + // We query exact matches from rows in the database + let databaseVectors = Array2d(data: database.rows.map { row in row.vector }) + let queryVectors = Array2d(data: database.rows.prefix(queryCount).map { row in row.vector }) + self.query = try client.generateQuery(for: queryVectors, using: secretKey) + + let response = try server.computeResponse(to: query, using: evaluationKey) + let decrypted = try client.decrypt(response: response, using: secretKey) + + // Validate correctness + let modulus = clientConfig.plaintextModuli.map { UInt64($0) }.reduce(1, *) + let expected = try databaseVectors.fixedPointCosineSimilarity( + queryVectors.transposed(), + modulus: modulus, + scalingFactor: Float(clientConfig.scalingFactor)) + precondition(decrypted.distances.data == expected.data, "Wrong response") + + self.evaluationKeySize = try evaluationKey.size() + self.evaluationKeyCount = evaluationKey.configuration.keyCount + self.querySize = try query.size() + self.queryCiphertextCount = query.ciphertextMatrices.map { matrix in matrix.ciphertexts.count }.sum() + self.responseSize = try response.size() + self.responseCiphertextCount = response.ciphertextMatrices + .map { matrix in matrix.ciphertexts.count }.sum() + self.noiseBudget = try response.scaledNoiseBudget(using: secretKey) + } +} + +func cosineSimilarityBenchmark(_: Scheme.Type) -> () -> Void { + { + let databaseConfig = DatabaseConfig( + rowCount: 4096, + vectorDimension: 128, + metadataCount: 0) + let encryptionConfig = EncryptionParametersConfig( + polyDegree: 4096, + // use plaintextModulusBits: [16, 17] for plaintext CRT + plaintextModulusBits: [17], + coefficientModulusBits: [27, 28, 28]) + let queryCount = 1 + + let benchmarkName = [ + "CosineSimilarityBenchmark", + String(describing: Scheme.self), + encryptionConfig.description, + "rowCount=\(databaseConfig.rowCount)", + "vectorDimension=\(databaseConfig.vectorDimension)", + "metadataCount=\(databaseConfig.metadataCount)", + "queryCount=\(queryCount)", + ].joined(separator: "/") + // swiftlint:disable closure_parameter_position + Benchmark(benchmarkName, configuration: benchmarkConfiguration) { ( + benchmark, + benchmarkContext: PnnsBenchmarkContext) in + for _ in benchmark.scaledIterations { + try blackHole( + benchmarkContext.server.computeResponse( + to: benchmarkContext.query, + using: benchmarkContext.evaluationKey)) + } + benchmark.measurement(.evaluationKeySize, benchmarkContext.evaluationKeySize) + benchmark.measurement(.evaluationKeyCount, benchmarkContext.evaluationKeyCount) + benchmark.measurement(.querySize, benchmarkContext.querySize) + benchmark.measurement(.queryCiphertextCount, benchmarkContext.queryCiphertextCount) + benchmark.measurement(.responseSize, benchmarkContext.responseSize) + benchmark.measurement(.responseCiphertextCount, benchmarkContext.responseCiphertextCount) + benchmark.measurement(.noiseBudget, benchmarkContext.noiseBudget) + } setup: { + try PnnsBenchmarkContext( + databaseConfig: databaseConfig, + parameterConfig: encryptionConfig, + queryCount: queryCount) + } + // swiftlint:enable closure_parameter_position + } +} + +extension BenchmarkMetric { + static var querySize: Self { .custom("Query byte size") } + static var queryCiphertextCount: Self { .custom("Query ciphertext count") } + static var evaluationKeySize: Self { .custom("Evaluation key byte size") } + static var evaluationKeyCount: Self { .custom("Evaluation key count") } + static var responseSize: Self { .custom("Response byte size") } + static var responseCiphertextCount: Self { .custom("Response ciphertext count") } + static var noiseBudget: Self { .custom("Noise budget x \(noiseBudgetScale)") } +} + +nonisolated(unsafe) let benchmarks: () -> Void = { + processBenchmark(Bfv.self)() + processBenchmark(Bfv.self)() + + cosineSimilarityBenchmark(Bfv.self)() + cosineSimilarityBenchmark(Bfv.self)() +} diff --git a/Package.swift b/Package.swift index e918d547..3c84a378 100644 --- a/Package.swift +++ b/Package.swift @@ -238,6 +238,20 @@ package.targets += [ plugins: [ .plugin(name: "BenchmarkPlugin", package: "package-benchmark"), ]), + .executableTarget( + name: "PNNSBenchmark", + dependencies: [ + .product(name: "Benchmark", package: "package-benchmark"), + "HomomorphicEncryption", + "HomomorphicEncryptionProtobuf", + "PrivateNearestNeighborsSearch", + "PrivateNearestNeighborsSearchProtobuf", + ], + path: "Benchmarks/PrivateNearestNeighborsSearchBenchmark", + swiftSettings: benchmarkSettings, + plugins: [ + .plugin(name: "BenchmarkPlugin", package: "package-benchmark"), + ]), ] // Set the minimum macOS version for the package