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

@@ -276,6 +276,19 @@ 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.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

lib/Dialect/LWE/IR/NewLWEAttributes.td

@@ -271,30 +271,23 @@ 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 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<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 = !mod_arith.mod_arith<modulus=298374 : i64>
+    #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.
+
+    // TODO(#1085): Validate syntax of polynomial ring after coefficientType changes.
+    ```
+    !ideal = !polynomial.int_polynomial<1 + x**1024>
+    !limb1 = !mod_arith.mod_arith<modulus=2251799814045697 : i64>
+    !limb2 = !mod_arith.mod_arith<modulus=65537 : i64>
+    #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", "2">:$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<slot_index = 2>
+#key_rlwe_2 = #lwe.key<slot_index = 0>
 
 // 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<slot_index = 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 = 3>
+#modulus_chain = #lwe.modulus_chain<elements = <7917 : i32>, 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<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