From 6d5c7ced173faabe7ec141a30ef3f3ac407fd5ff Mon Sep 17 00:00:00 2001 From: Fabian Boemer Date: Fri, 16 Aug 2024 07:20:03 -0700 Subject: [PATCH] Implement PlaintextMatrix.diagonal encoding --- Package.resolved | 11 +- Package.swift | 6 +- Sources/HomomorphicEncryption/Array2d.swift | 28 +++-- Sources/HomomorphicEncryption/Encoding.swift | 15 +-- Sources/HomomorphicEncryption/Scalar.swift | 28 ++++- .../DotProduct.swift | 41 +++++++ .../PlaintextMatrix.swift | 100 +++++++++++++++- .../Array2dTests.swift | 3 +- .../ScalarTests.swift | 10 +- .../PlaintextMatrixTests.swift | 113 +++++++++++++++++- 10 files changed, 322 insertions(+), 33 deletions(-) create mode 100644 Sources/PrivateNearestNeighborsSearch/DotProduct.swift diff --git a/Package.resolved b/Package.resolved index d05d705f..73c48ec8 100644 --- a/Package.resolved +++ b/Package.resolved @@ -1,5 +1,5 @@ { - "originHash" : "7ee1d955f789a2932c73eb859ad423a54265140cebf530f47f84bc523d26c86f", + "originHash" : "287ba43d965ab226c778e6332f9fbb9244920f2452fbe9bccef6f422c30d7e49", "pins" : [ { "identity" : "hdrhistogram-swift", @@ -28,6 +28,15 @@ "version" : "1.0.0" } }, + { + "identity" : "swift-algorithms", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-algorithms", + "state" : { + "revision" : "f6919dfc309e7f1b56224378b11e28bab5bccc42", + "version" : "1.2.0" + } + }, { "identity" : "swift-argument-parser", "kind" : "remoteSourceControl", diff --git a/Package.swift b/Package.swift index b5a01c95..ddbf78e4 100644 --- a/Package.swift +++ b/Package.swift @@ -51,6 +51,7 @@ let package = Package( .executable(name: "PIRShardDatabase", targets: ["PIRShardDatabase"]), ], dependencies: [ + .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-crypto.git", from: "3.4.0"), .package(url: "https://github.com/apple/swift-log.git", from: "1.0.0"), @@ -96,7 +97,10 @@ let package = Package( swiftSettings: librarySettings), .target( name: "PrivateNearestNeighborsSearch", - dependencies: ["HomomorphicEncryption"], + dependencies: [ + .product(name: "Algorithms", package: "swift-algorithms"), + "HomomorphicEncryption", + ], swiftSettings: librarySettings), .target( name: "TestUtilities", diff --git a/Sources/HomomorphicEncryption/Array2d.swift b/Sources/HomomorphicEncryption/Array2d.swift index ee829a8b..a2d139c8 100644 --- a/Sources/HomomorphicEncryption/Array2d.swift +++ b/Sources/HomomorphicEncryption/Array2d.swift @@ -15,11 +15,11 @@ /// Stores values in a 2 dimensional array. public struct Array2d: Equatable, Sendable { @usableFromInline package var data: [T] - @usableFromInline var rowCount: Int - @usableFromInline var columnCount: Int + @usableFromInline package var rowCount: Int + @usableFromInline package var columnCount: Int - @usableFromInline var shape: (Int, Int) { (rowCount, columnCount) } - @usableFromInline var count: Int { rowCount * columnCount } + @usableFromInline package var shape: (Int, Int) { (rowCount, columnCount) } + @usableFromInline package var count: Int { rowCount * columnCount } @inlinable package init(data: [T], rowCount: Int, columnCount: Int) { @@ -35,26 +35,34 @@ public struct Array2d: Equatable, self.rowCount = array.rowCount self.data = array.data.map { T($0) } } + + @inlinable + package static func zero(rowCount: Int, columnCount: Int) -> Self { + self.init( + data: [T](Array(repeating: T.zero, count: rowCount * columnCount)), + rowCount: rowCount, + columnCount: columnCount) + } } extension Array2d { @inlinable - func index(row: Int, column: Int) -> Int { + package func index(row: Int, column: Int) -> Int { row &* columnCount &+ column } @inlinable - func rowIndices(row: Int) -> Range { + package func rowIndices(row: Int) -> Range { index(row: row, column: 0).. StrideTo { + package func columnIndices(column: Int) -> StrideTo { stride(from: index(row: 0, column: column), to: index(row: rowCount, column: column), by: columnCount) } @inlinable - func row(row: Int) -> [T] { + package func row(row: Int) -> [T] { Array(data[rowIndices(row: row)]) } @@ -80,7 +88,7 @@ extension Array2d { } @inlinable - subscript(_ index: Int) -> T { + package subscript(_ index: Int) -> T { get { data[index] } @@ -90,7 +98,7 @@ extension Array2d { } @inlinable - subscript(_ row: Int, _ column: Int) -> T { + package subscript(_ row: Int, _ column: Int) -> T { get { data[index(row: row, column: column)] } diff --git a/Sources/HomomorphicEncryption/Encoding.swift b/Sources/HomomorphicEncryption/Encoding.swift index 5eaabe98..14ddaafd 100644 --- a/Sources/HomomorphicEncryption/Encoding.swift +++ b/Sources/HomomorphicEncryption/Encoding.swift @@ -169,16 +169,11 @@ extension Context { func encodeSimd(values: [some ScalarType]) throws -> Plaintext { guard !simdEncodingMatrix.isEmpty else { throw HeError.simdEncodingNotSupported(for: encryptionParameters) } let polyDegree = encryptionParameters.polyDegree - var array = Array2d( - data: [Scheme.Scalar](repeating: 0, - count: polyDegree), - rowCount: 1, - columnCount: polyDegree) + var array = Array2d.zero(rowCount: 1, columnCount: polyDegree) for index in 0..(context: plaintextContext, - data: array) + let poly = PolyRq<_, Eval>(context: plaintextContext, data: array) let coeffPoly = try poly.inverseNtt() return Plaintext(context: self, poly: coeffPoly) } @@ -189,10 +184,8 @@ extension Context { throw HeError.simdEncodingNotSupported(for: encryptionParameters) } let poly = try plaintext.poly.forwardNtt() - var values = [T](repeating: 0, count: encryptionParameters.polyDegree) - for index in 0.. Self { + public func nextMultiple(of rhs: Self, variableTime: Bool) -> Self { precondition(variableTime) precondition(self >= 0) - if factor == 0 { + if rhs == 0 { return 0 } - return dividingCeil(factor, variableTime: true) * factor + return dividingCeil(rhs, variableTime: true) * rhs + } + + /// Computes the largest value less than or equal to this value that is a multiple of `rhs`. + /// + /// This value must be non-negative. + /// - Parameters: + /// - rhs: Value of which the output is a multiple of. + /// - variableTime: Must be `true`, indicating this value and `other` are leaked through timing. + /// - Returns: the previous multiple of this value. + /// - Warning: Leaks this value and `other` through timing. + @inlinable + public func previousMultiple(of rhs: Self, variableTime: Bool) -> Self { + precondition(variableTime) + precondition(self >= 0) + if rhs == 0 { + return 0 + } + return (self / rhs) * rhs } } diff --git a/Sources/PrivateNearestNeighborsSearch/DotProduct.swift b/Sources/PrivateNearestNeighborsSearch/DotProduct.swift new file mode 100644 index 00000000..9c13701d --- /dev/null +++ b/Sources/PrivateNearestNeighborsSearch/DotProduct.swift @@ -0,0 +1,41 @@ +// 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. + +import Foundation + +/// Pre-computed values for matrix-vector multiplication using baby-step, giant-step algorithm. +/// +/// - seealso: Section 6.3 of . +struct BabyStepGiantStep: Codable, Equatable, Hashable, Sendable { + /// Dimension of the vector; "D" in the reference. + let vectorDimension: Int + /// Baby step; "g" in the reference. + let babyStep: Int + /// Giant step; "h" in the reference. + let giantStep: Int + + init(vectorDimension: Int, babyStep: Int, giantStep: Int) { + self.vectorDimension = vectorDimension + self.babyStep = babyStep + self.giantStep = giantStep + } + + init(vectorDimension: Int) { + let dimension = Int32(vectorDimension).nextPowerOfTwo + let babyStep = Int32(Double(dimension).squareRoot().rounded(.up)) + let giantStep = dimension.dividingCeil(babyStep, variableTime: true) + + self.init(vectorDimension: Int(dimension), babyStep: Int(babyStep), giantStep: Int(giantStep)) + } +} diff --git a/Sources/PrivateNearestNeighborsSearch/PlaintextMatrix.swift b/Sources/PrivateNearestNeighborsSearch/PlaintextMatrix.swift index 3c0b9907..5066d1e1 100644 --- a/Sources/PrivateNearestNeighborsSearch/PlaintextMatrix.swift +++ b/Sources/PrivateNearestNeighborsSearch/PlaintextMatrix.swift @@ -12,10 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. +import Algorithms import HomomorphicEncryption /// Different algorithms for packing a matrix of scalar values into plaintexts. -enum PlaintextMatrixPacking: Equatable, Sendable { +enum PlaintextMatrixPacking: Codable, Equatable, Hashable, Sendable { /// As many columns of data are packed sequentially into each plaintext SIMD row as possible, such that no SIMD row /// contains data from multiple columns. case denseColumn @@ -27,6 +28,11 @@ enum PlaintextMatrixPacking: Equatable, Sendable { /// with the constraint that either all or none of the entries stored within the last plaintext /// row are repeated. case denseRow + /// Packs the values using a generalized diagonal packing. + /// + /// Includes modifications for the baby-step, giant-step algorithm from Section 6.3 of + /// . + case diagonal(babyStepGiantStep: BabyStepGiantStep) } /// The dimensions of a matrix, a 2d array. @@ -156,6 +162,13 @@ struct PlaintextMatrix: Equatable, Sendabl dimensions: dimensions, values: values) try self.init(dimensions: dimensions, packing: packing, plaintexts: plaintexts) + case .diagonal: + let plaintexts = try PlaintextMatrix.diagonalPlaintexts( + context: context, + dimensions: dimensions, + packing: packing, + values: values) + try self.init(dimensions: dimensions, packing: packing, plaintexts: plaintexts) } } @@ -190,6 +203,11 @@ struct PlaintextMatrix: Equatable, Sendabl let rowsPerPlaintextCount = simdDimensions.rowCount * ( simdDimensions.columnCount / dimensions.columnCount.nextPowerOfTwo) return dimensions.rowCount.dividingCeil(rowsPerPlaintextCount, variableTime: true) + case .diagonal: + let plaintextsPerColumnCount = dimensions.rowCount.dividingCeil( + encryptionParameters.polyDegree, + variableTime: true) + return dimensions.columnCount.nextPowerOfTwo * plaintextsPerColumnCount } } @@ -322,6 +340,83 @@ struct PlaintextMatrix: Equatable, Sendabl return plaintexts } + /// Computes the plaintexts for diagonal packing. + /// - Parameters: + /// - context: Context for HE computation. + /// - dimensions: Plaintext matrix dimensions. + /// - packing: Plaintext packing; must be `.diagonal`. + /// - values: The data values to store in the plaintext matrix; stored in row-major format. + /// - Returns: The plaintexts for diagonal packing. + /// - Throws: Error upon failure to compute the plaintexts. + @inlinable + static func diagonalPlaintexts( + context: Context, + dimensions: Dimensions, + packing: PlaintextMatrixPacking, + values: [V]) throws -> [Scheme.CoeffPlaintext] + { + let encryptionParameters = context.encryptionParameters + guard let simdDimensions = context.simdDimensions else { + throw PNNSError.simdEncodingNotSupported(for: encryptionParameters) + } + let simdColumnCount = simdDimensions.columnCount + let simdRowCount = simdDimensions.rowCount + precondition(simdRowCount == 2, "simdRowCount must be 2") + guard dimensions.columnCount <= simdColumnCount else { + throw PNNSError.invalidMatrixDimensions(dimensions) + } + guard case let .diagonal(bsgs) = packing else { + let expectedBsgs = BabyStepGiantStep(vectorDimension: dimensions.columnCount) + throw PNNSError + .wrongPlaintextMatrixPacking(got: packing, expected: .diagonal(babyStepGiantStep: expectedBsgs)) + } + + let data = Array2d(data: values, rowCount: dimensions.rowCount, columnCount: dimensions.columnCount) + // Transposed from original shape, with extra zero columns. + // Encode diagonals + var packedValues = Array2d.zero( + rowCount: dimensions.columnCount.nextPowerOfTwo, + columnCount: dimensions.rowCount) + for rowIndex in 0... + let n = context.degree + for rowIndex in 0..: Equatable, Sendabl return try unpackDenseColumn() case .denseRow: return try unpackDenseRow() + case .diagonal: + // TODO: Implement + preconditionFailure("Unpacking diagonal plaintext matrix not supported") } } diff --git a/Tests/HomomorphicEncryptionTests/Array2dTests.swift b/Tests/HomomorphicEncryptionTests/Array2dTests.swift index 44ff0da2..59528a60 100644 --- a/Tests/HomomorphicEncryptionTests/Array2dTests.swift +++ b/Tests/HomomorphicEncryptionTests/Array2dTests.swift @@ -16,7 +16,7 @@ import XCTest class Array2dTests: XCTestCase { - func testZeroize() { + func testZeroAndZeroize() { func runTest(_: T.Type) { let data = [T](1...16) var array = Array2d(data: data, rowCount: 2, columnCount: 8) @@ -27,6 +27,7 @@ class Array2dTests: XCTestCase { rowCount: 2, columnCount: 8) XCTAssertEqual(array, zero) + XCTAssertEqual(array, Array2d.zero(rowCount: 2, columnCount: 8)) } runTest(Int.self) runTest(Int32.self) diff --git a/Tests/HomomorphicEncryptionTests/ScalarTests.swift b/Tests/HomomorphicEncryptionTests/ScalarTests.swift index 40bb6eec..0c2e0d7c 100644 --- a/Tests/HomomorphicEncryptionTests/ScalarTests.swift +++ b/Tests/HomomorphicEncryptionTests/ScalarTests.swift @@ -116,13 +116,21 @@ class ScalarTests: XCTestCase { } func testNextMultiple() { - XCTAssertEqual(0.nextMultiple(of: 7, variableTime: true), 0) + XCTAssertEqual(0.nextMultiple(of: 0, variableTime: true), 0) XCTAssertEqual(0.nextMultiple(of: 7, variableTime: true), 0) XCTAssertEqual(3.nextMultiple(of: 7, variableTime: true), 7) XCTAssertEqual(7.nextMultiple(of: 7, variableTime: true), 7) XCTAssertEqual(8.nextMultiple(of: 7, variableTime: true), 14) } + func testPreviousMultiple() { + XCTAssertEqual(0.previousMultiple(of: 0, variableTime: true), 0) + XCTAssertEqual(0.previousMultiple(of: 7, variableTime: true), 0) + XCTAssertEqual(3.previousMultiple(of: 7, variableTime: true), 0) + XCTAssertEqual(7.previousMultiple(of: 7, variableTime: true), 7) + XCTAssertEqual(8.previousMultiple(of: 7, variableTime: true), 7) + } + func testIsPrime() { let smallPrimes = [2, 3, 5, UInt32(1 << 14) - 65, UInt32(1 << 15) - 49, UInt32(1 << 16) - 17] let primes = smallPrimes + [(UInt32(1) << 28) - 183, (UInt32(1) << 29) - 3] diff --git a/Tests/PrivateNearestNeighborsSearchTests/PlaintextMatrixTests.swift b/Tests/PrivateNearestNeighborsSearchTests/PlaintextMatrixTests.swift index 25d6564a..f129e34d 100644 --- a/Tests/PrivateNearestNeighborsSearchTests/PlaintextMatrixTests.swift +++ b/Tests/PrivateNearestNeighborsSearchTests/PlaintextMatrixTests.swift @@ -29,7 +29,8 @@ final class PlaintextMatrixTests: XCTestCase { func testPlaintextMatrixError() throws { func runTest(rlweParams: PredefinedRlweParameters, _: Scheme.Type) throws { let encryptionParams = try EncryptionParameters(from: rlweParams) - guard encryptionParams.supportsSimdEncoding else { + // Parameters with large polyDegree are slow in debug mode + guard encryptionParams.supportsSimdEncoding, encryptionParams.polyDegree <= 16 else { return } let rowCount = encryptionParams.polyDegree @@ -141,7 +142,12 @@ final class PlaintextMatrixTests: XCTestCase { XCTAssertEqual(plaintextMatrix.packing, packing) XCTAssertEqual(plaintextMatrix.context, context) // Test round-trip - XCTAssertEqual(try plaintextMatrix.unpack(), encodeValues.flatMap { $0 }) + switch packing { + case .diagonal: // TODO: test .diagonal once implemented + break + default: + XCTAssertEqual(try plaintextMatrix.unpack(), encodeValues.flatMap { $0 }) + } // Test representation XCTAssertEqual(plaintextMatrix.plaintexts.count, expected.count) @@ -340,6 +346,109 @@ final class PlaintextMatrixTests: XCTestCase { try runTest(for: Bfv.self) } + func testPlaintextMatrixDiagonal() throws { + let kats: [((rowCount: Int, columnCount: Int), expected: [[Int]])] = [ + ((1, 3), [ + [1, 0, 0, 0, 0, 0, 0, 0], + [2, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 3, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0], + ]), + ((2, 3), [ + [1, 5, 0, 0, 0, 0, 0, 0], + [2, 6, 0, 0, 0, 0, 0, 0], + [0, 0, 3, 0, 0, 0, 0, 0], + [0, 0, 0, 4, 0, 0, 0, 0], + ]), + ((3, 3), [ + [1, 5, 9, 0, 0, 0, 0, 0], + [2, 6, 0, 0, 0, 0, 0, 0], + [7, 0, 3, 0, 0, 0, 0, 0], + [8, 0, 0, 4, 0, 0, 0, 0], + ]), + ((4, 3), [ + [1, 5, 9, 0, 0, 0, 0, 0], + [2, 6, 0, 10, 0, 0, 0, 0], + [7, 11, 3, 0, 0, 0, 0, 0], + [8, 12, 0, 4, 0, 0, 0, 0], + ]), + ((7, 3), [ + [1, 5, 9, 0, 13, 17, 21, 0], + [2, 6, 0, 10, 14, 18, 0, 0], + [7, 11, 3, 0, 19, 0, 15, 0], + [8, 12, 0, 4, 20, 0, 0, 16], + ]), + ((10, 3), [ + [1, 5, 9, 0, 13, 17, 21, 0], + [25, 29, 0, 0, 0, 0, 0, 0], + [2, 6, 0, 10, 14, 18, 0, 22], + [26, 30, 0, 0, 0, 0, 0, 0], + [7, 11, 3, 0, 19, 23, 15, 0], + [0, 0, 27, 0, 0, 0, 0, 0], + [8, 12, 0, 4, 20, 24, 0, 16], + [0, 0, 0, 28, 0, 0, 0, 0], + ]), + ((1, 4), [[1, 0, 0, 0, 0, 0, 0, 0], + [2, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 3, 0, 0, 0, 0, 0], + [0, 0, 4, 0, 0, 0, 0, 0]]), + ((2, 4), [[1, 6, 0, 0, 0, 0, 0, 0], + [2, 7, 0, 0, 0, 0, 0, 0], + [0, 0, 3, 8, 0, 0, 0, 0], + [0, 0, 4, 5, 0, 0, 0, 0]]), + ((3, 4), [[1, 6, 11, 0, 0, 0, 0, 0], + [2, 7, 12, 0, 0, 0, 0, 0], + [9, 0, 3, 8, 0, 0, 0, 0], + [10, 0, 4, 5, 0, 0, 0, 0]]), + ((4, 4), [[1, 6, 11, 16, 0, 0, 0, 0], + [2, 7, 12, 13, 0, 0, 0, 0], + [9, 14, 3, 8, 0, 0, 0, 0], + [10, 15, 4, 5, 0, 0, 0, 0]]), + ((7, 4), [[1, 6, 11, 16, 17, 22, 27, 0], + [2, 7, 12, 13, 18, 23, 28, 0], + [9, 14, 3, 8, 25, 0, 19, 24], + [10, 15, 4, 5, 26, 0, 20, 21]]), + ((8, 4), [[1, 6, 11, 16, 17, 22, 27, 32], + [2, 7, 12, 13, 18, 23, 28, 29], + [9, 14, 3, 8, 25, 30, 19, 24], + [10, 15, 4, 5, 26, 31, 20, 21]]), + ((9, 4), [[1, 6, 11, 16, 17, 22, 27, 32], + [33, 0, 0, 0, 0, 0, 0, 0], + [2, 7, 12, 13, 18, 23, 28, 29], + [34, 0, 0, 0, 0, 0, 0, 0], + [9, 14, 3, 8, 25, 30, 19, 24], + [0, 0, 35, 0, 0, 0, 0, 0], + [10, 15, 4, 5, 26, 31, 20, 21], + [0, 0, 36, 0, 0, 0, 0, 0]]), + ] + + func runTest(for _: Scheme.Type) throws { + let encryptionParams = try EncryptionParameters( + polyDegree: 8, + plaintextModulus: 1153, + coefficientModuli: Scheme.Scalar + .generatePrimes( + significantBitCounts: [25, 25], + preferringSmall: false, + nttDegree: 8), + errorStdDev: ErrorStdDev.stdDev32, + securityLevel: SecurityLevel.unchecked) + let context = try Context(encryptionParameters: encryptionParams) + for ((rowCount, columnCount), expected) in kats { + let dimensions = try MatrixDimensions(rowCount: rowCount, columnCount: columnCount) + let bsgs = BabyStepGiantStep(vectorDimension: dimensions.columnCount.nextPowerOfTwo) + try runPlaintextMatrixInitTest( + context: context, + dimensions: dimensions, + packing: .diagonal(babyStepGiantStep: bsgs), + expected: expected) + } + } + try runTest(for: NoOpScheme.self) + try runTest(for: Bfv.self) + try runTest(for: Bfv.self) + } + func testPlaintextMatrixConversion() throws { func runTest(for _: Scheme.Type) throws { let rlweParams = PredefinedRlweParameters.insecure_n_8_logq_5x18_logt_5