From f427c16bee3991946bac40ee4d763ba726b59d8d Mon Sep 17 00:00:00 2001 From: Zenithal Date: Mon, 10 Feb 2025 10:04:13 +0000 Subject: [PATCH] BGV: Add Variance-based Noise Model --- lib/Analysis/NoiseAnalysis/BGV/BUILD | 36 ++- .../BGV/NoiseByVarianceCoeffModel.cpp | 292 ++++++++++++++++++ .../BGV/NoiseByVarianceCoeffModel.h | 77 +++++ ...alysis.cpp => NoiseCoeffModelAnalysis.cpp} | 5 + lib/Transforms/ValidateNoise/BUILD | 2 + .../ValidateNoise/ValidateNoise.cpp | 5 + lib/Transforms/ValidateNoise/ValidateNoise.td | 12 +- 7 files changed, 422 insertions(+), 7 deletions(-) create mode 100644 lib/Analysis/NoiseAnalysis/BGV/NoiseByVarianceCoeffModel.cpp create mode 100644 lib/Analysis/NoiseAnalysis/BGV/NoiseByVarianceCoeffModel.h rename lib/Analysis/NoiseAnalysis/BGV/{NoiseByBoundCoeffModelAnalysis.cpp => NoiseCoeffModelAnalysis.cpp} (97%) diff --git a/lib/Analysis/NoiseAnalysis/BGV/BUILD b/lib/Analysis/NoiseAnalysis/BGV/BUILD index e626baa4e..ead4f845f 100644 --- a/lib/Analysis/NoiseAnalysis/BGV/BUILD +++ b/lib/Analysis/NoiseAnalysis/BGV/BUILD @@ -4,16 +4,16 @@ package( ) cc_library( - name = "NoiseByBoundCoeffModel", + name = "NoiseCoeffModelAnalysis", srcs = [ - "NoiseByBoundCoeffModel.cpp", - "NoiseByBoundCoeffModelAnalysis.cpp", + "NoiseCoeffModelAnalysis.cpp", ], hdrs = [ - "NoiseByBoundCoeffModel.h", ], deps = [ ":Noise", + ":NoiseByBoundCoeffModel", + ":NoiseByVarianceCoeffModel", "@heir//lib/Analysis:Utils", "@heir//lib/Analysis/DimensionAnalysis", "@heir//lib/Analysis/LevelAnalysis", @@ -31,6 +31,34 @@ cc_library( ], ) +cc_library( + name = "NoiseByBoundCoeffModel", + srcs = [ + "NoiseByBoundCoeffModel.cpp", + ], + hdrs = [ + "NoiseByBoundCoeffModel.h", + ], + deps = [ + ":Noise", + "@heir//lib/Parameters/BGV:Params", + ], +) + +cc_library( + name = "NoiseByVarianceCoeffModel", + srcs = [ + "NoiseByVarianceCoeffModel.cpp", + ], + hdrs = [ + "NoiseByVarianceCoeffModel.h", + ], + deps = [ + ":Noise", + "@heir//lib/Parameters/BGV:Params", + ], +) + cc_library( name = "Noise", srcs = ["Noise.cpp"], diff --git a/lib/Analysis/NoiseAnalysis/BGV/NoiseByVarianceCoeffModel.cpp b/lib/Analysis/NoiseAnalysis/BGV/NoiseByVarianceCoeffModel.cpp new file mode 100644 index 000000000..6b4a890b5 --- /dev/null +++ b/lib/Analysis/NoiseAnalysis/BGV/NoiseByVarianceCoeffModel.cpp @@ -0,0 +1,292 @@ +#include "lib/Analysis/NoiseAnalysis/BGV/NoiseByVarianceCoeffModel.h" + +#include +#include + +namespace mlir { +namespace heir { +namespace bgv { + +// the formulae below are mainly taken from MP24 +// "A Central Limit Framework for Ring-LWE Noise Analysis" +// https://ia.cr/2019/452 +// and CLP23 +// "Optimisations and tradeoffs for HElib" +// https://ia.cr/2023/104 + +template +using Model = NoiseByVarianceCoeffModel

; + +// https://stackoverflow.com/questions/27229371/inverse-error-function-in-c +static double erfinv(double a) { + double p, r, t; + t = fma(a, 0.0 - a, 1.0); + t = log(t); + if (fabs(t) > 6.125) { // maximum ulp error = 2.35793 + p = 3.03697567e-10; // 0x1.4deb44p-32 + p = fma(p, t, 2.93243101e-8); // 0x1.f7c9aep-26 + p = fma(p, t, 1.22150334e-6); // 0x1.47e512p-20 + p = fma(p, t, 2.84108955e-5); // 0x1.dca7dep-16 + p = fma(p, t, 3.93552968e-4); // 0x1.9cab92p-12 + p = fma(p, t, 3.02698812e-3); // 0x1.8cc0dep-9 + p = fma(p, t, 4.83185798e-3); // 0x1.3ca920p-8 + p = fma(p, t, -2.64646143e-1); // -0x1.0eff66p-2 + p = fma(p, t, 8.40016484e-1); // 0x1.ae16a4p-1 + } else { // maximum ulp error = 2.35002 + p = 5.43877832e-9; // 0x1.75c000p-28 + p = fma(p, t, 1.43285448e-7); // 0x1.33b402p-23 + p = fma(p, t, 1.22774793e-6); // 0x1.499232p-20 + p = fma(p, t, 1.12963626e-7); // 0x1.e52cd2p-24 + p = fma(p, t, -5.61530760e-5); // -0x1.d70bd0p-15 + p = fma(p, t, -1.47697632e-4); // -0x1.35be90p-13 + p = fma(p, t, 2.31468678e-3); // 0x1.2f6400p-9 + p = fma(p, t, 1.15392581e-2); // 0x1.7a1e50p-7 + p = fma(p, t, -2.32015476e-1); // -0x1.db2aeep-3 + p = fma(p, t, 8.86226892e-1); // 0x1.c5bf88p-1 + } + r = a * p; + return r; +} + +template +double Model

::toLogBound(const LocalParamType ¶m, + const StateType &noise) { + // error probability 0.1% + // though this only holds if every random variable is Gaussian + // or similar to Gaussian + // so this may give underestimation, see MP24 and CCH+23 + double alpha = 0.001; + auto ringDim = param.getSchemeParam()->getRingDim(); + double bound = + sqrt(2.0 * noise.getValue()) * erfinv(pow(1.0 - alpha, 1.0 / ringDim)); + return log2(bound); +} + +template +double Model

::toLogBudget(const LocalParamType ¶m, + const StateType &noise) { + return toLogTotal(param) - toLogBound(param, noise); +} + +template +double Model

::toLogTotal(const LocalParamType ¶m) { + double total = 0; + auto logqi = param.getSchemeParam()->getLogqi(); + for (auto i = 0; i <= param.getCurrentLevel(); ++i) { + total += logqi[i]; + } + return total - 1.0; +} + +template +std::string Model

::toLogBoundString(const LocalParamType ¶m, + const StateType &noise) { + auto logBound = toLogBound(param, noise); + std::stringstream stream; + stream << std::fixed << std::setprecision(2) << logBound; + return stream.str(); +} + +template +std::string Model

::toLogBudgetString(const LocalParamType ¶m, + const StateType &noise) { + auto logBudget = toLogBudget(param, noise); + std::stringstream stream; + stream << std::fixed << std::setprecision(2) << logBudget; + return stream.str(); +} + +template +std::string Model

::toLogTotalString(const LocalParamType ¶m) { + auto logTotal = toLogTotal(param); + std::stringstream stream; + stream << std::fixed << std::setprecision(2) << logTotal; + return stream.str(); +} + +template +double Model

::getVarianceErr(const LocalParamType ¶m) { + auto std0 = param.getSchemeParam()->getStd0(); + return std0 * std0; +} + +template +double Model

::getVarianceKey(const LocalParamType ¶m) { + // assume UNIFORM_TERNARY + return 2.0 / 3.0; +} + +template +typename Model

::StateType Model

::evalEncryptPk( + const LocalParamType ¶m) { + auto varianceError = getVarianceErr(param); + // uniform ternary + auto varianceKey = getVarianceKey(param); + auto t = param.getSchemeParam()->getPlaintextModulus(); + auto n = param.getSchemeParam()->getRingDim(); + // public key (-as + t * e, a) + // public key encryption (-aus + t(u * e + e_0) + m, au + e_1) + // v_fresh = m + t * (u * e + e_1 * s + e_0) + // var_fresh = t^2 * (2n * var_key + 1) * var_error + // for ringDim, see header comment for explanation + double fresh = t * t * varianceError * (2. * n * varianceKey + 1.); + return StateType::of(fresh); +} + +template +typename Model

::StateType Model

::evalEncryptSk( + const LocalParamType ¶m) { + auto t = param.getSchemeParam()->getPlaintextModulus(); + auto varianceError = getVarianceErr(param); + + // secret key s + // secret key encryption (-as + m + t * e, a) + // v_fresh = t * e + // var_fresh = t^2 * var_error + double fresh = t * t * varianceError; + return StateType::of(fresh); +} + +template +typename Model

::StateType Model

::evalEncrypt( + const LocalParamType ¶m) { + // P stands for public key encryption + if constexpr (P) { + return evalEncryptPk(param); + } else { + return evalEncryptSk(param); + } +} + +template +typename Model

::StateType Model

::evalConstant( + const LocalParamType ¶m) { + auto t = param.getSchemeParam()->getPlaintextModulus(); + // constant is v = m + t * 0 + // assume m is uniform from [-t/2, t/2] + // var_constant = t * t / 12 + return StateType::of(t * t / 12.0); +} + +template +typename Model

::StateType Model

::evalAdd(const StateType &lhs, + const StateType &rhs) { + // v_add = v_0 + v_1 + // assuming independent of course + return StateType::of(lhs.getValue() + rhs.getValue()); +} + +template +typename Model

::StateType Model

::evalMul( + const LocalParamType &resultParam, const StateType &lhs, + const StateType &rhs) { + auto ringDim = resultParam.getSchemeParam()->getRingDim(); + auto v0 = lhs.getValue(); + auto v1 = rhs.getValue(); + + // v_mul = v_0 * v_1 + // for ringDim, see header comment for explanation + return StateType::of(ringDim * v0 * v1); +} + +template +typename Model

::StateType Model

::evalModReduce( + const LocalParamType &inputParam, const StateType &input) { + auto cv = inputParam.getDimension(); + // for cv > 2 the rounding error term is different! + // like (tau_0, tau_1, tau_2) and the error becomes + // tau_0 + tau_1 s + tau_2 s^2 + assert(cv == 2); + + auto currentLogqi = + inputParam.getSchemeParam()->getLogqi()[inputParam.getCurrentLevel()]; + + double modulus = pow(2.0, currentLogqi); + + auto ringDim = inputParam.getSchemeParam()->getRingDim(); + auto varianceKey = getVarianceKey(inputParam); + + // modulus switching is essentially a scaling operation + // so the original error is scaled by the modulus + // v_scaled = v_input / modulus + // var_scaled = var_input / (modulus * modulus) + auto scaled = input.getValue() / (modulus * modulus); + // in the meantime, it will introduce an rounding error + // (tau_0, tau_1) to the (ct_0, ct_1) where ||tau_i|| < t / 2 + // so tau_0 + tau_1 * s has the variance + // var_added = var_const * (1.0 + var_key * ringDim) + auto varianceConst = evalConstant(inputParam).getValue(); + auto added = varianceConst * (1.0 + ringDim * varianceKey); + return StateType::of(scaled + added); +} + +template +typename Model

::StateType Model

::evalRelinearizeHYBRID( + const LocalParamType &inputParam, const StateType &input) { + // for v_input, after modup and moddown, it remains the same (with rounding). + // We only need to consider the error from key switching key + // and rounding error during moddown. + // Check the B.1.3 section of KPZ21 for more details. + + // also note that for cv > 3 (e.g. multiplication), we need to relinearize + // more terms like ct_3 and ct_4. + // this is a common path for mult relinearize and rotation relinearize + // so no assertion here for now. + + auto dnum = inputParam.getSchemeParam()->getDnum(); + auto varianceErr = getVarianceErr(inputParam); + auto varianceKey = getVarianceKey(inputParam); + auto ringDim = inputParam.getSchemeParam()->getRingDim(); + auto t = inputParam.getSchemeParam()->getPlaintextModulus(); + + auto currentLevel = inputParam.getCurrentLevel(); + // modup from Ql to QlP, so one more digit + auto currentNumDigit = ceil(static_cast(currentLevel + 1) / dnum) + 1; + + // log(qiq_{i+1}...), the digit size for a certain digit + // we use log(pip_{i+1}...) as an approximation, + // as we often choose P > each digit + auto logqi = inputParam.getSchemeParam()->getLogqi(); + auto logDigitSize = std::accumulate(logqi.begin(), logqi.end(), 0); + // omega in literature + auto digitSize = pow(2.0, logDigitSize); + + // the critical quantity for HYBRID key switching error is + // t * sum over all digit (ct_2 * e_ksk) + // there are "currentNumDigit" digits + // and c_2 uniformly from [-digitSize / 2, digitSize / 2] + // for ringDim, see header comment for explanation + auto varianceKeySwitch = t * t * currentNumDigit * + (digitSize * digitSize / 12.0) * ringDim * + varianceErr; + + // moddown by P + auto scaled = varianceKeySwitch / (digitSize * digitSize); + + // Some papers just say hey we mod down by P so the error added is just mod + // reduce, but the error for mod reduce is different for approximate mod down. + // Anyway, this term is not the major term. + auto varianceConst = evalConstant(inputParam).getValue(); + auto added = varianceConst * (1.0 + ringDim * varianceKey); + + // for relinearization after multiplication, often scaled + added is far less + // than input. + return StateType::of(input.getValue() + scaled + added); +} + +template +typename Model

::StateType Model

::evalRelinearize( + const LocalParamType &inputParam, const StateType &input) { + // assume HYBRID + // if we further introduce BV to SchemeParam we can have alternative + // implementation. + return evalRelinearizeHYBRID(inputParam, input); +} + +// instantiate template class +template class NoiseByVarianceCoeffModel; +template class NoiseByVarianceCoeffModel; + +} // namespace bgv +} // namespace heir +} // namespace mlir diff --git a/lib/Analysis/NoiseAnalysis/BGV/NoiseByVarianceCoeffModel.h b/lib/Analysis/NoiseAnalysis/BGV/NoiseByVarianceCoeffModel.h new file mode 100644 index 000000000..32dd8cf98 --- /dev/null +++ b/lib/Analysis/NoiseAnalysis/BGV/NoiseByVarianceCoeffModel.h @@ -0,0 +1,77 @@ +#ifndef INCLUDE_ANALYSIS_NOISEANALYSIS_BGV_NOISEBYVARIANCECOEFFMODEL_H_ +#define INCLUDE_ANALYSIS_NOISEANALYSIS_BGV_NOISEBYVARIANCECOEFFMODEL_H_ + +#include +#include +#include +#include +#include +#include + +#include "lib/Analysis/NoiseAnalysis/BGV/Noise.h" +#include "lib/Parameters/BGV/Params.h" + +namespace mlir { +namespace heir { +namespace bgv { + +// coefficient embedding noise model using variance +// use template here just for the sake of code reuse +// P for public key +template +class NoiseByVarianceCoeffModel { + public: + // for MP24, NoiseState stores the variance var for the one coefficient of + // critical quantity v = m + t * e, assuming coefficients are IID. + // + // MP24 states that for two polynomial multipication, the variance of one + // coefficient of the result can be approximated by ringDim * var_0 * var_1, + // because the polynomial multipication is a convolution. + using StateType = NoiseState; + using SchemeParamType = SchemeParam; + using LocalParamType = LocalParam; + + private: + static double getVarianceErr(const LocalParamType ¶m); + static double getVarianceKey(const LocalParamType ¶m); + + static StateType evalEncryptPk(const LocalParamType ¶m); + static StateType evalEncryptSk(const LocalParamType ¶m); + static StateType evalRelinearizeHYBRID(const LocalParamType &inputParam, + const StateType &input); + + public: + static StateType evalEncrypt(const LocalParamType ¶m); + static StateType evalConstant(const LocalParamType ¶m); + static StateType evalAdd(const StateType &lhs, const StateType &rhs); + static StateType evalMul(const LocalParamType &resultParam, + const StateType &lhs, const StateType &rhs); + static StateType evalRelinearize(const LocalParamType &inputParam, + const StateType &input); + static StateType evalModReduce(const LocalParamType &inputParam, + const StateType &input); + + // logTotal: log(Ql / 2) + // logBound: bound on ||m + t * e|| predicted by the model + // logBudget: logTotal - logBound + // as ||m + t * e|| < Ql / 2 for correct decryption + static double toLogBound(const LocalParamType ¶m, const StateType &noise); + static std::string toLogBoundString(const LocalParamType ¶m, + const StateType &noise); + static double toLogBudget(const LocalParamType ¶m, + const StateType &noise); + static std::string toLogBudgetString(const LocalParamType ¶m, + const StateType &noise); + static double toLogTotal(const LocalParamType ¶m); + static std::string toLogTotalString(const LocalParamType ¶m); +}; + +// user-facing typedefs +using NoiseByVarianceCoeffPkModel = NoiseByVarianceCoeffModel; +using NoiseByVarianceCoeffSkModel = NoiseByVarianceCoeffModel; + +} // namespace bgv +} // namespace heir +} // namespace mlir + +#endif // INCLUDE_ANALYSIS_NOISEANALYSIS_BGV_NOISEBYVARIANCECOEFFMODEL_H_ diff --git a/lib/Analysis/NoiseAnalysis/BGV/NoiseByBoundCoeffModelAnalysis.cpp b/lib/Analysis/NoiseAnalysis/BGV/NoiseCoeffModelAnalysis.cpp similarity index 97% rename from lib/Analysis/NoiseAnalysis/BGV/NoiseByBoundCoeffModelAnalysis.cpp rename to lib/Analysis/NoiseAnalysis/BGV/NoiseCoeffModelAnalysis.cpp index 677ca5c8d..a1eabce3f 100644 --- a/lib/Analysis/NoiseAnalysis/BGV/NoiseByBoundCoeffModelAnalysis.cpp +++ b/lib/Analysis/NoiseAnalysis/BGV/NoiseCoeffModelAnalysis.cpp @@ -3,6 +3,7 @@ #include "lib/Analysis/DimensionAnalysis/DimensionAnalysis.h" #include "lib/Analysis/LevelAnalysis/LevelAnalysis.h" #include "lib/Analysis/NoiseAnalysis/BGV/NoiseByBoundCoeffModel.h" +#include "lib/Analysis/NoiseAnalysis/BGV/NoiseByVarianceCoeffModel.h" #include "lib/Analysis/NoiseAnalysis/NoiseAnalysis.h" #include "lib/Analysis/Utils.h" #include "lib/Dialect/Mgmt/IR/MgmtOps.h" @@ -181,5 +182,9 @@ template class NoiseAnalysis; template class NoiseAnalysis; template class NoiseAnalysis; +// for by variance +template class NoiseAnalysis; +template class NoiseAnalysis; + } // namespace heir } // namespace mlir diff --git a/lib/Transforms/ValidateNoise/BUILD b/lib/Transforms/ValidateNoise/BUILD index be41d0f6b..2a60b5899 100644 --- a/lib/Transforms/ValidateNoise/BUILD +++ b/lib/Transforms/ValidateNoise/BUILD @@ -15,6 +15,8 @@ cc_library( "@heir//lib/Analysis/LevelAnalysis", "@heir//lib/Analysis/NoiseAnalysis", "@heir//lib/Analysis/NoiseAnalysis/BGV:NoiseByBoundCoeffModel", + "@heir//lib/Analysis/NoiseAnalysis/BGV:NoiseByVarianceCoeffModel", + "@heir//lib/Analysis/NoiseAnalysis/BGV:NoiseCoeffModelAnalysis", "@heir//lib/Analysis/SecretnessAnalysis", "@heir//lib/Dialect/Secret/IR:Dialect", "@llvm-project//llvm:Support", diff --git a/lib/Transforms/ValidateNoise/ValidateNoise.cpp b/lib/Transforms/ValidateNoise/ValidateNoise.cpp index 0e27f7615..eacdfb982 100644 --- a/lib/Transforms/ValidateNoise/ValidateNoise.cpp +++ b/lib/Transforms/ValidateNoise/ValidateNoise.cpp @@ -3,6 +3,7 @@ #include "lib/Analysis/DimensionAnalysis/DimensionAnalysis.h" #include "lib/Analysis/LevelAnalysis/LevelAnalysis.h" #include "lib/Analysis/NoiseAnalysis/BGV/NoiseByBoundCoeffModel.h" +#include "lib/Analysis/NoiseAnalysis/BGV/NoiseByVarianceCoeffModel.h" #include "lib/Analysis/NoiseAnalysis/NoiseAnalysis.h" #include "lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.h" #include "lib/Dialect/Secret/IR/SecretOps.h" @@ -161,6 +162,10 @@ struct ValidateNoise : impl::ValidateNoiseBase { run>(); } else if (model == "bgv-noise-by-bound-coeff-average-case-sk") { run>(); + } else if (model == "bgv-noise-by-variance-coeff-pk") { + run>(); + } else if (model == "bgv-noise-by-variance-coeff-sk") { + run>(); } else { getOperation()->emitOpError() << "Unknown noise model.\n"; signalPassFailure(); diff --git a/lib/Transforms/ValidateNoise/ValidateNoise.td b/lib/Transforms/ValidateNoise/ValidateNoise.td index 8a5a087fa..f0a6f4f39 100644 --- a/lib/Transforms/ValidateNoise/ValidateNoise.td +++ b/lib/Transforms/ValidateNoise/ValidateNoise.td @@ -9,16 +9,22 @@ def ValidateNoise : Pass<"validate-noise"> { This pass validates the noise of the HE circuit against a given noise model. Currently the pass works for BGV scheme, and there are two noise models - available: "bgv-noise-by-bound-coeff-average-case{-pk,-sk}" and - "bgv-noise-by-bound-coeff-worst-case{-pk,-sk}". + available: "bgv-noise-by-bound-coeff-average-case{-pk,-sk}", + "bgv-noise-by-bound-coeff-worst-case{-pk,-sk}" and + "bgv-noise-by-variance-coeff{-pk,-sk}". - The two models are taken from KPZ21, and they work by bounding + The first two models are taken from KPZ21, and they work by bounding the coefficient embedding of the ciphertexts. The difference of the two models is expansion factor used for multiplication of the coefficients, the first being `2\sqrt{N}` and the second being `N`. The `-pk`/`-sk` suffixes assume the input ciphertexts are encrypted using the public/secret key. + The third model is taken from MP24. It works by tracking the variance + of the coefficient embedding of the ciphertexts. This gives a more accurate + noise estimate, but it may give underestimates in some cases. See the paper + for more details. + This pass is experimental. The result should be observed using --debug-only=ValidateNoise.