[JaxiteWord] Update lowering and emission for CROSS polynomial APIs

lib/Dialect/JaxiteWord/IR/JaxiteWordOps.td

@@ -77,7 +77,6 @@ def GenParamsOp : JaxiteWord_Op<"gen_params"> {
     - batch: Batch size for TPU
     - r, c: Tiling dimensions (r * c = degree)
     - dnum: Key switching decomposition parameter
-    - numEvalMult: Number of evaluation multiplications
   }];
   let arguments = (ins
     // Scheme parameters
@@ -91,7 +90,6 @@ def GenParamsOp : JaxiteWord_Op<"gen_params"> {
     I32Attr:$r,
     I32Attr:$c,
     I32Attr:$dnum,
-    I32Attr:$numEvalMult,
     I32Attr:$compositeDegree
   );
   let results = (outs JaxiteWord_CryptoContext:$cryptoContext);
@@ -129,13 +127,17 @@ def ProgramInitializationOp : JaxiteWord_Op<"program_initialization"> {
   let summary = "Initialize CROSS program context with HEMul/HERot accessors";
   let description = [{
     Calls ctx.program_initialization() which pre-builds level-indexed
-    HEMul and HERot accessors and loads all keys onto the context.
+    HEMul and HERot accessors and loads all keys onto the context. The key
+    operands are attached to the CROSS context before calling
+    program_initialization, matching CROSS's params-dictionary API while keeping
+    gen_params key-independent.
     Replaces separate GenMulKeyOp + GenRotKeyOp for the new CROSS API.
   }];
   let arguments = (ins
     JaxiteWord_CryptoContext:$cryptoContext,
+    JaxiteWord_PublicKey:$publicKey,
     JaxiteWord_PrivateKey:$secretKey,
-    I64Attr:$totalHemulLevels,
+    JaxiteWord_EvalKey:$evaluationKey,
     DenseI64ArrayAttr:$totalRotationIndices,
     I32Attr:$dnum,
     I32Attr:$r,
@@ -186,6 +188,17 @@ def AddOp : JaxiteWord_BinaryOp<"add", [Commutative]> {
   let summary = "Homomorphic addition of two ciphertexts";
 }
 
+def AddPlainOp
+    : JaxiteWord_Op<"add_plain", [Pure, SameOperandsAndResultRings]> {
+  let summary = "Add a plaintext to a ciphertext";
+  let arguments = (ins
+    JaxiteWord_CryptoContext:$cryptoContext,
+    LWECiphertext:$ciphertext,
+    LWEPlaintext:$plaintext
+  );
+  let results = (outs LWECiphertext:$output);
+}
+
 def MulOp : JaxiteWord_Op<"mul", [Pure]> {
   let summary = "Homomorphic multiplication of two ciphertexts with relinearization";
   let arguments = (ins
@@ -236,10 +249,6 @@ def SubOp : JaxiteWord_BinaryOp<"sub", [SameOperandsAndResultRings]> {
   let summary = "Homomorphic subtraction of two ciphertexts";
 }
 
-def NegateOp : JaxiteWord_UnaryOp<"negate"> {
-  let summary = "Negate a ciphertext";
-}
-
 def SquareOp : JaxiteWord_UnaryOp<"square"> {
   let summary = "Square a ciphertext";
 }
@@ -252,26 +261,6 @@ def SubInPlaceOp : JaxiteWord_BinaryInPlaceOp<"sub_inplace"> {
   let summary = "In-place homomorphic subtraction";
 }
 
-def AddPlainOp : JaxiteWord_Op<"add_plain", [Pure, AllCiphertextTypesMatch]> {
-  let summary = "Add plaintext to ciphertext";
-  let arguments = (ins
-    JaxiteWord_CryptoContext:$cryptoContext,
-    LWEPlaintextOrCiphertext:$lhs,
-    LWEPlaintextOrCiphertext:$rhs
-  );
-  let results = (outs LWECiphertext:$output);
-}
-
-def SubPlainOp : JaxiteWord_Op<"sub_plain", [Pure, AllCiphertextTypesMatch]> {
-  let summary = "Subtract plaintext from ciphertext";
-  let arguments = (ins
-    JaxiteWord_CryptoContext:$cryptoContext,
-    LWEPlaintextOrCiphertext:$lhs,
-    LWEPlaintextOrCiphertext:$rhs
-  );
-  let results = (outs LWECiphertext:$output);
-}
-
 def MulPlainOp : JaxiteWord_Op<"mul_plain", [Pure]> {
   let summary = "Multiply ciphertext with plaintext";
   let arguments = (ins

lib/Dialect/JaxiteWord/Transforms/BUILD

@@ -22,15 +22,12 @@ cc_library(
     hdrs = ["ConfigureCryptoContext.h"],
     deps = [
         ":pass_inc_gen",
-        "@heir//lib/Analysis/MulDepthAnalysis",
-        "@heir//lib/Analysis/SecretnessAnalysis",
         "@heir//lib/Dialect:ModuleAttributes",
         "@heir//lib/Dialect/CKKS/IR:Dialect",
         "@heir//lib/Dialect/JaxiteWord/IR:Dialect",
         "@heir//lib/Utils",
         "@heir//lib/Utils:TransformUtils",
         "@llvm-project//llvm:Support",
-        "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",

lib/Dialect/JaxiteWord/Transforms/ConfigureCryptoContext.cpp

@@ -5,8 +5,6 @@
 #include <set>
 #include <string>
 
-#include "lib/Analysis/MulDepthAnalysis/MulDepthAnalysis.h"
-#include "lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.h"
 #include "lib/Dialect/CKKS/IR/CKKSAttributes.h"
 #include "lib/Dialect/CKKS/IR/CKKSDialect.h"
 #include "lib/Dialect/JaxiteWord/IR/JaxiteWordDialect.h"
@@ -15,32 +13,25 @@
 #include "lib/Dialect/ModuleAttributes.h"
 #include "lib/Utils/TransformUtils.h"
 #include "lib/Utils/Utils.h"
-#include "llvm/include/llvm/Support/Debug.h"               // from @llvm-project
-#include "llvm/include/llvm/Support/raw_ostream.h"         // from @llvm-project
-#include "mlir/include/mlir/Analysis/DataFlow/Utils.h"     // from @llvm-project
-#include "mlir/include/mlir/Analysis/DataFlowFramework.h"  // from @llvm-project
-#include "mlir/include/mlir/Dialect/Func/IR/FuncOps.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/BuiltinTypes.h"             // from @llvm-project
-#include "mlir/include/mlir/IR/ImplicitLocOpBuilder.h"     // from @llvm-project
-#include "mlir/include/mlir/IR/Operation.h"                // from @llvm-project
-#include "mlir/include/mlir/IR/Types.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/Support/LLVM.h"                // from @llvm-project
-#include "mlir/include/mlir/Support/LogicalResult.h"       // from @llvm-project
-#include "mlir/include/mlir/Support/WalkResult.h"          // from @llvm-project
-
-#define DEBUG_TYPE "jaxiteword-configure-crypto-context"
+#include "llvm/include/llvm/Support/raw_ostream.h"      // from @llvm-project
+#include "mlir/include/mlir/Dialect/Func/IR/FuncOps.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/BuiltinTypes.h"          // from @llvm-project
+#include "mlir/include/mlir/IR/ImplicitLocOpBuilder.h"  // from @llvm-project
+#include "mlir/include/mlir/IR/Operation.h"             // from @llvm-project
+#include "mlir/include/mlir/IR/Types.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/Support/LLVM.h"             // from @llvm-project
+#include "mlir/include/mlir/Support/LogicalResult.h"    // from @llvm-project
+#include "mlir/include/mlir/Support/WalkResult.h"       // from @llvm-project
 
 namespace mlir {
 namespace heir {
 namespace jaxiteword {
 
 struct Config {
-  int mulDepth;
-  bool hasRelinOp;
   SmallVector<int64_t> rotIndices;
   int64_t degree;
   int64_t numSlots;
@@ -63,18 +54,6 @@ struct ConfigureCryptoContext
  private:
   Config config;
 
-  bool checkHasRelinOp(func::FuncOp op) {
-    bool result = false;
-    walkFuncAndCallees(op, [&](Operation* op) {
-      if (isa<jaxiteword::MulOp, jaxiteword::RelinOp>(op)) {
-        result = true;
-        return WalkResult::interrupt();
-      }
-      return WalkResult::advance();
-    });
-    return result;
-  }
-
   SmallVector<int64_t> findAllRotIndices(func::FuncOp op) {
     std::set<int64_t> distinctRotIndices;
     walkFuncAndCallees(op, [&](Operation* op) {
@@ -90,10 +69,12 @@ struct ConfigureCryptoContext
   LogicalResult generateGenFunc(func::FuncOp op, const std::string& genFuncName,
                                 ImplicitLocOpBuilder& builder) {
     Type ccType = CryptoContextType::get(builder.getContext());
+
+    SmallVector<Type> funcArgTypes;
     SmallVector<Type> funcResultTypes = {ccType};
 
     FunctionType genFuncType =
-        FunctionType::get(builder.getContext(), {}, funcResultTypes);
+        FunctionType::get(builder.getContext(), funcArgTypes, funcResultTypes);
     auto genFuncOp = func::FuncOp::create(builder, genFuncName, genFuncType);
     builder.setInsertionPointToEnd(genFuncOp.addEntryBlock());
 
@@ -108,7 +89,6 @@ struct ConfigureCryptoContext
         /*r=*/static_cast<uint32_t>(config.r),
         /*c=*/static_cast<uint32_t>(config.c),
         /*dnum=*/static_cast<uint32_t>(config.dnum),
-        /*numEvalMult=*/static_cast<uint32_t>(config.mulDepth),
         /*compositeDegree=*/static_cast<uint32_t>(compositeDegree));
 
     func::ReturnOp::create(builder, cryptoContext);
@@ -119,9 +99,11 @@ struct ConfigureCryptoContext
                                    const std::string& configFuncName,
                                    ImplicitLocOpBuilder& builder) {
     Type ccType = CryptoContextType::get(builder.getContext());
+    Type pkType = PublicKeyType::get(builder.getContext());
     Type skType = PrivateKeyType::get(builder.getContext());
+    Type ekType = EvalKeyType::get(builder.getContext());
 
-    SmallVector<Type> funcArgTypes = {ccType, skType};
+    SmallVector<Type> funcArgTypes = {ccType, pkType, skType, ekType};
     SmallVector<Type> funcResultTypes;
 
     FunctionType configFuncType =
@@ -131,16 +113,17 @@ struct ConfigureCryptoContext
     builder.setInsertionPointToEnd(configFuncOp.addEntryBlock());
 
     Value cryptoContext = configFuncOp.getArgument(0);
-    Value secretKey = configFuncOp.getArgument(1);
-
-    ProgramInitializationOp::create(
-        builder, cryptoContext, secretKey,
-        /*totalHemulLevels=*/static_cast<int64_t>(config.mulDepth),
-        /*totalRotationIndices=*/config.rotIndices,
-        /*dnum=*/config.dnum,
-        /*r=*/config.r,
-        /*c=*/config.c,
-        /*batch=*/config.batch);
+    Value publicKey = configFuncOp.getArgument(1);
+    Value secretKey = configFuncOp.getArgument(2);
+    Value evaluationKey = configFuncOp.getArgument(3);
+
+    ProgramInitializationOp::create(builder, cryptoContext, publicKey,
+                                    secretKey, evaluationKey,
+                                    /*totalRotationIndices=*/config.rotIndices,
+                                    /*dnum=*/config.dnum,
+                                    /*r=*/config.r,
+                                    /*c=*/config.c,
+                                    /*batch=*/config.batch);
 
     func::ReturnOp::create(builder, ValueRange{});
     return success();
@@ -197,51 +180,6 @@ struct ConfigureCryptoContext
       module->removeAttr(ckks::CKKSDialect::kSchemeParamAttrName);
     }
 
-    LLVM_DEBUG(llvm::dbgs() << "Recomputing mul depth\n");
-    DataFlowSolver solver;
-    dataflow::loadBaselineAnalyses(solver);
-    solver.load<SecretnessAnalysis>();
-    solver.load<MulDepthAnalysis>();
-
-    if (failed(solver.initializeAndRun(module))) {
-      op->emitOpError() << "Failed to run mul depth analysis.\n";
-      return failure();
-    }
-
-    config.mulDepth = 0;
-    walkValues(op, [&](Value value) {
-      auto mulDepthState =
-          solver.lookupState<MulDepthLattice>(value)->getValue();
-      if (!mulDepthState.isInitialized()) {
-        LLVM_DEBUG(llvm::dbgs()
-                   << "mul depth uninitialized at " << value << "\n");
-        return;
-      }
-      auto depth = mulDepthState.getMulDepth();
-      if (depth > config.mulDepth) {
-        LLVM_DEBUG(llvm::dbgs() << "Found larger mul depth=" << depth << "\n");
-        config.mulDepth = depth;
-      }
-    });
-
-    if (mulDepth != 0) {
-      config.mulDepth = mulDepth;
-    }
-
-    if (config.mulDepth == 0) {
-      int mulCount = 0;
-      walkFuncAndCallees(op, [&](Operation* innerOp) {
-        if (isa<jaxiteword::MulOp>(innerOp)) mulCount++;
-        return WalkResult::advance();
-      });
-      if (mulCount > 0) {
-        config.mulDepth = mulCount;
-      } else if (!config.qTowers.empty()) {
-        config.mulDepth = 1;
-      }
-    }
-
-    config.hasRelinOp = checkHasRelinOp(op);
     config.rotIndices = findAllRotIndices(op);
 
     config.dnum = dnum;

lib/Dialect/JaxiteWord/Transforms/Passes.td

@@ -8,26 +8,25 @@ def ConfigureCryptoContext : Pass<"jaxiteword-configure-crypto-context"> {
   let description = [{
     This pass generates helper functions to generate and configure the CROSS
     crypto context for the given function. It analyzes the entry function to
-    determine the required CKKS scheme parameters, multiplication depth,
-    rotation indices, and relinearization needs, then synthesizes setup
-    functions using `jaxiteword.gen_params` and
+    determine the required CKKS scheme parameters and rotation indices, then
+    synthesizes setup functions using `jaxiteword.gen_params` and
     `jaxiteword.program_initialization`.
 
     For example, for an MLIR function `@my_func`, the generated helpers have
     the following signatures:
     ```mlir
-    func.func @my_func__generate_crypto_context() -> !jaxiteword.crypto_context
+    func.func @my_func__generate_crypto_context()
+        -> !jaxiteword.crypto_context
 
     func.func @my_func__configure_crypto_context(
-        !jaxiteword.crypto_context, !jaxiteword.private_key)
+        !jaxiteword.crypto_context, !jaxiteword.public_key,
+        !jaxiteword.private_key, !jaxiteword.eval_key)
     ```
   }];
   let dependentDialects = ["mlir::heir::jaxiteword::JaxiteWordDialect"];
   let options = [
     Option<"entryFunction", "entry-function", "std::string",
            /*default=*/"", "Name of entry function.">,
-    Option<"mulDepth", "mul-depth", "int",
-           /*default=*/"0", "Manually specify the mul depth (overrides analysis)">,
     Option<"dnum", "dnum", "int",
            /*default=*/"3", "Key-switching decomposition parameter">,
     Option<"r", "r", "int",

lib/Dialect/LWE/Conversions/LWEToJaxiteWord/BUILD

@@ -12,7 +12,6 @@ cc_library(
     hdrs = ["LWEToJaxiteWord.h"],
     deps = [
         ":pass_inc_gen",
-        "@heir//lib/Dialect/BGV/IR:Dialect",
         "@heir//lib/Dialect/CKKS/IR:Dialect",
         "@heir//lib/Dialect/JaxiteWord/IR:Dialect",
         "@heir//lib/Dialect/LWE/IR:Dialect",