Skip to content

Commit 0ba9fbf

Browse files
ZenithalHourlyRateZenithalHourlyRate
committed
BGV: Symbolic Noise Analysis for Dependency Tracking
1 parent dde839d commit 0ba9fbf

File tree

13 files changed

+709
-0
lines changed

13 files changed

+709
-0
lines changed

‎lib/Analysis/NoiseAnalysis/BGV/BUILD‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,3 +74,45 @@ cc_library(
7474
"@heir//lib/Utils:MathUtils",
7575
],
7676
)
77+
78+
cc_library(
79+
name = "NoiseBySymbolCoeffModelAnalysis",
80+
srcs = [
81+
"NoiseBySymbolCoeffModelAnalysis.cpp",
82+
],
83+
hdrs = [
84+
],
85+
deps = [
86+
":NoiseBySymbolCoeffModel",
87+
"@heir//lib/Analysis:Utils",
88+
"@heir//lib/Analysis/DimensionAnalysis",
89+
"@heir//lib/Analysis/LevelAnalysis",
90+
"@heir//lib/Analysis/NoiseAnalysis",
91+
"@heir//lib/Dialect/Mgmt/IR:Dialect",
92+
"@heir//lib/Dialect/Secret/IR:Dialect",
93+
"@heir//lib/Dialect/TensorExt/IR:Dialect",
94+
"@llvm-project//llvm:Support",
95+
"@llvm-project//mlir:ArithDialect",
96+
"@llvm-project//mlir:CallOpInterfaces",
97+
"@llvm-project//mlir:IR",
98+
"@llvm-project//mlir:Support",
99+
"@llvm-project//mlir:TensorDialect",
100+
],
101+
# required for gcc to link properly
102+
alwayslink = 1,
103+
)
104+
105+
cc_library(
106+
name = "NoiseBySymbolCoeffModel",
107+
srcs = [
108+
"NoiseBySymbolCoeffModel.cpp",
109+
],
110+
hdrs = [
111+
"NoiseBySymbolCoeffModel.h",
112+
],
113+
deps = [
114+
"@heir//lib/Analysis/NoiseAnalysis:Symbolic",
115+
"@heir//lib/Parameters/BGV:Params",
116+
"@heir//lib/Utils:MathUtils",
117+
],
118+
)

‎lib/Analysis/NoiseAnalysis/BGV/NoiseBySymbolCoeffModel.cpp‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
#include "lib/Analysis/NoiseAnalysis/BGV/NoiseBySymbolCoeffModel.h"
2+
3+
#include <cassert>
4+
#include <cmath>
5+
#include <iomanip>
6+
#include <ios>
7+
#include <limits>
8+
#include <numeric>
9+
#include <sstream>
10+
#include <stdexcept>
11+
#include <string>
12+
13+
#include "lib/Utils/MathUtils.h"
14+
15+
namespace mlir {
16+
namespace heir {
17+
namespace bgv {
18+
19+
using namespace mlir::heir::noise;
20+
21+
double NoiseBySymbolCoeffModel::toLogBound(const LocalParamType &param,
22+
const StateType &noise) const {
23+
double alpha = pow(2.0, -32);
24+
auto ringDim = param.getSchemeParam()->getRingDim();
25+
// variance is log2(Var(e))
26+
auto variance = noise.getVariance(ringDim);
27+
double bound = (1. / 2.) * (1 + variance) +
28+
log2(erfinv(pow(1.0 - alpha, 1.0 / ringDim)));
29+
return bound;
30+
}
31+
32+
double NoiseBySymbolCoeffModel::toLogBudget(const LocalParamType &param,
33+
const StateType &noise) const {
34+
return toLogTotal(param) - toLogBound(param, noise);
35+
}
36+
37+
double NoiseBySymbolCoeffModel::toLogTotal(const LocalParamType &param) const {
38+
double total = 0;
39+
auto logqi = param.getSchemeParam()->getLogqi();
40+
for (auto i = 0; i <= param.getCurrentLevel(); ++i) {
41+
total += logqi[i];
42+
}
43+
return total - 1.0;
44+
}
45+
46+
std::string NoiseBySymbolCoeffModel::toLogBoundString(
47+
const LocalParamType &param, const StateType &noise) const {
48+
auto logBound = toLogBound(param, noise);
49+
std::stringstream stream;
50+
stream << std::fixed << std::setprecision(2) << logBound;
51+
return stream.str();
52+
}
53+
54+
std::string NoiseBySymbolCoeffModel::toLogBudgetString(
55+
const LocalParamType &param, const StateType &noise) const {
56+
auto logBudget = toLogBudget(param, noise);
57+
std::stringstream stream;
58+
stream << std::fixed << std::setprecision(2) << logBudget;
59+
return stream.str();
60+
}
61+
62+
std::string NoiseBySymbolCoeffModel::toLogTotalString(
63+
const LocalParamType &param) const {
64+
auto logTotal = toLogTotal(param);
65+
std::stringstream stream;
66+
stream << std::fixed << std::setprecision(2) << logTotal;
67+
return stream.str();
68+
}
69+
70+
typename NoiseBySymbolCoeffModel::StateType
71+
NoiseBySymbolCoeffModel::evalEncryptPk(const LocalParamType &param,
72+
unsigned index) const {
73+
auto stderr = param.getSchemeParam()->getStd0();
74+
Expr ei = Symbol(Symbol::GAUSSIAN, "e" + std::to_string(index), stderr);
75+
Expr s = Symbol(Symbol::UNIFORM_TERNARY, "s", 0);
76+
Expr epk = Symbol(Symbol::GAUSSIAN, "epk", stderr);
77+
Expr ui = Symbol(Symbol::UNIFORM_TERNARY, "u" + std::to_string(index), 0);
78+
79+
Expr t = Symbol(Symbol::CONSTANT, "t",
80+
param.getSchemeParam()->getPlaintextModulus());
81+
return t * (ei * s + epk * ui);
82+
}
83+
84+
typename NoiseBySymbolCoeffModel::StateType
85+
NoiseBySymbolCoeffModel::evalEncryptSk(const LocalParamType &param,
86+
unsigned index) const {
87+
auto stderr = param.getSchemeParam()->getStd0();
88+
Expr ei = Symbol(Symbol::GAUSSIAN, "e" + std::to_string(index), stderr);
89+
90+
Expr t = Symbol(Symbol::CONSTANT, "t",
91+
param.getSchemeParam()->getPlaintextModulus());
92+
return t * ei;
93+
}
94+
95+
typename NoiseBySymbolCoeffModel::StateType
96+
NoiseBySymbolCoeffModel::evalEncrypt(const LocalParamType &param,
97+
unsigned index) const {
98+
auto usePublicKey = param.getSchemeParam()->getUsePublicKey();
99+
if (usePublicKey) {
100+
return evalEncryptPk(param, index);
101+
}
102+
return evalEncryptSk(param, index);
103+
}
104+
105+
typename NoiseBySymbolCoeffModel::StateType NoiseBySymbolCoeffModel::evalMul(
106+
const LocalParamType &resultParam, const StateType &lhs,
107+
const StateType &rhs) const {
108+
return lhs * rhs;
109+
}
110+
111+
typename NoiseBySymbolCoeffModel::StateType NoiseBySymbolCoeffModel::evalAdd(
112+
const StateType &lhs, const StateType &rhs) const {
113+
return lhs + rhs;
114+
}
115+
116+
} // namespace bgv
117+
} // namespace heir
118+
} // namespace mlir

‎lib/Analysis/NoiseAnalysis/BGV/NoiseBySymbolCoeffModel.h‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
#ifndef INCLUDE_ANALYSIS_NOISEANALYSIS_BGV_NOISEBYSYMBOLCOEFFMODEL_H_
2+
#define INCLUDE_ANALYSIS_NOISEANALYSIS_BGV_NOISEBYSYMBOLCOEFFMODEL_H_
3+
4+
#include <cassert>
5+
#include <string>
6+
7+
#include "lib/Analysis/NoiseAnalysis/Symbolic.h"
8+
#include "lib/Parameters/BGV/Params.h"
9+
10+
namespace mlir {
11+
namespace heir {
12+
namespace bgv {
13+
14+
// coefficient embedding noise model using variance
15+
class NoiseBySymbolCoeffModel {
16+
public:
17+
// for MP24, NoiseState stores the variance var for the one coefficient of
18+
// critical quantity v = m + t * e, assuming coefficients are IID.
19+
//
20+
// MP24 states that for two polynomial multipication, the variance of one
21+
// coefficient of the result can be approximated by ringDim * var_0 * var_1,
22+
// because the polynomial multipication is a convolution.
23+
using StateType = noise::Expression;
24+
using SchemeParamType = SchemeParam;
25+
using LocalParamType = LocalParam;
26+
27+
private:
28+
StateType evalEncryptPk(const LocalParamType &param, unsigned index) const;
29+
StateType evalEncryptSk(const LocalParamType &param, unsigned index) const;
30+
31+
public:
32+
StateType evalEncrypt(const LocalParamType &param, unsigned index) const;
33+
StateType evalAdd(const StateType &lhs, const StateType &rhs) const;
34+
StateType evalMul(const LocalParamType &resultParam, const StateType &lhs,
35+
const StateType &rhs) const;
36+
37+
// logTotal: log(Ql / 2)
38+
// logBound: bound on ||m + t * e|| predicted by the model
39+
// logBudget: logTotal - logBound
40+
// as ||m + t * e|| < Ql / 2 for correct decryption
41+
double toLogBound(const LocalParamType &param, const StateType &noise) const;
42+
std::string toLogBoundString(const LocalParamType &param,
43+
const StateType &noise) const;
44+
double toLogBudget(const LocalParamType &param, const StateType &noise) const;
45+
std::string toLogBudgetString(const LocalParamType &param,
46+
const StateType &noise) const;
47+
double toLogTotal(const LocalParamType &param) const;
48+
std::string toLogTotalString(const LocalParamType &param) const;
49+
};
50+
51+
} // namespace bgv
52+
} // namespace heir
53+
} // namespace mlir
54+
55+
#endif // INCLUDE_ANALYSIS_NOISEANALYSIS_BGV_NOISEBYSYMBOLCOEFFMODEL_H_

‎lib/Analysis/NoiseAnalysis/BGV/NoiseBySymbolCoeffModelAnalysis.cpp‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
#include <functional>
2+
#include <iomanip>
3+
4+
#include "lib/Analysis/DimensionAnalysis/DimensionAnalysis.h"
5+
#include "lib/Analysis/LevelAnalysis/LevelAnalysis.h"
6+
#include "lib/Analysis/NoiseAnalysis/BGV/NoiseBySymbolCoeffModel.h"
7+
#include "lib/Analysis/NoiseAnalysis/NoiseAnalysis.h"
8+
#include "lib/Analysis/Utils.h"
9+
#include "lib/Dialect/Mgmt/IR/MgmtOps.h"
10+
#include "lib/Dialect/Secret/IR/SecretOps.h"
11+
#include "lib/Dialect/TensorExt/IR/TensorExtOps.h"
12+
#include "llvm/include/llvm/ADT/TypeSwitch.h" // from @llvm-project
13+
#include "llvm/include/llvm/Support/Debug.h" // from @llvm-project
14+
#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
15+
#include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h" // from @llvm-project
16+
#include "mlir/include/mlir/IR/Operation.h" // from @llvm-project
17+
#include "mlir/include/mlir/IR/Value.h" // from @llvm-project
18+
#include "mlir/include/mlir/Interfaces/CallInterfaces.h" // from @llvm-project
19+
#include "mlir/include/mlir/Support/LLVM.h" // from @llvm-project
20+
21+
#define DEBUG_TYPE "NoiseAnalysis"
22+
23+
namespace mlir {
24+
namespace heir {
25+
26+
// explicit specialization of NoiseAnalysis for NoiseByBoundCoeffModel
27+
template <typename NoiseModel>
28+
void NoiseAnalysis<NoiseModel>::setToEntryState(LatticeType *lattice) {
29+
// At an entry point, we have no information about the noise.
30+
this->propagateIfChanged(lattice, lattice->join(NoiseState()));
31+
}
32+
33+
// explicit specialization of NoiseAnalysis for NoiseByBoundCoeffModel
34+
template <typename NoiseModel>
35+
void NoiseAnalysis<NoiseModel>::visitExternalCall(
36+
CallOpInterface call, ArrayRef<const LatticeType *> argumentLattices,
37+
ArrayRef<LatticeType *> resultLattices) {
38+
auto callback =
39+
std::bind(&NoiseAnalysis<NoiseModel>::propagateIfChangedWrapper, this,
40+
std::placeholders::_1, std::placeholders::_2);
41+
::mlir::heir::visitExternalCall<NoiseState, LatticeType>(
42+
call, argumentLattices, resultLattices, callback);
43+
}
44+
45+
// explicit specialization of NoiseAnalysis for NoiseByBoundCoeffModel
46+
template <typename NoiseModel>
47+
LogicalResult NoiseAnalysis<NoiseModel>::visitOperation(
48+
Operation *op, ArrayRef<const LatticeType *> operands,
49+
ArrayRef<LatticeType *> results) {
50+
auto getLocalParam = [&](Value value) {
51+
auto level = getLevelFromMgmtAttr(value);
52+
auto dimension = getDimensionFromMgmtAttr(value);
53+
return LocalParamType(&schemeParam, level, dimension);
54+
};
55+
56+
auto propagate = [&](Value value, NoiseState noise) {
57+
LLVM_DEBUG(llvm::dbgs()
58+
<< "Propagating "
59+
<< noiseModel.toLogBoundString(getLocalParam(value), noise)
60+
<< " " << noise.toString() << " to " << value << "\n");
61+
LatticeType *lattice = this->getLatticeElement(value);
62+
auto changeResult = lattice->join(noise);
63+
this->propagateIfChanged(lattice, changeResult);
64+
};
65+
66+
auto res =
67+
llvm::TypeSwitch<Operation &, LogicalResult>(*op)
68+
.Case<secret::GenericOp>([&](auto genericOp) {
69+
Block *body = genericOp.getBody();
70+
for (BlockArgument &arg : body->getArguments()) {
71+
auto localParam = getLocalParam(arg);
72+
NoiseState encrypted =
73+
noiseModel.evalEncrypt(localParam, arg.getArgNumber());
74+
propagate(arg, encrypted);
75+
}
76+
return success();
77+
})
78+
.template Case<arith::MulIOp>([&](auto mulOp) {
79+
SmallVector<OpResult> secretResults;
80+
this->getSecretResults(mulOp, secretResults);
81+
if (secretResults.empty()) {
82+
return success();
83+
}
84+
85+
auto localParam = getLocalParam(mulOp.getResult());
86+
NoiseState mult = noiseModel.evalMul(
87+
localParam, operands[0]->getValue(), operands[1]->getValue());
88+
propagate(mulOp.getResult(), mult);
89+
return success();
90+
})
91+
.template Case<arith::AddIOp>([&](auto addOp) {
92+
SmallVector<OpResult> secretResults;
93+
this->getSecretResults(addOp, secretResults);
94+
if (secretResults.empty()) {
95+
return success();
96+
}
97+
98+
NoiseState add = noiseModel.evalAdd(operands[0]->getValue(),
99+
operands[1]->getValue());
100+
propagate(addOp.getResult(), add);
101+
return success();
102+
})
103+
.template Case<mgmt::RelinearizeOp>([&](auto relinearizeOp) {
104+
propagate(relinearizeOp.getResult(), operands[0]->getValue());
105+
return success();
106+
})
107+
.Default([&](auto &op) { return success(); });
108+
return res;
109+
}
110+
111+
// template instantiation
112+
template class NoiseAnalysis<bgv::NoiseBySymbolCoeffModel>;
113+
114+
} // namespace heir
115+
} // namespace mlir

‎lib/Analysis/NoiseAnalysis/BUILD‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,3 +29,16 @@ cc_library(
2929
"@llvm-project//mlir:IR",
3030
],
3131
)
32+
33+
cc_library(
34+
name = "Symbolic",
35+
srcs = ["Symbolic.cpp"],
36+
hdrs = [
37+
"Symbolic.h",
38+
],
39+
deps = [
40+
":Noise",
41+
"@llvm-project//llvm:Support",
42+
"@llvm-project//mlir:Support",
43+
],
44+
)

0 commit comments

Comments
 (0)