diff --git a/lib/Target/OpenFhePke/BUILD b/lib/Target/OpenFhePke/BUILD index 24c1deca1e..76d4c6fa99 100644 --- a/lib/Target/OpenFhePke/BUILD +++ b/lib/Target/OpenFhePke/BUILD @@ -1,9 +1,7 @@ # OpenFhePke Emitter -load("@heir//bazel/openfhe:copts.bzl", "OPENMP_COPTS", "OPENMP_LINKOPTS") load("@heir//lib/Target/CompilationTarget:defs.bzl", "heir_backend_config") load("@rules_cc//cc:cc_library.bzl", "cc_library") -load("@rules_cc//cc:cc_test.bzl", "cc_test") package( default_applicable_licenses = ["@heir//:license"], @@ -131,38 +129,6 @@ cc_library( ], ) -cc_library( - name = "Interpreter", - srcs = ["Interpreter.cpp"], - hdrs = ["Interpreter.h"], - copts = select({ - "@heir//:config_openfhe_enable_timing": ["-DOPENFHE_ENABLE_TIMING"], - "//conditions:default": [], - }) + OPENMP_COPTS, - linkopts = OPENMP_LINKOPTS, - deps = [ - "@heir//lib/Dialect/Mgmt/IR:Dialect", - "@heir//lib/Dialect/ModArith/IR:Dialect", - "@heir//lib/Dialect/Openfhe/IR:Dialect", - "@heir//lib/Dialect/RNS/IR:Dialect", - "@heir//lib/Dialect/TensorExt/IR:Dialect", - "@hermetic_llvm//openmp:libomp", - "@llvm-project//llvm:Support", - "@llvm-project//mlir:AffineDialect", - "@llvm-project//mlir:Analysis", - "@llvm-project//mlir:ArithDialect", - "@llvm-project//mlir:FuncDialect", - "@llvm-project//mlir:IR", - "@llvm-project//mlir:LinalgDialect", - "@llvm-project//mlir:Parser", - "@llvm-project//mlir:SCFDialect", - "@llvm-project//mlir:Support", - "@llvm-project//mlir:TensorDialect", - "@openfhe//:core", - "@openfhe//:pke", - ], -) - cc_library( name = "OpenFhePkeDebugHeaderEmitter", srcs = ["OpenFhePkeDebugHeaderEmitter.cpp"], @@ -201,25 +167,6 @@ cc_library( ], ) -cc_test( - name = "InterpreterTest", - srcs = ["InterpreterTest.cpp"], - copts = OPENMP_COPTS, - linkopts = OPENMP_LINKOPTS, - tags = [ - "manual", - "notap", - ], - deps = [ - ":Interpreter", - "@googletest//:gtest_main", - "@llvm-project//mlir:IR", - "@llvm-project//mlir:Parser", - "@openfhe//:core", - "@openfhe//:pke", - ], -) - heir_backend_config( name = "openfhe_backend_config", srcs = ["TargetConfig.td"], diff --git a/lib/Target/OpenFhePke/Interpreter.cpp b/lib/Target/OpenFhePke/Interpreter.cpp deleted file mode 100644 index fb8ee128af..0000000000 --- a/lib/Target/OpenFhePke/Interpreter.cpp +++ /dev/null @@ -1,2176 +0,0 @@ -#include "lib/Target/OpenFhePke/Interpreter.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "lib/Dialect/Mgmt/IR/MgmtDialect.h" -#include "lib/Dialect/ModArith/IR/ModArithDialect.h" -#include "lib/Dialect/Openfhe/IR/OpenfheDialect.h" -#include "lib/Dialect/Openfhe/IR/OpenfheOps.h" -#include "lib/Dialect/Openfhe/IR/OpenfheTypes.h" -#include "lib/Dialect/RNS/IR/RNSDialect.h" -#include "lib/Dialect/TensorExt/IR/TensorExtDialect.h" -#include "llvm/include/llvm/ADT/STLExtras.h" // from @llvm-project -#include "llvm/include/llvm/ADT/TypeSwitch.h" // from @llvm-project -#include "llvm/include/llvm/Support/Casting.h" // from @llvm-project -#include "llvm/include/llvm/Support/raw_ostream.h" // from @llvm-project -#include "mlir/include/mlir/Analysis/Liveness.h" // from @llvm-project -#include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h" // from @llvm-project -#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project -#include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project -#include "mlir/include/mlir/Dialect/Linalg/IR/Linalg.h" // from @llvm-project -#include "mlir/include/mlir/Dialect/SCF/IR/SCF.h" // from @llvm-project -#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/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/MLIRContext.h" // from @llvm-project -#include "mlir/include/mlir/IR/Operation.h" // from @llvm-project -#include "mlir/include/mlir/IR/OperationSupport.h" // from @llvm-project -#include "mlir/include/mlir/IR/OwningOpRef.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 -#include "mlir/include/mlir/Parser/Parser.h" // from @llvm-project -#include "mlir/include/mlir/Support/LLVM.h" // from @llvm-project -#include "mlir/include/mlir/Support/TypeID.h" // from @llvm-project -#include "src/core/include/lattice/hal/lat-backend.h" // from @openfhe -#include "src/core/include/lattice/stdlatticeparms.h" // from @openfhe -#include "src/pke/include/ciphertext-fwd.h" // from @openfhe -#include "src/pke/include/constants-defs.h" // from @openfhe -#include "src/pke/include/cryptocontext-fwd.h" // from @openfhe -#include "src/pke/include/encoding/plaintext-fwd.h" // from @openfhe -#include "src/pke/include/gen-cryptocontext.h" // from @openfhe -#include "src/pke/include/key/evalkey-fwd.h" // from @openfhe -#include "src/pke/include/key/privatekey-fwd.h" // from @openfhe -#include "src/pke/include/key/publickey-fwd.h" // from @openfhe -#include "src/pke/include/scheme/ckksrns/gen-cryptocontext-ckksrns-params.h" // from @openfhe -#include "src/pke/include/scheme/ckksrns/gen-cryptocontext-ckksrns.h" // from @openfhe - -#ifdef OPENFHE_ENABLE_TIMING -#define TIME_OPERATION_VOID(op_name, code) \ - do { \ - auto start = std::chrono::high_resolution_clock::now(); \ - code; \ - auto end = std::chrono::high_resolution_clock::now(); \ - timingResults[op_name].totalTime += (end - start); \ - timingResults[op_name].count++; \ - } while (0) -#define TIME_OPERATION(op_name, result_val, code) \ - do { \ - auto start = std::chrono::high_resolution_clock::now(); \ - auto result = code; \ - auto end = std::chrono::high_resolution_clock::now(); \ - timingResults[op_name].totalTime += (end - start); \ - timingResults[op_name].count++; \ - ciphertexts.insert_or_assign(result_val, std::move(result)); \ - } while (0) -#define TIME_OPERATION_NONCT(op_name, result_val, code, map) \ - do { \ - auto start = std::chrono::high_resolution_clock::now(); \ - auto result = code; \ - auto end = std::chrono::high_resolution_clock::now(); \ - timingResults[op_name].totalTime += (end - start); \ - timingResults[op_name].count++; \ - (map).insert_or_assign(result_val, std::move(result)); \ - } while (0) -#else -#define TIME_OPERATION_VOID(op_name, code) \ - do { \ - code; \ - } while (0) -#define TIME_OPERATION(op_name, result_val, code) \ - do { \ - ciphertexts.insert_or_assign(result_val, code); \ - } while (0) -#define TIME_OPERATION_NONCT(op_name, result_val, code, map) \ - do { \ - (map).insert_or_assign(result_val, code); \ - } while (0) -#endif - -namespace mlir { -namespace heir { -namespace openfhe { - -using namespace lbcrypto; -using CiphertextT = Ciphertext; -using CCParamsT = CCParams; -using CryptoContextT = CryptoContext; -using EvalKeyT = EvalKey; -using PlaintextT = Plaintext; -using PrivateKeyT = PrivateKey; -using PublicKeyT = PublicKey; -using FastRotPrecompT = std::shared_ptr>; - -// Helper function for floor division on integers -static inline int floorDivInt(int lhs, int rhs) { - return static_cast(std::floor(static_cast(lhs) / rhs)); -} - -Interpreter::Interpreter(ModuleOp module) - : module(module), liveness(std::make_shared(module)) { - initializeDispatchTable(); -} - -// Static member initialization -llvm::DenseMap - Interpreter::operationDispatchTable; -bool Interpreter::dispatchTableInitialized = false; -MLIRContext* Interpreter::dispatchTableContext = nullptr; - -void Interpreter::initializeDispatchTable() { - auto* ctx = module.getContext(); - - // Reinitialize if context changed or not yet initialized - if (dispatchTableInitialized && dispatchTableContext == ctx) { - return; - } - - // Clear the dispatch table if context changed - if (dispatchTableContext != ctx) { - operationDispatchTable.clear(); - } - - dispatchTableContext = ctx; - -// Helper macro to register an operation type -#define REGISTER_OP(OpType) \ - operationDispatchTable[TypeID::get()] = [](Interpreter* interp, \ - Operation* op) { \ - interp->visit(llvm::cast(op)); \ - }; - - // Register all supported operations - REGISTER_OP(AddInPlaceOp); - REGISTER_OP(AddOp); - REGISTER_OP(AddPlainInPlaceOp); - REGISTER_OP(AddPlainOp); - REGISTER_OP(AutomorphOp); - REGISTER_OP(BootstrapOp); - REGISTER_OP(DecodeCKKSOp); - REGISTER_OP(DecodeOp); - REGISTER_OP(DecryptOp); - REGISTER_OP(EncryptOp); - REGISTER_OP(FastRotationOp); - REGISTER_OP(FastRotationExtOp); - REGISTER_OP(FastRotationPrecomputeOp); - REGISTER_OP(GenBootstrapKeyOp); - REGISTER_OP(GenContextOp); - REGISTER_OP(GenMulKeyOp); - REGISTER_OP(GenParamsOp); - REGISTER_OP(GenRotKeyOp); - REGISTER_OP(KeySwitchInPlaceOp); - REGISTER_OP(KeySwitchOp); - REGISTER_OP(KeySwitchDownOp); - REGISTER_OP(LevelReduceInPlaceOp); - REGISTER_OP(LevelReduceOp); - REGISTER_OP(MakeCKKSPackedPlaintextOp); - REGISTER_OP(MakePackedPlaintextOp); - REGISTER_OP(ModReduceInPlaceOp); - REGISTER_OP(ModReduceOp); - REGISTER_OP(MulConstInPlaceOp); - REGISTER_OP(MulConstOp); - REGISTER_OP(MulNoRelinOp); - REGISTER_OP(MulOp); - REGISTER_OP(MulPlainOp); - REGISTER_OP(NegateInPlaceOp); - REGISTER_OP(NegateOp); - REGISTER_OP(RelinInPlaceOp); - REGISTER_OP(RelinOp); - REGISTER_OP(RotOp); - REGISTER_OP(SetupBootstrapOp); - REGISTER_OP(SquareInPlaceOp); - REGISTER_OP(SquareOp); - REGISTER_OP(SubInPlaceOp); - REGISTER_OP(SubOp); - REGISTER_OP(SubPlainInPlaceOp); - REGISTER_OP(SubPlainOp); - REGISTER_OP(affine::AffineForOp); - REGISTER_OP(affine::AffineYieldOp); - REGISTER_OP(arith::AddFOp); - REGISTER_OP(arith::AddIOp); - REGISTER_OP(arith::AndIOp); - REGISTER_OP(arith::CmpIOp); - REGISTER_OP(arith::ConstantOp); - REGISTER_OP(arith::DivSIOp); - REGISTER_OP(arith::ExtFOp); - REGISTER_OP(arith::FloorDivSIOp); - REGISTER_OP(arith::MaxSIOp); - REGISTER_OP(arith::MinSIOp); - REGISTER_OP(arith::MulFOp); - REGISTER_OP(arith::MulIOp); - REGISTER_OP(arith::RemSIOp); - REGISTER_OP(arith::SelectOp); - REGISTER_OP(arith::SubFOp); - REGISTER_OP(arith::SubIOp); - REGISTER_OP(linalg::BroadcastOp); - REGISTER_OP(scf::ForOp); - REGISTER_OP(scf::ForallOp); - REGISTER_OP(scf::InParallelOp); - REGISTER_OP(scf::IfOp); - REGISTER_OP(scf::YieldOp); - REGISTER_OP(tensor::CollapseShapeOp); - REGISTER_OP(tensor::ConcatOp); - REGISTER_OP(tensor::EmptyOp); - REGISTER_OP(tensor::ExpandShapeOp); - REGISTER_OP(tensor::ExtractOp); - REGISTER_OP(tensor::ExtractSliceOp); - REGISTER_OP(tensor::FromElementsOp); - REGISTER_OP(tensor::InsertOp); - REGISTER_OP(tensor::InsertSliceOp); - REGISTER_OP(tensor::ParallelInsertSliceOp); - REGISTER_OP(tensor::SplatOp); - -#undef REGISTER_OP - - dispatchTableInitialized = true; -} - -void Interpreter::eraseValue(Value v) { - llvm::TypeSwitch(v.getType()) - .Case([&](auto ty) { plaintexts.erase(v); }) - .Case([&](auto ty) { ciphertexts.erase(v); }) - .Case([&](auto ty) { - auto elemType = ty.getElementType(); - if (elemType.isInteger() || elemType.isIndex()) { - intVectors.erase(v); - } else if (elemType.isF32()) { - floatVectors.erase(v); - } else if (elemType.isF64()) { - doubleVectors.erase(v); - } else if (isa(elemType)) { - plaintextVectors.erase(v); - } else if (isa(elemType)) { - ciphertextVectors.erase(v); - } else { - llvm::errs() << "Unsupported tensor element type " << elemType - << " in eraseValue\n"; - } - }) - .Case([&](auto ty) { - if (ty.isInteger(1)) { - boolValues.erase(v); - } else { - intValues.erase(v); - } - }) - .Case([&](auto ty) { - if (ty.isF32()) { - floatValues.erase(v); - } else { - doubleValues.erase(v); - } - }) - .Case( - [&](auto ty) { fastRotPrecomps.erase(v); }) - .Case([&](auto ty) { evalKeys.erase(v); }) - .Case([&](auto ty) { publicKeys.erase(v); }) - .Case([&](auto ty) { privateKeys.erase(v); }) - .Case( - [&](auto ty) { cryptoContexts.erase(v); }) - .Default([&](Type type) { - llvm::errs() << "Unsupported type " << type << " in eraseValue\n"; - }); -} - -// Only used for type-agnostic block arguments (func args, iter args, etc.) -void Interpreter::storeTypedValue(Value v, const TypedCppValue& typedVal) { - std::visit( - [&](auto&& arg) { - using T = std::decay_t; - if constexpr (std::is_same_v) { - // Nothing to store - } else if constexpr (std::is_same_v) { - boolValues[v] = arg; - } else if constexpr (std::is_same_v) { - intValues[v] = arg; - } else if constexpr (std::is_same_v) { - floatValues[v] = arg; - } else if constexpr (std::is_same_v) { - doubleValues[v] = arg; - } else if constexpr (std::is_same_v< - T, std::shared_ptr>>) { - intVectors[v] = arg; - } else if constexpr (std::is_same_v< - T, std::shared_ptr>>) { - floatVectors[v] = arg; - } else if constexpr (std::is_same_v< - T, std::shared_ptr>>) { - doubleVectors[v] = arg; - } else if constexpr (std::is_same_v) { - plaintexts[v] = arg; - } else if constexpr (std::is_same_v< - T, std::shared_ptr>>) { - plaintextVectors[v] = arg; - } else if constexpr (std::is_same_v) { - ciphertexts[v] = arg; - } else if constexpr (std::is_same_v>>) { - ciphertextVectors[v] = arg; - } else if constexpr (std::is_same_v) { - publicKeys[v] = arg; - } else if constexpr (std::is_same_v) { - privateKeys[v] = arg; - } else if constexpr (std::is_same_v) { - evalKeys[v] = arg; - } else if constexpr (std::is_same_v) { - cryptoContexts[v] = arg; - } else if constexpr (std::is_same_v) { - fastRotPrecomps[v] = arg; - } - }, - typedVal.value); -} - -// Only used for type-agnostic block terminators (return, yield) -TypedCppValue Interpreter::loadTypedValue(Value v) { - TypedCppValue result; - - llvm::TypeSwitch(v.getType()) - .Case([&](auto ty) { - if (auto it = plaintexts.find(v); it != plaintexts.end()) - result = TypedCppValue(it->second); - }) - .Case([&](auto ty) { - if (auto it = ciphertexts.find(v); it != ciphertexts.end()) - result = TypedCppValue(it->second); - }) - .Case([&](auto ty) { - auto elemType = ty.getElementType(); - if (elemType.isInteger() || elemType.isIndex()) { - if (auto it = intVectors.find(v); it != intVectors.end()) - result = TypedCppValue(it->second); - } else if (elemType.isF32()) { - if (auto it = floatVectors.find(v); it != floatVectors.end()) - result = TypedCppValue(it->second); - } else if (elemType.isF64()) { - if (auto it = doubleVectors.find(v); it != doubleVectors.end()) - result = TypedCppValue(it->second); - } else if (isa(elemType)) { - if (auto it = plaintextVectors.find(v); it != plaintextVectors.end()) - result = TypedCppValue(it->second); - } else if (isa(elemType)) { - if (auto it = ciphertextVectors.find(v); - it != ciphertextVectors.end()) - result = TypedCppValue(it->second); - } - }) - .Case([&](auto ty) { - if (ty.isInteger(1)) { - if (auto it = boolValues.find(v); it != boolValues.end()) - result = TypedCppValue(it->second); - } else { - if (auto it = intValues.find(v); it != intValues.end()) - result = TypedCppValue(it->second); - } - }) - .Case([&](auto ty) { - if (ty.isF32()) { - if (auto it = floatValues.find(v); it != floatValues.end()) - result = TypedCppValue(it->second); - } else { - if (auto it = doubleValues.find(v); it != doubleValues.end()) - result = TypedCppValue(it->second); - } - }) - .Case([&](auto ty) { - if (auto it = fastRotPrecomps.find(v); it != fastRotPrecomps.end()) - result = TypedCppValue(it->second); - }) - .Case([&](auto ty) { - if (auto it = evalKeys.find(v); it != evalKeys.end()) - result = TypedCppValue(it->second); - }) - .Case([&](auto ty) { - if (auto it = publicKeys.find(v); it != publicKeys.end()) - result = TypedCppValue(it->second); - }) - .Case([&](auto ty) { - if (auto it = privateKeys.find(v); it != privateKeys.end()) - result = TypedCppValue(it->second); - }) - .Case([&](auto ty) { - if (auto it = cryptoContexts.find(v); it != cryptoContexts.end()) - result = TypedCppValue(it->second); - }) - .Default([&](Type type) { - llvm::errs() << "Unsupported type " << type << " in loadTypedValue\n"; - }); - - return result; -} - -std::vector Interpreter::interpret( - const std::string& entryFunction, ArrayRef inputValues) { - if (!dispatchTableInitialized || - dispatchTableContext != module.getContext()) { - initializeDispatchTable(); - } - - // Clear all storage - boolValues.clear(); - intValues.clear(); - floatValues.clear(); - doubleValues.clear(); - intVectors.clear(); - floatVectors.clear(); - doubleVectors.clear(); - plaintexts.clear(); - plaintextVectors.clear(); - ciphertexts.clear(); - ciphertextVectors.clear(); - cryptoContexts.clear(); - publicKeys.clear(); - privateKeys.clear(); - evalKeys.clear(); - fastRotPrecomps.clear(); - func::FuncOp func = module.lookupSymbol(entryFunction); - std::vector results; - - FunctionType funcType = func.getFunctionType(); - SmallVector argTypes(funcType.getInputs()); - SmallVector returnTypes(funcType.getResults()); - - if (argTypes.size() != inputValues.size()) { - func->emitError() << "Input size does not match function signature"; - } - - for (const auto& [argIndex, argTy] : llvm::enumerate(argTypes)) { - storeTypedValue(func.getBody().getArgument(argIndex), - inputValues[argIndex]); - } - - llvm::outs() << "Interpreting function: " << func.getName() - << " with signature " << funcType << " and " - << inputValues.size() << " interpreted arguments\n"; - - // Walk only the operations in the entry block, not nested regions - // Nested regions (like loop bodies) will be handled by their parent ops - for (auto& op : func.getBody().front().getOperations()) { - if (auto returnOp = dyn_cast(&op)) { - results.reserve(returnOp.getOperands().size()); - for (auto retVal : returnOp.getOperands()) { - results.push_back(loadTypedValue(retVal)); - } - llvm::outs() << "Function returned " << results.size() << " values\n"; - } else { - visit(&op); - } - } - - return results; -} - -void Interpreter::visit(Operation* op) { - // Avoid dispatch overhead for trivial ops - if (isa(op)) { - intValues[op->getResult(0)] = intValues.at(op->getOperand(0)); - } else { - // Use jump table for faster dispatch - auto it = operationDispatchTable.find(op->getName().getTypeID()); - if (it != operationDispatchTable.end()) { - it->second(this, op); - } else { - op->emitError() << "Unsupported operation " << op->getName() - << " in interpreter\n"; - } - } - // If any of the operations op operands have no more uses, then remove them - // from storage. - if (!op->getParentOfType() && - !op->getParentOfType() && - !op->getParentOfType()) { - for (auto operand : op->getOperands()) { - if (liveness->isDeadAfter(operand, op)) { - eraseValue(operand); - } - } - } -} - -void Interpreter::visit(arith::ConstantOp op) { - auto valueAttr = op.getValue(); - if (auto intAttr = dyn_cast(valueAttr)) { - intValues[op.getResult()] = static_cast(intAttr.getInt()); - return; - } - - if (auto floatAttr = dyn_cast(valueAttr)) { - // Use float for 32-bit and smaller, double for 64-bit and larger - if (floatAttr.getType().isF32()) { - floatValues[op.getResult()] = - static_cast(floatAttr.getValueAsDouble()); - } else { - // F64 and larger types use double - doubleValues[op.getResult()] = floatAttr.getValueAsDouble(); - } - return; - } - - DenseElementsAttr denseElementsAttr = dyn_cast(valueAttr); - if (auto denseResourceAttr = dyn_cast(valueAttr)) { - const auto data = denseResourceAttr.getData(); - denseElementsAttr = - DenseElementsAttr::getFromRawBuffer(denseResourceAttr.getType(), data); - } - - if (denseElementsAttr) { - if (denseElementsAttr.getType().getElementType().isF32()) { - std::vector values; - for (auto val : denseElementsAttr.getValues()) { - values.push_back(static_cast(val.convertToFloat())); - } - floatVectors[op.getResult()] = - std::make_shared>(values); - return; - } - - if (denseElementsAttr.getType().getElementType().isF64()) { - std::vector values; - for (auto val : denseElementsAttr.getValues()) { - values.push_back(val.convertToDouble()); - } - doubleVectors[op.getResult()] = - std::make_shared>(values); - return; - } - - std::vector values; - for (auto val : denseElementsAttr.getValues()) { - values.push_back(val.getSExtValue()); - } - intVectors[op.getResult()] = std::make_shared>(values); - return; - } - - op->emitError() << "Unsupported constant attribute type " << valueAttr - << "\n"; -} - -// Macro for handling binary operations on integer types -#define HANDLE_BINARY_OP_INT(op, binop, opname) \ - do { \ - auto resultType = (op).getResult().getType(); \ - if (auto tensorType = dyn_cast(resultType)) { \ - auto elemType = tensorType.getElementType(); \ - auto numElements = tensorType.getNumElements(); \ - if (elemType.isInteger() || elemType.isIndex()) { \ - const auto& lhsVec = *intVectors.at((op).getLhs()); \ - const auto& rhsVec = *intVectors.at((op).getRhs()); \ - auto result = std::make_shared>(numElements); \ - for (int64_t i = 0; i < numElements; ++i) { \ - (*result)[i] = lhsVec[i] binop rhsVec[i]; \ - } \ - intVectors[(op).getResult()] = result; \ - } else { \ - (op)->emitError() << "Unsupported type for " opname ": " << resultType \ - << "\n"; \ - } \ - } else if (resultType.isInteger() || resultType.isIndex()) { \ - intValues[(op).getResult()] = \ - intValues.at((op).getLhs()) binop intValues.at((op).getRhs()); \ - } else { \ - (op)->emitError() << "Unsupported type for " opname ": " << resultType \ - << "\n"; \ - } \ - } while (0) - -// Macro for handling binary operations on floating-point types -#define HANDLE_BINARY_OP_FLOAT(op, binop, opname) \ - do { \ - auto resultType = (op).getResult().getType(); \ - if (auto tensorType = dyn_cast(resultType)) { \ - auto elemType = tensorType.getElementType(); \ - auto numElements = tensorType.getNumElements(); \ - if (elemType.isF32()) { \ - const auto& lhsVec = *floatVectors.at((op).getLhs()); \ - const auto& rhsVec = *floatVectors.at((op).getRhs()); \ - auto result = std::make_shared>(numElements); \ - for (int64_t i = 0; i < numElements; ++i) { \ - (*result)[i] = lhsVec[i] binop rhsVec[i]; \ - } \ - floatVectors[(op).getResult()] = result; \ - } else if (elemType.isF64()) { \ - const auto& lhsVec = *doubleVectors.at((op).getLhs()); \ - const auto& rhsVec = *doubleVectors.at((op).getRhs()); \ - auto result = std::make_shared>(numElements); \ - for (int64_t i = 0; i < numElements; ++i) { \ - (*result)[i] = lhsVec[i] binop rhsVec[i]; \ - } \ - doubleVectors[(op).getResult()] = result; \ - } else { \ - (op)->emitError() << "Unsupported type for " opname ": " << resultType \ - << "\n"; \ - } \ - } else if (resultType.isF32()) { \ - floatValues[(op).getResult()] = \ - floatValues.at((op).getLhs()) binop floatValues.at((op).getRhs()); \ - } else if (resultType.isF64()) { \ - doubleValues[(op).getResult()] = \ - doubleValues.at((op).getLhs()) binop doubleValues.at((op).getRhs()); \ - } else { \ - (op)->emitError() << "Unsupported type for " opname ": " << resultType \ - << "\n"; \ - } \ - } while (0) - -// Macro for handling binary operations on integer types using a function -#define HANDLE_BINARY_OP_INT_FUNC(op, func, opname) \ - do { \ - auto resultType = (op).getResult().getType(); \ - if (auto tensorType = dyn_cast(resultType)) { \ - auto elemType = tensorType.getElementType(); \ - auto numElements = tensorType.getNumElements(); \ - if (elemType.isInteger() || elemType.isIndex()) { \ - const auto& lhsVec = *intVectors.at((op).getLhs()); \ - const auto& rhsVec = *intVectors.at((op).getRhs()); \ - auto result = std::make_shared>(numElements); \ - for (int64_t i = 0; i < numElements; ++i) { \ - (*result)[i] = func(lhsVec[i], rhsVec[i]); \ - } \ - intVectors[(op).getResult()] = result; \ - } else { \ - (op)->emitError() << "Unsupported type for " opname ": " << resultType \ - << "\n"; \ - } \ - } else if (resultType.isInteger() || resultType.isIndex()) { \ - intValues[(op).getResult()] = \ - func(intValues.at((op).getLhs()), intValues.at((op).getRhs())); \ - } else { \ - (op)->emitError() << "Unsupported type for " opname ": " << resultType \ - << "\n"; \ - } \ - } while (0) - -void Interpreter::visit(arith::AddIOp op) { - HANDLE_BINARY_OP_INT(op, +, "arith.addi"); -} - -void Interpreter::visit(arith::AddFOp op) { - HANDLE_BINARY_OP_FLOAT(op, +, "arith.addf"); -} - -void Interpreter::visit(arith::SubIOp op) { - HANDLE_BINARY_OP_INT(op, -, "arith.subi"); -} - -void Interpreter::visit(arith::SubFOp op) { - HANDLE_BINARY_OP_FLOAT(op, -, "arith.subf"); -} - -void Interpreter::visit(arith::MulIOp op) { - HANDLE_BINARY_OP_INT(op, *, "arith.muli"); -} - -void Interpreter::visit(arith::MulFOp op) { - HANDLE_BINARY_OP_FLOAT(op, *, "arith.mulf"); -} - -void Interpreter::visit(arith::DivSIOp op) { - HANDLE_BINARY_OP_INT(op, /, "arith.divsi"); -} - -void Interpreter::visit(arith::RemSIOp op) { - HANDLE_BINARY_OP_INT(op, %, "arith.remsi"); -} - -void Interpreter::visit(arith::AndIOp op) { - HANDLE_BINARY_OP_INT(op, &, "arith.andi"); -} - -void Interpreter::visit(arith::MaxSIOp op) { - HANDLE_BINARY_OP_INT_FUNC(op, std::max, "arith.maxsi"); -} - -void Interpreter::visit(arith::MinSIOp op) { - HANDLE_BINARY_OP_INT_FUNC(op, std::min, "arith.minsi"); -} - -void Interpreter::visit(arith::CmpIOp op) { - auto lhsType = op.getLhs().getType(); - auto cmpFunc = [&](auto a, auto b) -> int { - switch (op.getPredicate()) { - case arith::CmpIPredicate::eq: - return a == b ? 1 : 0; - case arith::CmpIPredicate::ne: - return a != b ? 1 : 0; - case arith::CmpIPredicate::slt: - case arith::CmpIPredicate::ult: - return a < b ? 1 : 0; - case arith::CmpIPredicate::sle: - case arith::CmpIPredicate::ule: - return a <= b ? 1 : 0; - case arith::CmpIPredicate::sgt: - case arith::CmpIPredicate::ugt: - return a > b ? 1 : 0; - case arith::CmpIPredicate::sge: - case arith::CmpIPredicate::uge: - return a >= b ? 1 : 0; - } - return 0; - }; - - if (lhsType.isInteger()) { - intValues[op.getResult()] = - cmpFunc(intValues.at(op.getLhs()), intValues.at(op.getRhs())); - } else if (lhsType.isF32()) { - intValues[op.getResult()] = - cmpFunc(floatValues.at(op.getLhs()), floatValues.at(op.getRhs())); - } else if (lhsType.isF64()) { - intValues[op.getResult()] = - cmpFunc(doubleValues.at(op.getLhs()), doubleValues.at(op.getRhs())); - } else if (lhsType.isIndex()) { - intValues[op.getResult()] = - cmpFunc(intValues.at(op.getLhs()), intValues.at(op.getRhs())); - } else { - op->emitError() << "Unsupported type for arith.cmpi: " << lhsType << "\n"; - } -} - -void Interpreter::visit(arith::SelectOp op) { - int cond = intValues.at(op.getCondition()); - bool condBool = (cond != 0); - - auto resultType = op.getResult().getType(); - if (resultType.isInteger()) { - intValues[op.getResult()] = condBool ? intValues.at(op.getTrueValue()) - : intValues.at(op.getFalseValue()); - } else if (resultType.isF32()) { - floatValues[op.getResult()] = condBool ? floatValues.at(op.getTrueValue()) - : floatValues.at(op.getFalseValue()); - } else if (resultType.isF64()) { - doubleValues[op.getResult()] = condBool - ? doubleValues.at(op.getTrueValue()) - : doubleValues.at(op.getFalseValue()); - } else if (resultType.isIndex()) { - intValues[op.getResult()] = condBool ? intValues.at(op.getTrueValue()) - : intValues.at(op.getFalseValue()); - } else { - op->emitError() << "Unsupported type for arith.select: " << resultType - << "\n"; - } -} - -void Interpreter::visit(arith::FloorDivSIOp op) { - HANDLE_BINARY_OP_INT_FUNC(op, floorDivInt, "arith.floordivsi"); -} - -void Interpreter::visit(arith::ExtFOp op) { - auto inType = op.getIn().getType(); - auto outType = op.getOut().getType(); - - if (auto tensorType = dyn_cast(inType)) { - // Handle tensor case - auto elemType = tensorType.getElementType(); - auto outElemType = cast(outType).getElementType(); - - if (elemType.isF32() && outElemType.isF64()) { - // Convert float vector to double vector - auto inVec = floatVectors.at(op.getIn()); - auto outVec = - std::make_shared>(inVec->begin(), inVec->end()); - doubleVectors[op.getResult()] = outVec; - } else if (elemType.isF32()) { - floatVectors[op.getResult()] = floatVectors.at(op.getIn()); - } else { - doubleVectors[op.getResult()] = doubleVectors.at(op.getIn()); - } - } else { - // Scalar case - if (inType.isF32() && outType.isF64()) { - doubleValues[op.getResult()] = - static_cast(floatValues.at(op.getIn())); - } else if (inType.isF32()) { - floatValues[op.getResult()] = floatValues.at(op.getIn()); - } else { - doubleValues[op.getResult()] = doubleValues.at(op.getIn()); - } - } -} - -int Interpreter::getFlattenedTensorIndex(Value tensor, ValueRange indices) { - auto tensorType = cast(tensor.getType()); - auto shape = tensorType.getShape(); - int accum = intValues.at(indices[0]); - for (size_t i = 1; i < shape.size(); ++i) { - accum = accum * shape[i] + intValues.at(indices[i]); - } - return accum; -} - -void Interpreter::visit(tensor::EmptyOp op) { - auto tensorType = op.getResult().getType(); - auto numElements = tensorType.getNumElements(); - auto elementType = tensorType.getElementType(); - - if (elementType.isInteger()) { - intVectors[op.getResult()] = - std::make_shared>(numElements); - } else if (elementType.isF32()) { - floatVectors[op.getResult()] = - std::make_shared>(numElements); - } else if (elementType.isF64()) { - doubleVectors[op.getResult()] = - std::make_shared>(numElements); - } else if (isa(elementType)) { - plaintextVectors[op.getResult()] = - std::make_shared>(numElements); - } else { - ciphertextVectors[op.getResult()] = - std::make_shared>(numElements); - } -} - -void Interpreter::visit(tensor::ExtractOp op) { - int index = getFlattenedTensorIndex(op.getTensor(), op.getIndices()); - auto tensorType = cast(op.getTensor().getType()); - auto elemType = tensorType.getElementType(); - - if (elemType.isIndex() || elemType.isInteger()) { - intValues[op.getResult()] = (*intVectors.at(op.getTensor()))[index]; - } else if (elemType.isF32()) { - floatValues[op.getResult()] = (*floatVectors.at(op.getTensor()))[index]; - } else if (elemType.isF64()) { - doubleValues[op.getResult()] = (*doubleVectors.at(op.getTensor()))[index]; - } else if (isa(elemType)) { - plaintexts[op.getResult()] = (*plaintextVectors.at(op.getTensor()))[index]; - } else { - ciphertexts[op.getResult()] = - (*ciphertextVectors.at(op.getTensor()))[index]; - } -} - -void Interpreter::visit(tensor::InsertOp op) { - int index = getFlattenedTensorIndex(op.getDest(), op.getIndices()); - auto tensorType = cast(op.getDest().getType()); - auto elemType = tensorType.getElementType(); - - // Check if we can modify the tensor in-place (no copy needed) - bool canModifyInPlace = liveness->isDeadAfter(op.getDest(), op); - - if (elemType.isInteger(32) || elemType.isInteger(64)) { - auto srcVec = intVectors.at(op.getDest()); - auto vec = canModifyInPlace - ? srcVec - : std::make_shared>(*srcVec); - (*vec)[index] = intValues.at(op.getScalar()); - intVectors[op.getResult()] = vec; - } else if (elemType.isF32()) { - auto srcVec = floatVectors.at(op.getDest()); - auto vec = canModifyInPlace ? srcVec - : std::make_shared>(*srcVec); - (*vec)[index] = floatValues.at(op.getScalar()); - floatVectors[op.getResult()] = vec; - } else if (elemType.isF64()) { - auto srcVec = doubleVectors.at(op.getDest()); - auto vec = canModifyInPlace - ? srcVec - : std::make_shared>(*srcVec); - (*vec)[index] = doubleValues.at(op.getScalar()); - doubleVectors[op.getResult()] = vec; - } else if (isa(elemType)) { - auto srcVec = plaintextVectors.at(op.getDest()); - auto vec = canModifyInPlace - ? srcVec - : std::make_shared>(*srcVec); - (*vec)[index] = plaintexts.at(op.getScalar()); - plaintextVectors[op.getResult()] = vec; - } else { - auto srcVec = ciphertextVectors.at(op.getDest()); - auto vec = canModifyInPlace - ? srcVec - : std::make_shared>(*srcVec); - (*vec)[index] = ciphertexts.at(op.getScalar()); - ciphertextVectors[op.getResult()] = vec; - } -} - -void Interpreter::visit(tensor::SplatOp op) { - auto tensorType = op.getResult().getType(); - auto numElements = tensorType.getNumElements(); - auto elemType = tensorType.getElementType(); - - if (elemType.isInteger(32) || elemType.isInteger(64)) { - int val = intValues.at(op.getInput()); - intVectors[op.getResult()] = - std::make_shared>(numElements, val); - } else if (elemType.isF32()) { - float val = floatValues.at(op.getInput()); - floatVectors[op.getResult()] = - std::make_shared>(numElements, val); - } else if (elemType.isF64()) { - double val = doubleValues.at(op.getInput()); - doubleVectors[op.getResult()] = - std::make_shared>(numElements, val); - } -} - -void Interpreter::visit(tensor::FromElementsOp op) { - auto elements = op.getElements(); - if (elements.empty()) { - op.emitError("FromElementsOp requires at least one element"); - return; - } - - auto tensorType = op.getResult().getType(); - auto elemType = tensorType.getElementType(); - - if (elemType.isInteger(32) || elemType.isInteger(64)) { - auto result = std::make_shared>(elements.size()); - for (size_t i = 0; i < elements.size(); ++i) { - (*result)[i] = intValues.at(elements[i]); - } - intVectors[op.getResult()] = result; - } else if (elemType.isF32()) { - auto result = std::make_shared>(elements.size()); - for (size_t i = 0; i < elements.size(); ++i) { - (*result)[i] = floatValues.at(elements[i]); - } - floatVectors[op.getResult()] = result; - } else if (elemType.isF64()) { - auto result = std::make_shared>(elements.size()); - for (size_t i = 0; i < elements.size(); ++i) { - (*result)[i] = doubleValues.at(elements[i]); - } - doubleVectors[op.getResult()] = result; - } else if (isa(elemType)) { - auto result = std::make_shared>(elements.size()); - for (size_t i = 0; i < elements.size(); ++i) { - (*result)[i] = plaintexts.at(elements[i]); - } - plaintextVectors[op.getResult()] = result; - } else { - auto result = std::make_shared>(elements.size()); - for (size_t i = 0; i < elements.size(); ++i) { - (*result)[i] = ciphertexts.at(elements[i]); - } - ciphertextVectors[op.getResult()] = result; - } -} - -void Interpreter::visit(tensor::ConcatOp op) { - auto inputs = op.getInputs(); - auto tensorType = op.getResult().getType(); - auto elemType = tensorType.getElementType(); - - if (elemType.isInteger(32) || elemType.isInteger(64)) { - // Calculate total size and reserve - size_t totalSize = 0; - for (auto input : inputs) { - totalSize += intVectors.at(input)->size(); - } - auto result = std::make_shared>(); - result->reserve(totalSize); - for (auto input : inputs) { - const auto& vec = *intVectors.at(input); - result->insert(result->end(), vec.begin(), vec.end()); - } - intVectors[op.getResult()] = result; - } else if (elemType.isF32()) { - size_t totalSize = 0; - for (auto input : inputs) { - totalSize += floatVectors.at(input)->size(); - } - auto result = std::make_shared>(); - result->reserve(totalSize); - for (auto input : inputs) { - const auto& vec = *floatVectors.at(input); - result->insert(result->end(), vec.begin(), vec.end()); - } - floatVectors[op.getResult()] = result; - } else if (elemType.isF64()) { - size_t totalSize = 0; - for (auto input : inputs) { - totalSize += doubleVectors.at(input)->size(); - } - auto result = std::make_shared>(); - result->reserve(totalSize); - for (auto input : inputs) { - const auto& vec = *doubleVectors.at(input); - result->insert(result->end(), vec.begin(), vec.end()); - } - doubleVectors[op.getResult()] = result; - } else if (isa(elemType)) { - size_t totalSize = 0; - for (auto input : inputs) { - totalSize += plaintextVectors.at(input)->size(); - } - auto result = std::make_shared>(); - result->reserve(totalSize); - for (auto input : inputs) { - const auto& vec = *plaintextVectors.at(input); - result->insert(result->end(), vec.begin(), vec.end()); - } - plaintextVectors[op.getResult()] = result; - } else { - size_t totalSize = 0; - for (auto input : inputs) { - totalSize += ciphertextVectors.at(input)->size(); - } - auto result = std::make_shared>(); - result->reserve(totalSize); - for (auto input : inputs) { - const auto& vec = *ciphertextVectors.at(input); - result->insert(result->end(), vec.begin(), vec.end()); - } - ciphertextVectors[op.getResult()] = result; - } -} - -void Interpreter::visit(tensor::CollapseShapeOp op) { - // Just copy the shared_ptr (shape metadata doesn't affect our flat storage) - auto srcType = cast(op.getSrc().getType()); - auto elemType = srcType.getElementType(); - - if (elemType.isInteger(32) || elemType.isInteger(64)) { - intVectors[op.getResult()] = intVectors.at(op.getSrc()); - } else if (elemType.isF32()) { - floatVectors[op.getResult()] = floatVectors.at(op.getSrc()); - } else if (elemType.isF64()) { - doubleVectors[op.getResult()] = doubleVectors.at(op.getSrc()); - } else if (isa(elemType)) { - plaintextVectors[op.getResult()] = plaintextVectors.at(op.getSrc()); - } else { - ciphertextVectors[op.getResult()] = ciphertextVectors.at(op.getSrc()); - } -} - -void Interpreter::visit(tensor::ExpandShapeOp op) { - // Same as CollapseShapeOp - auto srcType = cast(op.getSrc().getType()); - auto elemType = srcType.getElementType(); - - if (elemType.isInteger(32) || elemType.isInteger(64)) { - intVectors[op.getResult()] = intVectors.at(op.getSrc()); - } else if (elemType.isF32()) { - floatVectors[op.getResult()] = floatVectors.at(op.getSrc()); - } else if (elemType.isF64()) { - doubleVectors[op.getResult()] = doubleVectors.at(op.getSrc()); - } else if (isa(elemType)) { - plaintextVectors[op.getResult()] = plaintextVectors.at(op.getSrc()); - } else { - ciphertextVectors[op.getResult()] = ciphertextVectors.at(op.getSrc()); - } -} - -void Interpreter::visit(tensor::ExtractSliceOp op) { - auto srcType = cast(op.getSource().getType()); - auto elemType = srcType.getElementType(); - - auto offsets = op.getStaticOffsets(); - auto sizes = op.getStaticSizes(); - auto strides = op.getStaticStrides(); - - auto sourceType = cast(op.getSource().getType()); - auto sourceShape = sourceType.getShape(); - - // Calculate total number of elements to extract - int64_t totalElements = 1; - for (int64_t size : sizes) { - totalElements *= size; - } - - // For multi-dimensional slices, we need to compute which elements to extract - // from the flattened source tensor - auto extractElement = [&](int64_t flatResultIndex) -> int64_t { - // Convert flat result index to multi-dimensional result indices - std::vector resultIndices(sizes.size()); - int64_t remaining = flatResultIndex; - for (int i = sizes.size() - 1; i >= 0; --i) { - resultIndices[i] = remaining % sizes[i]; - remaining /= sizes[i]; - } - - // Convert to source indices using offsets and strides - std::vector sourceIndices(offsets.size()); - for (size_t i = 0; i < offsets.size(); ++i) { - sourceIndices[i] = offsets[i] + resultIndices[i] * strides[i]; - } - - // Convert source indices to flat index - int64_t flatSourceIndex = sourceIndices[0]; - for (size_t i = 1; i < sourceIndices.size(); ++i) { - flatSourceIndex = flatSourceIndex * sourceShape[i] + sourceIndices[i]; - } - return flatSourceIndex; - }; - - if (elemType.isInteger()) { - auto result = std::vector(totalElements); - const auto& srcVec = *intVectors.at(op.getSource()); - for (int64_t i = 0; i < totalElements; ++i) { - result[i] = srcVec[extractElement(i)]; - } - intVectors[op.getResult()] = - std::make_shared>(std::move(result)); - } else if (elemType.isF32()) { - auto result = std::vector(totalElements); - const auto& srcVec = *floatVectors.at(op.getSource()); - for (int64_t i = 0; i < totalElements; ++i) { - result[i] = srcVec[extractElement(i)]; - } - floatVectors[op.getResult()] = - std::make_shared>(std::move(result)); - } else if (elemType.isF64()) { - auto result = std::vector(totalElements); - const auto& srcVec = *doubleVectors.at(op.getSource()); - for (int64_t i = 0; i < totalElements; ++i) { - result[i] = srcVec[extractElement(i)]; - } - doubleVectors[op.getResult()] = - std::make_shared>(std::move(result)); - } else if (isa(elemType)) { - auto result = std::vector(totalElements); - const auto& srcVec = *plaintextVectors.at(op.getSource()); - for (int64_t i = 0; i < totalElements; ++i) { - result[i] = srcVec[extractElement(i)]; - } - plaintextVectors[op.getResult()] = - std::make_shared>(std::move(result)); - } else if (isa(elemType)) { - auto result = std::vector(totalElements); - const auto& srcVec = *ciphertextVectors.at(op.getSource()); - for (int64_t i = 0; i < totalElements; ++i) { - result[i] = srcVec[extractElement(i)]; - } - ciphertextVectors[op.getResult()] = - std::make_shared>(std::move(result)); - } else { - op.emitError("Unsupported tensor type in ExtractSliceOp\n"); - } -} - -void Interpreter::visit(tensor::InsertSliceOp op) { - auto offsets = op.getStaticOffsets(); - auto sizes = op.getStaticSizes(); - auto strides = op.getStaticStrides(); - - auto destType = cast(op.getDest().getType()); - auto destShape = destType.getShape(); - - // Calculate total number of elements to insert - int64_t totalElements = 1; - for (int64_t size : sizes) { - totalElements *= size; - } - - // For multi-dimensional slices, we need to compute which dest elements to - // update - auto insertElement = [&](int64_t flatSourceIndex) -> int64_t { - // Convert flat source index to multi-dimensional source indices - std::vector sourceIndices(sizes.size()); - int64_t remaining = flatSourceIndex; - for (int i = sizes.size() - 1; i >= 0; --i) { - sourceIndices[i] = remaining % sizes[i]; - remaining /= sizes[i]; - } - - // Convert to dest indices using offsets and strides - std::vector destIndices(offsets.size()); - for (size_t i = 0; i < offsets.size(); ++i) { - destIndices[i] = offsets[i] + sourceIndices[i] * strides[i]; - } - - // Convert dest indices to flat index - int64_t flatDestIndex = destIndices[0]; - for (size_t i = 1; i < destIndices.size(); ++i) { - flatDestIndex = flatDestIndex * destShape[i] + destIndices[i]; - } - return flatDestIndex; - }; - - bool canModifyInPlace = liveness->isDeadAfter(op.getDest(), op); - - if (auto elemType = destType.getElementType(); elemType.isInteger()) { - auto srcDestVec = intVectors.at(op.getDest()); - auto destVec = canModifyInPlace - ? srcDestVec - : std::make_shared>(*srcDestVec); - const auto& srcVec = *intVectors.at(op.getSource()); - for (int64_t i = 0; i < totalElements; ++i) { - (*destVec)[insertElement(i)] = srcVec[i]; - } - intVectors[op.getResult()] = std::move(destVec); - } else if (elemType.isF32()) { - auto srcDestVec = floatVectors.at(op.getDest()); - auto destVec = canModifyInPlace - ? srcDestVec - : std::make_shared>(*srcDestVec); - const auto& srcVec = *floatVectors.at(op.getSource()); - for (int64_t i = 0; i < totalElements; ++i) { - (*destVec)[insertElement(i)] = srcVec[i]; - } - floatVectors[op.getResult()] = std::move(destVec); - } else if (elemType.isF64()) { - auto srcDestVec = doubleVectors.at(op.getDest()); - auto destVec = canModifyInPlace - ? srcDestVec - : std::make_shared>(*srcDestVec); - const auto& srcVec = *doubleVectors.at(op.getSource()); - for (int64_t i = 0; i < totalElements; ++i) { - (*destVec)[insertElement(i)] = srcVec[i]; - } - doubleVectors[op.getResult()] = std::move(destVec); - } else if (isa(elemType)) { - auto srcDestVec = plaintextVectors.at(op.getDest()); - auto destVec = canModifyInPlace - ? srcDestVec - : std::make_shared>(*srcDestVec); - const auto& srcVec = *plaintextVectors.at(op.getSource()); - for (int64_t i = 0; i < totalElements; ++i) { - (*destVec)[insertElement(i)] = srcVec[i]; - } - plaintextVectors[op.getResult()] = std::move(destVec); - } else if (isa(elemType)) { - auto srcDestVec = ciphertextVectors.at(op.getDest()); - auto destVec = - canModifyInPlace - ? srcDestVec - : std::make_shared>(*srcDestVec); - const auto& srcVec = *ciphertextVectors.at(op.getSource()); - for (int64_t i = 0; i < totalElements; ++i) { - (*destVec)[insertElement(i)] = srcVec[i]; - } - ciphertextVectors[op.getResult()] = std::move(destVec); - } else { - op.emitError("Unsupported tensor type in InsertSliceOp\n"); - } -} - -void Interpreter::visit(tensor::ParallelInsertSliceOp op) { - auto offsets = op.getOffsets(); - auto sizes = op.getStaticSizes(); - auto strides = op.getStaticStrides(); - - auto destType = cast(op.getDest().getType()); - auto destShape = destType.getShape(); - - // Calculate total number of elements to insert - int64_t totalElements = 1; - for (int64_t size : sizes) { - totalElements *= size; - } - - // For multi-dimensional slices, we need to compute which dest elements to - // update - auto insertElement = [&](int64_t flatSourceIndex) -> int64_t { - // Convert flat source index to multi-dimensional source indices - std::vector sourceIndices(sizes.size()); - int64_t remaining = flatSourceIndex; - for (int i = sizes.size() - 1; i >= 0; --i) { - sourceIndices[i] = remaining % sizes[i]; - remaining /= sizes[i]; - } - - // Convert to dest indices using offsets and strides - std::vector destIndices(offsets.size()); - for (size_t i = 0; i < offsets.size(); ++i) { - auto offsetVal = intValues.at(offsets[i]); - destIndices[i] = offsetVal + sourceIndices[i] * strides[i]; - } - - // Convert dest indices to flat index - int64_t flatDestIndex = destIndices[0]; - for (size_t i = 1; i < destIndices.size(); ++i) { - flatDestIndex = flatDestIndex * destShape[i] + destIndices[i]; - } - return flatDestIndex; - }; - - if (auto elemType = destType.getElementType(); elemType.isInteger()) { - auto destVec = intVectors.at(op.getDest()); - const auto& srcVec = *intVectors.at(op.getSource()); - for (int64_t i = 0; i < totalElements; ++i) { - (*destVec)[insertElement(i)] = srcVec[i]; - } - } else if (elemType.isF32()) { - auto destVec = floatVectors.at(op.getDest()); - const auto& srcVec = *floatVectors.at(op.getSource()); - for (int64_t i = 0; i < totalElements; ++i) { - (*destVec)[insertElement(i)] = srcVec[i]; - } - } else if (elemType.isF64()) { - auto destVec = doubleVectors.at(op.getDest()); - const auto& srcVec = *doubleVectors.at(op.getSource()); - for (int64_t i = 0; i < totalElements; ++i) { - (*destVec)[insertElement(i)] = srcVec[i]; - } - } else if (isa(elemType)) { - auto destVec = plaintextVectors.at(op.getDest()); - const auto& srcVec = *plaintextVectors.at(op.getSource()); - for (int64_t i = 0; i < totalElements; ++i) { - (*destVec)[insertElement(i)] = srcVec[i]; - } - } else if (isa(elemType)) { - auto destVec = ciphertextVectors.at(op.getDest()); - const auto& srcVec = *ciphertextVectors.at(op.getSource()); - for (int64_t i = 0; i < totalElements; ++i) { - (*destVec)[insertElement(i)] = srcVec[i]; - } - } else { - op.emitError("Unsupported tensor type in InsertSliceOp\n"); - } -} - -void Interpreter::visit(linalg::BroadcastOp op) { - // BroadcastOp copies the input op into a new tensor by adding the specified - // dims. - auto resultType = cast(op->getResults()[0].getType()); - auto resultShape = resultType.getShape(); - auto inputType = cast(op.getInput().getType()); - auto inputShape = inputType.getShape(); - auto numOutputElements = resultType.getNumElements(); - - auto calculate = [&](const auto& inputVec, auto& outputVec) { - // Create a map of the broadcast dimensions. - DenseMap broadcastDims; - for (const auto& dim : op.getDimensions()) { - broadcastDims[dim] = true; - } - - // Pre-allocate temporary index vectors outside the loop - std::vector outputIndices(resultShape.size()); - std::vector inputIndices(inputShape.size()); - - // Iterate over the output tensor's elements - for (int64_t i = 0; i < numOutputElements; ++i) { - // Calculate the multi-dimensional index in the output tensor by - // unflattening the index in the output shape. - int64_t temp = i; - for (int d = resultShape.size() - 1; d >= 0; --d) { - if (resultShape[d] > 0) { - outputIndices[d] = temp % resultShape[d]; - temp /= resultShape[d]; - } else { - outputIndices[d] = 0; - } - } - - // Calculate the multi-dimensional index in the input tensor. - int64_t inputDim = 0; - for (size_t d = 0; d < resultShape.size(); ++d) { - if (!broadcastDims.contains(d)) { - // If the output dimension is not broadcast, then the input and - // output dimensions are the same. - inputIndices[inputDim] = outputIndices[d]; - inputDim++; - } - } - - // Calculate the flattened index in the input tensor - int64_t inputFlatIndex = 0; - if (!inputShape.empty()) { - inputFlatIndex = inputIndices[0]; - for (size_t d = 1; d < inputShape.size(); ++d) { - inputFlatIndex = inputFlatIndex * inputShape[d] + inputIndices[d]; - } - } - - outputVec[i] = inputVec[inputFlatIndex]; - } - }; - - if (inputType.getElementType().isInteger()) { - auto inputVec = *intVectors.at(op.getInput()); - auto outputVec = std::vector(numOutputElements); - calculate(inputVec, outputVec); - intVectors[op->getResults()[0]] = - std::make_shared>(std::move(outputVec)); - } else if (inputType.getElementType().isF32()) { - auto inputVec = *floatVectors.at(op.getInput()); - auto outputVec = std::vector(numOutputElements); - calculate(inputVec, outputVec); - floatVectors[op->getResults()[0]] = - std::make_shared>(std::move(outputVec)); - } else if (inputType.getElementType().isF64()) { - auto inputVec = *doubleVectors.at(op.getInput()); - auto outputVec = std::vector(numOutputElements); - calculate(inputVec, outputVec); - doubleVectors[op->getResults()[0]] = - std::make_shared>(std::move(outputVec)); - } else { - op.emitError("Unsupported tensor type in BroadcastOp\n"); - } -} - -// SCF and Affine ops -void Interpreter::visit(scf::YieldOp op) { - // YieldOp is handled specially within loop bodies - // The operands are already in env and will be retrieved by the loop op -} - -void Interpreter::visit(affine::AffineYieldOp op) { - // AffineYieldOp is handled specially within loop bodies - // The operands are already in env and will be retrieved by the loop op -} - -void Interpreter::visit(scf::ForOp op) { - int lowerBound = intValues.at(op.getLowerBound()); - int upperBound = intValues.at(op.getUpperBound()); - int step = intValues.at(op.getStep()); - - // Initialize iter args with initial values - std::vector iterArgs; - iterArgs.reserve(op.getInitArgs().size()); - for (auto initArg : op.getInitArgs()) { - iterArgs.push_back(loadTypedValue(initArg)); - } - - // Cache operation visitors to avoid dispatch overhead in hot loop - std::vector> cachedOps; - scf::YieldOp yieldOp = nullptr; - for (auto& bodyOp : op.getBody()->getOperations()) { - if (auto yield = dyn_cast(&bodyOp)) { - yieldOp = yield; - } else { - Operation* opPtr = &bodyOp; // Capture pointer, not reference to loop var - cachedOps.push_back([this, opPtr]() { visit(opPtr); }); - } - } - - // Execute the loop - for (int i = lowerBound; i < upperBound; i += step) { - intValues[op.getInductionVar()] = i; - for (auto [blockArg, iterArg] : - llvm::zip(op.getRegionIterArgs(), iterArgs)) { - storeTypedValue(blockArg, iterArg); - } - - // Execute loop body - for (auto& cachedOp : cachedOps) { - cachedOp(); - } - - // Collect yield results - std::vector yieldResults; - yieldResults.reserve(iterArgs.size()); - if (yieldOp) { - for (auto yieldOperand : yieldOp.getOperands()) { - yieldResults.push_back(loadTypedValue(yieldOperand)); - } - } - - iterArgs = std::move(yieldResults); - } - - // Store final results - for (auto [result, finalValue] : llvm::zip(op.getResults(), iterArgs)) { - storeTypedValue(result, std::move(finalValue)); - } -} - -void Interpreter::visit(scf::InParallelOp op) { - for (auto& bodyOp : op.getBody()->getOperations()) { - visit(&bodyOp); - } -} - -void Interpreter::visit(scf::ForallOp op) { - if (op.getInductionVars().size() != 1) { - op.emitError("Only 1D forall loops are supported"); - } - - int lowerBound = op.getStaticLowerBound()[0]; - int upperBound = op.getStaticUpperBound()[0]; - int step = op.getStaticStep()[0]; - - // Assign the block arg values to the initial values. - for (auto& opOperand : op->getOpOperands()) { - BlockArgument blockArg = op.getTiedBlockArgument(&opOperand); - storeTypedValue(blockArg, loadTypedValue(opOperand.get())); - } - - // Cache operation visitors to avoids lookup overhead in the parallel loop. - std::vector> cachedOps; - for (auto& bodyOp : op.getBody()->getOperations()) { - Operation* opPtr = &bodyOp; - cachedOps.push_back([opPtr](Interpreter* interp) { interp->visit(opPtr); }); - } - -// Execute the loop in parallel using OpenMP. -#pragma omp parallel for - for (int i = lowerBound; i < upperBound; i += step) { - // Created a local interpreter for each thread. - Interpreter localInterpreter(*this); - localInterpreter.intValues[op.getInductionVar(0)] = i; - // Execute cached loop body operations using the local interpreter. - for (auto& cachedOp : cachedOps) { - cachedOp(&localInterpreter); - } - } - - // Assign the results to the final block arg values. - for (auto opResult : op->getOpResults()) { - OpOperand* opOperand = op.getTiedOpOperand(opResult); - BlockArgument blockArg = op.getTiedBlockArgument(opOperand); - storeTypedValue(opResult, loadTypedValue(blockArg)); - } -} - -void Interpreter::visit(scf::IfOp op) { - int condition = intValues.at(op.getCondition()); - bool condBool = (condition != 0); - std::vector results; - - if (condBool) { - for (auto& bodyOp : op.getThenRegion().front().getOperations()) { - if (auto yieldOp = dyn_cast(&bodyOp)) { - for (auto yieldOperand : yieldOp.getOperands()) { - results.push_back(loadTypedValue(yieldOperand)); - } - } else { - visit(&bodyOp); - } - } - } else if (!op.getElseRegion().empty()) { - for (auto& bodyOp : op.getElseRegion().front().getOperations()) { - if (auto yieldOp = dyn_cast(&bodyOp)) { - for (auto yieldOperand : yieldOp.getOperands()) { - results.push_back(loadTypedValue(yieldOperand)); - } - } else { - visit(&bodyOp); - } - } - } - - for (auto [result, value] : llvm::zip(op.getResults(), results)) { - storeTypedValue(result, std::move(value)); - } -} - -void Interpreter::visit(affine::AffineForOp op) { - // Get loop bounds from affine maps - auto lowerBoundMap = op.getLowerBoundMap(); - auto upperBoundMap = op.getUpperBoundMap(); - - // For simplicity, assume constant bounds (map with no inputs) - if (lowerBoundMap.getNumInputs() != 0 || upperBoundMap.getNumInputs() != 0) { - op.emitError("AffineForOp with non-constant bounds not yet supported"); - } - - int64_t lowerBound = lowerBoundMap.getSingleConstantResult(); - int64_t upperBound = upperBoundMap.getSingleConstantResult(); - int64_t step = op.getStep().getSExtValue(); - - // Initialize iter args with initial values - std::vector iterArgs; - iterArgs.reserve(op.getInits().size()); - for (auto initArg : op.getInits()) { - iterArgs.push_back(loadTypedValue(initArg)); - } - - // Cache operation visitors to avoid dispatch overhead in hot loop - std::vector> cachedOps; - affine::AffineYieldOp yieldOp = nullptr; - for (auto& bodyOp : op.getBody()->getOperations()) { - if (auto yield = dyn_cast(&bodyOp)) { - yieldOp = yield; - } else { - Operation* opPtr = &bodyOp; // Capture pointer, not reference to loop var - cachedOps.push_back([this, opPtr]() { visit(opPtr); }); - } - } - - // Execute the loop - for (int64_t i = lowerBound; i < upperBound; i += step) { - // Set up induction variable and iter args in env - eraseValue(op.getInductionVar()); - storeTypedValue(op.getInductionVar(), TypedCppValue(static_cast(i))); - - for (auto [blockArg, iterArg] : - llvm::zip(op.getRegionIterArgs(), iterArgs)) { - eraseValue(blockArg); - storeTypedValue(blockArg, iterArg); - } - - // Execute cached operations - for (auto& cachedOp : cachedOps) { - cachedOp(); - } - - // Collect yield results - std::vector yieldResults; - yieldResults.reserve(iterArgs.size()); - if (yieldOp) { - for (auto yieldOperand : yieldOp.getOperands()) { - yieldResults.push_back(loadTypedValue(yieldOperand)); - } - } - - // Update iter args for next iteration - iterArgs = std::move(yieldResults); - } - - // Store final results (move to avoid copying) - for (auto [result, finalValue] : llvm::zip(op.getResults(), iterArgs)) { - storeTypedValue(result, std::move(finalValue)); - } -} - -#ifdef OPENFHE_ENABLE_TIMING -#include -#include -#include - -void Interpreter::printTimingResults() { - double totalTime = 0; - for (const auto& entry : timingResults) { - totalTime += entry.second.totalTime.count(); - } - - struct TimingInfo { - std::string operation; - double totalTime; - int count; - double percentage; - }; - - std::vector sortedResults; - for (const auto& entry : timingResults) { - sortedResults.push_back( - {entry.first, entry.second.totalTime.count(), entry.second.count, - (entry.second.totalTime.count() / totalTime) * 100}); - } - - std::sort(sortedResults.begin(), sortedResults.end(), - [](const TimingInfo& a, const TimingInfo& b) { - return a.percentage > b.percentage; - }); - - std::cout << "--- Timing Results ---\n"; - std::cout << std::left << std::setw(30) << "Operation" << std::setw(20) - << "Total Time (s)" << std::setw(20) << "Total Time (%)" - << std::setw(20) << "Count" << "Average Latency (s)\n"; - for (const auto& entry : sortedResults) { - std::cout << std::left << std::setw(30) << entry.operation << std::setw(20) - << entry.totalTime << std::setw(20) << entry.percentage - << std::setw(20) << entry.count << (entry.totalTime / entry.count) - << "\n"; - } -} -#endif - -// Macro for handling binary operations on ciphertext types -#define HANDLE_CT_CT_BINOP(op, opName, evalMethod) \ - do { \ - auto cc = cryptoContexts.at((op).getCryptoContext()); \ - auto lhsCt = ciphertexts.at((op).getLhs()); \ - auto rhsCt = ciphertexts.at((op).getRhs()); \ - TIME_OPERATION(opName, (op).getOutput(), cc->evalMethod(lhsCt, rhsCt)); \ - } while (0) - -// OpenFHE ct-ct binary operations -void Interpreter::visit(AddOp op) { HANDLE_CT_CT_BINOP(op, "Add", EvalAdd); } -void Interpreter::visit(SubOp op) { HANDLE_CT_CT_BINOP(op, "Sub", EvalSub); } -void Interpreter::visit(MulOp op) { HANDLE_CT_CT_BINOP(op, "Mul", EvalMult); } -void Interpreter::visit(MulNoRelinOp op) { - HANDLE_CT_CT_BINOP(op, "MulNoRelin", EvalMultNoRelin); -} - -void Interpreter::visit(AddPlainOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto lhsVal = op.getLhs(); - auto rhsVal = op.getRhs(); - - if (isa(lhsVal.getType()) && - isa(rhsVal.getType())) { - auto lhsCt = ciphertexts.at(lhsVal); - auto rhsPt = plaintexts.at(rhsVal); - TIME_OPERATION("AddPlain", op.getOutput(), cc->EvalAdd(lhsCt, rhsPt)); - } else if (isa(lhsVal.getType()) && - isa(rhsVal.getType())) { - auto lhsPt = plaintexts.at(lhsVal); - auto rhsCt = ciphertexts.at(rhsVal); - TIME_OPERATION("AddPlain", op.getOutput(), cc->EvalAdd(lhsPt, rhsCt)); - } else { - op.emitError("AddPlainOp requires ciphertext and plaintext"); - } -} - -void Interpreter::visit(SubPlainOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto lhsVal = op.getLhs(); - auto rhsVal = op.getRhs(); - - // Check which is ciphertext and which is plaintext - if (isa(lhsVal.getType()) && - isa(rhsVal.getType())) { - auto lhsCt = ciphertexts.at(lhsVal); - auto rhsPt = plaintexts.at(rhsVal); - TIME_OPERATION("SubPlain", op.getOutput(), cc->EvalSub(lhsCt, rhsPt)); - } else if (isa(lhsVal.getType()) && - isa(rhsVal.getType())) { - auto lhsPt = plaintexts.at(lhsVal); - auto rhsCt = ciphertexts.at(rhsVal); - TIME_OPERATION("SubPlain", op.getOutput(), cc->EvalSub(lhsPt, rhsCt)); - } else { - op.emitError("SubPlainOp requires ciphertext and plaintext"); - } -} - -void Interpreter::visit(MulPlainOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto ct = ciphertexts.at(op.getCiphertext()); - auto pt = plaintexts.at(op.getPlaintext()); - TIME_OPERATION("MulPlain", op.getOutput(), cc->EvalMult(ct, pt)); -} - -void Interpreter::visit(MulConstOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto ct = ciphertexts.at(op.getCiphertext()); - auto constVal = intValues.at(op.getConstant()); - TIME_OPERATION("MulConst", op.getOutput(), cc->EvalMult(ct, constVal)); -} - -// OpenFHE unary operations -// Macro for unary operations on ciphertexts -#define HANDLE_CT_UNARY(op, opName, evalMethod) \ - do { \ - auto cc = cryptoContexts.at((op).getCryptoContext()); \ - auto inputCt = ciphertexts.at((op).getCiphertext()); \ - TIME_OPERATION(opName, (op).getOutput(), cc->evalMethod(inputCt)); \ - } while (0) - -void Interpreter::visit(NegateOp op) { - HANDLE_CT_UNARY(op, "Negate", EvalNegate); -} - -void Interpreter::visit(SquareOp op) { - HANDLE_CT_UNARY(op, "Square", EvalSquare); -} - -void Interpreter::visit(RelinOp op) { - HANDLE_CT_UNARY(op, "Relin", Relinearize); -} - -void Interpreter::visit(ModReduceOp op) { - HANDLE_CT_UNARY(op, "ModReduce", ModReduce); -} - -void Interpreter::visit(LevelReduceOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto ct = ciphertexts.at(op.getCiphertext()); - auto levelToDrop = op.getLevelToDrop(); - TIME_OPERATION("LevelReduce", op.getOutput(), - cc->LevelReduce(ct, nullptr, levelToDrop)); -} - -void Interpreter::visit(RotOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto ct = ciphertexts.at(op.getCiphertext()); - IntegerAttr staticShift = op.getStaticShiftAttr(); - Value dynamicShift = op.getDynamicShift(); - - if (staticShift) { - TIME_OPERATION("Rot", op.getOutput(), - cc->EvalRotate(ct, staticShift.getValue().getSExtValue())); - } else if (dynamicShift) { - TIME_OPERATION("Rot", op.getOutput(), - cc->EvalRotate(ct, intValues.at(dynamicShift))); - } else { - op.emitError("RotOp requires either static or dynamic shift"); - } -} - -void Interpreter::visit(AutomorphOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto ct = ciphertexts.at(op.getCiphertext()); - auto evalKey = evalKeys.at(op.getEvalKey()); - // Note: AutomorphOp requires building a map with the eval key - // For simplicity, we'll use index 0 as in the emitter - std::map evalKeyMap = {{0, evalKey}}; - TIME_OPERATION("Automorph", op.getOutput(), - cc->EvalAutomorphism(ct, 0, evalKeyMap)); -} - -void Interpreter::visit(KeySwitchOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto ct = ciphertexts.at(op.getCiphertext()); - auto key = evalKeys.at(op.getEvalKey()); - TIME_OPERATION("KeySwitch", op.getOutput(), cc->KeySwitch(ct, key)); -} - -void Interpreter::visit(BootstrapOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto ct = ciphertexts.at(op.getCiphertext()); - TIME_OPERATION("Bootstrap", op.getOutput(), cc->EvalBootstrap(ct)); -} - -// OpenFHE encryption/decryption operations -void Interpreter::visit(EncryptOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto pt = plaintexts.at(op.getPlaintext()); - // Note: EncryptOp takes an encryption key which could be public or private - // For now, we'll need to handle both cases - if (publicKeys.find(op.getEncryptionKey()) != publicKeys.end()) { - auto key = publicKeys.at(op.getEncryptionKey()); - TIME_OPERATION("Encrypt", op.getCiphertext(), cc->Encrypt(key, pt)); - } else { - auto key = privateKeys.at(op.getEncryptionKey()); - TIME_OPERATION("Encrypt", op.getCiphertext(), cc->Encrypt(key, pt)); - } -} - -void Interpreter::visit(DecryptOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto ciphertext = ciphertexts.at(op.getCiphertext()); - auto key = privateKeys.at(op.getPrivateKey()); - PlaintextT plaintext; - TIME_OPERATION_VOID("Decrypt", cc->Decrypt(key, ciphertext, &plaintext)); - plaintexts[op.getPlaintext()] = plaintext; -} - -void Interpreter::visit(MakePackedPlaintextOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - const auto& vec = *intVectors.at(op.getValue()); - std::vector vec64(vec.begin(), vec.end()); - TIME_OPERATION_NONCT("MakePackedPlaintext", op.getPlaintext(), - cc->MakePackedPlaintext(vec64), plaintexts); -} - -void Interpreter::visit(MakeCKKSPackedPlaintextOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto valType = op.getValue().getType(); - auto elemType = cast(valType).getElementType(); - - if (elemType.isF32()) { - const auto& vec = *floatVectors.at(op.getValue()); - std::vector vecDouble(vec.begin(), vec.end()); - TIME_OPERATION_NONCT("MakeCKKSPackedPlaintext", op.getPlaintext(), - cc->MakeCKKSPackedPlaintext(vecDouble), plaintexts); - } else if (elemType.isF64()) { - const auto& vec = *doubleVectors.at(op.getValue()); - TIME_OPERATION_NONCT("MakeCKKSPackedPlaintext", op.getPlaintext(), - cc->MakeCKKSPackedPlaintext(vec), plaintexts); - } else if (elemType.isInteger()) { - const auto& vec = *intVectors.at(op.getValue()); - std::vector vecDouble(vec.begin(), vec.end()); - TIME_OPERATION_NONCT("MakeCKKSPackedPlaintext", op.getPlaintext(), - cc->MakeCKKSPackedPlaintext(vecDouble), plaintexts); - } -} - -void Interpreter::visit(GenParamsOp op) { - int64_t mulDepth = op.getMulDepthAttr().getValue().getSExtValue(); - int64_t plainMod = op.getPlainModAttr().getValue().getSExtValue(); - int64_t evalAddCount = op.getEvalAddCountAttr().getValue().getSExtValue(); - int64_t keySwitchCount = op.getKeySwitchCountAttr().getValue().getSExtValue(); - - auto params = std::make_shared(); - params->SetMultiplicativeDepth(mulDepth); - if (plainMod > 0) params->SetPlaintextModulus(plainMod); - if (op.getRingDim() != 0) params->SetRingDim(op.getRingDim()); - if (op.getBatchSize() != 0) params->SetBatchSize(op.getBatchSize()); - if (op.getFirstModSize() != 0) params->SetFirstModSize(op.getFirstModSize()); - if (op.getScalingModSize() != 0) - params->SetScalingModSize(op.getScalingModSize()); - if (evalAddCount > 0) params->SetEvalAddCount(evalAddCount); - if (keySwitchCount > 0) params->SetKeySwitchCount(keySwitchCount); - if (op.getDigitSize() != 0) params->SetDigitSize(op.getDigitSize()); - if (op.getNumLargeDigits() != 0) - params->SetNumLargeDigits(op.getNumLargeDigits()); - if (op.getMaxRelinSkDeg() != 0) - params->SetMaxRelinSkDeg(op.getMaxRelinSkDeg()); - if (op.getInsecure()) params->SetSecurityLevel(HEStd_NotSet); - if (op.getEncryptionTechniqueExtended()) - params->SetEncryptionTechnique(EXTENDED); - if (!op.getKeySwitchingTechniqueBV()) - params->SetKeySwitchTechnique(HYBRID); - else - params->SetKeySwitchTechnique(BV); - if (op.getScalingTechniqueFixedManual()) - params->SetScalingTechnique(FIXEDMANUAL); - - params_.insert_or_assign(op.getResult(), std::move(params)); -} - -void Interpreter::visit(GenContextOp op) { - auto params = params_.at(op.getParams()); - CryptoContextT cc; - TIME_OPERATION_VOID("GenContext", cc = GenCryptoContext(*params)); - cc->Enable(PKE); - cc->Enable(KEYSWITCH); - cc->Enable(LEVELEDSHE); - if (op.getSupportFHE()) { - cc->Enable(ADVANCEDSHE); - cc->Enable(FHE); - } - cryptoContexts[op.getResult()] = cc; -} - -void Interpreter::visit(GenRotKeyOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto pk = privateKeys.at(op.getPrivateKey()); - std::vector rotIndices(op.getIndices().begin(), - op.getIndices().end()); - TIME_OPERATION_VOID("GenRotKey", cc->EvalRotateKeyGen(pk, rotIndices)); -} - -void Interpreter::visit(GenMulKeyOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto pk = privateKeys.at(op.getPrivateKey()); - TIME_OPERATION_VOID("GenMulKey", cc->EvalMultKeyGen(pk)); -} - -void Interpreter::visit(GenBootstrapKeyOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto pk = privateKeys.at(op.getPrivateKey()); - // Use full packing - ring dimension / 2 - auto numSlots = cc->GetRingDimension() / 2; - TIME_OPERATION_VOID("GenBootstrapKey", cc->EvalBootstrapKeyGen(pk, numSlots)); -} - -void Interpreter::visit(SetupBootstrapOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - std::vector levelBudget = { - static_cast( - op.getLevelBudgetEncode().getValue().getSExtValue()), - static_cast( - op.getLevelBudgetDecode().getValue().getSExtValue())}; - TIME_OPERATION_VOID("SetupBootstrap", cc->EvalBootstrapSetup(levelBudget)); -} - -void Interpreter::visit(FastRotationOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto ct = ciphertexts.at(op.getInput()); - auto index = intValues.at(op.getIndex()); - auto digits = fastRotPrecomps.at(op.getPrecomputedDigitDecomp()); - auto m = 2 * cc->GetRingDimension(); - TIME_OPERATION("FastRotation", op.getResult(), - cc->EvalFastRotation(ct, index, m, digits)); -} - -void Interpreter::visit(FastRotationPrecomputeOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto ct = ciphertexts.at(op.getInput()); - TIME_OPERATION_NONCT("FastRotationPrecompute", op.getResult(), - cc->EvalFastRotationPrecompute(ct), fastRotPrecomps); -} - -void Interpreter::visit(FastRotationExtOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto ct = ciphertexts.at(op.getInput()); - auto index = intValues.at(op.getIndex()); - auto digits = fastRotPrecomps.at(op.getPrecomputedDigitDecomp()); - bool addFirst = op.getAddFirst(); - TIME_OPERATION("FastRotationExt", op.getResult(), - cc->EvalFastRotationExt(ct, index, digits, addFirst)); -} - -void Interpreter::visit(KeySwitchDownOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto ct = ciphertexts.at(op.getCiphertext()); - TIME_OPERATION("KeySwitchDown", op.getResult(), cc->KeySwitchDown(ct)); -} - -void Interpreter::decodeCore(Operation* op, Value input, Value result, - bool isCKKS) { - auto plaintext = plaintexts.at(input); - - // Tensor case - if (auto tensorTy = dyn_cast(result.getType())) { - auto shape = tensorTy.getShape(); - auto nonUnitDims = llvm::count_if(shape, [](auto dim) { return dim != 1; }); - if (nonUnitDims != 1) { - op->emitError() - << "Only 1D tensors with one non-unit dimension supported"; - return; - } - - int64_t size = 0; - for (auto dim : shape) { - if (dim != 1) { - size = dim; - break; - } - } - plaintext->SetLength(size); - - if (isCKKS) { - auto ckksValues = plaintext->GetCKKSPackedValue(); - auto elemType = tensorTy.getElementType(); - - if (elemType.isF64()) { - auto res = std::make_shared>(); - res->reserve(ckksValues.size()); - for (const auto& val : ckksValues) { - res->push_back(val.real()); - } - doubleVectors[result] = res; - } else { - auto res = std::make_shared>(); - res->reserve(ckksValues.size()); - for (const auto& val : ckksValues) { - res->push_back(static_cast(val.real())); - } - floatVectors[result] = res; - } - return; - } - - auto packedValues = plaintext->GetPackedValue(); - auto res = std::make_shared>(); - res->reserve(packedValues.size()); - for (const auto& val : packedValues) { - res->push_back(val); - } - intVectors[result] = res; - return; - } - - // Scalar result - if (isCKKS) { - auto ckksValues = plaintext->GetCKKSPackedValue(); - auto elemType = result.getType(); - - if (elemType.isF64()) { - doubleValues[result] = ckksValues[0].real(); - } else { - floatValues[result] = static_cast(ckksValues[0].real()); - } - return; - } - - auto packedValues = plaintext->GetPackedValue(); - intValues[result] = static_cast(packedValues[0]); -} - -void Interpreter::visit(DecodeOp op) { - decodeCore(op, op.getInput(), op.getResult(), false); -} - -void Interpreter::visit(DecodeCKKSOp op) { - decodeCore(op, op.getInput(), op.getResult(), true); -} - -// In-place ops - -// Macro for handling binary in-place operations on ciphertext types -#define HANDLE_CT_CT_BINOP_INPLACE(op, opName, evalMethod) \ - do { \ - auto cc = cryptoContexts.at((op).getCryptoContext()); \ - auto lhsCt = ciphertexts.at((op).getLhs()); \ - auto rhsCt = ciphertexts.at((op).getRhs()); \ - TIME_OPERATION_VOID(opName, cc->evalMethod(lhsCt, rhsCt)); \ - ciphertexts[(op).getOutput()] = lhsCt; \ - } while (0) - -void Interpreter::visit(AddInPlaceOp op) { - HANDLE_CT_CT_BINOP_INPLACE(op, "AddInPlace", EvalAddInPlace); -} -void Interpreter::visit(SubInPlaceOp op) { - HANDLE_CT_CT_BINOP_INPLACE(op, "SubInPlace", EvalSubInPlace); -} - -// Macro for handling binary in-place ct-pt ops -#define HANDLE_CT_PT_BINOP_INPLACE(op, opName, evalMethod) \ - do { \ - auto cc = cryptoContexts.at(op.getCryptoContext()); \ - auto lhsVal = op.getLhs(); \ - auto rhsVal = op.getRhs(); \ - if (isa(lhsVal.getType()) && \ - isa(rhsVal.getType())) { \ - auto lhsCt = ciphertexts.at(lhsVal); \ - auto rhsPt = plaintexts.at(rhsVal); \ - TIME_OPERATION_VOID(opName, cc->evalMethod(lhsCt, rhsPt)); \ - ciphertexts[op.getOutput()] = lhsCt; \ - } else if (isa(lhsVal.getType()) && \ - isa(rhsVal.getType())) { \ - op.emitError(opName " LHS must be a ciphertext for in-place"); \ - } else { \ - op.emitError(opName " requires ciphertext and plaintext"); \ - } \ - } while (0) - -void Interpreter::visit(AddPlainInPlaceOp op) { - HANDLE_CT_PT_BINOP_INPLACE(op, "AddPlainInPlace", EvalAddInPlace); -} -void Interpreter::visit(SubPlainInPlaceOp op) { - HANDLE_CT_PT_BINOP_INPLACE(op, "SubPlainInPlace", EvalSubInPlace); -} - -void Interpreter::visit(MulConstInPlaceOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto ct = ciphertexts.at(op.getCiphertext()); - auto constVal = static_cast(intValues.at(op.getConstant())); - TIME_OPERATION_VOID("MulConstInPlace", cc->EvalMultInPlace(ct, constVal)); - ciphertexts[op.getOutput()] = ct; -} - -// Macro for handling in-place unary operations on integer types -#define HANDLE_CT_UNARY_INPLACE(op, opName, evalMethod) \ - do { \ - auto cc = cryptoContexts.at((op).getCryptoContext()); \ - auto inputCt = ciphertexts.at((op).getCiphertext()); \ - TIME_OPERATION_VOID(opName, cc->evalMethod(inputCt)); \ - ciphertexts[(op).getOutput()] = inputCt; \ - } while (0) - -void Interpreter::visit(ModReduceInPlaceOp op) { - HANDLE_CT_UNARY_INPLACE(op, "ModReduceInPlace", ModReduceInPlace); -} -void Interpreter::visit(NegateInPlaceOp op) { - HANDLE_CT_UNARY_INPLACE(op, "NegateInPlace", EvalNegateInPlace); -} -void Interpreter::visit(RelinInPlaceOp op) { - HANDLE_CT_UNARY_INPLACE(op, "RelinInPlace", RelinearizeInPlace); -} -void Interpreter::visit(SquareInPlaceOp op) { - HANDLE_CT_UNARY_INPLACE(op, "SquareInPlace", EvalSquareInPlace); -} - -void Interpreter::visit(KeySwitchInPlaceOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto ct = ciphertexts.at(op.getCiphertext()); - auto key = evalKeys.at(op.getEvalKey()); - TIME_OPERATION_VOID("KeySwitchInPlace", cc->KeySwitchInPlace(ct, key)); - ciphertexts[op.getOutput()] = ct; -} - -void Interpreter::visit(LevelReduceInPlaceOp op) { - auto cc = cryptoContexts.at(op.getCryptoContext()); - auto ct = ciphertexts.at(op.getCiphertext()); - auto levelToDrop = op.getLevelToDrop(); - TIME_OPERATION_VOID("LevelReduceInPlace", - cc->LevelReduceInPlace(ct, nullptr, levelToDrop)); - ciphertexts[op.getOutput()] = ct; -} - -void initContext(MLIRContext& context) { - mlir::DialectRegistry registry; - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - registry.insert(); - context.appendDialectRegistry(registry); - context.loadAllAvailableDialects(); -} - -OwningOpRef parse(MLIRContext* context, const std::string& mlirStr) { - return parseSourceString(mlirStr, context); -} - -} // namespace openfhe -} // namespace heir -} // namespace mlir diff --git a/lib/Target/OpenFhePke/Interpreter.h b/lib/Target/OpenFhePke/Interpreter.h deleted file mode 100644 index 61216012ed..0000000000 --- a/lib/Target/OpenFhePke/Interpreter.h +++ /dev/null @@ -1,277 +0,0 @@ -#ifndef LIB_TARGET_OPENFHEPKE_INTERPRETER_H_ -#define LIB_TARGET_OPENFHEPKE_INTERPRETER_H_ - -#include -#include -#include -#include -#include -#include - -#include "lib/Dialect/Openfhe/IR/OpenfheOps.h" -#include "llvm/include/llvm/ADT/DenseMap.h" // from @llvm-project -#include "mlir/include/mlir/Analysis/Liveness.h" // from @llvm-project -#include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h" // from @llvm-project -#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project -#include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project -#include "mlir/include/mlir/Dialect/Linalg/IR/Linalg.h" // from @llvm-project -#include "mlir/include/mlir/Dialect/SCF/IR/SCF.h" // from @llvm-project -#include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h" // from @llvm-project -#include "mlir/include/mlir/IR/BuiltinOps.h" // from @llvm-project -#include "mlir/include/mlir/IR/Operation.h" // from @llvm-project -#include "mlir/include/mlir/IR/Value.h" // from @llvm-project -#include "mlir/include/mlir/IR/ValueRange.h" // from @llvm-project -#include "src/pke/include/openfhe.h" // from @openfhe - -namespace mlir { -namespace heir { -namespace openfhe { - -using CCParamsT = lbcrypto::CCParams; -using CiphertextT = lbcrypto::Ciphertext; -using CryptoContextT = lbcrypto::CryptoContext; -using EvalKeyT = lbcrypto::EvalKey; -using PlaintextT = lbcrypto::Plaintext; -using PrivateKeyT = lbcrypto::PrivateKey; -using PublicKeyT = lbcrypto::PublicKey; -using FastRotPrecompT = std::shared_ptr>; - -struct TypedCppValue { - using Variant = std::variant< - std::monostate, // NULL_TY - bool, // BOOL - int, // INT - float, // FLOAT - double, // DOUBLE - std::shared_ptr>, // INT_VECTOR - std::shared_ptr>, // FLOAT_VECTOR - std::shared_ptr>, // DOUBLE_VECTOR - PlaintextT, // PLAINTEXT - std::shared_ptr>, // PLAINTEXT_VECTOR - CiphertextT, // CIPHERTEXT - std::shared_ptr>, // CIPHERTEXT_VECTOR - PublicKeyT, // PUBLIC_KEY - PrivateKeyT, // PRIVATE_KEY - EvalKeyT, // EVAL_KEY - CryptoContextT, // CRYPTO_CONTEXT - FastRotPrecompT // FAST_ROTATION_PRECOMP - >; - - Variant value; - - TypedCppValue() = default; - - TypedCppValue(bool v) : value(v) {} - TypedCppValue(int v) : value(v) {} - TypedCppValue(float v) : value(v) {} - TypedCppValue(double v) : value(v) {} - TypedCppValue(int64_t v) : value(static_cast(v)) {} - - // Constructors for shared_ptr (preferred - no copy) - TypedCppValue(std::shared_ptr> v) - : value(std::move(v)) {} - TypedCppValue(std::shared_ptr> v) : value(std::move(v)) {} - TypedCppValue(std::shared_ptr> v) : value(std::move(v)) {} - TypedCppValue(std::shared_ptr> v) - : value(std::move(v)) {} - TypedCppValue(std::shared_ptr> v) - : value(std::move(v)) {} - - // Convenience constructors for raw vectors (wraps in shared_ptr) - TypedCppValue(const std::vector& v) - : value(std::make_shared>(v)) {} - TypedCppValue(std::vector&& v) - : value(std::make_shared>(std::move(v))) {} - TypedCppValue(const std::vector& v) - : value(std::make_shared>(v)) {} - TypedCppValue(std::vector&& v) - : value(std::make_shared>(std::move(v))) {} - TypedCppValue(const std::vector& v) - : value(std::make_shared>(v)) {} - TypedCppValue(std::vector&& v) - : value(std::make_shared>(std::move(v))) {} - TypedCppValue(const std::vector& v) - : value(std::make_shared>(v)) {} - TypedCppValue(std::vector&& v) - : value(std::make_shared>(std::move(v))) {} - TypedCppValue(const std::vector& v) - : value(std::make_shared>(v)) {} - TypedCppValue(std::vector&& v) - : value(std::make_shared>(std::move(v))) {} - - TypedCppValue(PlaintextT v) : value(std::move(v)) {} - TypedCppValue(CiphertextT v) : value(std::move(v)) {} - TypedCppValue(PublicKeyT v) : value(std::move(v)) {} - TypedCppValue(PrivateKeyT v) : value(std::move(v)) {} - TypedCppValue(EvalKeyT v) : value(std::move(v)) {} - TypedCppValue(CryptoContextT v) : value(std::move(v)) {} - TypedCppValue(FastRotPrecompT v) : value(std::move(v)) {} -}; - -class Interpreter { - public: - Interpreter(ModuleOp module); - - std::vector interpret(const std::string& entryFunction, - ArrayRef inputValues); - -#ifdef OPENFHE_ENABLE_TIMING - void printTimingResults(); -#endif - - void visit(Operation* op); - - // Upstream ops - void visit(arith::AddIOp op); - void visit(arith::AddFOp op); - void visit(arith::AndIOp op); - void visit(arith::CmpIOp op); - void visit(arith::ConstantOp op); - void visit(arith::DivSIOp op); - void visit(arith::ExtFOp op); - void visit(arith::FloorDivSIOp op); - void visit(arith::MulIOp op); - void visit(arith::MulFOp op); - void visit(arith::MinSIOp op); - void visit(arith::MaxSIOp op); - void visit(arith::RemSIOp op); - void visit(arith::SelectOp op); - void visit(arith::SubIOp op); - void visit(arith::SubFOp op); - void visit(linalg::BroadcastOp op); - void visit(tensor::CollapseShapeOp op); - void visit(tensor::ConcatOp op); - void visit(tensor::EmptyOp op); - void visit(tensor::ExpandShapeOp op); - void visit(tensor::ExtractOp op); - void visit(tensor::ExtractSliceOp op); - void visit(tensor::FromElementsOp op); - void visit(tensor::InsertOp op); - void visit(tensor::ParallelInsertSliceOp op); - void visit(tensor::InsertSliceOp op); - void visit(tensor::SplatOp op); - - // SCF and Affine ops - void visit(scf::IfOp op); - void visit(scf::ForOp op); - void visit(scf::ForallOp op); - void visit(scf::InParallelOp op); - void visit(scf::YieldOp op); - void visit(affine::AffineForOp op); - void visit(affine::AffineYieldOp op); - - // OpenFHE ops - void visit(AddInPlaceOp op); - void visit(AddOp op); - void visit(AddPlainInPlaceOp op); - void visit(AddPlainOp op); - void visit(AutomorphOp op); - void visit(BootstrapOp op); - void visit(DecodeCKKSOp op); - void visit(DecodeOp op); - void visit(DecryptOp op); - void visit(EncryptOp op); - void visit(FastRotationOp op); - void visit(FastRotationExtOp op); - void visit(FastRotationPrecomputeOp op); - void visit(GenBootstrapKeyOp op); - void visit(GenContextOp op); - void visit(GenMulKeyOp op); - void visit(GenParamsOp op); - void visit(GenRotKeyOp op); - void visit(KeySwitchInPlaceOp op); - void visit(KeySwitchOp op); - void visit(KeySwitchDownOp op); - void visit(LevelReduceInPlaceOp op); - void visit(LevelReduceOp op); - void visit(MakeCKKSPackedPlaintextOp op); - void visit(MakePackedPlaintextOp op); - void visit(ModReduceInPlaceOp op); - void visit(ModReduceOp op); - void visit(MulConstInPlaceOp op); - void visit(MulConstOp op); - void visit(MulNoRelinOp op); - void visit(MulOp op); - void visit(MulPlainOp op); - void visit(NegateInPlaceOp op); - void visit(NegateOp op); - void visit(RelinInPlaceOp op); - void visit(RelinOp op); - void visit(RotOp op); - void visit(SetupBootstrapOp op); - void visit(SquareInPlaceOp op); - void visit(SquareOp op); - void visit(SubInPlaceOp op); - void visit(SubOp op); - void visit(SubPlainInPlaceOp op); - void visit(SubPlainOp op); - - int getFlattenedTensorIndex(Value tensor, ValueRange indices); - - private: - // Helper to erase a value from all storage maps (for liveness) - void eraseValue(Value v); - - // Helper to convert TypedCppValue to type-specific storage (for inputs) - void storeTypedValue(Value v, const TypedCppValue& typedVal); - - // Helper for decoding - void decodeCore(Operation* op, Value input, Value result, bool isCKKS); - - // Helper to convert from type-specific storage to TypedCppValue (for outputs) - TypedCppValue loadTypedValue(Value v); - ModuleOp module; - - // Type-specific storage - zero variant overhead! - llvm::DenseMap boolValues; - llvm::DenseMap intValues; - llvm::DenseMap floatValues; - llvm::DenseMap doubleValues; - - // Vectors stored as shared_ptr to avoid expensive copying - llvm::DenseMap>> intVectors; - llvm::DenseMap>> floatVectors; - llvm::DenseMap>> doubleVectors; - - // OpenFHE types (already shared_ptr internally) - llvm::DenseMap plaintexts; - llvm::DenseMap>> - plaintextVectors; - llvm::DenseMap ciphertexts; - llvm::DenseMap>> - ciphertextVectors; - llvm::DenseMap cryptoContexts; - llvm::DenseMap publicKeys; - llvm::DenseMap privateKeys; - llvm::DenseMap evalKeys; - llvm::DenseMap fastRotPrecomps; - - // liveness as shared_ptr to avoid expensive copying - std::shared_ptr liveness; - llvm::DenseMap> params_; - - // Jump table for fast operation dispatch - using OperationVisitor = std::function; - - static llvm::DenseMap - operationDispatchTable; - static bool dispatchTableInitialized; - static MLIRContext* dispatchTableContext; - void initializeDispatchTable(); - - struct TimingData { - std::chrono::duration totalTime{0}; - int count{0}; - }; - std::map timingResults; -}; - -void initContext(MLIRContext& context); - -OwningOpRef parse(MLIRContext* context, const std::string& mlirStr); - -} // namespace openfhe -} // namespace heir -} // namespace mlir - -#endif // LIB_TARGET_OPENFHEPKE_INTERPRIl INTERPRETER_H_ diff --git a/lib/Target/OpenFhePke/InterpreterTest.cpp b/lib/Target/OpenFhePke/InterpreterTest.cpp deleted file mode 100644 index 37ef6c0ddb..0000000000 --- a/lib/Target/OpenFhePke/InterpreterTest.cpp +++ /dev/null @@ -1,1500 +0,0 @@ -#include -#include -#include -#include -#include - -#include "gtest/gtest.h" // from @googletest -#include "lib/Target/OpenFhePke/Interpreter.h" -#include "mlir/include/mlir/IR/BuiltinOps.h" // from @llvm-project -#include "mlir/include/mlir/IR/MLIRContext.h" // from @llvm-project -#include "mlir/include/mlir/IR/OwningOpRef.h" // from @llvm-project -#include "mlir/include/mlir/Parser/Parser.h" // from @llvm-project -#include "src/core/include/lattice/hal/lat-backend.h" // from @openfhe -#include "src/pke/include/constants-defs.h" // from @openfhe -#include "src/pke/include/cryptocontext-fwd.h" // from @openfhe -#include "src/pke/include/encoding/plaintext-fwd.h" // from @openfhe -#include "src/pke/include/gen-cryptocontext.h" // from @openfhe -#include "src/pke/include/key/keypair.h" // from @openfhe -#include "src/pke/include/openfhe.h" // from @openfhe -#include "src/pke/include/scheme/bgvrns/gen-cryptocontext-bgvrns-params.h" // from @openfhe -#include "src/pke/include/scheme/bgvrns/gen-cryptocontext-bgvrns.h" // from @openfhe -#include "src/pke/include/scheme/ckksrns/gen-cryptocontext-ckksrns-params.h" // from @openfhe -#include "src/pke/include/scheme/ckksrns/gen-cryptocontext-ckksrns.h" // from @openfhe - -namespace mlir { -namespace heir { -namespace openfhe { - -using namespace lbcrypto; - -OwningOpRef parseTest(MLIRContext* context, - const std::string& mlirStr) { - return parseSourceString(mlirStr, context); -} - -TEST(InterpreterTest, TestTrivial) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main() { - return - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector results = interpreter.interpret(entryFunction, {}); - EXPECT_TRUE(results.empty()); -} - -TEST(InterpreterTest, TestAdd) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%a: i32, %b: i32) -> i32 { - %c = arith.addi %a, %b : i32 - return %c : i32 - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(3), TypedCppValue(4)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - EXPECT_EQ(std::get(results[0].value), 7); -} - -TEST(InterpreterTest, TestAddFloat) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%a: f32, %b: f32) -> f32 { - %c = arith.addf %a, %b : f32 - return %c : f32 - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(3.0f), - TypedCppValue(4.0f)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - EXPECT_EQ(std::get(results[0].value), 7.0f); -} - -TEST(InterpreterTest, TestFloorDivSI) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%a: i32, %b: i32) -> i32 { - %c = arith.floordivsi %a, %b : i32 - return %c : i32 - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(-7), TypedCppValue(3)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - EXPECT_EQ(std::get(results[0].value), -3); -} - -TEST(InterpreterTest, TestElementwiseAdd) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%a: tensor<3xi32>, %b: tensor<3xi32>) -> tensor<3xi32> { - %c = arith.addi %a, %b : tensor<3xi32> - return %c : tensor<3xi32> - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector a = {1, 2, 3}; - std::vector b = {2, 3, 4}; - std::vector expected = {3, 5, 7}; - std::vector inputs = {TypedCppValue(a), TypedCppValue(b)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - EXPECT_EQ(*std::get>>(results[0].value), - expected); -} - -TEST(InterpreterTest, TestMul) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%a: i32, %b: i32) -> i32 { - %c = arith.muli %a, %b : i32 - return %c : i32 - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(3), TypedCppValue(4)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - EXPECT_EQ(std::get(results[0].value), 12); -} - -TEST(InterpreterTest, TestDiv) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%a: i32, %b: i32) -> i32 { - %c = arith.divsi %a, %b : i32 - return %c : i32 - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(12), TypedCppValue(3)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - EXPECT_EQ(std::get(results[0].value), 4); -} - -TEST(InterpreterTest, TestRem) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%a: i32, %b: i32) -> i32 { - %c = arith.remsi %a, %b : i32 - return %c : i32 - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(10), TypedCppValue(3)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - EXPECT_EQ(std::get(results[0].value), 1); -} - -TEST(InterpreterTest, TestAnd) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%a: i32, %b: i32) -> i32 { - %c = arith.andi %a, %b : i32 - return %c : i32 - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(1), TypedCppValue(1)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - EXPECT_EQ(std::get(results[0].value), 1); -} - -TEST(InterpreterTest, TestCmpILt) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%a: i32, %b: i32) -> i32 { - %cmp = arith.cmpi slt, %a, %b : i32 - %c = arith.extui %cmp : i1 to i32 - return %c : i32 - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(3), TypedCppValue(5)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - EXPECT_EQ(std::get(results[0].value), 1); -} - -TEST(InterpreterTest, TestCmpIEq) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%a: i32, %b: i32) -> i32 { - %cmp = arith.cmpi eq, %a, %b : i32 - %c = arith.extui %cmp : i1 to i32 - return %c : i32 - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(5), TypedCppValue(5)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - EXPECT_EQ(std::get(results[0].value), 1); -} - -TEST(InterpreterTest, TestSelect) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%a: i32, %b: i32) -> i32 { - %cmp = arith.cmpi slt, %a, %b : i32 - %c = arith.select %cmp, %a, %b : i32 - return %c : i32 - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(10), TypedCppValue(20)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - EXPECT_EQ(std::get(results[0].value), 10); -} - -TEST(InterpreterTest, TestSelectFalse) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%a: i32, %b: i32) -> i32 { - %cmp = arith.cmpi slt, %a, %b : i32 - %c = arith.select %cmp, %a, %b : i32 - return %c : i32 - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(20), TypedCppValue(10)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - EXPECT_EQ(std::get(results[0].value), 10); -} - -TEST(InterpreterTest, TestTensorSplat) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%val: i32) -> tensor<4xi32> { - %t = tensor.splat %val : tensor<4xi32> - return %t : tensor<4xi32> - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(42)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - auto vec = *std::get>>(results[0].value); - EXPECT_EQ(vec.size(), 4); - EXPECT_EQ(vec[0], 42); - EXPECT_EQ(vec[1], 42); - EXPECT_EQ(vec[2], 42); - EXPECT_EQ(vec[3], 42); -} - -TEST(InterpreterTest, TestTensorFromElements) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%a: i32, %b: i32, %c: i32) -> tensor<3xi32> { - %t = tensor.from_elements %a, %b, %c : tensor<3xi32> - return %t : tensor<3xi32> - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(10), TypedCppValue(20), - TypedCppValue(30)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - auto vec = *std::get>>(results[0].value); - EXPECT_EQ(vec.size(), 3); - EXPECT_EQ(vec[0], 10); - EXPECT_EQ(vec[1], 20); - EXPECT_EQ(vec[2], 30); -} - -TEST(InterpreterTest, TestTensorExtract) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%a: i32, %b: i32, %c: i32) -> i32 { - %t = tensor.from_elements %a, %b, %c : tensor<3xi32> - %idx = arith.constant 1 : index - %val = tensor.extract %t[%idx] : tensor<3xi32> - return %val : i32 - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(10), TypedCppValue(20), - TypedCppValue(30)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - EXPECT_EQ(std::get(results[0].value), 20); -} - -TEST(InterpreterTest, TestTensorInsert) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%val: i32) -> tensor<3xi32> { - %t = tensor.empty() : tensor<3xi32> - %idx = arith.constant 1 : index - %t2 = tensor.insert %val into %t[%idx] : tensor<3xi32> - return %t2 : tensor<3xi32> - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(99)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - auto vec = *std::get>>(results[0].value); - EXPECT_EQ(vec.size(), 3); - EXPECT_EQ(vec[1], 99); -} - -TEST(InterpreterTest, TestLoop) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main() -> tensor<6xi32> { - %0 = arith.constant dense<0> : tensor<6xi32> - %1 = arith.constant dense<1> : tensor<6xi32> - %2 = affine.for %arg1 = 0 to 6 iter_args(%arg2 = %0) -> (tensor<6xi32>) { - %extracted = tensor.extract %1[%arg1] : tensor<6xi32> - %inserted = tensor.insert %extracted into %arg2[%arg1] : tensor<6xi32> - affine.yield %inserted : tensor<6xi32> - } - return %2 : tensor<6xi32> - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - auto vec = *std::get>>(results[0].value); - EXPECT_EQ(vec.size(), 6); - EXPECT_EQ(vec[1], 1); -} - -TEST(InterpreterTest, TestLinalgBroadcast) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main() -> tensor<1x512xf32> { - %cst_0 = arith.constant dense<1.000000e+00> : tensor - %1 = tensor.empty() : tensor<1x512xf32> - %broadcasted = linalg.broadcast ins(%cst_0 : tensor) outs(%1 : tensor<1x512xf32>) dimensions = [0, 1] - func.return %broadcasted : tensor<1x512xf32> - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - auto resultVec = - *std::get>>(results[0].value); - EXPECT_EQ(resultVec.size(), 512); // 1x512 - EXPECT_EQ(resultVec[0], 1.000000f); -} - -TEST(InterpreterTest, TestLinalgBroadcastMultiDim) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main() -> tensor<2x3x4x1x5xf32> { - %cst_0 = arith.constant dense<[[1.0], [2.0], [3.0]]> : tensor<3x1xf32> - %1 = tensor.empty() : tensor<2x3x4x1x5xf32> - %broadcasted = linalg.broadcast ins(%cst_0 : tensor<3x1xf32>) outs(%1 : tensor<2x3x4x1x5xf32>) dimensions = [0, 2, 4] - func.return %broadcasted : tensor<2x3x4x1x5xf32> - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - auto resultVec = - *std::get>>(results[0].value); - EXPECT_EQ(resultVec.size(), 120); // 2x3x4x1x5 - std::vector> inputVec = {{1.0f}, {2.0f}, {3.0f}}; - for (size_t i = 0; i < 2; ++i) { - for (size_t j = 0; j < 3; ++j) { - for (size_t k = 0; k < 4; ++k) { - for (size_t l = 0; l < 5; ++l) { - EXPECT_EQ(resultVec[i * 60 + j * 20 + k * 5 + l], inputVec[j][0]); - } - } - } - } -} - -// Helper function to set up a basic BGV crypto context for testing -struct CryptoSetup { - CryptoContext cc; - KeyPair keyPair; - - CryptoSetup(uint32_t multDepth = 2) { - CCParams parameters; - parameters.SetPlaintextModulus(65537); - parameters.SetMultiplicativeDepth(multDepth); - - cc = GenCryptoContext(parameters); - cc->Enable(PKE); - cc->Enable(KEYSWITCH); - cc->Enable(LEVELEDSHE); - - keyPair = cc->KeyGen(); - cc->EvalMultKeyGen(keyPair.secretKey); - } -}; - -// Common LWE type definitions header for MLIR tests -static const char* kLWETypesHeader = R"mlir( -!ct = !openfhe.ciphertext -!pt = !openfhe.plaintext -)mlir"; - -TEST(InterpreterTest, TestOpenfheAdd) { - CryptoSetup setup; - - // Create plaintexts - std::vector vec1 = {1, 2, 3, 4}; - std::vector vec2 = {5, 6, 7, 8}; - auto pt1 = setup.cc->MakePackedPlaintext(vec1); - auto pt2 = setup.cc->MakePackedPlaintext(vec2); - - // Encrypt - auto ct1 = setup.cc->Encrypt(setup.keyPair.publicKey, pt1); - auto ct2 = setup.cc->Encrypt(setup.keyPair.publicKey, pt2); - - // Test via interpreter - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -module { - func.func @main(%cc: !openfhe.crypto_context, %ct1: !ct, %ct2: !ct) -> !ct { - %result = openfhe.add %cc, %ct1, %ct2 : (!openfhe.crypto_context, !ct, !ct) -> !ct - return %result : !ct - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(setup.cc), - TypedCppValue(ct1), TypedCppValue(ct2)}; - std::vector results = interpreter.interpret("main", inputs); - - EXPECT_EQ(results.size(), 1); - auto resultCt = std::get(results[0].value); - - // Decrypt and verify - Plaintext resultPt; - setup.cc->Decrypt(setup.keyPair.secretKey, resultCt, &resultPt); - resultPt->SetLength(vec1.size()); - - auto resultVec = resultPt->GetPackedValue(); - EXPECT_EQ(resultVec.size(), vec1.size()); - for (size_t i = 0; i < vec1.size(); i++) { - EXPECT_EQ(resultVec[i], vec1[i] + vec2[i]); - } -} - -TEST(InterpreterTest, TestOpenfheAddInPlace) { - CryptoSetup setup; - - // Create plaintexts - std::vector vec1 = {1, 2, 3, 4}; - std::vector vec2 = {5, 6, 7, 8}; - auto pt1 = setup.cc->MakePackedPlaintext(vec1); - auto pt2 = setup.cc->MakePackedPlaintext(vec2); - - // Encrypt - auto ct1 = setup.cc->Encrypt(setup.keyPair.publicKey, pt1); - auto ct2 = setup.cc->Encrypt(setup.keyPair.publicKey, pt2); - - // Test via interpreter - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -module { - func.func @main(%cc: !openfhe.crypto_context, %ct1: !ct, %ct2: !ct) -> !ct { - %result = openfhe.add_inplace %cc, %ct1, %ct2 : (!openfhe.crypto_context, !ct, !ct) -> !ct - return %result : !ct - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(setup.cc), - TypedCppValue(ct1), TypedCppValue(ct2)}; - std::vector results = interpreter.interpret("main", inputs); - - EXPECT_EQ(results.size(), 1); - auto resultCt = std::get(results[0].value); - - // Decrypt and verify - Plaintext resultPt; - setup.cc->Decrypt(setup.keyPair.secretKey, resultCt, &resultPt); - resultPt->SetLength(vec1.size()); - - auto resultVec = resultPt->GetPackedValue(); - EXPECT_EQ(resultVec.size(), vec1.size()); - for (size_t i = 0; i < vec1.size(); i++) { - EXPECT_EQ(resultVec[i], vec1[i] + vec2[i]); - } -} - -TEST(InterpreterTest, TestOpenfheSubInPlace) { - CryptoSetup setup; - - std::vector vec1 = {10, 20, 30, 40}; - std::vector vec2 = {3, 5, 7, 9}; - auto pt1 = setup.cc->MakePackedPlaintext(vec1); - auto pt2 = setup.cc->MakePackedPlaintext(vec2); - - auto ct1 = setup.cc->Encrypt(setup.keyPair.publicKey, pt1); - auto ct2 = setup.cc->Encrypt(setup.keyPair.publicKey, pt2); - - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -module { - func.func @main(%cc: !openfhe.crypto_context, %ct1: !ct, %ct2: !ct) -> !ct { - %result = openfhe.sub_inplace %cc, %ct1, %ct2 : (!openfhe.crypto_context, !ct, !ct) -> !ct - return %result : !ct - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(setup.cc), - TypedCppValue(ct1), TypedCppValue(ct2)}; - std::vector results = interpreter.interpret("main", inputs); - - auto resultCt = std::get(results[0].value); - Plaintext resultPt; - setup.cc->Decrypt(setup.keyPair.secretKey, resultCt, &resultPt); - resultPt->SetLength(vec1.size()); - - auto resultVec = resultPt->GetPackedValue(); - for (size_t i = 0; i < vec1.size(); i++) { - EXPECT_EQ(resultVec[i], vec1[i] - vec2[i]); - } -} - -TEST(InterpreterTest, TestOpenfheNegateInPlace) { - CryptoSetup setup; - - std::vector vec = {5, 10, 15, 20}; - auto pt = setup.cc->MakePackedPlaintext(vec); - auto ct = setup.cc->Encrypt(setup.keyPair.publicKey, pt); - - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -module { - func.func @main(%cc: !openfhe.crypto_context, %ct: !ct) -> !ct { - %result = openfhe.negate_inplace %cc, %ct : (!openfhe.crypto_context, !ct) -> !ct - return %result : !ct - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(setup.cc), - TypedCppValue(ct)}; - std::vector results = interpreter.interpret("main", inputs); - - auto resultCt = std::get(results[0].value); - Plaintext resultPt; - setup.cc->Decrypt(setup.keyPair.secretKey, resultCt, &resultPt); - resultPt->SetLength(vec.size()); - - auto resultVec = resultPt->GetPackedValue(); - for (size_t i = 0; i < vec.size(); i++) { - EXPECT_EQ(resultVec[i], -vec[i]); - } -} - -TEST(InterpreterTest, TestOpenfheSquareInPlace) { - CryptoSetup setup(2); // Need depth for squaring - - std::vector vec = {2, 3, 4, 5}; - auto pt = setup.cc->MakePackedPlaintext(vec); - auto ct = setup.cc->Encrypt(setup.keyPair.publicKey, pt); - - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -module { - func.func @main(%cc: !openfhe.crypto_context, %ct: !ct) -> !ct { - %result = openfhe.square_inplace %cc, %ct : (!openfhe.crypto_context, !ct) -> !ct - return %result : !ct - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(setup.cc), - TypedCppValue(ct)}; - std::vector results = interpreter.interpret("main", inputs); - - auto resultCt = std::get(results[0].value); - Plaintext resultPt; - setup.cc->Decrypt(setup.keyPair.secretKey, resultCt, &resultPt); - resultPt->SetLength(vec.size()); - - auto resultVec = resultPt->GetPackedValue(); - for (size_t i = 0; i < vec.size(); i++) { - EXPECT_EQ(resultVec[i], vec[i] * vec[i]); - } -} - -TEST(InterpreterTest, TestOpenfheSub) { - CryptoSetup setup; - - std::vector vec1 = {10, 20, 30, 40}; - std::vector vec2 = {3, 5, 7, 9}; - auto pt1 = setup.cc->MakePackedPlaintext(vec1); - auto pt2 = setup.cc->MakePackedPlaintext(vec2); - - auto ct1 = setup.cc->Encrypt(setup.keyPair.publicKey, pt1); - auto ct2 = setup.cc->Encrypt(setup.keyPair.publicKey, pt2); - - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -module { - func.func @main(%cc: !openfhe.crypto_context, %ct1: !ct, %ct2: !ct) -> !ct { - %result = openfhe.sub %cc, %ct1, %ct2 : (!openfhe.crypto_context, !ct, !ct) -> !ct - return %result : !ct - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(setup.cc), - TypedCppValue(ct1), TypedCppValue(ct2)}; - std::vector results = interpreter.interpret("main", inputs); - - auto resultCt = std::get(results[0].value); - Plaintext resultPt; - setup.cc->Decrypt(setup.keyPair.secretKey, resultCt, &resultPt); - resultPt->SetLength(vec1.size()); - - auto resultVec = resultPt->GetPackedValue(); - for (size_t i = 0; i < vec1.size(); i++) { - EXPECT_EQ(resultVec[i], vec1[i] - vec2[i]); - } -} - -TEST(InterpreterTest, TestOpenfheMul) { - CryptoSetup setup(2); // Need depth 2 for multiplication - - std::vector vec1 = {2, 3, 4, 5}; - std::vector vec2 = {3, 4, 5, 6}; - auto pt1 = setup.cc->MakePackedPlaintext(vec1); - auto pt2 = setup.cc->MakePackedPlaintext(vec2); - - auto ct1 = setup.cc->Encrypt(setup.keyPair.publicKey, pt1); - auto ct2 = setup.cc->Encrypt(setup.keyPair.publicKey, pt2); - - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -module { - func.func @main(%cc: !openfhe.crypto_context, %ct1: !ct, %ct2: !ct) -> !ct { - %result = openfhe.mul %cc, %ct1, %ct2 : (!openfhe.crypto_context, !ct, !ct) -> !ct - return %result : !ct - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(setup.cc), - TypedCppValue(ct1), TypedCppValue(ct2)}; - std::vector results = interpreter.interpret("main", inputs); - - auto resultCt = std::get(results[0].value); - Plaintext resultPt; - setup.cc->Decrypt(setup.keyPair.secretKey, resultCt, &resultPt); - resultPt->SetLength(vec1.size()); - - auto resultVec = resultPt->GetPackedValue(); - for (size_t i = 0; i < vec1.size(); i++) { - EXPECT_EQ(resultVec[i], vec1[i] * vec2[i]); - } -} - -TEST(InterpreterTest, TestOpenfheMulPlain) { - CryptoSetup setup; - - std::vector vec1 = {2, 3, 4, 5}; - std::vector vec2 = {10, 10, 10, 10}; - auto pt1 = setup.cc->MakePackedPlaintext(vec1); - auto pt2 = setup.cc->MakePackedPlaintext(vec2); - - auto ct1 = setup.cc->Encrypt(setup.keyPair.publicKey, pt1); - - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -module { - func.func @main(%cc: !openfhe.crypto_context, %ct: !ct, %pt: !pt) -> !ct { - %result = openfhe.mul_plain %cc, %ct, %pt : (!openfhe.crypto_context, !ct, !pt) -> !ct - return %result : !ct - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(setup.cc), - TypedCppValue(ct1), TypedCppValue(pt2)}; - std::vector results = interpreter.interpret("main", inputs); - - auto resultCt = std::get(results[0].value); - Plaintext resultPt; - setup.cc->Decrypt(setup.keyPair.secretKey, resultCt, &resultPt); - resultPt->SetLength(vec1.size()); - - auto resultVec = resultPt->GetPackedValue(); - for (size_t i = 0; i < vec1.size(); i++) { - EXPECT_EQ(resultVec[i], vec1[i] * vec2[i]); - } -} - -TEST(InterpreterTest, TestOpenfheNegate) { - CryptoSetup setup; - - std::vector vec = {5, 10, 15, 20}; - auto pt = setup.cc->MakePackedPlaintext(vec); - auto ct = setup.cc->Encrypt(setup.keyPair.publicKey, pt); - - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -module { - func.func @main(%cc: !openfhe.crypto_context, %ct: !ct) -> !ct { - %result = openfhe.negate %cc, %ct : (!openfhe.crypto_context, !ct) -> !ct - return %result : !ct - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(setup.cc), - TypedCppValue(ct)}; - std::vector results = interpreter.interpret("main", inputs); - - auto resultCt = std::get(results[0].value); - Plaintext resultPt; - setup.cc->Decrypt(setup.keyPair.secretKey, resultCt, &resultPt); - resultPt->SetLength(vec.size()); - - auto resultVec = resultPt->GetPackedValue(); - for (size_t i = 0; i < vec.size(); i++) { - EXPECT_EQ(resultVec[i], -vec[i]); - } -} - -TEST(InterpreterTest, TestOpenfheSquare) { - CryptoSetup setup(2); // Need depth for squaring - - std::vector vec = {2, 3, 4, 5}; - auto pt = setup.cc->MakePackedPlaintext(vec); - auto ct = setup.cc->Encrypt(setup.keyPair.publicKey, pt); - - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -module { - func.func @main(%cc: !openfhe.crypto_context, %ct: !ct) -> !ct { - %result = openfhe.square %cc, %ct : (!openfhe.crypto_context, !ct) -> !ct - return %result : !ct - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(setup.cc), - TypedCppValue(ct)}; - std::vector results = interpreter.interpret("main", inputs); - - auto resultCt = std::get(results[0].value); - Plaintext resultPt; - setup.cc->Decrypt(setup.keyPair.secretKey, resultCt, &resultPt); - resultPt->SetLength(vec.size()); - - auto resultVec = resultPt->GetPackedValue(); - for (size_t i = 0; i < vec.size(); i++) { - EXPECT_EQ(resultVec[i], vec[i] * vec[i]); - } -} - -TEST(InterpreterTest, TestOpenfheRot) { - CryptoSetup setup; - - // Generate rotation keys - setup.cc->EvalRotateKeyGen(setup.keyPair.secretKey, {2}); - - std::vector vec = {1, 2, 3, 4, 5, 6, 7, 8}; - // Cyclically replicate the vector to fill the slots - std::vector replicatedVec; - int64_t numSlots = setup.cc->GetRingDimension() / 2; - replicatedVec.reserve(numSlots); - for (int i = 0; i < numSlots; i++) - replicatedVec.push_back(vec[i % vec.size()]); - - auto pt = setup.cc->MakePackedPlaintext(replicatedVec); - auto ct = setup.cc->Encrypt(setup.keyPair.publicKey, pt); - - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -module { - func.func @main(%cc: !openfhe.crypto_context, %ct: !ct) -> !ct { - %result = openfhe.rot %cc, %ct {static_shift = 2 : i32} : (!openfhe.crypto_context, !ct) -> !ct - return %result : !ct - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(setup.cc), - TypedCppValue(ct)}; - std::vector results = interpreter.interpret("main", inputs); - - auto resultCt = std::get(results[0].value); - Plaintext resultPt; - setup.cc->Decrypt(setup.keyPair.secretKey, resultCt, &resultPt); - resultPt->SetLength(vec.size()); - - auto resultVec = resultPt->GetPackedValue(); - // Rotation by 2 should shift: [1,2,3,4,5,6,7,8] -> [3,4,5,6,7,8,1,2] - EXPECT_EQ(resultVec[0], 3); - EXPECT_EQ(resultVec[1], 4); - EXPECT_EQ(resultVec[6], 1); - EXPECT_EQ(resultVec[7], 2); -} - -TEST(InterpreterTest, TestOpenfheRotDynamicShift) { - CryptoSetup setup; - - // Generate rotation keys - setup.cc->EvalRotateKeyGen(setup.keyPair.secretKey, {2}); - - std::vector vec = {1, 2, 3, 4, 5, 6, 7, 8}; - // Cyclically replicate the vector to fill the slots - std::vector replicatedVec; - int64_t numSlots = setup.cc->GetRingDimension() / 2; - replicatedVec.reserve(numSlots); - for (int i = 0; i < numSlots; i++) - replicatedVec.push_back(vec[i % vec.size()]); - - auto pt = setup.cc->MakePackedPlaintext(replicatedVec); - auto ct = setup.cc->Encrypt(setup.keyPair.publicKey, pt); - - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -module { - func.func @main(%cc: !openfhe.crypto_context, %ct: !ct) -> !ct { - %0 = arith.constant 2 : i32 - %result = openfhe.rot %cc, %ct, %0 : (!openfhe.crypto_context, !ct, i32) -> !ct - return %result : !ct - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(setup.cc), - TypedCppValue(ct)}; - std::vector results = interpreter.interpret("main", inputs); - - auto resultCt = std::get(results[0].value); - Plaintext resultPt; - setup.cc->Decrypt(setup.keyPair.secretKey, resultCt, &resultPt); - resultPt->SetLength(vec.size()); - - auto resultVec = resultPt->GetPackedValue(); - // Rotation by 2 should shift: [1,2,3,4,5,6,7,8] -> [3,4,5,6,7,8,1,2] - EXPECT_EQ(resultVec[0], 3); - EXPECT_EQ(resultVec[1], 4); - EXPECT_EQ(resultVec[6], 1); - EXPECT_EQ(resultVec[7], 2); -} - -TEST(InterpreterTest, TestOpenfheMakePackedPlaintext) { - CryptoSetup setup; - - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -module { - func.func @main(%cc: !openfhe.crypto_context, %vec: tensor<4xi32>) -> !pt { - %result = openfhe.make_packed_plaintext %cc, %vec : (!openfhe.crypto_context, tensor<4xi32>) -> !pt - return %result : !pt - } -} -)mlir"; - auto module = parseTest(&context, mlirStr); - - Interpreter interpreter(module.get()); - - std::vector vec = {10, 20, 30, 40}; - std::vector inputs = {TypedCppValue(setup.cc), - TypedCppValue(vec)}; - std::vector results = interpreter.interpret("main", inputs); - - EXPECT_EQ(results.size(), 1); - auto resultPt = std::get(results[0].value); - resultPt->SetLength(vec.size()); - - auto resultVec = resultPt->GetPackedValue(); - EXPECT_EQ(resultVec.size(), vec.size()); - for (size_t i = 0; i < vec.size(); i++) { - EXPECT_EQ(resultVec[i], vec[i]); - } -} - -TEST(InterpreterTest, TestOpenfheEncryptDecrypt) { - CryptoSetup setup; - - std::vector vec = {7, 14, 21, 28}; - auto pt = setup.cc->MakePackedPlaintext(vec); - - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -module { - func.func @main(%cc: !openfhe.crypto_context, %pt: !pt, %pk: !openfhe.public_key, %sk: !openfhe.private_key) -> !pt { - %ct = openfhe.encrypt %cc, %pt, %pk : (!openfhe.crypto_context, !pt, !openfhe.public_key) -> !ct - %result = openfhe.decrypt %cc, %ct, %sk : (!openfhe.crypto_context, !ct, !openfhe.private_key) -> !pt - return %result : !pt - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(setup.cc), - TypedCppValue(pt), - TypedCppValue(setup.keyPair.publicKey), - TypedCppValue(setup.keyPair.secretKey)}; - std::vector results = interpreter.interpret("main", inputs); - - auto resultPt = std::get(results[0].value); - resultPt->SetLength(vec.size()); - - auto resultVec = resultPt->GetPackedValue(); - for (size_t i = 0; i < vec.size(); i++) { - EXPECT_EQ(resultVec[i], vec[i]); - } -} - -TEST(InterpreterTest, TestOpenfheRLWEDecodeBGVScalar) { - CryptoSetup setup; - - std::vector vec = {42}; - auto pt = setup.cc->MakePackedPlaintext(vec); - - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -module attributes {scheme.bgv} { - func.func @main(%pt: !pt) -> i32 { - %result = openfhe.decode %pt : !pt -> i32 - return %result : i32 - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(pt)}; - std::vector results = interpreter.interpret("main", inputs); - - EXPECT_EQ(results.size(), 1); - EXPECT_EQ(std::get(results[0].value), 42); -} - -TEST(InterpreterTest, TestOpenfheRLWEDecodeBGVTensor) { - CryptoSetup setup; - - std::vector vec = {1, 2, 3, 4, 5, 6, 7, 8}; - auto pt = setup.cc->MakePackedPlaintext(vec); - - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -module attributes {scheme.bgv} { - func.func @main(%pt: !pt) -> tensor<8xi32> { - %result = openfhe.decode %pt : !pt -> tensor<8xi32> - return %result : tensor<8xi32> - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(pt)}; - std::vector results = interpreter.interpret("main", inputs); - - EXPECT_EQ(results.size(), 1); - auto resultVec = - *std::get>>(results[0].value); - EXPECT_EQ(resultVec.size(), vec.size()); - for (size_t i = 0; i < vec.size(); i++) { - EXPECT_EQ(resultVec[i], vec[i]); - } -} - -// Helper function to set up a basic CKKS crypto context for testing -struct CKKSSetup { - CryptoContext cc; - KeyPair keyPair; - - CKKSSetup(uint32_t multDepth = 2) { - CCParams parameters; - parameters.SetMultiplicativeDepth(multDepth); - parameters.SetScalingModSize(50); - parameters.SetBatchSize(8); - - cc = GenCryptoContext(parameters); - cc->Enable(PKE); - cc->Enable(KEYSWITCH); - cc->Enable(LEVELEDSHE); - - keyPair = cc->KeyGen(); - cc->EvalMultKeyGen(keyPair.secretKey); - } -}; - -TEST(InterpreterTest, TestOpenfheRLWEDecodeCKKSScalar) { - CKKSSetup setup; - - std::vector vec = {3.14}; - auto pt = setup.cc->MakeCKKSPackedPlaintext(vec); - - MLIRContext context; - initContext(context); - std::string mlirStr = R"mlir( -!pt = !openfhe.plaintext - -module attributes {scheme.ckks} { - func.func @main(%pt: !pt) -> f32 { - %result = openfhe.decode_ckks %pt : !pt -> f32 - return %result : f32 - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(pt)}; - std::vector results = interpreter.interpret("main", inputs); - - EXPECT_EQ(results.size(), 1); - // Use approximate comparison for floating point - EXPECT_NEAR(std::get(results[0].value), 3.14f, 0.01f); -} - -TEST(InterpreterTest, TestOpenfheRLWEDecodeCKKSTensor) { - CKKSSetup setup; - - std::vector vec = {1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8}; - auto pt = setup.cc->MakeCKKSPackedPlaintext(vec); - - MLIRContext context; - initContext(context); - std::string mlirStr = R"mlir( -!pt = !openfhe.plaintext - -module attributes {scheme.ckks} { - func.func @main(%pt: !pt) -> tensor<8xf32> { - %result = openfhe.decode_ckks %pt : !pt -> tensor<8xf32> - return %result : tensor<8xf32> - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(pt)}; - std::vector results = interpreter.interpret("main", inputs); - - EXPECT_EQ(results.size(), 1); - auto resultVec = - *std::get>>(results[0].value); - EXPECT_EQ(resultVec.size(), vec.size()); - for (size_t i = 0; i < vec.size(); i++) { - EXPECT_NEAR(resultVec[i], static_cast(vec[i]), 0.01f); - } -} - -TEST(InterpreterTest, TestDoubleConstant) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main() -> f64 { - %c = arith.constant 3.14159265358979 : f64 - return %c : f64 - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector results = interpreter.interpret(entryFunction, {}); - EXPECT_EQ(results.size(), 1); - EXPECT_NEAR(std::get(results[0].value), 3.14159265358979, 1e-10); -} - -TEST(InterpreterTest, TestFloatConstant) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main() -> f32 { - %c = arith.constant 3.14 : f32 - return %c : f32 - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector results = interpreter.interpret(entryFunction, {}); - EXPECT_EQ(results.size(), 1); - EXPECT_NEAR(std::get(results[0].value), 3.14f, 0.01f); -} - -TEST(InterpreterTest, TestDoubleTensorConstant) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main() -> tensor<3xf64> { - %c = arith.constant dense<[1.1, 2.2, 3.3]> : tensor<3xf64> - return %c : tensor<3xf64> - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector results = interpreter.interpret(entryFunction, {}); - EXPECT_EQ(results.size(), 1); - auto vec = *std::get>>(results[0].value); - EXPECT_EQ(vec.size(), 3); - EXPECT_NEAR(vec[0], 1.1, 0.01); - EXPECT_NEAR(vec[1], 2.2, 0.01); - EXPECT_NEAR(vec[2], 3.3, 0.01); -} - -TEST(InterpreterTest, TestDenseResourceConstant) { - MLIRContext context; - initContext(context); - std::vector expected = {1.1f, 2.2f, 3.3f, 0.0f}; - - auto module = parseTest(&context, R"mlir( - module { - func.func @main() -> tensor<4xf32> { - %c = arith.constant dense_resource : tensor<4xf32> - return %c : tensor<4xf32> - } - } - {-# - dialect_resources: { - builtin: { - dense_elements_f32: "0x40000000CDCC8C3FCDCC0C403333534000000000" - } - } - #-} - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector results = interpreter.interpret(entryFunction, {}); - EXPECT_EQ(results.size(), 1); - auto vec = *std::get>>(results[0].value); - EXPECT_EQ(vec.size(), 4); - for (size_t i = 0; i < 4; i++) { - EXPECT_NEAR(vec[i], expected[i], 1e-10); - } -} - -TEST(InterpreterTest, TestDoubleTensorSplat) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%val: f64) -> tensor<4xf64> { - %t = tensor.splat %val : tensor<4xf64> - return %t : tensor<4xf64> - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(3.14)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - auto vec = *std::get>>(results[0].value); - EXPECT_EQ(vec.size(), 4); - for (size_t i = 0; i < 4; ++i) { - EXPECT_NEAR(vec[i], 3.14, 0.01); - } -} - -TEST(InterpreterTest, TestExtFOpFloatToDouble) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%a: f32) -> f64 { - %c = arith.extf %a : f32 to f64 - return %c : f64 - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector inputs = {TypedCppValue(3.14f)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - EXPECT_NEAR(std::get(results[0].value), 3.14, 0.01); -} - -TEST(InterpreterTest, TestExtFOpFloatVectorToDoubleVector) { - MLIRContext context; - initContext(context); - auto module = parseTest(&context, R"mlir( - module { - func.func @main(%a: tensor<3xf32>) -> tensor<3xf64> { - %c = arith.extf %a : tensor<3xf32> to tensor<3xf64> - return %c : tensor<3xf64> - } - } - )mlir"); - Interpreter interpreter(module.get()); - std::string entryFunction = "main"; - std::vector a = {1.1f, 2.2f, 3.3f}; - std::vector inputs = {TypedCppValue(a)}; - std::vector results = - interpreter.interpret(entryFunction, inputs); - EXPECT_EQ(results.size(), 1); - auto vec = *std::get>>(results[0].value); - EXPECT_EQ(vec.size(), 3); - for (size_t i = 0; i < 3; ++i) { - EXPECT_NEAR(vec[i], a[i], 0.01); - } -} - -TEST(InterpreterTest, TestOpenfheRLWEDecodeCKKSScalarDouble) { - CKKSSetup setup; - - std::vector vec = {3.14159265358979}; - auto pt = setup.cc->MakeCKKSPackedPlaintext(vec); - - MLIRContext context; - initContext(context); - std::string mlirStr = R"mlir( -!pt = !openfhe.plaintext - -module attributes {scheme.ckks} { - func.func @main(%pt: !pt) -> f64 { - %result = openfhe.decode_ckks %pt : !pt -> f64 - return %result : f64 - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(pt)}; - std::vector results = interpreter.interpret("main", inputs); - - EXPECT_EQ(results.size(), 1); - // Use approximate comparison for floating point - EXPECT_NEAR(std::get(results[0].value), 3.14159265358979, 0.01); -} - -TEST(InterpreterTest, TestOpenfheRLWEDecodeCKKSTensorDouble) { - CKKSSetup setup; - - std::vector vec = {1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8}; - auto pt = setup.cc->MakeCKKSPackedPlaintext(vec); - - MLIRContext context; - initContext(context); - std::string mlirStr = R"mlir( -!pt = !openfhe.plaintext - -module attributes {scheme.ckks} { - func.func @main(%pt: !pt) -> tensor<8xf64> { - %result = openfhe.decode_ckks %pt : !pt -> tensor<8xf64> - return %result : tensor<8xf64> - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = {TypedCppValue(pt)}; - std::vector results = interpreter.interpret("main", inputs); - - EXPECT_EQ(results.size(), 1); - auto resultVec = - *std::get>>(results[0].value); - EXPECT_EQ(resultVec.size(), vec.size()); - for (size_t i = 0; i < vec.size(); i++) { - EXPECT_NEAR(resultVec[i], vec[i], 0.01); - } -} - -TEST(InterpreterTest, TestOpenfheParallelFastRotation) { - CryptoSetup setup; - - // Generate rotation keys - setup.cc->EvalRotateKeyGen(setup.keyPair.secretKey, {1, 2, 3, 4}); - - std::vector vec = {1, 2, 3, 4, 5, 6, 7, 8}; - // Cyclically replicate the vector to fill the slots - std::vector replicatedVec; - int64_t numSlots = setup.cc->GetRingDimension() / 2; - replicatedVec.reserve(numSlots); - for (int i = 0; i < numSlots; i++) - replicatedVec.push_back(vec[i % vec.size()]); - - auto pt = setup.cc->MakePackedPlaintext(replicatedVec); - auto ct = setup.cc->Encrypt(setup.keyPair.publicKey, pt); - - std::vector indices = {1, 2, 3, 4}; - - MLIRContext context; - initContext(context); - std::string mlirStr = std::string(kLWETypesHeader) + R"mlir( -!cc = !openfhe.crypto_context -!digit_decomp = !openfhe.digit_decomp - -module attributes {scheme.ckks} { - func.func @main(%cc: !cc, %ct: !ct, %indices: tensor<4xindex>) -> tensor<4x!ct> { - %digit_decomp = openfhe.fast_rotation_precompute %cc, %ct : (!cc, !ct) -> !digit_decomp - %0 = tensor.empty() : tensor<4x!ct> - %1 = scf.forall (%arg0) in (4) shared_outs(%arg1 = %0) -> (tensor<4x!ct>) { - %extracted_9 = tensor.extract %indices[%arg0] : tensor<4xindex> - %ct_10 = openfhe.fast_rotation %cc, %ct, %extracted_9, %digit_decomp {cyclotomicOrder = 64 : index} : (!cc, !ct, index, !digit_decomp) -> !ct - %from_elements_11 = tensor.from_elements %ct_10 : tensor<1x!ct> - scf.forall.in_parallel { - tensor.parallel_insert_slice %from_elements_11 into %arg1[%arg0] [1] [1] : tensor<1x!ct> into tensor<4x!ct> - } - } - %c0 = arith.constant 0 : index - %extracted = tensor.extract %1[%c0] : tensor<4x!ct> - %c1_0 = arith.constant 1 : index - %extracted_1 = tensor.extract %1[%c1_0] : tensor<4x!ct> - %c2_2 = arith.constant 2 : index - %extracted_3 = tensor.extract %1[%c2_2] : tensor<4x!ct> - %c3_4 = arith.constant 3 : index - %extracted_5 = tensor.extract %1[%c3_4] : tensor<4x!ct> - %2 = tensor.from_elements %extracted, %extracted_1, %extracted_3, %extracted_5 : tensor<4x!ct> - return %2 : tensor<4x!ct> - } -} -)mlir"; - auto module = parse(&context, mlirStr); - - Interpreter interpreter(module.get()); - std::vector inputs = { - TypedCppValue(setup.cc), TypedCppValue(ct), TypedCppValue(indices)}; - std::vector results = interpreter.interpret("main", inputs); - - EXPECT_EQ(results.size(), 1); - auto resultVec = - *std::get>>(results[0].value); - EXPECT_EQ(resultVec.size(), 4); - - std::vector> expected = {{2, 3, 4, 5, 6, 7, 8, 1}, - {3, 4, 5, 6, 7, 8, 1, 2}, - {4, 5, 6, 7, 8, 1, 2, 3}, - {5, 6, 7, 8, 1, 2, 3, 4}}; - for (int i = 0; i < 4; i++) { - // Decrypt and verify - Plaintext resultPt; - setup.cc->Decrypt(setup.keyPair.secretKey, resultVec[i], &resultPt); - resultPt->SetLength(vec.size()); - auto resultVec = resultPt->GetPackedValue(); - EXPECT_EQ(resultVec.size(), vec.size()); - EXPECT_EQ(resultVec, expected[i]); - } -} - -} // namespace openfhe -} // namespace heir -} // namespace mlir diff --git a/tests/Examples/openfhe/ckks/lenet/BUILD b/tests/Examples/openfhe/ckks/lenet/BUILD index 107547d3ee..786e28a916 100644 --- a/tests/Examples/openfhe/ckks/lenet/BUILD +++ b/tests/Examples/openfhe/ckks/lenet/BUILD @@ -1,100 +1,43 @@ -# load("@heir//tools:heir-opt.bzl", "heir_opt") -load("@pybind11_bazel//:build_defs.bzl", "pybind_extension") -load("@rules_cc//cc:cc_binary.bzl", "cc_binary") -load("@rules_cc//cc:cc_library.bzl", "cc_library") -load("@rules_python//python:py_test.bzl", "py_test") +# load("@rules_python//python:py_test.bzl", "py_test") +# load("@heir//tools:heir-openfhe.bzl", "openfhe_lib") package(default_applicable_licenses = ["@heir//:license"]) -# A binary that lets you pass in a lenet.mlir in openfhe dialect -# (pre-compiled) and runs it with timing. For performance profiling, -# use -# -# bazel build -c opt --copt=-g --linkopt=-lprofiler \ -# tests/Examples/openfhe/ckks/lenet:lenet_binary -# -# CPUPROFILE=prof.out bazel-bin/tests/Examples/openfhe/ckks/lenet/lenet_binary \ -# tests/Examples/openfhe/ckks/lenet/pre_compiled_lenet.openfhe.mlir -# -# Then a pprof invocation such as -# -# pprof --text --lines --focus=mlir::heir::openfhe::Interpreter \ -# ./bazel-bin/tests/Examples/openfhe/ckks/lenet/lenet_binary \ -# prof.out > interpreter_focus.txt -cc_binary( - name = "lenet_binary", - srcs = ["lenet_main.cpp"], - tags = [ - "manual", - "nofastbuild", - "notap", - ], - deps = [ - "@heir//lib/Target/OpenFhePke:Interpreter", - "@llvm-project//mlir:IR", - "@llvm-project//mlir:Parser", - "@openfhe//:core", - "@openfhe//:pke", - ], -) - -# TODO(#2702): Re-enable once pooling is fully supported. -# heir_opt( -# name = "lenet_mlir_opt", -# src = "@heir//tests/Examples/common/lenet:lenet.mlir", -# generated_filename = "lenet.openfhe.mlir", -# pass_flags = [ +# openfhe_lib( +# name = "lenet_openfhe", +# generated_lib_header = "lenet_openfhe_lib.inc.h", +# heir_opt_flags = [ # "--annotate-module=backend=openfhe scheme=ckks", -# "--torch-linalg-to-ckks=min-slot-count=1024", +# "--torch-linalg-to-ckks=min-slot-count=1024 greedy-level-budget=15 greedy-modulus-switch-after-mul=true experimental-disable-loop-unroll=true first-mod-bits=30 scaling-mod-bits=24", # "--scheme-to-openfhe", # ], +# mlir_src = "lenet.mlir", +# pybind_target_name = "lenet_openfhe_pybind", # tags = [ # "nofastbuild", # "requires-mem:28g", # ], # ) - -cc_library( - name = "interpreter_shim", - srcs = ["interpreter_shim.cpp"], - hdrs = ["interpreter_shim.h"], - deps = [ - "@heir//lib/Target/OpenFhePke:Interpreter", - "@llvm-project//mlir:IR", - "@llvm-project//mlir:Support", - "@openfhe//:core", - "@openfhe//:pke", - ], -) - -pybind_extension( - name = "lenet_interpreter", - srcs = ["interpreter_bindings.cpp"], - deps = [ - ":interpreter_shim", - ], -) - -py_test( - name = "lenet_test", - size = "enormous", - srcs = ["lenet_test.py"], - data = [ - ":lenet.openfhe.mlir", - "@heir//tests/Examples/common/mnist/data:t10k-images-idx3-ubyte", - "@heir//tests/Examples/common/mnist/data:t10k-labels-idx1-ubyte", - ], - main = "lenet_test.py", - tags = [ - "manual", - "nofastbuild", # openfhe is slow unless -c opt - "notap", - "requires-mem:28g", - ], - deps = [ - ":lenet_interpreter", - "@abseil-py//absl/testing:absltest", - "@heir_pip_deps//numpy", - "@heir_pip_deps//torch", - ], -) +# +# py_test( +# name = "lenet_test", +# size = "large", +# srcs = ["lenet_test.py"], +# data = [ +# ":lenet.openfhe.mlir", +# "@heir//tests/Examples/common/mnist/data:t10k-images-idx3-ubyte", +# "@heir//tests/Examples/common/mnist/data:t10k-labels-idx1-ubyte", +# ], +# main = "lenet_test.py", +# strict_deps = False, +# tags = [ +# "nofastbuild", # openfhe is slow unless -c opt +# "requires-mem:28g", +# ], +# deps = [ +# ":lenet_openfhe_pybind", +# "@abseil-py//absl/testing:absltest", +# "@heir_pip_deps//numpy", +# "@heir_pip_deps//torch", +# ], +# ) diff --git a/tests/Examples/openfhe/ckks/lenet/interpreter_bindings.cpp b/tests/Examples/openfhe/ckks/lenet/interpreter_bindings.cpp deleted file mode 100644 index 64baa48084..0000000000 --- a/tests/Examples/openfhe/ckks/lenet/interpreter_bindings.cpp +++ /dev/null @@ -1,12 +0,0 @@ -// A shim around the HEIR OpenFHE interpreter to provide pybind11 bindings. -#include -#include - -#include "tests/Examples/openfhe/ckks/lenet/interpreter_shim.h" - -namespace py = pybind11; - -PYBIND11_MODULE(lenet_interpreter, m) { - m.def("lenet_interpreter", &lenet_interpreter, - py::call_guard()); -} diff --git a/tests/Examples/openfhe/ckks/lenet/interpreter_shim.cpp b/tests/Examples/openfhe/ckks/lenet/interpreter_shim.cpp deleted file mode 100644 index cc854e25cf..0000000000 --- a/tests/Examples/openfhe/ckks/lenet/interpreter_shim.cpp +++ /dev/null @@ -1,64 +0,0 @@ -#include -#include -#include -#include -#include -#include - -#include "lib/Target/OpenFhePke/Interpreter.h" -#include "mlir/include/mlir/IR/BuiltinOps.h" // from @llvm-project -#include "mlir/include/mlir/IR/MLIRContext.h" // from @llvm-project -#include "mlir/include/mlir/IR/OwningOpRef.h" // from @llvm-project -#include "mlir/include/mlir/Support/LLVM.h" // from @llvm-project -#include "src/core/include/lattice/hal/lat-backend.h" // from @openfhe -#include "src/pke/include/cryptocontext-fwd.h" // from @openfhe - -using namespace lbcrypto; -using namespace mlir::heir::openfhe; -using CryptoContextT = CryptoContext; -using floatvec = std::shared_ptr>; -using mlir::MLIRContext; -using mlir::ModuleOp; - -std::pair, double> lenet_interpreter( - const std::string& mlirSrc, const std::vector& input) { - // Load the MLIR module from a file - MLIRContext context; - initContext(context); - mlir::OwningOpRef module = parse(&context, mlirSrc); - Interpreter interpreter(module.get()); - - std::cout << "Generating crypto context" << std::endl; - TypedCppValue ccInitial = - interpreter.interpret("lenet__generate_crypto_context", {})[0]; - - auto keyPair = std::get(ccInitial.value)->KeyGen(); - auto publicKey = keyPair.publicKey; - auto secretKey = keyPair.secretKey; - std::vector args = {ccInitial, TypedCppValue(secretKey)}; - std::cout << "Configuring crypto context" << std::endl; - TypedCppValue cc = std::move( - interpreter.interpret("lenet__configure_crypto_context", args)[0]); - - std::cout << "Encrypting input" << std::endl; - TypedCppValue arg0Enc = interpreter.interpret( - "lenet__encrypt__arg0", - {cc, TypedCppValue(input), TypedCppValue(publicKey)})[0]; - - std::cout << "Running module" << std::endl; - auto start = std::chrono::high_resolution_clock::now(); - TypedCppValue outputEncrypted = - interpreter.interpret("lenet", {cc, arg0Enc})[0]; - auto end = std::chrono::high_resolution_clock::now(); - - std::cout << "Decrypting output" << std::endl; - std::vector actualVal = - interpreter.interpret("lenet__decrypt__result0", - {cc, outputEncrypted, TypedCppValue(secretKey)}); - - auto resultVec = std::get(actualVal[0].value); - double duration = - std::chrono::duration_cast(end - start) - .count(); - return {*resultVec, duration}; -} diff --git a/tests/Examples/openfhe/ckks/lenet/interpreter_shim.h b/tests/Examples/openfhe/ckks/lenet/interpreter_shim.h deleted file mode 100644 index c222756549..0000000000 --- a/tests/Examples/openfhe/ckks/lenet/interpreter_shim.h +++ /dev/null @@ -1,15 +0,0 @@ -// A wrapper around an invocation of the C++ interpreter for OpenFHE customized -// for the lenet model test. This is a workaround due to long compilation times -// of generated C++ functions with >> 50k lines. - -#include -#include -#include - -/// Run the lenet model on the given input data, returning the output vector -/// and the time taken in seconds. -/// -/// @param mlirSrc The MLIR OpenFHE dialect IR to run, as a string -/// @param inputData The input vector to pass to the model. -std::pair, double> lenet_interpreter( - const std::string& mlirSrc, const std::vector& inputData); diff --git a/tests/Examples/openfhe/ckks/lenet/lenet_main.cpp b/tests/Examples/openfhe/ckks/lenet/lenet_main.cpp deleted file mode 100644 index 846b0dcae7..0000000000 --- a/tests/Examples/openfhe/ckks/lenet/lenet_main.cpp +++ /dev/null @@ -1,90 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include - -#include "lib/Target/OpenFhePke/Interpreter.h" -#include "mlir/include/mlir/IR/BuiltinOps.h" // from @llvm-project -#include "mlir/include/mlir/IR/MLIRContext.h" // from @llvm-project -#include "mlir/include/mlir/IR/OwningOpRef.h" // from @llvm-project -#include "mlir/include/mlir/Parser/Parser.h" // from @llvm-project -#include "src/core/include/lattice/hal/lat-backend.h" // from @openfhe -#include "src/pke/include/cryptocontext-fwd.h" // from @openfhe - -namespace mlir { -namespace heir { -namespace openfhe { - -using namespace lbcrypto; -using CryptoContextT = CryptoContext; - -} // namespace openfhe -} // namespace heir -} // namespace mlir - -void reportTime(const std::string& operation, - const std::chrono::high_resolution_clock::time_point& start, - const std::chrono::high_resolution_clock::time_point& end) { - auto duration = std::chrono::duration_cast(end - start); - std::cout << operation << " time: " << duration.count() << " seconds\n"; -} - -int main(int argc, char* argv[]) { - // Generated by the bazel rule - std::string generatedModuleFile = argv[1]; - - // Load the MLIR module from a file - mlir::MLIRContext context; - mlir::heir::openfhe::initContext(context); - mlir::OwningOpRef module = - parseSourceFile(generatedModuleFile, &context); - mlir::heir::openfhe::Interpreter interpreter(module.get()); - - auto start = std::chrono::high_resolution_clock::now(); - mlir::heir::openfhe::TypedCppValue ccInitial = - interpreter.interpret("lenet__generate_crypto_context", {})[0]; - auto end = std::chrono::high_resolution_clock::now(); - reportTime("lenet__generate_crypto_context", start, end); - - auto keyPair = - std::get(ccInitial.value)->KeyGen(); - auto publicKey = keyPair.publicKey; - auto secretKey = keyPair.secretKey; - std::vector args = { - ccInitial, mlir::heir::openfhe::TypedCppValue(secretKey)}; - - start = std::chrono::high_resolution_clock::now(); - mlir::heir::openfhe::TypedCppValue cc = std::move( - interpreter.interpret("lenet__configure_crypto_context", args)[0]); - end = std::chrono::high_resolution_clock::now(); - reportTime("lenet__configure_crypto_context", start, end); - - // Run one inference - std::vector inputVector(32 * 32, 0.0f); - start = std::chrono::high_resolution_clock::now(); - mlir::heir::openfhe::TypedCppValue arg0Enc = interpreter.interpret( - "lenet__encrypt__arg0", - {cc, mlir::heir::openfhe::TypedCppValue(inputVector), - mlir::heir::openfhe::TypedCppValue(publicKey)})[0]; - end = std::chrono::high_resolution_clock::now(); - reportTime("lenet__encrypt__arg0", start, end); - - start = std::chrono::high_resolution_clock::now(); - mlir::heir::openfhe::TypedCppValue outputEncrypted = - interpreter.interpret("lenet", {cc, arg0Enc})[0]; - end = std::chrono::high_resolution_clock::now(); - reportTime("lenet", start, end); - - start = std::chrono::high_resolution_clock::now(); - std::vector actualVal = - interpreter.interpret( - "lenet__decrypt__result0", - {cc, outputEncrypted, mlir::heir::openfhe::TypedCppValue(secretKey)}); - end = std::chrono::high_resolution_clock::now(); - reportTime("lenet__decrypt__result0", start, end); - - return 0; -} diff --git a/tests/Examples/openfhe/ckks/lenet/lenet_test.py b/tests/Examples/openfhe/ckks/lenet/lenet_test.py index b293150f50..d50e6e5f48 100644 --- a/tests/Examples/openfhe/ckks/lenet/lenet_test.py +++ b/tests/Examples/openfhe/ckks/lenet/lenet_test.py @@ -1,11 +1,11 @@ import os - +import time import numpy as np import torch from torch.utils.data import DataLoader, Dataset from absl.testing import absltest -from tests.Examples.openfhe.ckks.lenet import lenet_interpreter as lenet +import tests.Examples.openfhe.ckks.lenet_openfhe_pybind as lenet # fmt: off DATA_PATH = "tests/Examples/openfhe/ckks/mnist/data" @@ -86,6 +86,14 @@ def test_run_test(self): shuffle=False, # SequentialSampler equivalent ) + crypto_context = lenet.lenet__generate_crypto_context() + key_pair = crypto_context.KeyGen() + public_key = key_pair.publicKey + secret_key = key_pair.secretKey + crypto_context = lenet.lenet__configure_crypto_context( + crypto_context, secret_key + ) + total = 4 correct = 0 samples_processed = 0 @@ -96,11 +104,22 @@ def test_run_test(self): input_tensor = batch_data.contiguous() # (1, 1, 28, 28) input_vector = input_tensor.flatten().tolist() - (output, runtime_ms) = lenet.lenet_interpreter(mlir_src, input_vector) - print(f"runtime (ms): {runtime_ms}") + input_encrypted = lenet.lenet__encrypt__arg1( + crypto_context, input_vector, public_key + ) + + start_time = time.time() + output_encrypted = lenet.lenet(crypto_context, input_encrypted) + end_time = time.time() + + time_elapsed_ms = (end_time - start_time) * 1000.0 + print(f"CPU time used: {time_elapsed_ms:.2f} ms") label = batch_target.item() - max_id = max(range(len(output)), key=lambda index: output[index]) + output = lenet.lenet__decrypt__result0( + crypto_context, output_encrypted, secret_key + ) + max_id = max(range(len(output)), key=lambda i: output[i]) if max_id == label: correct += 1 diff --git a/tests/Examples/openfhe/ckks/mnist/mnist_test.py b/tests/Examples/openfhe/ckks/mnist/mnist_test.py index e37cded22b..edc08ac93c 100644 --- a/tests/Examples/openfhe/ckks/mnist/mnist_test.py +++ b/tests/Examples/openfhe/ckks/mnist/mnist_test.py @@ -145,12 +145,11 @@ def test_run_test(self): time_elapsed_ms = (end_time - start_time) * 1000.0 print(f"CPU time used: {time_elapsed_ms:.2f} ms") - output = [0.0] * 10 output = mnist.mnist__decrypt__result0( crypto_context, output_encrypted, secret_key ) label = batch_target.item() - max_id = max(range(len(output)), key=lambda index: output[index]) + max_id = max(range(len(output)), key=lambda i: output[i]) # NOTE: For the test to pass with the default placeholder 'output', # the `mnist__decrypt__result0` function must be fixed to return an diff --git a/tests/Examples/openfhe/ckks/parallelization/BUILD b/tests/Examples/openfhe/ckks/parallelization/BUILD index f536cc23a2..1878efc229 100644 --- a/tests/Examples/openfhe/ckks/parallelization/BUILD +++ b/tests/Examples/openfhe/ckks/parallelization/BUILD @@ -1,4 +1,4 @@ -load("@heir//tests/Examples/openfhe:test.bzl", "openfhe_end_to_end_test", "openfhe_interpreter_test") +load("@heir//tests/Examples/openfhe:test.bzl", "openfhe_end_to_end_test") package(default_applicable_licenses = ["@heir//:license"]) @@ -36,21 +36,3 @@ openfhe_end_to_end_test( tags = ["notap"], test_src = "parallelization_skip_test.cpp", ) - -openfhe_interpreter_test( - name = "parallelization_interpreter_test", - generated_heir_opt_filename = "module.openfhe.mlir", - heir_opt_flags = [ - "--openfhe-fast-rotation-precompute", - "--boolean-vectorize", - "--canonicalize", - "--cse", - "--openfhe-alloc-to-inplace", - ], - mlir_src = "@heir//tests/Examples/openfhe/ckks/parallelization:parallelization.mlir", - test_src = "parallelization_interpreter_test.cpp", - deps = [ - "@llvm-project//mlir:IR", - "@llvm-project//mlir:Parser", - ], -) diff --git a/tests/Examples/openfhe/ckks/parallelization/parallelization_interpreter_test.cpp b/tests/Examples/openfhe/ckks/parallelization/parallelization_interpreter_test.cpp deleted file mode 100644 index b6519eaa1f..0000000000 --- a/tests/Examples/openfhe/ckks/parallelization/parallelization_interpreter_test.cpp +++ /dev/null @@ -1,65 +0,0 @@ -#include -#include -#include - -#include "gtest/gtest.h" // from @googletest -#include "lib/Target/OpenFhePke/Interpreter.h" -#include "mlir/include/mlir/IR/BuiltinOps.h" // from @llvm-project -#include "mlir/include/mlir/IR/MLIRContext.h" // from @llvm-project -#include "mlir/include/mlir/IR/OwningOpRef.h" // from @llvm-project -#include "mlir/include/mlir/Parser/Parser.h" // from @llvm-project -#include "mlir/include/mlir/Support/LLVM.h" // from @llvm-project -#include "src/core/include/lattice/hal/lat-backend.h" // from @openfhe -#include "src/pke/include/ciphertext-fwd.h" // from @openfhe -#include "src/pke/include/cryptocontext-fwd.h" // from @openfhe -#include "src/pke/include/encoding/plaintext-fwd.h" // from @openfhe - -namespace mlir { -namespace heir { -namespace openfhe { - -using namespace lbcrypto; -using CryptoContextT = CryptoContext; - -// Copybara manages this declaration via regex -const std::string WORKSPACE_BASE = ""; - -TEST(ParallelizationInterpreterTest, RunTest) { - // Generated by the bazel rule - std::string generatedModuleFile = - WORKSPACE_BASE + - "tests/Examples/openfhe/ckks/parallelization/module.openfhe.mlir"; - - // Load the MLIR module from a file - MLIRContext context; - initContext(context); - OwningOpRef module = - parseSourceFile(generatedModuleFile, &context); - Interpreter interpreter(module.get()); - - TypedCppValue ccInitial = - interpreter.interpret("rotations__generate_crypto_context", {})[0]; - - auto keyPair = std::get(ccInitial.value)->KeyGen(); - auto publicKey = keyPair.publicKey; - auto secretKey = keyPair.secretKey; - std::vector args = {ccInitial, TypedCppValue(secretKey)}; - TypedCppValue cc = std::move( - interpreter.interpret("rotations__configure_crypto_context", args)[0]); - - std::vector arg0Vals(1024, 1.0); - CryptoContextT cryptoContext = std::get(cc.value); - lbcrypto::Plaintext ptxt = - cryptoContext->MakeCKKSPackedPlaintext(arg0Vals, 1); - ptxt->SetLength(1024); - Ciphertext arg0Enc = - cryptoContext->Encrypt(keyPair.publicKey, ptxt); - std::vector arg0Vec = {arg0Enc}; - - TypedCppValue outputEncrypted = - interpreter.interpret("rotations", {cc, TypedCppValue(arg0Vec)})[0]; -} - -} // namespace openfhe -} // namespace heir -} // namespace mlir diff --git a/tests/Examples/openfhe/test.bzl b/tests/Examples/openfhe/test.bzl index 90e387420b..024575c3bf 100644 --- a/tests/Examples/openfhe/test.bzl +++ b/tests/Examples/openfhe/test.bzl @@ -2,7 +2,6 @@ load("@heir//bazel/openfhe:copts.bzl", "OPENMP_COPTS", "OPENMP_LINKOPTS") load("@heir//tools:heir-openfhe.bzl", "openfhe_lib") -load("@heir//tools:heir-opt.bzl", "heir_opt") load("@rules_cc//cc:cc_test.bzl", "cc_test") def openfhe_end_to_end_test(name, mlir_src, test_src, generated_lib_header, heir_opt_flags = [], heir_translate_flags = [], externalize_constants = True, ext_const_output_dir = "", data = [], size = "small", tags = [], deps = [], generate_debug_helper = False, **kwargs): @@ -43,50 +42,3 @@ def openfhe_end_to_end_test(name, mlir_src, test_src, generated_lib_header, heir linkopts = OPENMP_LINKOPTS, **kwargs ) - -def openfhe_interpreter_test(name, mlir_src, test_src, generated_heir_opt_filename = "", heir_opt_flags = [], data = [], tags = [], deps = [], copts = [], timeout = "moderate", **kwargs): - """A rule for running generating OpenFHE dialect and exposing it to an interpreter. - - Args: - name: The name of the cc_test target and the generated .cc file basename. - mlir_src: The source mlir file to run through heir-translate - test_src: The C++ test harness source file. - heir_opt_flags: Flags to pass to heir-opt before heir-translate - generated_heir_opt_filename: The filename of the file output by heir-opt - data: Data dependencies to be passed to cc_test/heir_opt - tags: Tags to pass to cc_test - deps: Deps to pass to cc_test - copts: Additional copts to pass to cc_test - timeout: Timeout to pass to cc_test - **kwargs: Keyword arguments to pass to cc_test. - """ - heir_opt_name = "%s_heir_opt" % name - if not generated_heir_opt_filename: - generated_heir_opt_filename = "%s_heir_opt.mlir" % name - - heir_opt( - name = heir_opt_name, - src = mlir_src, - pass_flags = heir_opt_flags, - generated_filename = generated_heir_opt_filename, - data = data, - ) - - cc_test( - name = name, - srcs = [test_src], - deps = deps + [ - "@heir//lib/Target/OpenFhePke:Interpreter", - "@openfhe//:pke", - "@openfhe//:core", - "@googletest//:gtest_main", - # for mlir source file parsing - "@llvm-project//mlir:Support", - ], - timeout = timeout, - tags = tags + ["manual", "notap"], - data = data + [":" + generated_heir_opt_filename], - copts = OPENMP_COPTS + copts, - linkopts = OPENMP_LINKOPTS, - **kwargs - )