diff --git a/lib/Dialect/JaxiteWord/IR/JaxiteWordOps.td b/lib/Dialect/JaxiteWord/IR/JaxiteWordOps.td index a5a8c8dffb..bf273799cc 100644 --- a/lib/Dialect/JaxiteWord/IR/JaxiteWordOps.td +++ b/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 diff --git a/lib/Dialect/JaxiteWord/Transforms/BUILD b/lib/Dialect/JaxiteWord/Transforms/BUILD index 3d1c04675b..a82056a426 100644 --- a/lib/Dialect/JaxiteWord/Transforms/BUILD +++ b/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", diff --git a/lib/Dialect/JaxiteWord/Transforms/ConfigureCryptoContext.cpp b/lib/Dialect/JaxiteWord/Transforms/ConfigureCryptoContext.cpp index 367cee033b..5a18d28992 100644 --- a/lib/Dialect/JaxiteWord/Transforms/ConfigureCryptoContext.cpp +++ b/lib/Dialect/JaxiteWord/Transforms/ConfigureCryptoContext.cpp @@ -5,8 +5,6 @@ #include #include -#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 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(op)) { - result = true; - return WalkResult::interrupt(); - } - return WalkResult::advance(); - }); - return result; - } - SmallVector findAllRotIndices(func::FuncOp op) { std::set 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 funcArgTypes; SmallVector 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(config.r), /*c=*/static_cast(config.c), /*dnum=*/static_cast(config.dnum), - /*numEvalMult=*/static_cast(config.mulDepth), /*compositeDegree=*/static_cast(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 funcArgTypes = {ccType, skType}; + SmallVector funcArgTypes = {ccType, pkType, skType, ekType}; SmallVector 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(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(); - solver.load(); - - 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(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(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; diff --git a/lib/Dialect/JaxiteWord/Transforms/Passes.td b/lib/Dialect/JaxiteWord/Transforms/Passes.td index f9ab1ebdba..a81c72810d 100644 --- a/lib/Dialect/JaxiteWord/Transforms/Passes.td +++ b/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", diff --git a/lib/Dialect/LWE/Conversions/LWEToJaxiteWord/BUILD b/lib/Dialect/LWE/Conversions/LWEToJaxiteWord/BUILD index 49ee33ca34..cd7166f8a1 100644 --- a/lib/Dialect/LWE/Conversions/LWEToJaxiteWord/BUILD +++ b/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", diff --git a/lib/Dialect/LWE/Conversions/LWEToJaxiteWord/LWEToJaxiteWord.cpp b/lib/Dialect/LWE/Conversions/LWEToJaxiteWord/LWEToJaxiteWord.cpp index 8949062aa8..35f4f6f485 100644 --- a/lib/Dialect/LWE/Conversions/LWEToJaxiteWord/LWEToJaxiteWord.cpp +++ b/lib/Dialect/LWE/Conversions/LWEToJaxiteWord/LWEToJaxiteWord.cpp @@ -3,8 +3,6 @@ #include #include -#include "lib/Dialect/BGV/IR/BGVDialect.h" -#include "lib/Dialect/BGV/IR/BGVOps.h" #include "lib/Dialect/CKKS/IR/CKKSDialect.h" #include "lib/Dialect/CKKS/IR/CKKSOps.h" #include "lib/Dialect/JaxiteWord/IR/JaxiteWordDialect.h" @@ -20,6 +18,7 @@ #include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h" // from @llvm-project #include "mlir/include/mlir/IR/BuiltinAttributes.h" // from @llvm-project #include "mlir/include/mlir/IR/BuiltinTypes.h" // from @llvm-project +#include "mlir/include/mlir/IR/Matchers.h" // from @llvm-project #include "mlir/include/mlir/IR/PatternMatch.h" // from @llvm-project #include "mlir/include/mlir/IR/TypeUtilities.h" // from @llvm-project #include "mlir/include/mlir/Pass/Pass.h" // from @llvm-project @@ -53,6 +52,59 @@ class JaxiteWordTypeConverter : public TypeConverter { namespace { +bool containsCryptoArgument(func::FuncOp funcOp) { + return llvm::any_of(funcOp.getArgumentTypes(), [&](Type argType) { + return DialectEqual()( + &getElementTypeOrSelf(argType).getDialect()); + }); +} + +bool funcNeedsCryptoContextAndKeys(func::FuncOp funcOp) { + return containsDialects(funcOp) || + containsCryptoArgument(funcOp); +} + +void insertCryptoContextAndKeys(func::FuncOp funcOp) { + if (!funcNeedsCryptoContextAndKeys(funcOp)) return; + if (funcOp.getFunctionType().getNumInputs() >= 2 && + mlir::isa( + funcOp.getFunctionType().getInput(0)) && + mlir::isa( + funcOp.getFunctionType().getInput(1))) { + return; + } + auto cryptoContextType = + jaxiteword::CryptoContextType::get(funcOp.getContext()); + auto evalKeyType = jaxiteword::EvalKeyType::get(funcOp.getContext()); + (void)funcOp.insertArgument(0, evalKeyType, nullptr, funcOp.getLoc()); + (void)funcOp.insertArgument(0, cryptoContextType, nullptr, funcOp.getLoc()); +} + +void updateCryptoFuncCalls(Operation* op) { + op->walk([&](func::CallOp callOp) { + auto callee = getCalledFunction(callOp); + if (failed(callee) || !funcNeedsCryptoContextAndKeys(callee.value())) { + return; + } + if (callOp.getNumOperands() == callee.value().getNumArguments()) { + return; + } + auto caller = callOp->getParentOfType(); + if (!caller || caller.getNumArguments() < 2 || + !mlir::isa( + caller.getArgument(0).getType()) || + !mlir::isa(caller.getArgument(1).getType())) { + return; + } + SmallVector newOperands; + newOperands.push_back(caller.getArgument(0)); + newOperands.push_back(caller.getArgument(1)); + newOperands.append(callOp.getOperands().begin(), + callOp.getOperands().end()); + callOp->setOperands(newOperands); + }); +} + FailureOr getContextualCryptoContextForJaxiteWord(Operation* op) { auto funcOp = op->getParentOfType(); if (!funcOp) return failure(); @@ -75,38 +127,21 @@ FailureOr getContextualEvalKeyForJaxiteWord(Operation* op) { return funcOp.getArgument(1); } -struct AddCryptoContextAndKeys : public OpConversionPattern { - AddCryptoContextAndKeys(mlir::MLIRContext* context) - : OpConversionPattern(context, /* benefit= */ 2) {} - - using OpConversionPattern::OpConversionPattern; - - LogicalResult matchAndRewrite( - func::FuncOp op, OpAdaptor adaptor, - ConversionPatternRewriter& rewriter) const override { - auto containsCryptoOps = - ::mlir::heir::containsDialects(op); - if (!containsCryptoOps) return failure(); - - auto cryptoContextType = jaxiteword::CryptoContextType::get(getContext()); - auto evalKeyType = jaxiteword::EvalKeyType::get(getContext()); - - rewriter.startOpModification(op); - bool hasContext = op.getFunctionType().getNumInputs() > 0 && - mlir::isa( - op.getFunctionType().getInput(0)); - if (!hasContext) { - if (failed(op.insertArgument(0, evalKeyType, nullptr, op.getLoc()))) - return failure(); - if (failed(op.insertArgument(0, cryptoContextType, nullptr, op.getLoc()))) - return failure(); - } - rewriter.finalizeOpModification(op); - - return success(); +static FailureOr getStaticRotationIndex(ckks::RotateOp op, + Value dynamicShift) { + auto i64Type = IntegerType::get(op.getContext(), 64); + if (IntegerAttr staticShift = op.getStaticShiftAttr()) { + return IntegerAttr::get(i64Type, staticShift.getValue().getSExtValue()); } -}; + if (!dynamicShift) { + return failure(); + } + IntegerAttr intAttr; + if (matchPattern(dynamicShift, m_Constant(&intAttr))) { + return IntegerAttr::get(i64Type, intAttr.getValue().getSExtValue()); + } + return failure(); +} template struct ConvertBinOp : public OpConversionPattern { @@ -125,27 +160,10 @@ struct ConvertBinOp : public OpConversionPattern { } }; -struct ConvertMulOp : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - LogicalResult matchAndRewrite( - ckks::MulOp op, OpAdaptor adaptor, - ConversionPatternRewriter& rewriter) const override { - FailureOr ctx = getContextualCryptoContextForJaxiteWord(op); - if (failed(ctx)) return failure(); - - FailureOr evalKey = getContextualEvalKeyForJaxiteWord(op); - if (failed(evalKey)) return failure(); - - rewriter.replaceOpWithNewOp( - op, this->getTypeConverter()->convertType(op.getOutput().getType()), - ctx.value(), adaptor.getLhs(), adaptor.getRhs(), evalKey.value()); - return success(); - } -}; - -template -struct ConvertNegateOp : public OpConversionPattern { +// The JaxiteWord API requires ciphertext-plaintext operand ordering even for +// commutative source operations. +template +struct ConvertCommutativePlainOp : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; LogicalResult matchAndRewrite( @@ -154,9 +172,23 @@ struct ConvertNegateOp : public OpConversionPattern { FailureOr ctx = getContextualCryptoContextForJaxiteWord(op); if (failed(ctx)) return failure(); - rewriter.replaceOpWithNewOp( + Value ciphertext; + Value plaintext; + if (isa( + getElementTypeOrSelf(op.getLhs().getType()))) { + ciphertext = adaptor.getLhs(); + plaintext = adaptor.getRhs(); + } else if (isa( + getElementTypeOrSelf(op.getRhs().getType()))) { + ciphertext = adaptor.getRhs(); + plaintext = adaptor.getLhs(); + } else { + return rewriter.notifyMatchFailure(op, "expected one ciphertext operand"); + } + + rewriter.replaceOpWithNewOp( op, this->getTypeConverter()->convertType(op.getOutput().getType()), - ctx.value(), adaptor.getInput()); + ctx.value(), ciphertext, plaintext); return success(); } }; @@ -173,20 +205,16 @@ struct ConvertRotateOp : public OpConversionPattern { FailureOr evalKey = getContextualEvalKeyForJaxiteWord(op); if (failed(evalKey)) return failure(); - Value dynamicShift = adaptor.getDynamicShift(); - IntegerAttr staticShift = op.getStaticShiftAttr(); - if (!staticShift && !dynamicShift) { + FailureOr indexAttr = + getStaticRotationIndex(op, adaptor.getDynamicShift()); + if (failed(indexAttr)) { return rewriter.notifyMatchFailure( - op, "rotate op must have either static or dynamic shift"); - } - if (dynamicShift) { - return rewriter.notifyMatchFailure( - op, "jaxiteword rotation requires static shift"); + op, "jaxiteword rotation requires statically known shift"); } rewriter.replaceOpWithNewOp( op, this->getTypeConverter()->convertType(op.getOutput().getType()), - ctx.value(), adaptor.getInput(), evalKey.value(), staticShift); + ctx.value(), adaptor.getInput(), evalKey.value(), indexAttr.value()); return success(); } }; @@ -236,9 +264,11 @@ struct ConvertEncodeOp : public OpConversionPattern { FailureOr ctx = getContextualCryptoContextForJaxiteWord(op); if (failed(ctx)) return failure(); - rewriter.replaceOpWithNewOp( - op, this->getTypeConverter()->convertType(op.getOutput().getType()), + auto newOp = jaxiteword::EncodeOp::create( + rewriter, op.getLoc(), + this->getTypeConverter()->convertType(op.getOutput().getType()), ctx.value(), adaptor.getInput()); + rewriter.replaceOp(op, newOp.getResult()); return success(); } }; @@ -298,35 +328,44 @@ struct LWEToJaxiteWord : public impl::LWEToJaxiteWordBase { MLIRContext* context = &getContext(); Operation* op = getOperation(); + op->walk([&](func::FuncOp funcOp) { insertCryptoContextAndKeys(funcOp); }); + updateCryptoFuncCalls(op); + RewritePatternSet patterns(context); ConversionTarget target(*context); target.addLegalDialect(); - target.addIllegalDialect(); + target.addIllegalDialect(); target.addLegalOp(); JaxiteWordTypeConverter typeConverter(context); - target.addDynamicallyLegalOp([&](func::FuncOp op) { - auto containsCryptoOps = - ::mlir::heir::containsDialects(op); - if (!containsCryptoOps) return true; - bool hasArgs = op.getFunctionType().getNumInputs() >= 2; - return typeConverter.isSignatureLegal(op.getFunctionType()) && hasArgs && + target.addDynamicallyLegalOp([&](func::FuncOp funcOp) { + if (!funcNeedsCryptoContextAndKeys(funcOp)) return true; + bool hasArgs = funcOp.getFunctionType().getNumInputs() >= 2; + return typeConverter.isSignatureLegal(funcOp.getFunctionType()) && + hasArgs && mlir::isa( - op.getFunctionType().getInput(0)) && + funcOp.getFunctionType().getInput(0)) && mlir::isa( - op.getFunctionType().getInput(1)); + funcOp.getFunctionType().getInput(1)); + }); + + target.addDynamicallyLegalOp([&](func::CallOp callOp) { + if (auto callee = getCalledFunction(callOp); succeeded(callee)) { + if (funcNeedsCryptoContextAndKeys(callee.value())) { + return callOp.getNumOperands() == + callOp.getCalleeType().getNumInputs(); + } + } + return true; }); populateFunctionOpInterfaceTypeConversionPattern( patterns, typeConverter); addTensorConversionPatterns(typeConverter, patterns, target); - patterns.add(typeConverter, context); patterns.add>(typeConverter, context); patterns.add>(typeConverter, @@ -335,17 +374,20 @@ struct LWEToJaxiteWord : public impl::LWEToJaxiteWordBase { context); patterns.add>(typeConverter, context); + patterns.add< + ConvertCommutativePlainOp, + ConvertCommutativePlainOp, + ConvertCommutativePlainOp, + ConvertCommutativePlainOp>( + typeConverter, context); patterns.add>(typeConverter, context); - patterns.add(typeConverter, context); + patterns.add>( + typeConverter, context); patterns.add>( typeConverter, context); - patterns.add>(typeConverter, context); - patterns.add>(typeConverter, context); patterns.add(typeConverter, context); patterns.add(typeConverter, context); - patterns.add>(typeConverter, - context); patterns.add>(typeConverter, context); patterns.add>(typeConverter, context); @@ -353,19 +395,6 @@ struct LWEToJaxiteWord : public impl::LWEToJaxiteWordBase { patterns.add(typeConverter, context); patterns.add(typeConverter, context); patterns.add(typeConverter, context); - patterns.add>( - typeConverter, context); - patterns.add>( - typeConverter, context); - patterns.add>( - typeConverter, context); - patterns.add>( - typeConverter, context); - patterns.add>( - typeConverter, context); - patterns.add>( - typeConverter, context); - if (failed(applyPartialConversion(op, target, std::move(patterns)))) { return signalPassFailure(); } diff --git a/lib/Target/JaxiteWord/JaxiteWordEmitter.cpp b/lib/Target/JaxiteWord/JaxiteWordEmitter.cpp index 0122cb4517..e7b5c8e5a9 100644 --- a/lib/Target/JaxiteWord/JaxiteWordEmitter.cpp +++ b/lib/Target/JaxiteWord/JaxiteWordEmitter.cpp @@ -30,6 +30,7 @@ #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/DialectRegistry.h" // from @llvm-project +#include "mlir/include/mlir/IR/TypeUtilities.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/ValueRange.h" // from @llvm-project @@ -115,12 +116,12 @@ LogicalResult JaxiteWordEmitter::translate(Operation& op) { LogicalResult status = llvm::TypeSwitch(op) .Case([&](auto op) { return printOperation(op); }) - .Case( + .Case( [&](auto op) { return printOperation(op); }) - .Case( [&](auto op) { return printOperation(op); }) .Case([&](auto op) { return printOperation(op); }) .Case( - [&](auto op) { return printOperation(op); }) + arith::DivSIOp, arith::FloorDivSIOp, arith::RemSIOp, + arith::CmpIOp, arith::SelectOp, arith::ExtSIOp, arith::ExtUIOp, + arith::TruncIOp>([&](auto op) { return printOperation(op); }) .Default([&](Operation&) { return op.emitOpError("unable to find printer for op"); }); @@ -168,9 +169,6 @@ LogicalResult JaxiteWordEmitter::printOperation(func::FuncOp funcOp) { << "Failed to emit JaxiteWord type " << arg.getType(); } os << ",\n"; - if (isa(arg.getType())) { - CiphertextArg_ = argName; - } } os.unindent(); os << ")"; @@ -213,6 +211,26 @@ LogicalResult JaxiteWordEmitter::printOperation(func::FuncOp funcOp) { return success(); } +LogicalResult JaxiteWordEmitter::printOperation(func::CallOp op) { + if (op.getNumResults() == 1) { + emitAssignPrefix(op.getResult(0)); + } else if (op.getNumResults() > 1) { + os << "("; + for (auto [idx, result] : llvm::enumerate(op.getResults())) { + if (idx > 0) os << ", "; + os << variableNames->getNameForValue(result); + } + os << ") = "; + } + + os << op.getCallee().str() << "("; + os << commaSeparatedValues(op.getOperands(), [&](Value value) { + return variableNames->getNameForValue(value); + }); + os << ")\n"; + return success(); +} + LogicalResult JaxiteWordEmitter::printOperation(func::ReturnOp op) { std::function resultValue = [&](Value value) { if (isa(value)) { @@ -236,38 +254,35 @@ LogicalResult JaxiteWordEmitter::printOperation(func::ReturnOp op) { } LogicalResult JaxiteWordEmitter::printOperation(AddOp op) { - return printBinaryOpHelper( - op.getResult(), op.getLhs(), op.getRhs(), - [&](StringRef lhs, StringRef rhs, StringRef result) { - os << lhs << "\n"; - os << llvm::formatv(kAddCoreTemplate.data(), result, rhs); - }); + return printBinaryOpHelper(op.getResult(), op.getLhs(), op.getRhs(), + op.getCryptoContext(), op, "he_add", "add"); } -LogicalResult JaxiteWordEmitter::printOperation(SubOp op) { - return printBinaryOpHelper( - op.getResult(), op.getLhs(), op.getRhs(), - [&](StringRef lhs, StringRef rhs, StringRef result) { - os << lhs << "\n"; - std::string rhsCiphertext = (rhs + ".ciphertext").str(); - os << llvm::formatv(kSubTemplate.data(), result, rhsCiphertext); - }); +LogicalResult JaxiteWordEmitter::printOperation(AddPlainOp op) { + auto ctx = variableNames->getNameForValue(op.getCryptoContext()); + auto ct = variableNames->getNameForValue(op.getCiphertext()); + auto pt = variableNames->getNameForValue(op.getPlaintext()); + auto result = variableNames->getNameForValue(op.getOutput()); + auto level = + getCrossLevelExpr(op.getCiphertext(), ctx, op, /*extraOffset=*/0); + os << result << " = " << ctx << ".he_add[" << level << "].add_plain(" << ct + << ", " << pt << ")\n"; + return success(); } -LogicalResult JaxiteWordEmitter::printOperation(NegateOp op) { - emitAssignPrefix(op.getResult()); - os << variableNames->getNameForValue(op.getCiphertext()) << ".mul(-1)\n"; - return success(); +LogicalResult JaxiteWordEmitter::printOperation(SubOp op) { + return printBinaryOpHelper(op.getResult(), op.getLhs(), op.getRhs(), + op.getCryptoContext(), op, "he_sub", "sub"); } LogicalResult JaxiteWordEmitter::printOperation(SquareOp op) { auto ct = variableNames->getNameForValue(op.getCiphertext()); auto ctx = variableNames->getNameForValue(op.getCryptoContext()); + auto result = variableNames->getNameForValue(op.getResult()); auto level = getCrossLevelExpr(op.getCiphertext(), ctx, op, /*extraOffset=*/1); - - emitAssignPrefix(op.getResult()); - os << ctx << ".he_mul[" << level << "].mul(" << ct << ", " << ct << ")\n"; + os << result << " = " << ctx << ".he_mul[" << level << "].mul(" << ct << ", " + << ct << ")\n"; return success(); } @@ -275,25 +290,19 @@ void JaxiteWordEmitter::emitAssignPrefix(Value result) { os << variableNames->getNameForValue(result) << " = "; } -LogicalResult JaxiteWordEmitter::printBinaryOpHelper( - Value result, Value lhs, Value rhs, - llvm::function_ref callback) { - auto lhsName = variableNames->getNameForValue(lhs); - auto rhsName = variableNames->getNameForValue(rhs); +LogicalResult JaxiteWordEmitter::printBinaryOpHelper(Value result, Value lhs, + Value rhs, Value ctx, + Operation* op, + StringRef accessor, + StringRef method) { auto resultName = variableNames->getNameForValue(result); - - emitAssignPrefix(result); - callback(lhsName, rhsName, resultName); - return success(); -} - -LogicalResult JaxiteWordEmitter::printInPlaceBinaryOpHelper( - Value lhs, Value rhs, - llvm::function_ref callback) { auto lhsName = variableNames->getNameForValue(lhs); auto rhsName = variableNames->getNameForValue(rhs); + auto ctxName = variableNames->getNameForValue(ctx); + auto level = getCrossLevelExpr(lhs, ctxName, op, /*extraOffset=*/0); - callback(lhsName, rhsName); + os << resultName << " = " << ctxName << "." << accessor << "[" << level + << "]." << method << "(" << lhsName << ", " << rhsName << ")\n"; return success(); } @@ -308,7 +317,6 @@ LogicalResult JaxiteWordEmitter::printMulOpHelper( auto resultName = variableNames->getNameForValue(result); auto level = getCrossLevelExpr(lhs, ctxName, op, /*extraOffset=*/1); - emitAssignPrefix(result); callback(lhsName, rhsName, ctxName, resultName, level); return success(); } @@ -323,11 +331,10 @@ LogicalResult JaxiteWordEmitter::printOperation(EncodeOp op) { LogicalResult JaxiteWordEmitter::printOperation(EncryptOp op) { auto ctx = variableNames->getNameForValue(op.getCryptoContext()); auto pk = variableNames->getNameForValue(op.getPublicKey()); + auto pt = variableNames->getNameForValue(op.getPlaintext()); + auto result = variableNames->getNameForValue(op.getResult()); os << ctx << ".public_key = " << pk << "\n"; - - emitAssignPrefix(op.getResult()); - os << ctx << ".encrypt(" << variableNames->getNameForValue(op.getPlaintext()) - << ")\n"; + os << result << " = " << ctx << ".encrypt(" << pt << ")\n"; return success(); } @@ -336,8 +343,8 @@ LogicalResult JaxiteWordEmitter::printOperation(MulOp op) { op.getCryptoContext(), op.getOperation(), [&](StringRef lhs, StringRef rhs, StringRef ctx, StringRef result, StringRef level) { - os << ctx << ".he_mul[" << level << "].hemul(" - << lhs << ", " << rhs << ")\n"; + os << result << " = " << ctx << ".he_mul[" << level + << "].mul(" << lhs << ", " << rhs << ")\n"; }); } @@ -346,7 +353,7 @@ LogicalResult JaxiteWordEmitter::printOperation(MulNoRelinOp op) { op.getCryptoContext(), op.getOperation(), [&](StringRef lhs, StringRef rhs, StringRef ctx, StringRef result, StringRef level) { - os << ctx << ".he_mul[" << level + os << result << " = " << ctx << ".he_mul[" << level << "].hemul_no_relin(" << lhs << ", " << rhs << ")\n"; }); @@ -355,32 +362,39 @@ LogicalResult JaxiteWordEmitter::printOperation(MulNoRelinOp op) { LogicalResult JaxiteWordEmitter::printOperation(RelinOp op) { auto ct = variableNames->getNameForValue(op.getCiphertext()); auto ctx = variableNames->getNameForValue(op.getCryptoContext()); + auto result = variableNames->getNameForValue(op.getOutput()); auto level = getCrossLevelExpr(op.getCiphertext(), ctx, op, /*extraOffset=*/1); - - auto result = variableNames->getNameForValue(op.getOutput()); - os << llvm::formatv(kRelinTemplate.data(), result, ctx, level, ct); - + os << result << " = " << ctx << ".he_mul[" << level << "].relinearize(" << ct + << ")\n"; return success(); } LogicalResult JaxiteWordEmitter::printOperation(ModReduceOp op) { + auto ctx = variableNames->getNameForValue(op.getCryptoContext()); auto ct = variableNames->getNameForValue(op.getCiphertext()); - emitAssignPrefix(op.getResult()); - os << ct << "\n"; + auto result = variableNames->getNameForValue(op.getResult()); + auto srcLevel = + getCrossLevelExpr(op.getCiphertext(), ctx, op, /*extraOffset=*/0); + auto dstLevel = getCrossLevelExpr(op.getResult(), ctx, op, /*extraOffset=*/0); + if (srcLevel == dstLevel) { + os << result << " = " << ct << "\n"; + return success(); + } + os << result << " = " << ctx << ".he_rescale[" << srcLevel << ", " << dstLevel + << "].rescale(" << ct << ")\n"; return success(); } LogicalResult JaxiteWordEmitter::printOperation(RotOp op) { auto ct = variableNames->getNameForValue(op.getCiphertext()); auto ctx = variableNames->getNameForValue(op.getCryptoContext()); + auto result = variableNames->getNameForValue(op.getResult()); auto rotIndex = op.getIndex(); auto level = getCrossLevelExpr(op.getCiphertext(), ctx, op, /*extraOffset=*/0); - - emitAssignPrefix(op.getResult()); - os << ctx << ".he_rot[" << level << ", " << rotIndex << "].rotate(" << ct - << ")\n"; + os << result << " = " << ctx << ".he_rot[" << level << ", " << rotIndex + << "].rotate(" << ct << ")\n"; return success(); } @@ -388,16 +402,10 @@ LogicalResult JaxiteWordEmitter::printOperation(DecryptOp op) { auto ctx = variableNames->getNameForValue(op.getCryptoContext()); auto sk = variableNames->getNameForValue(op.getSecretKey()); auto ct = variableNames->getNameForValue(op.getCiphertext()); - - auto ctType = cast(op.getCiphertext().getType()); - int current = ctType.getModulusChain().getCurrent(); - int maxCurrent = getMaxCurrentInModule(op); - int rescales = maxCurrent - current; - - os << llvm::formatv(kDecryptTemplate.data(), ctx, sk, rescales, ct); - + auto result = variableNames->getNameForValue(op.getResult()); + os << ctx << ".secret_key = " << sk << "\n"; emitAssignPrefix(op.getResult()); - os << ctx << ".decrypt(_ct_for_dec)\n"; + os << ctx << ".decrypt(" << ct << ")\n"; return success(); } @@ -424,50 +432,34 @@ LogicalResult JaxiteWordEmitter::printOperation(MulPlainOp op) { auto ctx = variableNames->getNameForValue(op.getCryptoContext()); auto ct = variableNames->getNameForValue(op.getCiphertext()); auto pt = variableNames->getNameForValue(op.getPlaintext()); + auto result = variableNames->getNameForValue(op.getResult()); auto level = getCrossLevelExpr(op.getCiphertext(), ctx, op, /*extraOffset=*/0); - - os << ctx << ".ptct_mul[" << level << "].set_plaintext(" << pt << ")\n"; - emitAssignPrefix(op.getResult()); - os << ctx << ".ptct_mul[" << level << "].mul(" << ct << ")\n"; + os << result << " = " << ctx << ".ptct_mul[" << level << "].mul(" << ct + << ", " << pt << ")\n"; return success(); } -LogicalResult JaxiteWordEmitter::printOperation(AddPlainOp op) { +LogicalResult JaxiteWordEmitter::printOperation(AddInPlaceOp op) { + auto ctx = variableNames->getNameForValue(op.getCryptoContext()); auto lhs = variableNames->getNameForValue(op.getLhs()); auto rhs = variableNames->getNameForValue(op.getRhs()); - os << lhs << ".ciphertext = " << lhs << ".ciphertext + " << rhs << "\n"; - os << llvm::formatv(kAddModReduceTemplate.data(), lhs); - - emitAssignPrefix(op.getResult()); - os << lhs << "\n"; + auto level = getCrossLevelExpr(op.getLhs(), ctx, op, /*extraOffset=*/0); + os << lhs << " = " << ctx << ".he_add[" << level << "].add(" << lhs << ", " + << rhs << ")\n"; return success(); } -LogicalResult JaxiteWordEmitter::printOperation(SubPlainOp op) { +LogicalResult JaxiteWordEmitter::printOperation(SubInPlaceOp op) { + auto ctx = variableNames->getNameForValue(op.getCryptoContext()); auto lhs = variableNames->getNameForValue(op.getLhs()); auto rhs = variableNames->getNameForValue(op.getRhs()); - os << llvm::formatv(kSubTemplate.data(), lhs, rhs); - emitAssignPrefix(op.getResult()); - os << lhs << "\n"; + auto level = getCrossLevelExpr(op.getLhs(), ctx, op, /*extraOffset=*/0); + os << lhs << " = " << ctx << ".he_sub[" << level << "].sub(" << lhs << ", " + << rhs << ")\n"; return success(); } -LogicalResult JaxiteWordEmitter::printOperation(AddInPlaceOp op) { - return printInPlaceBinaryOpHelper( - op.getLhs(), op.getRhs(), [&](StringRef lhs, StringRef rhs) { - os << llvm::formatv(kAddCoreTemplate.data(), lhs, rhs); - }); -} - -LogicalResult JaxiteWordEmitter::printOperation(SubInPlaceOp op) { - return printInPlaceBinaryOpHelper( - op.getLhs(), op.getRhs(), [&](StringRef lhs, StringRef rhs) { - std::string rhsCiphertext = (rhs + ".ciphertext").str(); - os << llvm::formatv(kSubTemplate.data(), lhs, rhsCiphertext); - }); -} - LogicalResult JaxiteWordEmitter::printOperation(GenKeyPairOp op) { auto pk = variableNames->getNameForValue(op.getPublicKey()); auto sk = variableNames->getNameForValue(op.getSecretKey()); @@ -485,12 +477,17 @@ LogicalResult JaxiteWordEmitter::printOperation(GenMulKeyOp op) { auto ctx = variableNames->getNameForValue(op.getCryptoContext()); auto sk = variableNames->getNameForValue(op.getSecretKey()); - os << ek << " = key_gen.gen_evaluation_key(" << sk << ", " << "q=" << ctx + os << ek << "_raw = key_gen.gen_evaluation_key(" << sk << ", " << "q=" << ctx << ".q_towers, " << "P=" << ctx << ".p_towers, " << "dnum=" << ctx << ".parameters.get('dnum', 3)" << ")\n"; - - heMulVarName_ = "he_mul"; - os << llvm::formatv(kGenMulKeyTemplate.data(), ctx, heMulVarName_, ek); + os << ek << " = [\n"; + os.indent(); + os << "jnp.array(" << ek + << "_raw[\"a\"], dtype=jnp.uint32).transpose(0, 2, 1),\n"; + os << "jnp.array(" << ek + << "_raw[\"b\"], dtype=jnp.uint32).transpose(0, 2, 1),\n"; + os.unindent(); + os << "]\n"; return success(); } @@ -500,20 +497,21 @@ LogicalResult JaxiteWordEmitter::printOperation(GenRotKeyOp op) { auto ctx = variableNames->getNameForValue(op.getCryptoContext()); auto sk = variableNames->getNameForValue(op.getSecretKey()); - rotKeysDictVarName_ = rk + "_dict"; - std::string indicesStr; - llvm::raw_string_ostream indicesOs(indicesStr); - llvm::interleaveComma(op.getIndices(), indicesOs); - - os << llvm::formatv(kGenRotKeyTemplate.data(), rotKeysDictVarName_, - indicesStr, sk, ctx, heRotVarName_, rk); + os << rk << " = {}\n"; + os << "for _rot_idx in ["; + llvm::interleaveComma(op.getIndices(), os); + os << "]:\n"; + os.indent(); + os << rk << "[_rot_idx] = key_gen.gen_rotation_key(" << sk << ", " << ctx + << ".q_towers, " << ctx << ".p_towers, rot_index=_rot_idx, dnum=" << ctx + << ".parameters.get('dnum', 3))[_rot_idx]\n"; + os.unindent(); return success(); } LogicalResult JaxiteWordEmitter::printOperation(GenParamsOp op) { auto ctx = variableNames->getNameForValue(op.getCryptoContext()); - cryptoContextVarName_ = ctx; os << "params = {\n"; os.indent(); @@ -527,8 +525,8 @@ LogicalResult JaxiteWordEmitter::printOperation(GenParamsOp op) { os << "\"r\": " << op.getR() << ",\n"; os << "\"c\": " << op.getC() << ",\n"; os << "\"dnum\": " << op.getDnum() << ",\n"; - os << "\"numEvalMult\": " << op.getNumEvalMult() << ",\n"; os << "\"scaling_factor\": " << op.getScalingFactor() << ",\n"; + os << "\"output_scale\": " << op.getScalingFactor() << ",\n"; os << "\"q_towers\": ["; llvm::interleaveComma(qTowers, os); @@ -553,12 +551,18 @@ LogicalResult JaxiteWordEmitter::printOperation(GenParamsOp op) { LogicalResult JaxiteWordEmitter::printOperation(ProgramInitializationOp op) { auto ctx = variableNames->getNameForValue(op.getCryptoContext()); + auto pk = variableNames->getNameForValue(op.getPublicKey()); auto sk = variableNames->getNameForValue(op.getSecretKey()); + auto ek = variableNames->getNameForValue(op.getEvaluationKey()); + os << ctx << ".public_key = " << pk << "\n"; os << ctx << ".secret_key = " << sk << "\n"; - os << ctx << ".program_initialization("; - os << "total_hemul_levels=" << op.getTotalHemulLevels() << ", "; + os << ctx << ".evaluation_key = " << ek << "\n"; + os << ctx << ".parameters[\"public_key\"] = " << pk << "\n"; + os << ctx << ".parameters[\"secret_key\"] = " << sk << "\n"; + os << ctx << ".parameters[\"evaluation_key\"] = " << ek << "\n"; + os << ctx << ".program_initialization("; os << "total_rotation_indices=["; auto rotIndices = op.getTotalRotationIndices(); llvm::interleaveComma(rotIndices, os); @@ -682,7 +686,8 @@ LogicalResult JaxiteWordEmitter::printOperation(tensor::EmptyOp op) { return success(); } - os << " = np.zeros(("; + emitAssignPrefix(op.getResult()); + os << "np.zeros(("; for (size_t i = 0; i < shape.size(); ++i) { if (i > 0) os << ", "; os << shape[i]; @@ -1028,6 +1033,13 @@ LogicalResult JaxiteWordEmitter::printOperation(arith::DivSIOp op) { return success(); } +LogicalResult JaxiteWordEmitter::printOperation(arith::FloorDivSIOp op) { + os << variableNames->getNameForValue(op.getResult()) << " = " + << variableNames->getNameForValue(op.getLhs()) << " // " + << variableNames->getNameForValue(op.getRhs()) << "\n"; + return success(); +} + LogicalResult JaxiteWordEmitter::printOperation(arith::RemSIOp op) { os << variableNames->getNameForValue(op.getResult()) << " = " << variableNames->getNameForValue(op.getLhs()) << " % " @@ -1119,14 +1131,14 @@ FailureOr JaxiteWordEmitter::convertType(Type type) { return llvm::TypeSwitch>(type) .Case( - [&](auto) { return std::string("Ciphertext"); }) + [&](auto) { return std::string("Polynomial"); }) .Case([&](auto) { return std::string("np.ndarray"); }) .Case([&](auto) { return std::string("np.ndarray"); }) .Case([&](auto) { return std::string("dict"); }) .Case( [&](auto) { return std::string("ckks.CKKSContext"); }) .Case( - [&](auto) { return std::string("Ciphertext"); }) + [&](auto) { return std::string("Polynomial"); }) .Default([&](Type) { return failure(); }); } diff --git a/lib/Target/JaxiteWord/JaxiteWordEmitter.h b/lib/Target/JaxiteWord/JaxiteWordEmitter.h index c744b6205f..a8fbcf1994 100644 --- a/lib/Target/JaxiteWord/JaxiteWordEmitter.h +++ b/lib/Target/JaxiteWord/JaxiteWordEmitter.h @@ -43,28 +43,13 @@ class JaxiteWordEmitter { // values. SelectVariableNames* variableNames; - // ciphertext arg. - std::string CiphertextArg_; - - // A list of modulus to be used for the add operation. - std::string ModulusListArg_; - - // Crypto context variable name (set by GenParamsOp, used by accessor calls) - std::string cryptoContextVarName_; - - // Legacy member variables kept for backward compatibility with old pipeline - // (GenMulKeyOp / GenRotKeyOp path). Not used by the new - // ProgramInitializationOp path. - std::string heMulVarName_; - std::string heRotVarName_; - std::string rotKeysDictVarName_; - LogicalResult printOperation(ModuleOp moduleOp); LogicalResult printOperation(func::FuncOp funcOp); + LogicalResult printOperation(func::CallOp op); LogicalResult printOperation(func::ReturnOp returnOp); LogicalResult printOperation(AddOp op); + LogicalResult printOperation(AddPlainOp op); LogicalResult printOperation(SubOp op); - LogicalResult printOperation(NegateOp op); LogicalResult printOperation(SquareOp op); LogicalResult printOperation(MulOp op); LogicalResult printOperation(MulNoRelinOp op); @@ -72,8 +57,6 @@ class JaxiteWordEmitter { LogicalResult printOperation(RotOp op); LogicalResult printOperation(RelinOp op); - LogicalResult printOperation(AddPlainOp op); - LogicalResult printOperation(SubPlainOp op); LogicalResult printOperation(MulPlainOp op); LogicalResult printOperation(AddInPlaceOp op); LogicalResult printOperation(SubInPlaceOp op); @@ -109,6 +92,7 @@ class JaxiteWordEmitter { LogicalResult printOperation(arith::SubIOp op); LogicalResult printOperation(arith::MulIOp op); LogicalResult printOperation(arith::DivSIOp op); + LogicalResult printOperation(arith::FloorDivSIOp op); LogicalResult printOperation(arith::RemSIOp op); LogicalResult printOperation(arith::CmpIOp op); LogicalResult printOperation(arith::SelectOp op); @@ -126,13 +110,9 @@ class JaxiteWordEmitter { void emitAssignPrefix(Value result); - LogicalResult printBinaryOpHelper( - Value result, Value lhs, Value rhs, - llvm::function_ref callback); - - LogicalResult printInPlaceBinaryOpHelper( - Value lhs, Value rhs, - llvm::function_ref callback); + LogicalResult printBinaryOpHelper(Value result, Value lhs, Value rhs, + Value ctx, Operation* op, + StringRef accessor, StringRef method); LogicalResult printMulOpHelper( Value result, Value lhs, Value rhs, Value ctx, Operation* op, diff --git a/lib/Target/JaxiteWord/JaxiteWordTemplates.h b/lib/Target/JaxiteWord/JaxiteWordTemplates.h index d979ef6c05..fb21879a5b 100644 --- a/lib/Target/JaxiteWord/JaxiteWordTemplates.h +++ b/lib/Target/JaxiteWord/JaxiteWordTemplates.h @@ -10,120 +10,13 @@ namespace jaxiteword { constexpr std::string_view kModulePrelude = R"python( import jax import jax.numpy as jnp +import key_gen import numpy as np -from ciphertext import Ciphertext from polynomial import Polynomial import ckks_ctx as ckks )python"; -// Template for GenMulKeyOp -// This template initializes HEMul for homomorphic multiplication and sets up -// relinearization. It computes r and c from the degree if they are not provided -// in the parameters. -constexpr std::string_view kGenMulKeyTemplate = R"python( -_degree = {0}.parameters.get('degree') -if _degree is not None: - _log_degree = int(math.log2(_degree)) - _half_k = _log_degree // 2 - _default_r = 2 ** _half_k - _default_c = _degree // _default_r -else: - _default_r = 4 - _default_c = 4 -{1} = HEMul( - batch={0}.parameters.get('batch', 1), - r={0}.parameters.get('r', _default_r), - c={0}.parameters.get('c', _default_c), - dnum={0}.parameters.get('dnum', 3), - num_eval_mult={0}.parameters.get('numEvalMult', 1), - original_moduli={0}.q_towers, - extend_moduli={0}.p_towers -) -{1}.control_gen(degree_layout=({0}.parameters.get('r', _default_r), {0}.parameters.get('c', _default_c))) -{1}.setup_relinearization(jnp.array({2}["a"], dtype=jnp.uint32).transpose(0,2,1), jnp.array({2}["b"], dtype=jnp.uint32).transpose(0,2,1)) -)python"; - -// Template for GenRotKeyOp -// This template generates rotation keys for power-of-2 indices and initializes -// HERot. It computes r and c from the degree if they are not provided in the -// parameters. -constexpr std::string_view kGenRotKeyTemplate = R"python( -{0} = {{}} -_all_indices = [{1}] -_max_abs_rot_idx = max([abs(idx) for idx in _all_indices]) if _all_indices else 1 -_power_of_2_indices = [] -_pow2 = 1 -while _pow2 <= _max_abs_rot_idx: - _power_of_2_indices.append(_pow2) - _pow2 <<= 1 -_neg_power_of_2_indices = [-idx for idx in _power_of_2_indices] -_all_pow2_indices = _power_of_2_indices + _neg_power_of_2_indices - -for _rot_idx in _all_pow2_indices: - {0}[_rot_idx] = key_gen.gen_rotation_key({2}, {3}.q_towers, {3}.p_towers, rot_index=_rot_idx, dnum={3}.parameters.get('dnum', 3)) - -_degree_rot = {3}.parameters.get('degree') -if _degree_rot is not None: - _log_degree_rot = int(math.log2(_degree_rot)) - _default_r_rot = 1 << (_log_degree_rot // 2) - _default_c_rot = _degree_rot // _default_r_rot -else: - _default_r_rot = 4 - _default_c_rot = 4 -{4} = HERot( - r={3}.parameters.get('r', _default_r_rot), - c={3}.parameters.get('c', _default_c_rot), - dnum={3}.parameters.get('dnum', 3), - rotate_in_ciphertext_moduli={3}.q_towers, - extend_moduli={3}.p_towers -) -{4}.control_gen(batch=1, degree_layout=({3}.parameters.get('r', _default_r_rot), {3}.parameters.get('c', _default_c_rot))) -{5} = {0} -)python"; - -// Template for DecryptOp -// This template prepares the ciphertext for decryption by extracting the -// required moduli. -constexpr std::string_view kDecryptTemplate = R"python( -{0}.secret_key = {1} -_rescales = {2} -_num_moduli = len({0}.q_towers) - _rescales * {0}.composite_degree -_q_sub = {0}.q_towers[:_num_moduli] -_ct_for_dec = Polynomial( - {{'batch': 1, 'num_elements': 2, 'degree': {0}.degree, - 'precision': 32, 'num_moduli': _num_moduli, - 'degree_layout': ({0}.degree,)}}, - {{'moduli': _q_sub}}) -_ct_for_dec.set_batch_polynomial({3}.polynomial.reshape(1, 2, {0}.degree, _num_moduli)) -)python"; - -// Template for AddOp and AddInPlaceOp -// This template performs addition and modular reduction. -constexpr std::string_view kAddCoreTemplate = R"python( -{0}.add({1}) -{0}.ciphertext = jnp.where({0}.ciphertext >= {0}.moduli_array, {0}.ciphertext - {0}.moduli_array, {0}.ciphertext) -)python"; - -// Template for SubOp, SubInPlaceOp, and SubPlainOp -// This template performs subtraction and modular reduction. -// {1} should be rhs.ciphertext for Sub/SubInPlace and just rhs for SubPlain. -constexpr std::string_view kSubTemplate = R"python( -{0}.ciphertext = jnp.where({0}.ciphertext < {1}, {0}.ciphertext + {0}.moduli_array - {1}, {0}.ciphertext - {1}) -)python"; - -// Template for AddPlainOp modular reduction -constexpr std::string_view kAddModReduceTemplate = R"python( -{0}.ciphertext = jnp.where({0}.ciphertext >= {0}.moduli_array, {0}.ciphertext - {0}.moduli_array, {0}.ciphertext) -)python"; - -// Template for RelinOp -constexpr std::string_view kRelinTemplate = R"python( -{0} = {1}.he_mul[{2}].relinearize({3}) -_s = {0}.polynomial.shape -{0}.polynomial = {0}.polynomial.reshape(_s[0], _s[1], {0}.degree, _s[-1]) -)python"; - } // namespace jaxiteword } // namespace heir } // namespace mlir diff --git a/tests/Dialect/JaxiteWord/IR/ops.mlir b/tests/Dialect/JaxiteWord/IR/ops.mlir index 263ce206cf..f1ecd9e2e6 100644 --- a/tests/Dialect/JaxiteWord/IR/ops.mlir +++ b/tests/Dialect/JaxiteWord/IR/ops.mlir @@ -24,6 +24,13 @@ !ct_L1_D3 = !lwe.lwe_ciphertext, ciphertext_space = #ciphertext_space_L1_D3, key = #key, modulus_chain = #modulus_chain_L7_C1> module attributes {scheme.ckks} { + // CHECK: func.func @add_plain( + func.func @add_plain(%ctx: !jaxiteword.crypto_context<>, %ct: !ct_L1, %pt: !pt) -> !ct_L1 { + // CHECK: jaxiteword.add_plain + %0 = jaxiteword.add_plain %ctx, %ct, %pt : (!jaxiteword.crypto_context<>, !ct_L1, !pt) -> !ct_L1 + return %0 : !ct_L1 + } + // CHECK: func.func @simple_mul( func.func @simple_mul(%arg0: !jaxiteword.crypto_context<>, %arg1: !jaxiteword.eval_key<>, %arg2: tensor<1x!ct_L1> {tensor_ext.original_type = #original_type}, %arg3: tensor<1x!ct_L1> {tensor_ext.original_type = #original_type}) -> (tensor<1x!ct_L0> {tensor_ext.original_type = #original_type}) { %c0 = arith.constant 0 : index diff --git a/tests/Dialect/JaxiteWord/IR/verifier.mlir b/tests/Dialect/JaxiteWord/IR/verifier.mlir new file mode 100644 index 0000000000..5d33dc7228 --- /dev/null +++ b/tests/Dialect/JaxiteWord/IR/verifier.mlir @@ -0,0 +1,27 @@ +// RUN: heir-opt --verify-diagnostics %s + +!Z1073741441_i64 = !mod_arith.int<1073741441 : i64> +!Z536870273_i64 = !mod_arith.int<536870273 : i64> +#encoding = #lwe.inverse_canonical_encoding +#key = #lwe.key<> +#modulus_chain_L0 = #lwe.modulus_chain, current = 0> +#modulus_chain_L1 = #lwe.modulus_chain, current = 1> +#ring_f64_1_x8 = #polynomial.ring> +!rns_L0 = !rns.rns +!rns_L1 = !rns.rns +#ring_rns_L0_1_x8 = #polynomial.ring> +#ring_rns_L1_1_x8 = #polynomial.ring> +#ciphertext_space_L0 = #lwe.ciphertext_space +#ciphertext_space_L1 = #lwe.ciphertext_space +!pt = !lwe.lwe_plaintext> +!ct_L0 = !lwe.lwe_ciphertext, ciphertext_space = #ciphertext_space_L0, key = #key, modulus_chain = #modulus_chain_L0> +!ct_L1 = !lwe.lwe_ciphertext, ciphertext_space = #ciphertext_space_L1, key = #key, modulus_chain = #modulus_chain_L1> + +module { + func.func @add_plain_mismatched_rings( + %ctx: !jaxiteword.crypto_context<>, %ct: !ct_L1, %pt: !pt) -> !ct_L0 { + // expected-error@+1 {{requires all operands and results to have the same rings}} + %0 = jaxiteword.add_plain %ctx, %ct, %pt : (!jaxiteword.crypto_context<>, !ct_L1, !pt) -> !ct_L0 + return %0 : !ct_L0 + } +} diff --git a/tests/Dialect/JaxiteWord/Transforms/configure_crypto_context.mlir b/tests/Dialect/JaxiteWord/Transforms/configure_crypto_context.mlir index 6ba1dcd271..3bd5beb7bb 100644 --- a/tests/Dialect/JaxiteWord/Transforms/configure_crypto_context.mlir +++ b/tests/Dialect/JaxiteWord/Transforms/configure_crypto_context.mlir @@ -33,10 +33,15 @@ module attributes {ckks.schemeParam = #ckks.scheme_param !jaxiteword.crypto_context // CHECK: jaxiteword.gen_params // CHECK-SAME: degree = 8192 // CHECK-SAME: numSlots = 4096 // CHECK: @simple_mul__configure_crypto_context +// CHECK-SAME: !jaxiteword.crypto_context +// CHECK-SAME: !jaxiteword.public_key +// CHECK-SAME: !jaxiteword.private_key +// CHECK-SAME: !jaxiteword.eval_key // CHECK: jaxiteword.program_initialization // CHECK-SAME: totalRotationIndices = array diff --git a/tests/Dialect/JaxiteWord/Transforms/configure_crypto_context_defaults.mlir b/tests/Dialect/JaxiteWord/Transforms/configure_crypto_context_defaults.mlir index e0d1df3ceb..c4fce70b89 100644 --- a/tests/Dialect/JaxiteWord/Transforms/configure_crypto_context_defaults.mlir +++ b/tests/Dialect/JaxiteWord/Transforms/configure_crypto_context_defaults.mlir @@ -32,6 +32,7 @@ module { } // CHECK: @simple_mul__generate_crypto_context +// CHECK-SAME: () -> !jaxiteword.crypto_context // CHECK: jaxiteword.gen_params // CHECK-SAME: batch = 1 : i32 // CHECK-SAME: c = 4 : i32 diff --git a/tests/Dialect/LWE/Conversions/lwe_to_jaxiteword/nonconstant_rotate.mlir b/tests/Dialect/LWE/Conversions/lwe_to_jaxiteword/nonconstant_rotate.mlir new file mode 100644 index 0000000000..8d7c76977a --- /dev/null +++ b/tests/Dialect/LWE/Conversions/lwe_to_jaxiteword/nonconstant_rotate.mlir @@ -0,0 +1,22 @@ +// RUN: not heir-opt --lwe-to-jaxiteword %s 2>&1 | FileCheck %s + +!Z1032955396097_i64_ = !mod_arith.int<1032955396097 : i64> +!Z1095233372161_i64_ = !mod_arith.int<1095233372161 : i64> +!Z65537_i64_ = !mod_arith.int<65537 : i64> +#full_crt_packing_encoding = #lwe.full_crt_packing_encoding +#key = #lwe.key<> +#modulus_chain_L5_C1_ = #lwe.modulus_chain, current = 1> +!rns_L1_ = !rns.rns +#ring_Z65537_i64_1_x1024_ = #polynomial.ring> +#ring_rns_L1_1_x1024_ = #polynomial.ring> +#ciphertext_space_L1_ = #lwe.ciphertext_space +!ct_L1_ = !lwe.lwe_ciphertext, ciphertext_space = #ciphertext_space_L1_, key = #key, modulus_chain = #modulus_chain_L5_C1_> + +module { + func.func @test_nonconstant_rotate(%ct: !ct_L1_, %shift: index) -> !ct_L1_ { + %rotated = ckks.rotate %ct, %shift : index : !ct_L1_ + return %rotated : !ct_L1_ + } +} + +// CHECK: failed to legalize operation 'ckks.rotate' diff --git a/tests/Dialect/LWE/Conversions/lwe_to_jaxiteword/radd.mlir b/tests/Dialect/LWE/Conversions/lwe_to_jaxiteword/radd.mlir index 923f410c00..9f8d3658fc 100644 --- a/tests/Dialect/LWE/Conversions/lwe_to_jaxiteword/radd.mlir +++ b/tests/Dialect/LWE/Conversions/lwe_to_jaxiteword/radd.mlir @@ -11,6 +11,7 @@ #ring_rns_L1_1_x1024_ = #polynomial.ring> #ciphertext_space_L1_ = #lwe.ciphertext_space !ct_L1_ = !lwe.lwe_ciphertext, ciphertext_space = #ciphertext_space_L1_, key = #key, modulus_chain = #modulus_chain_L5_C1_> +!pt_ = !lwe.lwe_plaintext> module { // CHECK: func.func @test_radd @@ -20,4 +21,32 @@ module { %sum = lwe.radd %ct, %ct_0 : (!ct_L1_, !ct_L1_) -> !ct_L1_ return %sum : !ct_L1_ } + + // CHECK: func.func @test_radd_plain + func.func @test_radd_plain(%ct: !ct_L1_, %pt: !pt_) -> !ct_L1_ { + // CHECK: jaxiteword.add_plain {{.*}}, %{{.*}}, %{{.*}} : (!jaxiteword.crypto_context<>, !ct_L1, !pt) -> !ct_L1 + %sum = lwe.radd_plain %ct, %pt : (!ct_L1_, !pt_) -> !ct_L1_ + return %sum : !ct_L1_ + } + + // CHECK: func.func @test_radd_plain_reversed + func.func @test_radd_plain_reversed(%pt: !pt_, %ct: !ct_L1_) -> !ct_L1_ { + // CHECK: jaxiteword.add_plain {{.*}}, %{{.*}}, %{{.*}} : (!jaxiteword.crypto_context<>, !ct_L1, !pt) -> !ct_L1 + %sum = lwe.radd_plain %pt, %ct : (!pt_, !ct_L1_) -> !ct_L1_ + return %sum : !ct_L1_ + } + + // CHECK: func.func @test_ckks_add_plain + func.func @test_ckks_add_plain(%ct: !ct_L1_, %pt: !pt_) -> !ct_L1_ { + // CHECK: jaxiteword.add_plain + %sum = ckks.add_plain %ct, %pt : (!ct_L1_, !pt_) -> !ct_L1_ + return %sum : !ct_L1_ + } + + // CHECK: func.func @test_rmul_plain_reversed + func.func @test_rmul_plain_reversed(%pt: !pt_, %ct: !ct_L1_) -> !ct_L1_ { + // CHECK: jaxiteword.mul_plain {{.*}}, %{{.*}}, %{{.*}} : (!jaxiteword.crypto_context<>, !ct_L1, !pt) -> !ct_L1 + %product = lwe.rmul_plain %pt, %ct : (!pt_, !ct_L1_) -> !ct_L1_ + return %product : !ct_L1_ + } } diff --git a/tests/Dialect/LWE/Conversions/lwe_to_jaxiteword/smoke_test.mlir b/tests/Dialect/LWE/Conversions/lwe_to_jaxiteword/smoke_test.mlir index 29c336675e..d35d993790 100644 --- a/tests/Dialect/LWE/Conversions/lwe_to_jaxiteword/smoke_test.mlir +++ b/tests/Dialect/LWE/Conversions/lwe_to_jaxiteword/smoke_test.mlir @@ -10,27 +10,47 @@ #ring_Z65537_i64_1_x1024_ = #polynomial.ring> #ring_rns_L1_1_x1024_ = #polynomial.ring> #ciphertext_space_L1_ = #lwe.ciphertext_space +#ciphertext_space_L1_D3_ = #lwe.ciphertext_space !ct_L1_ = !lwe.lwe_ciphertext, ciphertext_space = #ciphertext_space_L1_, key = #key, modulus_chain = #modulus_chain_L5_C1_> +!ct_L1_D3_ = !lwe.lwe_ciphertext, ciphertext_space = #ciphertext_space_L1_D3_, key = #key, modulus_chain = #modulus_chain_L5_C1_> !pt_ = !lwe.lwe_plaintext> !pkey_L1_ = !lwe.lwe_public_key module { // CHECK: @test_ops // CHECK-SAME: (%{{[^:]*}}: !jaxiteword.crypto_context<>, %{{[^:]*}}: !jaxiteword.eval_key<> - func.func @test_ops(%ct: !ct_L1_, %ct2: !ct_L1_, %pt: !pt_) -> (!ct_L1_, !ct_L1_, !ct_L1_, !ct_L1_, !ct_L1_, !ct_L1_) { + func.func @test_ops(%ct: !ct_L1_, %ct2: !ct_L1_, %pt: !pt_) -> (!ct_L1_, !ct_L1_, !ct_L1_) { // CHECK: jaxiteword.add %add = lwe.radd %ct, %ct2 : (!ct_L1_, !ct_L1_) -> !ct_L1_ // CHECK: jaxiteword.sub %sub = lwe.rsub %ct, %ct2 : (!ct_L1_, !ct_L1_) -> !ct_L1_ - // CHECK: jaxiteword.negate - %neg = lwe.rnegate %ct : !ct_L1_ - // CHECK: jaxiteword.add_plain - %add_plain = lwe.radd_plain %ct, %pt : (!ct_L1_, !pt_) -> !ct_L1_ - // CHECK: jaxiteword.sub_plain - %sub_plain = lwe.rsub_plain %ct, %pt : (!ct_L1_, !pt_) -> !ct_L1_ // CHECK: jaxiteword.mul_plain %mul_plain = lwe.rmul_plain %ct, %pt : (!ct_L1_, !pt_) -> !ct_L1_ - return %add, %sub, %neg, %add_plain, %sub_plain, %mul_plain : !ct_L1_, !ct_L1_, !ct_L1_, !ct_L1_, !ct_L1_, !ct_L1_ + return %add, %sub, %mul_plain : !ct_L1_, !ct_L1_, !ct_L1_ + } + + // CHECK: @test_ckks_mul + func.func @test_ckks_mul(%ct: !ct_L1_, %ct2: !ct_L1_) -> !ct_L1_D3_ { + // CHECK: jaxiteword.mul_no_relin + %mul = ckks.mul %ct, %ct2 : (!ct_L1_, !ct_L1_) -> !ct_L1_D3_ + return %mul : !ct_L1_D3_ + } + + // CHECK: @test_static_rotate + func.func @test_static_rotate(%ct: !ct_L1_) -> !ct_L1_ { + // CHECK: jaxiteword.rot + // CHECK-SAME: index = 3 : i64 + %rotated = ckks.rotate %ct {static_shift = 3 : index} : !ct_L1_ + return %rotated : !ct_L1_ + } + + // CHECK: @test_constant_dynamic_rotate + func.func @test_constant_dynamic_rotate(%ct: !ct_L1_) -> !ct_L1_ { + %shift = arith.constant -2 : index + // CHECK: jaxiteword.rot + // CHECK-SAME: index = -2 : i64 + %rotated = ckks.rotate %ct, %shift : index : !ct_L1_ + return %rotated : !ct_L1_ } // CHECK: @test_encode_encrypt diff --git a/tests/Emitter/JaxiteWord/emit_jaxiteword.mlir b/tests/Emitter/JaxiteWord/emit_jaxiteword.mlir index cf1a999229..95d948339b 100644 --- a/tests/Emitter/JaxiteWord/emit_jaxiteword.mlir +++ b/tests/Emitter/JaxiteWord/emit_jaxiteword.mlir @@ -5,17 +5,44 @@ #ring_f64_1_x8 = #polynomial.ring> #ring_i32_1_x8 = #polynomial.ring> #ciphertext_space = #lwe.ciphertext_space +#ciphertext_space_D3 = #lwe.ciphertext_space #modulus_chain = #lwe.modulus_chain, current = 0> +#modulus_chain_L2 = #lwe.modulus_chain, current = 1> !ct_L1 = !lwe.lwe_ciphertext, ciphertext_space = #ciphertext_space, key = #key, modulus_chain = #modulus_chain> +!ct_L2 = !lwe.lwe_ciphertext, ciphertext_space = #ciphertext_space, key = #key, modulus_chain = #modulus_chain_L2> +!ct_L2_D3 = !lwe.lwe_ciphertext, ciphertext_space = #ciphertext_space_D3, key = #key, modulus_chain = #modulus_chain_L2> +!pt = !lwe.lwe_plaintext> // CHECK: def test_add( +// CHECK: {{.*}}: Polynomial, +// CHECK: .he_add[ +// CHECK-SAME: ].add( func.func @test_add(%ctx: !jaxiteword.crypto_context<>, %ct1 : !ct_L1, %ct2 : !ct_L1) -> !ct_L1 { %out = jaxiteword.add %ctx, %ct1, %ct2 : (!jaxiteword.crypto_context<>, !ct_L1, !ct_L1) -> !ct_L1 return %out : !ct_L1 } +// CHECK: def test_add_plain( +// CHECK: {{.*}} = {{.*}}.he_add[{{.*}}.max_level].add_plain({{.*}}, {{.*}}) +func.func @test_add_plain(%ctx: !jaxiteword.crypto_context<>, %ct : !ct_L2, %pt : !pt) -> !ct_L2 { + %out = jaxiteword.add_plain %ctx, %ct, %pt : (!jaxiteword.crypto_context<>, !ct_L2, !pt) -> !ct_L2 + return %out : !ct_L2 +} + +// CHECK: def test_sub( +// CHECK: .he_sub[ +// CHECK-SAME: ].sub( +func.func @test_sub(%ctx: !jaxiteword.crypto_context<>, %ct1 : !ct_L1, %ct2 : !ct_L1) -> !ct_L1 { + %out = jaxiteword.sub %ctx, %ct1, %ct2 : (!jaxiteword.crypto_context<>, !ct_L1, !ct_L1) -> !ct_L1 + return %out : !ct_L1 +} + // CHECK: def test_mul( -// CHECK: hemul( +// CHECK: {{.*}}_raw = key_gen.gen_evaluation_key +// CHECK: {{.*}} = [ +// CHECK: jnp.array({{.*}}_raw["a"], dtype=jnp.uint32).transpose(0, 2, 1), +// CHECK: jnp.array({{.*}}_raw["b"], dtype=jnp.uint32).transpose(0, 2, 1), +// CHECK: .he_mul[ func.func @test_mul(%ctx: !jaxiteword.crypto_context<>, %ct1 : !ct_L1, %ct2 : !ct_L1) -> !ct_L1 { %pk, %sk = jaxiteword.gen_keypair %ctx : (!jaxiteword.crypto_context<>) -> (!jaxiteword.public_key<>, !jaxiteword.private_key<>) %ek = jaxiteword.gen_mulkey %ctx, %sk : (!jaxiteword.crypto_context<>, !jaxiteword.private_key<>) -> !jaxiteword.eval_key<> @@ -24,14 +51,58 @@ func.func @test_mul(%ctx: !jaxiteword.crypto_context<>, %ct1 : !ct_L1, %ct2 : !c } // CHECK: def test_mul_no_relin( -// CHECK: hemul_no_relin( +// CHECK: .hemul_no_relin( func.func @test_mul_no_relin(%ctx: !jaxiteword.crypto_context<>, %ct1 : !ct_L1, %ct2 : !ct_L1) -> !ct_L1 { %out = jaxiteword.mul_no_relin %ctx, %ct1, %ct2 : (!jaxiteword.crypto_context<>, !ct_L1, !ct_L1) -> !ct_L1 return %out : !ct_L1 } +// CHECK: def test_relin( +// CHECK: .relinearize( +func.func @test_relin(%ctx: !jaxiteword.crypto_context<>, %ct: !ct_L2_D3, %ek: !jaxiteword.eval_key<>) -> !ct_L2 { + %out = jaxiteword.relin %ctx, %ct, %ek : (!jaxiteword.crypto_context<>, !ct_L2_D3, !jaxiteword.eval_key<>) -> !ct_L2 + return %out : !ct_L2 +} + +// CHECK: def test_rescale( +// CHECK: .he_rescale[ +// CHECK-SAME: ].rescale( +func.func @test_rescale(%ctx: !jaxiteword.crypto_context<>, %ct: !ct_L2) -> !ct_L1 { + %out = jaxiteword.mod_reduce %ctx, %ct : (!jaxiteword.crypto_context<>, !ct_L2) -> !ct_L1 + return %out : !ct_L1 +} + +// CHECK: def test_rotate( +// CHECK: .he_rot[ +// CHECK-SAME: ].rotate( +func.func @test_rotate(%ctx: !jaxiteword.crypto_context<>, %ct: !ct_L1, %ek: !jaxiteword.eval_key<>) -> !ct_L1 { + %out = jaxiteword.rot %ctx, %ct, %ek {index = 2 : i64} : (!jaxiteword.crypto_context<>, !ct_L1, !jaxiteword.eval_key<>) -> !ct_L1 + return %out : !ct_L1 +} + +// CHECK: def test_mul_plain( +// CHECK: .ptct_mul[ +// CHECK-SAME: ].mul({{.*}}, {{.*}}) +func.func @test_mul_plain(%ctx: !jaxiteword.crypto_context<>, %ct: !ct_L1, %pt: !pt) -> !ct_L1 { + %out = jaxiteword.mul_plain %ctx, %ct, %pt : (!jaxiteword.crypto_context<>, !ct_L1, !pt) -> !ct_L1 + return %out : !ct_L1 +} + +// CHECK: def test_floor_div_si( +// CHECK: {{.*}} = {{.*}} // {{.*}} +func.func @test_floor_div_si(%lhs: i32, %rhs: i32) -> i32 { + %out = arith.floordivsi %lhs, %rhs : i32 + return %out : i32 +} + // CHECK: def test_gen_params( +// CHECK: params = { // CHECK: "scaling_factor": 563019763943521 +// CHECK: "output_scale": 563019763943521 +// CHECK-NOT: "public_key": +// CHECK-NOT: "secret_key": +// CHECK-NOT: "evaluation_key": +// CHECK: {{.*}} = ckks.CKKSContext(params) func.func @test_gen_params() -> !jaxiteword.crypto_context<> { %ctx = jaxiteword.gen_params { degree = 8192 : i64, @@ -43,8 +114,30 @@ func.func @test_gen_params() -> !jaxiteword.crypto_context<> { r = 4 : i32, c = 4 : i32, dnum = 3 : i32, - numEvalMult = 2 : i32, compositeDegree = 1 : i32 } : () -> !jaxiteword.crypto_context<> return %ctx : !jaxiteword.crypto_context<> } + +// CHECK: def test_program_initialization( +// CHECK: {{.*}}.public_key = {{.*}} +// CHECK: {{.*}}.secret_key = {{.*}} +// CHECK: {{.*}}.evaluation_key = {{.*}} +// CHECK: {{.*}}.parameters["public_key"] = {{.*}} +// CHECK: {{.*}}.parameters["secret_key"] = {{.*}} +// CHECK: {{.*}}.parameters["evaluation_key"] = {{.*}} +// CHECK: {{.*}}.program_initialization(total_rotation_indices=[1, 2], dnum=3, r=4, c=4, batch=1) +func.func @test_program_initialization( + %ctx: !jaxiteword.crypto_context<>, + %pk: !jaxiteword.public_key<>, + %sk: !jaxiteword.private_key<>, + %ek: !jaxiteword.eval_key<>) { + jaxiteword.program_initialization %ctx, %pk, %sk, %ek { + totalRotationIndices = array, + dnum = 3 : i32, + r = 4 : i32, + c = 4 : i32, + batch = 1 : i32 + } : (!jaxiteword.crypto_context<>, !jaxiteword.public_key<>, !jaxiteword.private_key<>, !jaxiteword.eval_key<>) -> () + return +}