From f24037c18a974e40a417226e13e1b08c4d3e50e0 Mon Sep 17 00:00:00 2001 From: Asra Ali Date: Tue, 19 Nov 2024 12:30:10 -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: 698106948 --- lib/Dialect/LWE/IR/LWEDialect.cpp | 13 ++++ lib/Dialect/LWE/IR/NewLWEAttributes.td | 84 ++++++++++++++------- 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, 93 insertions(+), 31 deletions(-) diff --git a/lib/Dialect/LWE/IR/LWEDialect.cpp b/lib/Dialect/LWE/IR/LWEDialect.cpp index 74594b2d7..759f90a11 100644 --- a/lib/Dialect/LWE/IR/LWEDialect.cpp +++ b/lib/Dialect/LWE/IR/LWEDialect.cpp @@ -276,6 +276,19 @@ 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.getSlotIndex() != 0 && (ciphertextSpace.getSize() != 2)) { + return emitError() << "a ciphertext with nontrivial slot rotation must " + "have size 2, but found size " + << ciphertextSpace.getSize(); + } + 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..ce79d6471 100644 --- a/lib/Dialect/LWE/IR/NewLWEAttributes.td +++ b/lib/Dialect/LWE/IR/NewLWEAttributes.td @@ -271,30 +271,23 @@ 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 with which the message is currently + encrypted. + + The key attribute describes the key with which the message is currently + encrypted and decryption can be 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 key is $(1, s(x))$. + + The `slot_index` describes the key after using a Galois automorphism to + rotate the plaintext slots by `slot_index`. This will correspond to an + action $\phi_k: x \rightarrow x^k$ for some `k` that depends on the + structure of the Galois group for the chosen scheme parameters. The + corresponding key will have a new basis $(1, s(x^(k)))$. }]; let parameters = (ins - "::mlir::StringAttr":$id, - DefaultValuedParameter<"unsigned", "1">:$size, - OptionalArrayRefParameter<"unsigned int">:$basis + DefaultValuedParameter<"int", "0">:$slot_index ); let assemblyFormat = "`<` struct(params) `>`"; @@ -314,19 +307,56 @@ 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 = !mod_arith.mod_arith + #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. + + // TODO(#1085): Validate syntax of polynomial ring after coefficientType changes. + ``` + !ideal = !polynomial.int_polynomial<1 + x**1024> + !limb1 = !mod_arith.mod_arith + !limb2 = !mod_arith.mod_arith + #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", "2">:$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..4fd6d9a40 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..95b49c8ae 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 {{a ciphertext with nontrivial slot rotation must have size 2, but found 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