Async ciphertext APIs for rotation & mod-switching

Author: fboemer

Sources/HomomorphicEncryption/Ciphertext.swift

@@ -937,4 +937,90 @@ extension Ciphertext {
         }
         throw HeError.errorCastingPolyFormat(from: Format.self, to: Scheme.CanonicalCiphertextFormat.self)
     }
+
+    // MARK: Async rotations
+
+    /// Asynchronously rotates the columns of a ciphertext.
+    ///
+    /// - Parameters:
+    ///   - step: Number of slots to rotate. Negative values indicate a left rotation, and positive values indicate a
+    /// right rotation. Must have absolute value in `[1, N / 2 - 1]` where `N` is the RLWE ring dimension, given by
+    /// ``EncryptionParameters/polyDegree``.
+    ///   - evaluationKey: Evaluation key to use in the HE computation. Must contain the Galois element associated with
+    /// `step`, see ``GaloisElement/rotatingColumns(by:degree:)``.
+    /// - Throws: failure to rotate ciphertext's columns.
+    /// - seealso: ``HeScheme/rotateColumns(of:by:using:)-7h3fz`` for an alternate API and more information.
+    @inlinable
+    public mutating func rotateColumns(by step: Int,
+                                       using evaluationKey: EvaluationKey<Scheme>) async throws
+        where Format == Scheme.CanonicalCiphertextFormat
+    {
+        try await Scheme.rotateColumnsAsync(of: &self, by: step, using: evaluationKey)
+    }
+
+    /// Asynchronously swaps the rows of a ciphertext.
+    ///
+    /// A plaintext in ``EncodeFormat/simd`` format can be viewed a `2 x (N / 2)` matrix of coefficients.
+    /// For instance, for `N = 8`, given a ciphertext encrypting a plaintext with values
+    /// ```
+    /// [1, 2, 3, 4, 5, 6, 7, 8]
+    /// ```
+    /// calling ``HeScheme/swapRows(of:using:)`` with `step: 1` will yield a ciphertext decrypting to
+    /// ```
+    /// [5, 6, 7, 8, 1, 2, 3, 4]
+    /// ```
+    /// - Parameter evaluationKey: Evaluation key to use in the HE computation. Must contain the Galois element
+    /// associated with `step`, see ``GaloisElement/rotatingColumns(by:degree:)``.
+    /// - Throws: error upon failure to swap the ciphertext's rows.
+    /// - seealso: ``HeScheme/swapRows(of:using:)-50tac`` for an alternate API.
+    @inlinable
+    public mutating func swapRows(using evaluationKey: EvaluationKey<Scheme>) async throws
+        where Format == Scheme.CanonicalCiphertextFormat
+    {
+        try await Scheme.swapRowsAsync(of: &self, using: evaluationKey)
+    }
+
+    /// Asynchronously performs modulus switching on the ciphertext.
+    ///
+    /// - Throws: Error upon failure to mod-switch.
+    /// - seealso: ``HeScheme/modSwitchDown(_:)`` for an alternative API and more information.
+    @inlinable
+    public mutating func modSwitchDown() async throws where Format == Scheme.CanonicalCiphertextFormat {
+        try await Scheme.modSwitchDownAsync(&self)
+    }
+
+    /// Asynchronously performs modulus switching to a single modulus.
+    ///
+    /// If the ciphertext already has a single modulus, this is a no-op.
+    /// - Throws: Error upon failure to modulus switch.
+    /// - seealso: ``Ciphertext/modSwitchDown()`` for more information and an alternative API.
+    @inlinable
+    public mutating func modSwitchDownToSingle() async throws where Format == Scheme.CanonicalCiphertextFormat {
+        try await Scheme.modSwitchDownToSingleAsync(&self)
+    }
+}
+
+extension Ciphertext where Format == Scheme.CanonicalCiphertextFormat {
+    /// Asynchronously applies a Galois transformation.
+    ///
+    /// - Parameters:
+    ///   - 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``.
+    ///   - key: Evaluation key. Must contain Galois element `element`.
+    /// - Throws: Error upon failure to apply the Galois transformation.
+    /// - seealso: ``HeScheme/applyGalois(ciphertext:element:using:)`` for an alternative API and more information.
+    @inlinable
+    public mutating func applyGalois(element: Int, using key: EvaluationKey<Scheme>) async throws {
+        try await Scheme.applyGaloisAsync(ciphertext: &self, element: element, using: key)
+    }
+
+    /// Asynchronously Relinearizes the ciphertext.
+    ///
+    /// - Parameter key: Evaluation key to relinearize with. Must contain a `RelinearizationKey`.
+    /// - Throws: Error upon failure to relinearize.
+    /// - seealso: ``HeScheme/relinearize(_:using:)`` for an alternative API and more information.
+    @inlinable
+    public mutating func relinearize(using key: EvaluationKey<Scheme>) async throws {
+        try await Scheme.relinearizeAsync(&self, using: key)
+    }
 }

Sources/PrivateNearestNeighborSearch/CiphertextMatrix.swift

@@ -337,10 +337,7 @@ extension CiphertextMatrix {
         let rotateCount = simdColumnCount / (copiesInMask * columnCountPowerOfTwo) - 1
         var ciphertextCopyRight = ciphertext
         for await _ in (0..<rotateCount).async {
-            try await Scheme.rotateColumnsAsync(
-                of: &ciphertextCopyRight,
-                by: columnCountPowerOfTwo,
-                using: evaluationKey)
+            try await ciphertextCopyRight.rotateColumns(by: columnCountPowerOfTwo, using: evaluationKey)
             try await ciphertext += ciphertextCopyRight
         }
         // e.g., `ciphertext` now encrypts
@@ -349,7 +346,7 @@ extension CiphertextMatrix {
 
         // Duplicate values to both SIMD rows
         var ciphertextCopy = ciphertext
-        try await Scheme.swapRowsAsync(of: &ciphertextCopy, using: evaluationKey)
+        try await ciphertextCopy.swapRows(using: evaluationKey)
         try await ciphertext += ciphertextCopy
         // e.g., `ciphertext` now encrypts
         // [[3, 4, 3, 4, 3, 4, 3, 4],

Sources/PrivateNearestNeighborSearch/MatrixMultiplication.swift

@@ -176,7 +176,7 @@ extension PlaintextMatrix {
         for step in 0..<babyStepGiantStep.babyStep {
             rotatedStates.append(state)
             if step != babyStepGiantStep.babyStep - 1 {
-                try await Scheme.rotateColumnsAsync(of: &state, by: -1, using: evaluationKey)
+                try await state.rotateColumns(by: -1, using: evaluationKey)
             }
         }
         let rotatedCiphertexts: [Scheme.EvalCiphertext] = try await .init(

Sources/_HomomorphicEncryptionExtras/Ciphertext.swift

@@ -32,4 +32,23 @@ extension Ciphertext {
     {
         try Scheme.rotateColumnsMultiStep(of: &self, by: step, using: evaluationKey)
     }
+
+    /// Asynchronously rotates the columns of the ciphertext by combining multiple rotation steps corresponding to
+    /// Galois elements
+    /// available in the `evaluationKey`.
+    ///
+    /// - Parameters:
+    ///   - step: Number of slots to rotate. Negative values indicate a left rotation, and positive values indicate a
+    /// right rotation. Must have absolute value in `[1, N / 2 - 1]` where `N` is the RLWE ring dimension, given by
+    /// `EncryptionParameters/polyDegree`.
+    ///   - evaluationKey: Evaluation key to use in the HE computation. Must contain Galois elements which can be
+    /// combined for the desired rotation step.
+    /// - Throws: Error upon failure to rotate ciphertext's columns.
+    /// - seealso: `HeScheme/_rotateColumnsMultiStep(of:by:using:)` for an alternative API and more information.
+    @inlinable
+    public mutating func rotateColumnsMultiStep(by step: Int, using evaluationKey: EvaluationKey<Scheme>) async throws
+        where Format == Scheme.CanonicalCiphertextFormat
+    {
+        try await Scheme.rotateColumnsMultiStepAsync(of: &self, by: step, using: evaluationKey)
+    }
 }

Sources/_TestUtilities/HeApiTestUtils.swift

@@ -567,7 +567,7 @@ public enum HeAPITestHelpers {
             try await ciphertextProduct2 *= ciphertext2
 
             var relinearizedProd = ciphertextProduct
-            try relinearizedProd.relinearize(using: #require(testEnv.evaluationKey))
+            try await relinearizedProd.relinearize(using: #require(testEnv.evaluationKey))
             #expect(relinearizedProd.polys.count == Scheme.freshCiphertextPolyCount)
 
             let coeffCiphertext = try await ciphertextProduct.convertToCoeffFormat()
@@ -1233,7 +1233,7 @@ public enum HeAPITestHelpers {
             // with mod-switch down
             if context.coefficientModuli.count > 2 {
                 var ciphertext = testEnv.ciphertext1
-                try ciphertext.modSwitchDown()
+                try await ciphertext.modSwitchDown()
                 let evalCiphertext = try await ciphertext.convertToEvalFormat()
                 let moduliCount = evalCiphertext.moduli.count
                 let evalPlaintext = try testEnv.context.encode(values: data2, format: .simd, moduliCount: moduliCount)
@@ -1253,84 +1253,108 @@ public enum HeAPITestHelpers {
         context: Scheme.Context,
         scheme _: Scheme.Type) async throws
     {
-        func runRotationTest(context: Scheme.Context, galoisElements: [Int], multiStep: Bool) async throws {
-            let degree = context.degree
+        func runRotationTestSync(context: Scheme.Context, galoisElements: [Int], multiStep: Bool) throws {
             let testEnv = try TestEnv<Scheme>(context: context, format: .simd, galoisElements: galoisElements)
             let evaluationKey = try #require(testEnv.evaluationKey)
+            let degree = context.degree
+
             for step in 1..<min(8, degree / 2) {
                 let expectedData = Array(testEnv.data1[degree / 2 - step..<degree / 2] + testEnv
                     .data1[0..<degree / 2 - step] + testEnv
                     .data1[degree - step..<degree] + testEnv.data1[degree / 2..<degree - step])
                 var rotatedCiphertext = testEnv.ciphertext1
-                var rotatedCiphertextAsync = testEnv.ciphertext1
                 if multiStep {
                     try rotatedCiphertext.rotateColumnsMultiStep(by: step, using: evaluationKey)
-                    try await Scheme.rotateColumnsMultiStepAsync(
-                        of: &rotatedCiphertextAsync,
-                        by: step,
-                        using: evaluationKey)
                 } else {
                     try rotatedCiphertext.rotateColumns(by: step, using: evaluationKey)
-                    try await Scheme.rotateColumnsAsync(of: &rotatedCiphertextAsync, by: step, using: evaluationKey)
                 }
                 try testEnv.checkDecryptsDecodes(ciphertext: rotatedCiphertext, format: .simd, expected: expectedData)
-                try testEnv.checkDecryptsDecodes(
-                    ciphertext: rotatedCiphertextAsync,
-                    format: .simd,
-                    expected: expectedData)
 
                 if multiStep {
                     try rotatedCiphertext.rotateColumnsMultiStep(by: -step, using: evaluationKey)
-                    try await Scheme.rotateColumnsMultiStepAsync(
-                        of: &rotatedCiphertextAsync,
-                        by: -step,
-                        using: evaluationKey)
                 } else {
                     try rotatedCiphertext.rotateColumns(by: -step, using: evaluationKey)
-                    try await Scheme.rotateColumnsAsync(of: &rotatedCiphertextAsync, by: -step, using: evaluationKey)
                 }
-                try testEnv.checkDecryptsDecodes(ciphertext: rotatedCiphertext,
-                                                 format: .simd,
-                                                 expected: testEnv.data1)
-                try testEnv.checkDecryptsDecodes(ciphertext: rotatedCiphertextAsync,
-                                                 format: .simd,
-                                                 expected: testEnv.data1)
+                try testEnv.checkDecryptsDecodes(ciphertext: rotatedCiphertext, format: .simd, expected: testEnv.data1)
+            }
+        }
+
+        func runRotationTestAsync(context: Scheme.Context, galoisElements: [Int], multiStep: Bool) async throws {
+            let testEnv = try TestEnv<Scheme>(context: context, format: .simd, galoisElements: galoisElements)
+            let evaluationKey = try #require(testEnv.evaluationKey)
+            let degree = context.degree
+
+            for step in 1..<min(8, degree / 2) {
+                let expectedData = Array(testEnv.data1[degree / 2 - step..<degree / 2] + testEnv
+                    .data1[0..<degree / 2 - step] + testEnv
+                    .data1[degree - step..<degree] + testEnv.data1[degree / 2..<degree - step])
+                var rotatedCiphertext = testEnv.ciphertext1
+                if multiStep {
+                    try await rotatedCiphertext.rotateColumnsMultiStep(by: step, using: evaluationKey)
+                } else {
+                    try await rotatedCiphertext.rotateColumns(by: step, using: evaluationKey)
+                }
+                try testEnv.checkDecryptsDecodes(ciphertext: rotatedCiphertext, format: .simd, expected: expectedData)
+
+                if multiStep {
+                    try await rotatedCiphertext.rotateColumnsMultiStep(by: -step, using: evaluationKey)
+                } else {
+                    try await rotatedCiphertext.rotateColumns(by: -step, using: evaluationKey)
+                }
+                try testEnv.checkDecryptsDecodes(ciphertext: rotatedCiphertext, format: .simd, expected: testEnv.data1)
             }
         }
 
         guard context.supportsSimdEncoding, context.supportsEvaluationKey else {
             return
         }
 
+        let degree = context.degree
+        let galoisElementsRotate = try (1..<(degree >> 1)).flatMap { step in
+            try [
+                GaloisElement.rotatingColumns(by: step, degree: degree),
+                GaloisElement.rotatingColumns(by: -step, degree: degree),
+            ]
+        }
+        let galoisElementsMultiStep = try GaloisElement.rotatingColumnsMultiStep(degree: degree)
+
+        try runRotationTestSync(context: context, galoisElements: galoisElementsRotate, multiStep: false)
+        try await runRotationTestAsync(context: context, galoisElements: galoisElementsMultiStep, multiStep: true)
+    }
+
+    /// Testing ciphertext rotation of the scheme.
+    @inlinable
+    public static func schemeSwapRowsTest<Scheme: HeScheme>(
+        context: Scheme.Context,
+        scheme _: Scheme.Type) async throws
+    {
+        guard context.supportsSimdEncoding, context.supportsEvaluationKey else {
+            return
+        }
         let degree = context.degree
         let galoisElementsSwap = [GaloisElement.swappingRows(degree: degree)]
         let testEnv = try TestEnv<Scheme>(context: context, format: .simd, galoisElements: galoisElementsSwap)
         let evaluationKey = try #require(testEnv.evaluationKey)
         let expectedData = Array(testEnv.data1[degree / 2..<degree] + testEnv.data1[0..<degree / 2])
         var ciphertext = testEnv.ciphertext1
-        var ciphertextAsync = ciphertext
 
-        try ciphertext.swapRows(using: evaluationKey)
-        try await Scheme.swapRowsAsync(of: &ciphertextAsync, using: evaluationKey)
+        func swapRowsSyncTest() throws {
+            try ciphertext.swapRows(using: evaluationKey)
+            try testEnv.checkDecryptsDecodes(ciphertext: ciphertext, format: .simd, expected: expectedData)
 
-        try testEnv.checkDecryptsDecodes(ciphertext: ciphertext, format: .simd, expected: expectedData)
-        try testEnv.checkDecryptsDecodes(ciphertext: ciphertextAsync, format: .simd, expected: expectedData)
+            try ciphertext.swapRows(using: evaluationKey)
+            try testEnv.checkDecryptsDecodes(ciphertext: ciphertext, format: .simd, expected: testEnv.data1)
+        }
+        try swapRowsSyncTest()
 
-        try ciphertext.swapRows(using: evaluationKey)
-        try await Scheme.swapRowsAsync(of: &ciphertextAsync, using: evaluationKey)
-        try testEnv.checkDecryptsDecodes(ciphertext: ciphertext, format: .simd, expected: testEnv.data1)
-        try testEnv.checkDecryptsDecodes(ciphertext: ciphertextAsync, format: .simd, expected: testEnv.data1)
+        func swapRowsAsyncTest() async throws {
+            try await ciphertext.swapRows(using: evaluationKey)
+            try testEnv.checkDecryptsDecodes(ciphertext: ciphertext, format: .simd, expected: expectedData)
 
-        let galoisElementsRotate = try (1..<(degree >> 1)).flatMap { step in
-            try [
-                GaloisElement.rotatingColumns(by: step, degree: degree),
-                GaloisElement.rotatingColumns(by: -step, degree: degree),
-            ]
+            try await ciphertext.swapRows(using: evaluationKey)
+            try testEnv.checkDecryptsDecodes(ciphertext: ciphertext, format: .simd, expected: testEnv.data1)
         }
-        let galoisElementsMultiStep = try GaloisElement.rotatingColumnsMultiStep(degree: degree)
-
-        try await runRotationTest(context: context, galoisElements: galoisElementsRotate, multiStep: false)
-        try await runRotationTest(context: context, galoisElements: galoisElementsMultiStep, multiStep: true)
+        try await swapRowsAsyncTest()
     }
 
     /// Testing apply Galois element of the scheme.
@@ -1358,25 +1382,37 @@ public enum HeAPITestHelpers {
 
         let dataCount = testEnv.data1.count
         let halfDataCount = dataCount / 2
-        for (step, element) in elements.enumerated() {
-            for modSwitchCount in 0...max(0, context.coefficientModuli.count - 2) {
-                var rotatedCiphertext = testEnv.ciphertext1
-                var rotatedCiphertextAsync = testEnv.ciphertext1
+        func syncTest() throws {
+            for (step, element) in elements.enumerated() {
+                for modSwitchCount in 0...max(0, context.coefficientModuli.count - 2) {
+                    var rotatedCiphertext = testEnv.ciphertext1
+
+                    for _ in 0..<modSwitchCount {
+                        try rotatedCiphertext.modSwitchDown()
+                    }
+                    try rotatedCiphertext.applyGalois(element: element, using: evaluationKey)
+                    let expected = rotate(original: testEnv.data1, halfDataCount: halfDataCount, step: step + 1)
+                    try testEnv.checkDecryptsDecodes(ciphertext: rotatedCiphertext, format: .simd, expected: expected)
+                }
+            }
+        }
+        try syncTest()
 
-                for _ in 0..<modSwitchCount {
-                    try rotatedCiphertext.modSwitchDown()
-                    try await Scheme.modSwitchDownAsync(&rotatedCiphertextAsync)
+        func asyncTest() async throws {
+            for (step, element) in elements.enumerated() {
+                for modSwitchCount in 0...max(0, context.coefficientModuli.count - 2) {
+                    var rotatedCiphertext = testEnv.ciphertext1
+
+                    for _ in 0..<modSwitchCount {
+                        try await rotatedCiphertext.modSwitchDown()
+                    }
+                    try await rotatedCiphertext.applyGalois(element: element, using: evaluationKey)
+                    let expected = rotate(original: testEnv.data1, halfDataCount: halfDataCount, step: step + 1)
+                    try testEnv.checkDecryptsDecodes(ciphertext: rotatedCiphertext, format: .simd, expected: expected)
                 }
-                try rotatedCiphertext.applyGalois(element: element, using: evaluationKey)
-                try await Scheme.applyGaloisAsync(
-                    ciphertext: &rotatedCiphertextAsync,
-                    element: element,
-                    using: evaluationKey)
-                let expected = rotate(original: testEnv.data1, halfDataCount: halfDataCount, step: step + 1)
-                try testEnv.checkDecryptsDecodes(ciphertext: rotatedCiphertext, format: .simd, expected: expected)
-                try testEnv.checkDecryptsDecodes(ciphertext: rotatedCiphertextAsync, format: .simd, expected: expected)
             }
         }
+        try await asyncTest()
     }
 
     /// Testing noise budget estimation.