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

@@ -276,6 +276,20 @@ LogicalResult PlaintextSpaceAttr::verify(
   return success();
 }
 
+LogicalResult NewLWECiphertextType::verify(
+    llvm::function_ref<mlir::InFlightDiagnostic()> 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

lib/Dialect/LWE/IR/NewLWEAttributes.td

@@ -271,30 +271,27 @@ def LWE_PlaintextSpaceAttr : AttrDef<LWE_Dialect, "PlaintextSpace"> {
 def LWE_KeyAttr : AttrDef<LWE_Dialect, "Key"> {
   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<LWE_Dialect, "CiphertextSpace"> {
     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<storage=i64, modulus=298374>
+    #ring = #polynomial.ring<coefficientType = !cmod, modulus = !ideal>
+    #ciphertext_space = #lwe.ciphertext_space<ring = #ring, encryption_type = #encryption_type>
+    ```
+
+    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<storage=i64, modulus=2251799814045697>
+    !limb2 = modarith<storage=i64, modulus=65537>
+    #rns_mod = !rns.rns<!limb1, !limb2>
+    #ring = #polynomial.ring<coefficientType = #rns_mod, modulus = #ideal>
+    #ciphertext_space = #lwe.ciphertext_space<ring = #ring, encryption_type = #encryption_type>
+    ```
+
+    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<coefficientType = i32, modulus = #ideal>
+    #ciphertext_space = #lwe.ciphertext_space<ring = #ring, encryption_type = #encryption_type>
+    ```
 
     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) `>`";

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<NewLWECiphertext, "new-lwe-ciphertext-like">;

tests/Dialect/LWE/IR/attributes.mlir

@@ -146,9 +146,9 @@ func.func @test_fn() {
 
 // -----
 
-#key = #lwe.key<id = "1234">
-#key_rlwe_rotate = #lwe.key<id = "1234", basis = 0, 2>
-#key_rlwe_2 = #lwe.key<id = "1234", size = 2>
+#key = #lwe.key<>
+#key_rlwe_rotate = #lwe.key<rotation_basis = 0, 2>
+#key_rlwe_2 = #lwe.key<rotation_basis = 0, 1, 1>
 
 // CHECK-LABEL: test_fn
 func.func @test_fn() {

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<scaling_factor = 10000>
 // expected-error@below {{modulus must be 1 mod n for full CRT packing}}
 #plaintext_space = #lwe.plaintext_space<ring = #ring, encoding = #crt>
+
+// -----
+
+#key = #lwe.key<rotation_basis = 0, 1, 1>
+#my_poly = #polynomial.int_polynomial<1 + x**1024>
+#ring = #polynomial.ring<coefficientType = i32, coefficientModulus = 7917 : i32, polynomialModulus=#my_poly>
+!public_key = !lwe.new_lwe_public_key<key = #key, ring = #ring>
+
+#preserve_overflow = #lwe.preserve_overflow<>
+#application_data = #lwe.application_data<message_type = i1, overflow = #preserve_overflow>
+#inverse_canonical_enc = #lwe.inverse_canonical_encoding<scaling_factor = 10000>
+#plaintext_space = #lwe.plaintext_space<ring = #ring, encoding = #inverse_canonical_enc>
+#ciphertext_space = #lwe.ciphertext_space<ring = #ring, encryption_type = msb, size = 2>
+#modulus_chain = #lwe.modulus_chain<elements = <7917 : i32>, current = 0>
+
+// expected-error@below {{ciphertext space size 2 does not match rotation basis size 3}}
+!new_lwe_ciphertext = !lwe.new_lwe_ciphertext<application_data = #application_data, plaintext_space = #plaintext_space, key = #key, ciphertext_space = #ciphertext_space, modulus_chain = #modulus_chain>

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<id = "1234", size = 1>
+#key = #lwe.key<>
 !secret_key = !lwe.new_lwe_secret_key<key = #key, ring = #ring>
 
 // CHECK-LABEL: test_new_lwe_secret_key