From 8e94c4518fbd577ac3b9e16ecb0acca07cf4d80a Mon Sep 17 00:00:00 2001 From: Asra Ali Date: Mon, 11 Nov 2024 11:32:11 -0800 Subject: [PATCH] lwe: add ciphertext size and clarify modulus types on new lwe types This PR: * Removes key identifier until we have the need for >1 "keyholder" * Adds a ciphertext size descriping the "vectorspace" dimension of the ciphertexts (number of polynomials) * Adds a description of how rings with RNS types look PiperOrigin-RevId: 695422841 --- lib/Dialect/LWE/IR/LWEDialect.cpp | 14 ++++ lib/Dialect/LWE/IR/NewLWEAttributes.td | 87 ++++++++++++++------- lib/Dialect/LWE/IR/NewLWETypes.td | 2 + tests/Dialect/LWE/IR/attributes.mlir | 6 +- tests/Dialect/LWE/IR/attributes_errors.mlir | 17 ++++ tests/Dialect/LWE/IR/types.mlir | 2 +- 6 files changed, 97 insertions(+), 31 deletions(-) diff --git a/lib/Dialect/LWE/IR/LWEDialect.cpp b/lib/Dialect/LWE/IR/LWEDialect.cpp index 665b78c5b..0053079ea 100644 --- a/lib/Dialect/LWE/IR/LWEDialect.cpp +++ b/lib/Dialect/LWE/IR/LWEDialect.cpp @@ -276,6 +276,20 @@ LogicalResult PlaintextSpaceAttr::verify( return success(); } +LogicalResult NewLWECiphertextType::verify( + llvm::function_ref emitError, + mlir::heir::lwe::ApplicationDataAttr, mlir::heir::lwe::PlaintextSpaceAttr, + mlir::heir::lwe::CiphertextSpaceAttr ciphertextSpace, + mlir::heir::lwe::KeyAttr keyAttr, mlir::heir::lwe::ModulusChainAttr) { + if (!keyAttr.getRotationBasis().empty() && + (ciphertextSpace.getSize() != keyAttr.getRotationBasis().size())) { + return emitError() << "ciphertext space size " << ciphertextSpace.getSize() + << " does not match rotation basis size " + << keyAttr.getRotationBasis().size(); + } + return success(); +} + } // namespace lwe } // namespace heir } // namespace mlir diff --git a/lib/Dialect/LWE/IR/NewLWEAttributes.td b/lib/Dialect/LWE/IR/NewLWEAttributes.td index 6870f1f7d..123104c98 100644 --- a/lib/Dialect/LWE/IR/NewLWEAttributes.td +++ b/lib/Dialect/LWE/IR/NewLWEAttributes.td @@ -271,30 +271,27 @@ def LWE_PlaintextSpaceAttr : AttrDef { def LWE_KeyAttr : AttrDef { let mnemonic = "key"; let description = [{ - An attribute describing the key used for encrypting the ciphertext. - - This attribute includes a key identifier for the original key used to - encrypt the secret key. - - The `key_size` parameter is used to describe the number of polynomials of - the secret key. This is typically $1$ for RLWE ciphertexts and greater than - $1$ for LWE instances. A ciphertext encrypted with a `key_size` of $k$ will - have size $k+1$. - - The key basis describes the inner product used in the phase calculation in - decryption. This attribute is only supported for RLWE ciphertexts whose - `key_size` is $1$. An RLWE ciphertext is canonically encrypted against key - basis `(1, s)`. After a multiplication, its size will increase and the basis - will be `(1, s, s^2)`. The array that represents the key basis is - constructed by listing the powers of `s` at each position of the array. For - example, `(1, s, s^2)` corresponds to `[0, 1, 2]`, while `(1, s^2)` - corresponds to `[0, 2]`. + An attribute describing the key used for decrypting the ciphertext. + + The key attribute describes the key with which decryption is performed. For + example, if the decryption of a ciphertext $c = (c_0(x), c_1(x))$ is + performed by computing the inner product $(c_0(x), c_1(x)) \cdot (1, s(x))$ + then the decryption key has a basis $(1, s(x))$. + + The `rotation_basis` describes the powers of `x` at each position of the + array, so $(1, s(x))$ becomes $[0, 1]$. After a Galois rotation mapping $x + -> x^i$, the key will have a new basis $(1, s(x^i))$, which will be + represented by $[0, i]$. + + Note that the ciphertext size must match the size of the key's basis. + + By default, the rotation basis is assumed to be $[0, 1, ..., 1]$ so that + after a multiplication of two ciphertexts, the `rotation_basis` is $[0, 1, + 1]$. }]; let parameters = (ins - "::mlir::StringAttr":$id, - DefaultValuedParameter<"unsigned", "1">:$size, - OptionalArrayRefParameter<"unsigned int">:$basis + OptionalArrayRefParameter<"unsigned int">:$rotation_basis ); let assemblyFormat = "`<` struct(params) `>`"; @@ -314,19 +311,55 @@ def LWE_CiphertextSpaceAttr : AttrDef { An attribute describing the ciphertext space and the transformation from plaintext space to ciphertext space of an FHE scheme. - The ciphertext space information includes the ring structure, which contains - the ciphertext modulus $q$. Ciphertexts using an RNS representation for $q$ - will represent their ciphertext components in the ring attribute. Scalar LWE - ciphertexts (as opposed to RLWE) use an ideal polynomial of degree 1, $x$. - CGGI ciphertexts will typically use a power of two modulus. + The ciphertext space information includes the ring attribute, describing the + space that the ciphertext elements belong to. The ring attribute contains a + coefficient type attribute that describes the semantics of the coefficient. + For example, a ring modulo $1 + x^1024$ with coefficients modulo $q = + 298374$ will be described as + + ``` + !ideal = !polynomial.int_polynomial<1 + x**1024> + !cmod = !modarith + #ring = #polynomial.ring + #ciphertext_space = #lwe.ciphertext_space + ``` + + Ciphertexts using an RNS representation for $q$ will use an RNS type in + their ring's coefficient type attribute. + + ``` + !ideal = !polynomial.int_polynomial<1 + x**1024> + !limb1 = modarith + !limb2 = modarith + #rns_mod = !rns.rns + #ring = #polynomial.ring + #ciphertext_space = #lwe.ciphertext_space + ``` + + Scalar LWE ciphertexts (like those used in CGGI) use an ideal polynomial of + degree 1, $x$. CGGI ciphertexts will typically use a power of two modulus + and may use a native integer type for its coefficient modulus. + + ``` + !ideal = !polynomial.int_polynomial<1 + x**1024> + #ring = #polynomial.ring + #ciphertext_space = #lwe.ciphertext_space + ``` The ciphertext encoding info is used to describe the way the plaintext data is encoded into the ciphertext (in the MSB, LSB, or mixed). + + The `size` parameter is used to describe the number of polynomials + comprising the ciphertext. This is typically $2$ for RLWE ciphertexts that + are made up of an $(a, b)$ pair and greater than $2$ for LWE instances. For + example, after an RLWE multiplication of two size $2$ ciphertexts, + the ciphertext's size will be $3$. }]; let parameters = (ins "::mlir::polynomial::RingAttr":$ring, - "::mlir::heir::lwe::LweEncryptionType":$encryption_type + "::mlir::heir::lwe::LweEncryptionType":$encryption_type, + DefaultValuedParameter<"unsigned", "1">:$size ); let assemblyFormat = "`<` struct(params) `>`"; diff --git a/lib/Dialect/LWE/IR/NewLWETypes.td b/lib/Dialect/LWE/IR/NewLWETypes.td index ce576d9b6..0edad9516 100644 --- a/lib/Dialect/LWE/IR/NewLWETypes.td +++ b/lib/Dialect/LWE/IR/NewLWETypes.td @@ -66,6 +66,8 @@ def NewLWECiphertext : LWE_Type<"NewLWECiphertext", "new_lwe_ciphertext"> { "KeyAttr":$key, OptionalParameter<"ModulusChainAttr">:$modulus_chain ); + + let genVerifyDecl = 1; } def NewLWECiphertextLike : TypeOrContainer; diff --git a/tests/Dialect/LWE/IR/attributes.mlir b/tests/Dialect/LWE/IR/attributes.mlir index 73afee962..63e8628c2 100644 --- a/tests/Dialect/LWE/IR/attributes.mlir +++ b/tests/Dialect/LWE/IR/attributes.mlir @@ -146,9 +146,9 @@ func.func @test_fn() { // ----- -#key = #lwe.key -#key_rlwe_rotate = #lwe.key -#key_rlwe_2 = #lwe.key +#key = #lwe.key<> +#key_rlwe_rotate = #lwe.key +#key_rlwe_2 = #lwe.key // CHECK-LABEL: test_fn func.func @test_fn() { diff --git a/tests/Dialect/LWE/IR/attributes_errors.mlir b/tests/Dialect/LWE/IR/attributes_errors.mlir index b8215be55..1b4967a1d 100644 --- a/tests/Dialect/LWE/IR/attributes_errors.mlir +++ b/tests/Dialect/LWE/IR/attributes_errors.mlir @@ -85,3 +85,20 @@ func.func @test_invalid_inverse_canonical_embedding_encoding() { #crt = #lwe.full_crt_packing_encoding // expected-error@below {{modulus must be 1 mod n for full CRT packing}} #plaintext_space = #lwe.plaintext_space + +// ----- + +#key = #lwe.key +#my_poly = #polynomial.int_polynomial<1 + x**1024> +#ring = #polynomial.ring +!public_key = !lwe.new_lwe_public_key + +#preserve_overflow = #lwe.preserve_overflow<> +#application_data = #lwe.application_data +#inverse_canonical_enc = #lwe.inverse_canonical_encoding +#plaintext_space = #lwe.plaintext_space +#ciphertext_space = #lwe.ciphertext_space +#modulus_chain = #lwe.modulus_chain, current = 0> + +// expected-error@below {{ciphertext space size 2 does not match rotation basis size 3}} +!new_lwe_ciphertext = !lwe.new_lwe_ciphertext diff --git a/tests/Dialect/LWE/IR/types.mlir b/tests/Dialect/LWE/IR/types.mlir index 73e1ed6ad..e72ef14bd 100644 --- a/tests/Dialect/LWE/IR/types.mlir +++ b/tests/Dialect/LWE/IR/types.mlir @@ -33,7 +33,7 @@ func.func @test_valid_rlwe_ciphertext(%arg0 : !ciphertext_rlwe) -> !ciphertext_r return %arg0 : !ciphertext_rlwe } -#key = #lwe.key +#key = #lwe.key<> !secret_key = !lwe.new_lwe_secret_key // CHECK-LABEL: test_new_lwe_secret_key