encode plaintext at highest used level instead of max level

Author: mdgrs

lib/Dialect/LWE/Conversions/LWEToLattigo/LWEToLattigo.cpp

@@ -541,9 +541,12 @@ struct ConvertRlweEncodeOp : public OpConversionPattern<EncodeOp> {
     Value params = result2.value();
 
     Value input = adaptor.getInput();
+    // A missing level attribute leaves the allocation at
+    // params.MaxLevel().
     auto alloc = AllocOp::create(
         rewriter, op.getLoc(),
-        this->typeConverter->convertType(op.getOutput().getType()), params);
+        this->typeConverter->convertType(op.getOutput().getType()), params,
+        op.getLevelAttr());
 
     auto encoding = op.getEncoding();
     int64_t scale = lwe::getScalingFactorFromEncodingAttr(encoding);

lib/Dialect/LWE/IR/LWEOps.td

@@ -201,6 +201,15 @@ def LWE_RLWEEncodeOp : LWE_Op<"rlwe_encode", [
     be floating points, and a scaling factor described by the encoding will be
     applied.
 
+    The optional `level` attribute records the level of the ciphertext this
+    plaintext will be combined with, i.e., the `current` index of that
+    ciphertext's modulus chain. It is
+    recorded on the op rather than in `!lwe.lwe_plaintext` because that type is
+    not RNS-aware (see #1643), and it must be attached where the plaintext is
+    created, since passes such as `--split-preprocessing` later sever the
+    def-use chain to the consuming ciphertext op. When absent, backends fall
+    back to the top level of the modulus chain.
+
     Examples:
 
     ```
@@ -212,7 +221,7 @@ def LWE_RLWEEncodeOp : LWE_Op<"rlwe_encode", [
     SignlessIntegerOrFloatLike:$input,
     AnyPlaintextEncodingAttr:$encoding,
     Polynomial_RingAttr:$ring,
-    OptionalAttr<I64Attr>:$level,
+    OptionalAttr<ConfinedAttr<I64Attr, [IntNonNegative]>>:$level,
     OptionalAttr<I64Attr>:$scale
   );
 

lib/Dialect/LWE/IR/LWETypes.cpp

@@ -1,6 +1,7 @@
 #include "lib/Dialect/LWE/IR/LWETypes.h"
 
 #include <cstdint>
+#include <optional>
 
 #include "lib/Dialect/LWE/IR/LWEAttributes.h"
 #include "lib/Dialect/Polynomial/IR/PolynomialAttributes.h"
@@ -111,6 +112,12 @@ LWECiphertextType cloneAtLevel(LWECiphertextType inputType, int64_t level) {
       lwe::KeyAttr::get(ctx, 0), newChain);
 }
 
+std::optional<int64_t> getLevel(LWECiphertextType ctType) {
+  auto modulusChain = ctType.getModulusChain();
+  if (!modulusChain) return std::nullopt;
+  return modulusChain.getCurrent();
+}
+
 }  // namespace lwe
 }  // namespace heir
 }  // namespace mlir

lib/Dialect/LWE/IR/LWETypes.h

@@ -2,6 +2,7 @@
 #define LIB_DIALECT_LWE_IR_LWETYPES_H_
 
 #include <cstdint>
+#include <optional>
 
 #include "lib/Dialect/LWE/IR/LWEAttributes.h"
 #include "mlir/include/mlir/IR/MLIRContext.h"         // from @llvm-project
@@ -44,6 +45,11 @@ FailureOr<LWECiphertextType> applyModReduce(LWECiphertextType inputType);
 // value.
 LWECiphertextType cloneAtLevel(LWECiphertextType inputType, int64_t level);
 
+// Return the level of a ciphertext type, i.e., the index of the top modulus of
+// its modulus chain. Returns nullopt for schemes that don't use a modulus
+// chain
+std::optional<int64_t> getLevel(LWECiphertextType ctType);
+
 }  // namespace lwe
 }  // namespace heir
 }  // namespace mlir

lib/Dialect/LWE/Transforms/AnnotatePlaintextLevel.cpp

@@ -0,0 +1,154 @@
+#include "lib/Dialect/LWE/Transforms/AnnotatePlaintextLevel.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <optional>
+
+#include "lib/Dialect/LWE/IR/LWEOps.h"
+#include "lib/Dialect/LWE/IR/LWETraits.h"
+#include "lib/Dialect/LWE/IR/LWETypes.h"
+#include "lib/Dialect/ModuleAttributes.h"
+#include "llvm/include/llvm/ADT/DenseSet.h"               // from @llvm-project
+#include "llvm/include/llvm/ADT/STLExtras.h"              // from @llvm-project
+#include "llvm/include/llvm/ADT/SmallVector.h"            // from @llvm-project
+#include "mlir/include/mlir/IR/BuiltinOps.h"              // from @llvm-project
+#include "mlir/include/mlir/IR/Operation.h"               // from @llvm-project
+#include "mlir/include/mlir/IR/TypeUtilities.h"           // from @llvm-project
+#include "mlir/include/mlir/IR/Value.h"                   // from @llvm-project
+#include "mlir/include/mlir/IR/Visitors.h"                // from @llvm-project
+#include "mlir/include/mlir/Interfaces/CallInterfaces.h"  // from @llvm-project
+#include "mlir/include/mlir/Support/LLVM.h"               // from @llvm-project
+
+namespace mlir {
+namespace heir {
+namespace lwe {
+
+#define GEN_PASS_DEF_ANNOTATEPLAINTEXTLEVEL
+#include "lib/Dialect/LWE/Transforms/Passes.h.inc"
+
+namespace {
+
+// Return the highest level among an op's ciphertext operands and results, or
+// nullopt if it touches no ciphertext carrying a modulus chain.
+std::optional<int64_t> highestCiphertextLevel(Operation* op) {
+  std::optional<int64_t> result;
+  auto join = [&](Value value) {
+    auto ctType =
+        dyn_cast<LWECiphertextType>(getElementTypeOrSelf(value.getType()));
+    if (!ctType) return;
+    std::optional<int64_t> level = getLevel(ctType);
+    if (!level) return;
+    result = result ? std::max(*result, *level) : *level;
+  };
+  for (Value operand : op->getOperands()) join(operand);
+  for (Value opResult : op->getResults()) join(opResult);
+  return result;
+}
+
+// Return true if an op's results account for every place a value it consumes
+// can end up, so the walk can follow them. Ops with regions route operands into
+// block arguments (`scf.for`'s iter_args) and calls route them into a callee,
+// neither of which the walk models; a terminator has no results to follow.
+bool forwardsThroughResults(Operation* op) {
+  return op->getNumResults() > 0 && op->getNumRegions() == 0 &&
+         !isa<CallOpInterface>(op) &&
+         llvm::all_of(op->getResults(), [](Value result) {
+           return isa<LWEPlaintextType>(getElementTypeOrSelf(result.getType()));
+         });
+}
+
+// Return the level to encode the plaintext at, or nullopt to leave it at the
+// top of the modulus chain. Plaintexts are not always consumed directly: the
+// encode op may feed a `tensor.from_elements` or similar first, so keep walking
+// forward through the ops that forward it until one that consumes it is
+// reached.
+//
+// The two kinds of consumer constrain the level differently:
+//
+//   - A ct-pt op sets a *lower bound*. The backend combines at the lower of the
+//     two levels, so encoding above the ciphertext costs limbs but nothing
+//     else, while encoding below it silently truncates that ciphertext. The
+//     highest such use is therefore the smallest level that serves them all.
+//
+//   - An encryption sets an *exact* level. It has no ciphertext operand to be
+//     clamped against: the backend builds the ciphertext at the plaintext's own
+//     level. Both `lwe.rlwe_encrypt` and `lwe.trivial_encrypt` work
+//     this way.
+//
+// One encoding can serve both only when the encryption's level also covers
+// every ct-pt use. When it does not, or when two encryptions want different
+// levels, no annotation is correct and the fallback takes over.
+std::optional<int64_t> findUseLevel(RLWEEncodeOp encodeOp) {
+  std::optional<int64_t> combinedLevel;
+  std::optional<int64_t> encryptedLevel;
+  SmallVector<Value> worklist = {encodeOp.getOutput()};
+  DenseSet<Operation*> visited;
+
+  while (!worklist.empty()) {
+    Value value = worklist.pop_back_val();
+    for (Operation* user : value.getUsers()) {
+      if (!visited.insert(user).second) continue;
+
+      if (isa<RLWEEncryptOp, TrivialEncryptOp>(user)) {
+        // A consumer whose ciphertexts carry no modulus chain constrains
+        // nothing, and leaves the plaintext to whatever its other uses need.
+        std::optional<int64_t> level = highestCiphertextLevel(user);
+        if (!level) continue;
+        if (encryptedLevel && *encryptedLevel != *level) return std::nullopt;
+        encryptedLevel = *level;
+        continue;
+      }
+      if (user->hasTrait<IsCiphertextPlaintextOp>()) {
+        // Likewise for a ct-pt op whose ciphertexts carry no modulus chain.
+        if (std::optional<int64_t> level = highestCiphertextLevel(user)) {
+          combinedLevel =
+              combinedLevel ? std::max(*combinedLevel, *level) : *level;
+        }
+        continue;
+      }
+      if (!forwardsThroughResults(user)) return std::nullopt;
+      for (Value userResult : user->getResults()) {
+        worklist.push_back(userResult);
+      }
+    }
+  }
+
+  if (!encryptedLevel) return combinedLevel;
+  if (combinedLevel && *combinedLevel > *encryptedLevel) return std::nullopt;
+  return encryptedLevel;
+}
+
+struct AnnotatePlaintextLevel
+    : impl::AnnotatePlaintextLevelBase<AnnotatePlaintextLevel> {
+  using AnnotatePlaintextLevelBase::AnnotatePlaintextLevelBase;
+
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+
+    // BFV does no level management: every ciphertext sits at the bottom of the
+    // modulus chain, while the backend's ciphertexts still span the whole chain
+    // (the extra modulus is an artifact of the encryption technique). So the
+    // chain's `current` is not an encoding level here, and there is nothing to
+    // save by lowering one.
+    if (moduleIsBFV(module)) {
+      module->walk([&](RLWEEncodeOp encodeOp) { encodeOp.removeLevelAttr(); });
+      return;
+    }
+
+    module->walk([&](RLWEEncodeOp encodeOp) {
+      std::optional<int64_t> level = findUseLevel(encodeOp);
+      // Drop a level that no longer has a use to justify it
+      if (!level) {
+        encodeOp.removeLevelAttr();
+        return;
+      }
+      encodeOp.setLevel(*level);
+    });
+  }
+};
+
+}  // namespace
+
+}  // namespace lwe
+}  // namespace heir
+}  // namespace mlir

lib/Dialect/LWE/Transforms/AnnotatePlaintextLevel.h

@@ -0,0 +1,19 @@
+#ifndef LIB_DIALECT_LWE_TRANSFORMS_ANNOTATEPLAINTEXTLEVEL_H_
+#define LIB_DIALECT_LWE_TRANSFORMS_ANNOTATEPLAINTEXTLEVEL_H_
+
+// IWYU pragma: begin_keep
+#include "mlir/include/mlir/Pass/Pass.h"  // from @llvm-project
+// IWYU pragma: end_keep
+
+namespace mlir {
+namespace heir {
+namespace lwe {
+
+#define GEN_PASS_DECL_ANNOTATEPLAINTEXTLEVEL
+#include "lib/Dialect/LWE/Transforms/Passes.h.inc"
+
+}  // namespace lwe
+}  // namespace heir
+}  // namespace mlir
+
+#endif  // LIB_DIALECT_LWE_TRANSFORMS_ANNOTATEPLAINTEXTLEVEL_H_

lib/Dialect/LWE/Transforms/BUILD

@@ -13,13 +13,30 @@ cc_library(
     ],
     deps = [
         ":AddDebugPort",
+        ":AnnotatePlaintextLevel",
         ":DecomposeLWEOps",
         ":ImplementTrivialEncryptionAsAddition",
         ":pass_inc_gen",
         "@heir//lib/Dialect/LWE/IR:Dialect",
     ],
 )
 
+cc_library(
+    name = "AnnotatePlaintextLevel",
+    srcs = ["AnnotatePlaintextLevel.cpp"],
+    hdrs = ["AnnotatePlaintextLevel.h"],
+    deps = [
+        ":pass_inc_gen",
+        "@heir//lib/Dialect:ModuleAttributes",
+        "@heir//lib/Dialect/LWE/IR:Dialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:CallOpInterfaces",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
 cc_library(
     name = "AddDebugPort",
     srcs = ["AddDebugPort.cpp"],

lib/Dialect/LWE/Transforms/Passes.h

@@ -4,6 +4,7 @@
 // IWYU pragma: begin_keep
 #include "lib/Dialect/LWE/IR/LWEDialect.h"
 #include "lib/Dialect/LWE/Transforms/AddDebugPort.h"
+#include "lib/Dialect/LWE/Transforms/AnnotatePlaintextLevel.h"
 #include "lib/Dialect/LWE/Transforms/DecomposeLWEOps.h"
 #include "lib/Dialect/LWE/Transforms/ImplementTrivialEncryptionAsAddition.h"
 // IWYU pragma: end_keep

lib/Dialect/LWE/Transforms/Passes.td

@@ -69,6 +69,42 @@ def ImplementTrivialEncryptionAsAddition : Pass<"implement-trivial-encryption-as
   let options = [];
 }
 
+def AnnotatePlaintextLevel : Pass<"lwe-annotate-plaintext-level", "mlir::ModuleOp"> {
+  let summary = "Annotate `lwe.rlwe_encode` ops with the level they are used at";
+  let description = [{
+    Sets the `level` attribute of each `lwe.rlwe_encode` op to the level its
+    plaintext is used at, found by walking forward from the encode op through
+    plaintext-typed values.
+
+    Backends allocate a plaintext at the top of the modulus chain unless told
+    otherwise, but encoding cost is dominated by per-RNS-limb work, so a
+    plaintext only ever combined with a level 3 ciphertext is cheaper to encode
+    with 4 limbs than with all of them. For a ct-pt op the result is unchanged
+    either way: combining a ciphertext and a plaintext yields the lower of the
+    two levels, so annotating with the *highest* level among such uses is the
+    smallest choice that serves them all.
+
+    `lwe.rlwe_encrypt` is not a ct-pt op and does not follow that rule. It has
+    no ciphertext operand to be clamped against, so the backend builds the
+    ciphertext at the plaintext's own level and the annotation is exact rather
+    than a lower bound. Annotating an encryption therefore changes what the
+    client emits, not just how many limbs it encodes: where the entry level sits
+    below the top of the chain, the ciphertext now enters at the level its type
+    declares instead of at `MaxLevel`. A plaintext that both feeds an encryption
+    and is combined with a higher-level ciphertext has no correct level, and is
+    left to the fallback.
+
+    Run this after CSE. Identical encode ops that CSE has merged are cheaper as
+    one plaintext at the highest level than as several at their own levels,
+    since each encode pays a fixed cost on top of the per-limb work.
+
+    (* example filepath=tests/Dialect/LWE/Transforms/annotate_plaintext_level.mlir *)
+  }];
+  let dependentDialects = [
+    "mlir::heir::lwe::LWEDialect",
+  ];
+}
+
 def DecomposeLWEOps : Pass<"lwe-decompose"> {
   let summary = "Decomposes complex LWE ops into primitive LWE ops";
   let description = [{

lib/Dialect/Lattigo/IR/LattigoBGVOps.td

@@ -17,9 +17,14 @@ def Lattigo_BGVNewPlaintextOp : Lattigo_BGVOp<"new_plaintext"> {
   let summary = "Create a new plaintext in the Lattigo BGV dialect";
   let description = [{
     This operation creates a new plaintext value in the Lattigo BGV dialect.
+
+    The optional `level` attribute is the level to allocate the plaintext at.
+    When absent, the
+    plaintext is allocated at `params.MaxLevel()`.
   }];
   let arguments = (ins
-    Lattigo_BGVParameter:$params
+    Lattigo_BGVParameter:$params,
+    OptionalAttr<ConfinedAttr<I64Attr, [IntNonNegative]>>:$level
   );
   let results = (outs Lattigo_RLWEPlaintext:$plaintext);
 }