Bump params for lattigo when bootstrapping is present

PiperOrigin-RevId: 895960176

Author: j2kun

lib/Dialect/Lattigo/IR/LattigoCKKSOps.td

@@ -215,8 +215,8 @@ def Lattigo_CKKSMulOp : Lattigo_CKKSBinaryInPlaceOp<"mul", [IncreasesMulDepthOpI
   }];
 }
 
-class Lattigo_CKKSUnaryOp<string mnemonic> :
-        Lattigo_CKKSOp<mnemonic> {
+class Lattigo_CKKSUnaryOp<string mnemonic, list<Trait> traits = []> :
+        Lattigo_CKKSOp<mnemonic, traits> {
   let arguments = (ins
     Lattigo_CKKSEvaluator:$evaluator,
     Lattigo_RLWECiphertext:$input
@@ -322,7 +322,7 @@ def Lattigo_CKKSRotateOp : Lattigo_CKKSUnaryInPlaceOp<"rotate", [
   let hasVerifier = 1;
 }
 
-def Lattigo_CKKSBootstrapOp : Lattigo_CKKSUnaryOp<"bootstrap"> {
+def Lattigo_CKKSBootstrapOp : Lattigo_CKKSUnaryOp<"bootstrap", [ResetsMulDepthOpInterface]> {
   let summary = "Bootstrap a ciphertext in the Lattigo CKKS dialect";
   let description = [{
     Bootstraps a ciphertext value in the Lattigo CKKS dialect.

lib/Transforms/GenerateParam/BUILD

@@ -28,6 +28,7 @@ cc_library(
         "@heir//lib/Analysis/NoiseAnalysis/BGV:NoiseCanEmbModel",
         "@heir//lib/Analysis/RangeAnalysis",
         "@heir//lib/Analysis/SecretnessAnalysis",
+        "@heir//lib/Dialect:HEIRInterfaces",
         "@heir//lib/Dialect:ModuleAttributes",
         "@heir//lib/Dialect/BGV/IR:Dialect",
         "@heir//lib/Dialect/CKKS/IR:Dialect",
@@ -37,6 +38,7 @@ cc_library(
         "@heir//lib/Dialect/Secret/IR:Dialect",
         "@heir//lib/Parameters/BGV:Params",
         "@heir//lib/Parameters/CKKS:Params",
+        "@heir//lib/Utils",
         "@heir//lib/Utils:LogArithmetic",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:Analysis",

lib/Transforms/GenerateParam/GenerateParamCKKS.cpp

@@ -7,6 +7,7 @@
 #include "lib/Dialect/CKKS/IR/CKKSAttributes.h"
 #include "lib/Dialect/CKKS/IR/CKKSDialect.h"
 #include "lib/Dialect/CKKS/IR/CKKSEnums.h"
+#include "lib/Dialect/HEIRInterfaces.h"
 #include "lib/Dialect/Mgmt/IR/MgmtAttributes.h"
 #include "lib/Dialect/ModuleAttributes.h"
 #include "lib/Parameters/CKKS/Params.h"
@@ -17,10 +18,12 @@
 #include "mlir/include/mlir/Analysis/DataFlowFramework.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/Builders.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/Diagnostics.h"              // from @llvm-project
 #include "mlir/include/mlir/IR/Operation.h"                // from @llvm-project
 #include "mlir/include/mlir/IR/Value.h"                    // from @llvm-project
 #include "mlir/include/mlir/Support/LLVM.h"                // from @llvm-project
+#include "mlir/include/mlir/Support/WalkResult.h"          // from @llvm-project
 
 // IWYU pragma: begin_keep
 #include "lib/Transforms/GenerateParam/GenerateParam.h"
@@ -36,6 +39,18 @@ namespace heir {
 #define GEN_PASS_DEF_GENERATEPARAMCKKS
 #include "lib/Transforms/GenerateParam/GenerateParam.h.inc"
 
+namespace {
+bool containsBootstrap(Operation* op) {
+  auto result = op->walk([&](Operation* walkOp) {
+    if (isa<ResetsMulDepthOpInterface>(walkOp)) {
+      return WalkResult::interrupt();
+    }
+    return WalkResult::advance();
+  });
+  return result.wasInterrupted();
+}
+}  // namespace
+
 struct GenerateParamCKKS : impl::GenerateParamCKKSBase<GenerateParamCKKS> {
   using GenerateParamCKKSBase::GenerateParamCKKSBase;
 
@@ -132,6 +147,17 @@ struct GenerateParamCKKS : impl::GenerateParamCKKSBase<GenerateParamCKKS> {
     if (moduleIsLattigo(getOperation())) {
       encryptionTechniqueExtended = true;
       LDBG() << "For lattigo, fixing extended encryption technique";
+
+      // Lattigo bootstrapping requires LogN >= 14, i.e., ringDim >= 16384.
+      // Since ringDim is computed from slotNumber (minRingDim = 2 *
+      // slotNumber), we bump slotNumber to 8192 if bootstrapping is present.
+      if (containsBootstrap(getOperation())) {
+        if (slotNumber < 8192) {
+          LDBG() << "Lattigo bootstrapping detected, bumping slotNumber from "
+                 << slotNumber << " to 8192";
+          slotNumber = 8192;
+        }
+      }
     }
 
     auto schemeParam = ckks::SchemeParam::getConcreteSchemeParam(

tests/Transforms/generate_param_ckks/rolled_matvec.mlir

@@ -6,7 +6,10 @@
 
 // CHECK: module
 // CHECK-SAME: #ckks.scheme_param
-// CHECK-SAME: logN = 13
+
+// because the pass has bootstrapping, lattigo bumps the parameters to 2**14 at a minimum
+// CHECK-SAME: logN = 14
+
 // CHECK-SAME: Q = [{{[0-9]*}}, {{[0-9]*}}]
 // CHECK-SAME: P = [{{[0-9]*}}]
 // CHECK-SAME: logDefaultScale = 45