BGV: Symbolic Noise Analysis for Dependency Tracking

lib/Analysis/NoiseAnalysis/BGV/BUILD

@@ -78,3 +78,45 @@ cc_library(
         "@heir//lib/Utils:MathUtils",
     ],
 )
+
+cc_library(
+    name = "NoiseBySymbolCoeffModelAnalysis",
+    srcs = [
+        "NoiseBySymbolCoeffModelAnalysis.cpp",
+    ],
+    hdrs = [
+    ],
+    deps = [
+        ":NoiseBySymbolCoeffModel",
+        "@heir//lib/Analysis:Utils",
+        "@heir//lib/Analysis/DimensionAnalysis",
+        "@heir//lib/Analysis/LevelAnalysis",
+        "@heir//lib/Analysis/NoiseAnalysis",
+        "@heir//lib/Dialect/Mgmt/IR:Dialect",
+        "@heir//lib/Dialect/Secret/IR:Dialect",
+        "@heir//lib/Dialect/TensorExt/IR:Dialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:CallOpInterfaces",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+    ],
+    # required for gcc to link properly
+    alwayslink = 1,
+)
+
+cc_library(
+    name = "NoiseBySymbolCoeffModel",
+    srcs = [
+        "NoiseBySymbolCoeffModel.cpp",
+    ],
+    hdrs = [
+        "NoiseBySymbolCoeffModel.h",
+    ],
+    deps = [
+        "@heir//lib/Analysis/NoiseAnalysis:Symbolic",
+        "@heir//lib/Parameters/BGV:Params",
+        "@heir//lib/Utils:MathUtils",
+    ],
+)

lib/Analysis/NoiseAnalysis/BGV/NoiseBySymbolCoeffModel.cpp

@@ -0,0 +1,118 @@
+#include "lib/Analysis/NoiseAnalysis/BGV/NoiseBySymbolCoeffModel.h"
+
+#include <cassert>
+#include <cmath>
+#include <iomanip>
+#include <ios>
+#include <limits>
+#include <numeric>
+#include <sstream>
+#include <stdexcept>
+#include <string>
+
+#include "lib/Utils/MathUtils.h"
+
+namespace mlir {
+namespace heir {
+namespace bgv {
+
+using namespace mlir::heir::noise;
+
+double NoiseBySymbolCoeffModel::toLogBound(const LocalParamType &param,
+                                           const StateType &noise) const {
+  double alpha = pow(2.0, -32);
+  auto ringDim = param.getSchemeParam()->getRingDim();
+  // variance is log2(Var(e))
+  auto variance = noise.getVariance(ringDim);
+  double bound = (1. / 2.) * (1 + variance) +
+                 log2(erfinv(pow(1.0 - alpha, 1.0 / ringDim)));
+  return bound;
+}
+
+double NoiseBySymbolCoeffModel::toLogBudget(const LocalParamType &param,
+                                            const StateType &noise) const {
+  return toLogTotal(param) - toLogBound(param, noise);
+}
+
+double NoiseBySymbolCoeffModel::toLogTotal(const LocalParamType &param) const {
+  double total = 0;
+  auto logqi = param.getSchemeParam()->getLogqi();
+  for (auto i = 0; i <= param.getCurrentLevel(); ++i) {
+    total += logqi[i];
+  }
+  return total - 1.0;
+}
+
+std::string NoiseBySymbolCoeffModel::toLogBoundString(
+    const LocalParamType &param, const StateType &noise) const {
+  auto logBound = toLogBound(param, noise);
+  std::stringstream stream;
+  stream << std::fixed << std::setprecision(2) << logBound;
+  return stream.str();
+}
+
+std::string NoiseBySymbolCoeffModel::toLogBudgetString(
+    const LocalParamType &param, const StateType &noise) const {
+  auto logBudget = toLogBudget(param, noise);
+  std::stringstream stream;
+  stream << std::fixed << std::setprecision(2) << logBudget;
+  return stream.str();
+}
+
+std::string NoiseBySymbolCoeffModel::toLogTotalString(
+    const LocalParamType &param) const {
+  auto logTotal = toLogTotal(param);
+  std::stringstream stream;
+  stream << std::fixed << std::setprecision(2) << logTotal;
+  return stream.str();
+}
+
+typename NoiseBySymbolCoeffModel::StateType
+NoiseBySymbolCoeffModel::evalEncryptPk(const LocalParamType &param,
+                                       unsigned index) const {
+  auto stderr = param.getSchemeParam()->getStd0();
+  Expr ei = Symbol(Symbol::GAUSSIAN, "e" + std::to_string(index), stderr);
+  Expr s = Symbol(Symbol::UNIFORM_TERNARY, "s", 0);
+  Expr epk = Symbol(Symbol::GAUSSIAN, "epk", stderr);
+  Expr ui = Symbol(Symbol::UNIFORM_TERNARY, "u" + std::to_string(index), 0);
+
+  Expr t = Symbol(Symbol::CONSTANT, "t",
+                  param.getSchemeParam()->getPlaintextModulus());
+  return t * (ei * s + epk * ui);
+}
+
+typename NoiseBySymbolCoeffModel::StateType
+NoiseBySymbolCoeffModel::evalEncryptSk(const LocalParamType &param,
+                                       unsigned index) const {
+  auto stderr = param.getSchemeParam()->getStd0();
+  Expr ei = Symbol(Symbol::GAUSSIAN, "e" + std::to_string(index), stderr);
+
+  Expr t = Symbol(Symbol::CONSTANT, "t",
+                  param.getSchemeParam()->getPlaintextModulus());
+  return t * ei;
+}
+
+typename NoiseBySymbolCoeffModel::StateType
+NoiseBySymbolCoeffModel::evalEncrypt(const LocalParamType &param,
+                                     unsigned index) const {
+  auto usePublicKey = param.getSchemeParam()->getUsePublicKey();
+  if (usePublicKey) {
+    return evalEncryptPk(param, index);
+  }
+  return evalEncryptSk(param, index);
+}
+
+typename NoiseBySymbolCoeffModel::StateType NoiseBySymbolCoeffModel::evalMul(
+    const LocalParamType &resultParam, const StateType &lhs,
+    const StateType &rhs) const {
+  return lhs * rhs;
+}
+
+typename NoiseBySymbolCoeffModel::StateType NoiseBySymbolCoeffModel::evalAdd(
+    const StateType &lhs, const StateType &rhs) const {
+  return lhs + rhs;
+}
+
+}  // namespace bgv
+}  // namespace heir
+}  // namespace mlir

lib/Analysis/NoiseAnalysis/BGV/NoiseBySymbolCoeffModel.h

@@ -0,0 +1,55 @@
+#ifndef INCLUDE_ANALYSIS_NOISEANALYSIS_BGV_NOISEBYSYMBOLCOEFFMODEL_H_
+#define INCLUDE_ANALYSIS_NOISEANALYSIS_BGV_NOISEBYSYMBOLCOEFFMODEL_H_
+
+#include <cassert>
+#include <string>
+
+#include "lib/Analysis/NoiseAnalysis/Symbolic.h"
+#include "lib/Parameters/BGV/Params.h"
+
+namespace mlir {
+namespace heir {
+namespace bgv {
+
+// coefficient embedding noise model using variance
+class NoiseBySymbolCoeffModel {
+ 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 = noise::Expression;
+  using SchemeParamType = SchemeParam;
+  using LocalParamType = LocalParam;
+
+ private:
+  StateType evalEncryptPk(const LocalParamType &param, unsigned index) const;
+  StateType evalEncryptSk(const LocalParamType &param, unsigned index) const;
+
+ public:
+  StateType evalEncrypt(const LocalParamType &param, unsigned index) const;
+  StateType evalAdd(const StateType &lhs, const StateType &rhs) const;
+  StateType evalMul(const LocalParamType &resultParam, const StateType &lhs,
+                    const StateType &rhs) const;
+
+  // 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
+  double toLogBound(const LocalParamType &param, const StateType &noise) const;
+  std::string toLogBoundString(const LocalParamType &param,
+                               const StateType &noise) const;
+  double toLogBudget(const LocalParamType &param, const StateType &noise) const;
+  std::string toLogBudgetString(const LocalParamType &param,
+                                const StateType &noise) const;
+  double toLogTotal(const LocalParamType &param) const;
+  std::string toLogTotalString(const LocalParamType &param) const;
+};
+
+}  // namespace bgv
+}  // namespace heir
+}  // namespace mlir
+
+#endif  // INCLUDE_ANALYSIS_NOISEANALYSIS_BGV_NOISEBYSYMBOLCOEFFMODEL_H_

lib/Analysis/NoiseAnalysis/BGV/NoiseBySymbolCoeffModelAnalysis.cpp

@@ -0,0 +1,115 @@
+#include <functional>
+#include <iomanip>
+
+#include "lib/Analysis/DimensionAnalysis/DimensionAnalysis.h"
+#include "lib/Analysis/LevelAnalysis/LevelAnalysis.h"
+#include "lib/Analysis/NoiseAnalysis/BGV/NoiseBySymbolCoeffModel.h"
+#include "lib/Analysis/NoiseAnalysis/NoiseAnalysis.h"
+#include "lib/Analysis/Utils.h"
+#include "lib/Dialect/Mgmt/IR/MgmtOps.h"
+#include "lib/Dialect/Secret/IR/SecretOps.h"
+#include "lib/Dialect/TensorExt/IR/TensorExtOps.h"
+#include "llvm/include/llvm/ADT/TypeSwitch.h"             // from @llvm-project
+#include "llvm/include/llvm/Support/Debug.h"              // from @llvm-project
+#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"     // from @llvm-project
+#include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.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/Interfaces/CallInterfaces.h"  // from @llvm-project
+#include "mlir/include/mlir/Support/LLVM.h"               // from @llvm-project
+
+#define DEBUG_TYPE "NoiseAnalysis"
+
+namespace mlir {
+namespace heir {
+
+// explicit specialization of NoiseAnalysis for NoiseByBoundCoeffModel
+template <typename NoiseModel>
+void NoiseAnalysis<NoiseModel>::setToEntryState(LatticeType *lattice) {
+  // At an entry point, we have no information about the noise.
+  this->propagateIfChanged(lattice, lattice->join(NoiseState()));
+}
+
+// explicit specialization of NoiseAnalysis for NoiseByBoundCoeffModel
+template <typename NoiseModel>
+void NoiseAnalysis<NoiseModel>::visitExternalCall(
+    CallOpInterface call, ArrayRef<const LatticeType *> argumentLattices,
+    ArrayRef<LatticeType *> resultLattices) {
+  auto callback =
+      std::bind(&NoiseAnalysis<NoiseModel>::propagateIfChangedWrapper, this,
+                std::placeholders::_1, std::placeholders::_2);
+  ::mlir::heir::visitExternalCall<NoiseState, LatticeType>(
+      call, argumentLattices, resultLattices, callback);
+}
+
+// explicit specialization of NoiseAnalysis for NoiseByBoundCoeffModel
+template <typename NoiseModel>
+LogicalResult NoiseAnalysis<NoiseModel>::visitOperation(
+    Operation *op, ArrayRef<const LatticeType *> operands,
+    ArrayRef<LatticeType *> results) {
+  auto getLocalParam = [&](Value value) {
+    auto level = getLevelFromMgmtAttr(value);
+    auto dimension = getDimensionFromMgmtAttr(value);
+    return LocalParamType(&schemeParam, level, dimension);
+  };
+
+  auto propagate = [&](Value value, NoiseState noise) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Propagating "
+               << noiseModel.toLogBoundString(getLocalParam(value), noise)
+               << " " << noise.toString() << " to " << value << "\n");
+    LatticeType *lattice = this->getLatticeElement(value);
+    auto changeResult = lattice->join(noise);
+    this->propagateIfChanged(lattice, changeResult);
+  };
+
+  auto res =
+      llvm::TypeSwitch<Operation &, LogicalResult>(*op)
+          .Case<secret::GenericOp>([&](auto genericOp) {
+            Block *body = genericOp.getBody();
+            for (BlockArgument &arg : body->getArguments()) {
+              auto localParam = getLocalParam(arg);
+              NoiseState encrypted =
+                  noiseModel.evalEncrypt(localParam, arg.getArgNumber());
+              propagate(arg, encrypted);
+            }
+            return success();
+          })
+          .template Case<arith::MulIOp>([&](auto mulOp) {
+            SmallVector<OpResult> secretResults;
+            this->getSecretResults(mulOp, secretResults);
+            if (secretResults.empty()) {
+              return success();
+            }
+
+            auto localParam = getLocalParam(mulOp.getResult());
+            NoiseState mult = noiseModel.evalMul(
+                localParam, operands[0]->getValue(), operands[1]->getValue());
+            propagate(mulOp.getResult(), mult);
+            return success();
+          })
+          .template Case<arith::AddIOp>([&](auto addOp) {
+            SmallVector<OpResult> secretResults;
+            this->getSecretResults(addOp, secretResults);
+            if (secretResults.empty()) {
+              return success();
+            }
+
+            NoiseState add = noiseModel.evalAdd(operands[0]->getValue(),
+                                                operands[1]->getValue());
+            propagate(addOp.getResult(), add);
+            return success();
+          })
+          .template Case<mgmt::RelinearizeOp>([&](auto relinearizeOp) {
+            propagate(relinearizeOp.getResult(), operands[0]->getValue());
+            return success();
+          })
+          .Default([&](auto &op) { return success(); });
+  return res;
+}
+
+// template instantiation
+template class NoiseAnalysis<bgv::NoiseBySymbolCoeffModel>;
+
+}  // namespace heir
+}  // namespace mlir

lib/Analysis/NoiseAnalysis/BUILD

@@ -32,3 +32,16 @@ cc_library(
         "@llvm-project//mlir:IR",
     ],
 )
+
+cc_library(
+    name = "Symbolic",
+    srcs = ["Symbolic.cpp"],
+    hdrs = [
+        "Symbolic.h",
+    ],
+    deps = [
+        ":Noise",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:Support",
+    ],
+)