ckks: add optional targetLevel attribute to ckks.bootstrap

See the issue I filed #2436

PiperOrigin-RevId: 838756394

Author: asraa

lib/Dialect/CKKS/IR/CKKSOps.cpp

@@ -1,9 +1,11 @@
 #include "lib/Dialect/CKKS/IR/CKKSOps.h"
 
+#include <cstdint>
 #include <optional>
 
 #include "lib/Dialect/LWE/IR/LWEOps.h"
 #include "lib/Dialect/LWE/IR/LWEPatterns.h"
+#include "lib/Dialect/LWE/IR/LWETypes.h"
 #include "mlir/include/mlir/IR/Location.h"            // from @llvm-project
 #include "mlir/include/mlir/IR/MLIRContext.h"         // from @llvm-project
 #include "mlir/include/mlir/IR/PatternMatch.h"        // from @llvm-project
@@ -33,6 +35,21 @@ LogicalResult RescaleOp::verify() {
 
 LogicalResult LevelReduceOp::verify() { return lwe::verifyLevelReduceOp(this); }
 
+LogicalResult BootstrapOp::verify() {
+  std::optional<int64_t> targetLevel = getTargetLevel();
+  if (targetLevel.has_value()) {
+    // If a target level is specified, then the result ciphertext must have that
+    // many levels.
+    lwe::LWECiphertextType outputType = lwe::getCtTy(getOutput());
+    if (outputType.getModulusChain().getCurrent() != targetLevel.value()) {
+      return emitOpError() << "output ciphertext must have "
+                           << targetLevel.value() << " levels but has "
+                           << outputType.getModulusChain().getCurrent();
+    }
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Op type inference.
 //===----------------------------------------------------------------------===//

lib/Dialect/CKKS/IR/CKKSOps.td

@@ -213,13 +213,16 @@ def CKKS_BootstrapOp : CKKS_Op<"bootstrap", [ElementwiseMappable]> {
   }];
 
   let arguments = (ins
-    LWECiphertextLike:$input
+    LWECiphertextLike:$input,
+    OptionalAttr<I64Attr>:$targetLevel
   );
 
   let results = (outs
     LWECiphertextLike:$output
   );
 
+  let hasVerifier = 1;
+
   let assemblyFormat = "operands attr-dict `:` qualified(type($input)) `->` qualified(type($output))" ;
 }
 

lib/Dialect/LWE/Conversions/LWEToOpenfhe/BUILD

@@ -18,6 +18,7 @@ cc_library(
         "@heir//lib/Dialect/CKKS/IR:Dialect",
         "@heir//lib/Dialect/LWE/IR:Dialect",
         "@heir//lib/Dialect/Openfhe/IR:Dialect",
+        "@heir//lib/Parameters/CKKS:Params",
         "@heir//lib/Utils",
         "@heir//lib/Utils:ConversionUtils",
         "@llvm-project//llvm:Support",

lib/Dialect/LWE/Conversions/LWEToOpenfhe/LWEToOpenfhe.cpp

@@ -4,6 +4,7 @@
 
 #include "lib/Dialect/BGV/IR/BGVDialect.h"
 #include "lib/Dialect/BGV/IR/BGVOps.h"
+#include "lib/Dialect/CKKS/IR/CKKSAttributes.h"
 #include "lib/Dialect/CKKS/IR/CKKSDialect.h"
 #include "lib/Dialect/CKKS/IR/CKKSOps.h"
 #include "lib/Dialect/LWE/Conversions/LWEToOpenfhe/LWEToOpenfhe.h"
@@ -14,6 +15,7 @@
 #include "lib/Dialect/Openfhe/IR/OpenfheDialect.h"
 #include "lib/Dialect/Openfhe/IR/OpenfheOps.h"
 #include "lib/Dialect/Openfhe/IR/OpenfheTypes.h"
+#include "lib/Parameters/CKKS/Params.h"
 #include "lib/Utils/ConversionUtils.h"
 #include "lib/Utils/Utils.h"
 #include "llvm/include/llvm/ADT/STLExtras.h"             // from @llvm-project
@@ -22,6 +24,7 @@
 #include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"   // from @llvm-project
 #include "mlir/include/mlir/IR/Attributes.h"             // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinAttributes.h"      // from @llvm-project
+#include "mlir/include/mlir/IR/BuiltinOps.h"             // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinTypes.h"           // from @llvm-project
 #include "mlir/include/mlir/IR/PatternMatch.h"           // from @llvm-project
@@ -275,6 +278,20 @@ struct ConvertBootstrapOp : public OpConversionPattern<ckks::BootstrapOp> {
       return rewriter.notifyMatchFailure(op, "No crypto context arg");
     }
 
+    // TODO(#2436): Support bootstrap target level in OpenFHE
+    if (op.getTargetLevel().has_value()) {
+      // Right now we don't support any bootstrap ops with a target level.
+      // Ideally, we would want to check that the target level is equal to the
+      // number of Qis available in the scheme parameters (max levels) minus the
+      // levels consumed by bootstrapping to emit a full bootstrap op. The
+      // latter info is not persisted in the IR. So we simply rely on higher
+      // level passes with access to the bootstrap waterline to remove the
+      // target level attribute.
+      // TODO(#1207): Persist the number of consumed levels from bootstrapping
+      return rewriter.notifyMatchFailure(
+          op, "variadic bootstrapping is not supported in OpenFHE");
+    }
+
     Value cryptoContext = result.value();
     rewriter.replaceOpWithNewOp<openfhe::BootstrapOp>(
         op, op.getOutput().getType(), cryptoContext, adaptor.getInput());

tests/Dialect/CKKS/IR/verifier.mlir

@@ -88,3 +88,30 @@ module attributes {ckks.schemeParam = #ckks.scheme_param<logN = 13, Q = [3602879
     return %ct_2 : !ct_L0_
   }
 }
+
+// -----
+
+!Z35184372121601_i64_ = !mod_arith.int<35184372121601 : i64>
+!Z36028797019389953_i64_ = !mod_arith.int<36028797019389953 : i64>
+// note the scaling factor is 45
+// after mul it should be 90
+#inverse_canonical_encoding = #lwe.inverse_canonical_encoding<scaling_factor = 45>
+#key = #lwe.key<>
+#modulus_chain_L1_C0_ = #lwe.modulus_chain<elements = <36028797019389953 : i64, 35184372121601 : i64>, current = 0>
+#modulus_chain_L1_C1_ = #lwe.modulus_chain<elements = <36028797019389953 : i64, 35184372121601 : i64>, current = 1>
+#ring_f64_1_x1024_ = #polynomial.ring<coefficientType = f64, polynomialModulus = <1 + x**1024>>
+!rns_L0_ = !rns.rns<!Z36028797019389953_i64_>
+!rns_L1_ = !rns.rns<!Z36028797019389953_i64_, !Z35184372121601_i64_>
+#ring_rns_L0_1_x1024_ = #polynomial.ring<coefficientType = !rns_L0_, polynomialModulus = <1 + x**1024>>
+#ring_rns_L1_1_x1024_ = #polynomial.ring<coefficientType = !rns_L1_, polynomialModulus = <1 + x**1024>>
+#ciphertext_space_L0_ = #lwe.ciphertext_space<ring = #ring_rns_L0_1_x1024_, encryption_type = lsb>
+#ciphertext_space_L1_ = #lwe.ciphertext_space<ring = #ring_rns_L1_1_x1024_, encryption_type = lsb>
+!ct_L0_ = !lwe.lwe_ciphertext<application_data = <message_type = i16>, plaintext_space = <ring = #ring_f64_1_x1024_, encoding = #inverse_canonical_encoding>, ciphertext_space = #ciphertext_space_L0_, key = #key, modulus_chain = #modulus_chain_L1_C0_>
+!ct_L1_ = !lwe.lwe_ciphertext<application_data = <message_type = i16>, plaintext_space = <ring = #ring_f64_1_x1024_, encoding = #inverse_canonical_encoding>, ciphertext_space = #ciphertext_space_L1_, key = #key, modulus_chain = #modulus_chain_L1_C1_>
+module attributes {ckks.schemeParam = #ckks.scheme_param<logN = 13, Q = [36028797019389953, 35184372121601], P = [36028797019488257], logDefaultScale = 45>, scheme.ckks} {
+  func.func @bootstrap(%ct: !ct_L0_) -> !ct_L1_ {
+    // expected-error@+1 {{'ckks.bootstrap' op output ciphertext must have 2 levels but has 1}}
+    %ct_2 = ckks.bootstrap %ct {targetLevel = 2} : !ct_L0_ -> !ct_L1_
+    return %ct_2 : !ct_L1_
+  }
+}