lwe: per-level encode level attr on LWE_RLWEEncodeOp

Thread the modulus-chain level onto LWE_RLWEEncodeOp (as an optional attr,
set from the ciphertext type's modulus chain in the secret->ckks encode
lowering and the client-side trivial encryption) so backends with canonical
per-level scales can encode at the scale of the target level (e.g.
GetScale(level)). TODO(#1643): LWEPlaintextType should carry RNS ring info
instead.

Author: AlexanderViand

lib/Dialect/LWE/IR/LWEOps.td

@@ -211,7 +211,8 @@ def LWE_RLWEEncodeOp : LWE_Op<"rlwe_encode", [
   let arguments = (ins
     SignlessIntegerOrFloatLike:$input,
     AnyPlaintextEncodingAttr:$encoding,
-    Polynomial_RingAttr:$ring
+    Polynomial_RingAttr:$ring,
+    OptionalAttr<I64Attr>:$level
   );
 
   let results = (outs LWEPlaintextLike:$output);

lib/Dialect/LWE/Transforms/ImplementTrivialEncryptionAsAddition.cpp

@@ -163,7 +163,8 @@ func::FuncOp getOrCreateEncryptionOfZerosFunc(func::FuncOp parentFunc,
   auto plaintextSpace = plaintextType.getPlaintextSpace();
   auto encodeOp = RLWEEncodeOp::create(
       builder, plaintextType, constantOp.getResult(),
-      plaintextSpace.getEncoding(), plaintextSpace.getRing());
+      plaintextSpace.getEncoding(), plaintextSpace.getRing(),
+      /*level=*/nullptr);
   auto encrypted = RLWEEncryptOp::create(
       builder, ciphertextType, encodeOp.getResult(), encFuncOp.getArgument(0));
   encrypted->setAttrs(originalOp->getAttrs());

lib/Dialect/Secret/Conversions/Patterns.cpp

@@ -86,10 +86,13 @@ LogicalResult ConvertClientConceal::lowerToTrivialEncryption(
 
   // Intentionally use op.getCleartext() because we don't want to type-convert
   // the input to a ciphertext.
+  IntegerAttr levelAttr;
+  if (auto mc = ctTy.getModulusChain())
+    levelAttr = rewriter.getI64IntegerAttr(mc.getCurrent());
   auto encoded = lwe::RLWEEncodeOp::create(
       rewriter, op.getLoc(), encodeOpResultTy, op.getCleartext(),
       ctTy.getPlaintextSpace().getEncoding(),
-      ctTy.getPlaintextSpace().getRing());
+      ctTy.getPlaintextSpace().getRing(), levelAttr);
   auto newOp =
       lwe::TrivialEncryptOp::create(rewriter, op.getLoc(), resultTy, encoded);
   newOp->setAttrs(op->getAttrs());
@@ -145,11 +148,15 @@ LogicalResult ConvertClientConceal::matchAndRewrite(
   auto plaintextTy = lwe::LWEPlaintextType::get(op.getContext(),
                                                 resultCtTy.getPlaintextSpace());
 
+  IntegerAttr encLevelAttr;
+  if (auto mc = resultCtTy.getModulusChain())
+    encLevelAttr = rewriter.getI64IntegerAttr(mc.getCurrent());
+
   auto encryptFn = [&](Value cleartext) -> lwe::RLWEEncryptOp {
-    auto encoded =
-        lwe::RLWEEncodeOp::create(rewriter, op.getLoc(), plaintextTy, cleartext,
-                                  resultCtTy.getPlaintextSpace().getEncoding(),
-                                  resultCtTy.getPlaintextSpace().getRing());
+    auto encoded = lwe::RLWEEncodeOp::create(
+        rewriter, op.getLoc(), plaintextTy, cleartext,
+        resultCtTy.getPlaintextSpace().getEncoding(),
+        resultCtTy.getPlaintextSpace().getRing(), encLevelAttr);
     auto encryptOp = lwe::RLWEEncryptOp::create(
         rewriter, op.getLoc(), resultCtTy, encoded.getResult(), keyBlockArg);
     // Copy attributes from the original op to preserve any mgmt attrs needed by

lib/Utils/ContextAwareConversionUtils.cpp

@@ -53,23 +53,27 @@ FailureOr<Value> encodeCleartextAsPlaintext(
     return failure();
   }
 
-  // TODO(#1643): inherit level information to plaintext type from init-op
-  // mgmt attr. This actually needs to make LWEPlaintextType RNS aware.
+  // TODO(#1643): LWEPlaintextType should carry RNS ring info. For now,
+  // the level is passed as an attribute on the encode op.
   auto plaintextTy = lwe::LWEPlaintextType::get(
       ctx, lwe::PlaintextSpaceAttr::get(ctx, plaintextSpace.getRing(),
                                         plaintextEncoding));
 
+  IntegerAttr levelAttr;
+  if (auto mc = ciphertextElementType.getModulusChain())
+    levelAttr = builder.getI64IntegerAttr(mc.getCurrent());
+
   // cleartext is a ciphertext-semantic tensor, so it
   // could be a tensor<Nxty>, tensor<k x N x ty>, (where k=1 is possible).
   // For rank 1, it's a single ciphertext and can be encoded directly.
   auto cleartextTensorTy = cast<RankedTensorType>(cleartext.getType());
   int64_t numSlots =
       cleartextTensorTy.getDimSize(cleartextTensorTy.getRank() - 1);
   if (cleartextTensorTy.getRank() == 1) {
-    Value encodeOp =
-        lwe::RLWEEncodeOp::create(builder, plaintextTy, cleartext,
-                                  plaintextEncoding, plaintextSpace.getRing())
-            .getResult();
+    Value encodeOp = lwe::RLWEEncodeOp::create(
+                         builder, plaintextTy, cleartext, plaintextEncoding,
+                         plaintextSpace.getRing(), levelAttr)
+                         .getResult();
     return encodeOp;
   }
 
@@ -87,10 +91,10 @@ FailureOr<Value> encodeCleartextAsPlaintext(
     SmallVector<OpFoldResult> strides(2, builder.getIndexAttr(1));
     auto slice = tensor::ExtractSliceOp::create(builder, sliceTy, cleartext,
                                                 offsets, sizes, strides);
-    Value encodedSlice =
-        lwe::RLWEEncodeOp::create(builder, plaintextTy, slice,
-                                  plaintextEncoding, plaintextSpace.getRing())
-            .getResult();
+    Value encodedSlice = lwe::RLWEEncodeOp::create(
+                             builder, plaintextTy, slice, plaintextEncoding,
+                             plaintextSpace.getRing(), levelAttr)
+                             .getResult();
     encodedSlices.push_back(encodedSlice);
   }