Skip to content

Commit 0d623e9

Browse files
copybara-githubcopybara-github
committed
Merge pull request #2402 from ZenithalHourlyRate:ckks-param-fixedmanual
PiperOrigin-RevId: 833850099
2 parents f2c6f94 + 2ce97bc commit 0d623e9

7 files changed

Lines changed: 252 additions & 18 deletions

File tree

‎lib/Parameters/CKKS/BUILD‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ cc_library(
1212
deps = [
1313
"@heir//lib/Dialect/CKKS/IR:Dialect",
1414
"@heir//lib/Parameters:RLWEParams",
15+
"@heir//lib/Parameters:RLWESecurityParams",
1516
"@llvm-project//llvm:Support",
17+
"@openfhe//:core",
1618
],
1719
)

‎lib/Parameters/CKKS/Params.cpp‎

Lines changed: 222 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,240 @@
11
#include "lib/Parameters/CKKS/Params.h"
22

3+
#include <algorithm>
34
#include <cassert>
5+
#include <cmath>
6+
#include <cstddef>
7+
#include <cstdint>
48
#include <utility>
59
#include <vector>
610

11+
#include "lib/Dialect/CKKS/IR/CKKSAttributes.h"
12+
#include "lib/Dialect/CKKS/IR/CKKSEnums.h"
713
#include "lib/Parameters/RLWEParams.h"
14+
#include "lib/Parameters/RLWESecurityParams.h"
815
#include "llvm/include/llvm/Support/raw_ostream.h" // from @llvm-project
16+
#include "src/core/include/openfhecore.h" // from @openfhe
917

1018
namespace mlir {
1119
namespace heir {
1220
namespace ckks {
1321

22+
/// By original we mean the method in RNS-CKKS implementation
23+
/// Corresponds to FIXED* in OpenFHE
24+
static std::vector<int64_t> moduliQGenerationOpenFHEFixed(int logFirstMod,
25+
int logDefaultScale,
26+
int numLevel,
27+
int ringDim) {
28+
auto cyclOrder = ringDim * 2;
29+
std::vector<int64_t> moduliQ(numLevel);
30+
lbcrypto::NativeInteger q =
31+
lbcrypto::FirstPrime<NativeInteger>(logDefaultScale, cyclOrder);
32+
moduliQ[numLevel - 1] = q.ConvertToInt();
33+
34+
auto maxPrime{q};
35+
auto minPrime{q};
36+
37+
auto qPrev = q;
38+
auto qNext = q;
39+
if (numLevel > 2) {
40+
for (size_t i = numLevel - 2, cnt = 0; i >= 1; --i, ++cnt) {
41+
if ((cnt % 2) == 0) {
42+
qPrev = PreviousPrime(qPrev, cyclOrder);
43+
moduliQ[i] = qPrev.ConvertToInt();
44+
} else {
45+
qNext = NextPrime(qNext, cyclOrder);
46+
moduliQ[i] = qNext.ConvertToInt();
47+
}
48+
49+
if (moduliQ[i] > maxPrime)
50+
maxPrime = moduliQ[i];
51+
else if (moduliQ[i] < minPrime)
52+
minPrime = moduliQ[i];
53+
}
54+
}
55+
56+
if (logFirstMod == logDefaultScale) { // this requires dcrtBits < 60
57+
moduliQ[0] =
58+
lbcrypto::NextPrime<NativeInteger>(maxPrime, cyclOrder).ConvertToInt();
59+
} else {
60+
moduliQ[0] = lbcrypto::LastPrime<NativeInteger>(logFirstMod, cyclOrder)
61+
.ConvertToInt();
62+
63+
// find if the value of moduliQ[0] is already in the vector starting with
64+
// moduliQ[1] and if there is, then get another prime for moduliQ[0]
65+
const auto pos = std::find(moduliQ.begin() + 1, moduliQ.end(), moduliQ[0]);
66+
if (pos != moduliQ.end()) {
67+
moduliQ[0] = lbcrypto::NextPrime<NativeInteger>(maxPrime, cyclOrder)
68+
.ConvertToInt();
69+
}
70+
}
71+
return moduliQ;
72+
}
73+
74+
/// See "Reduced Error" paper https://eprint.iacr.org/2020/1118
75+
/// Corresponds to FLEXIBLE* in OpenFHE
76+
static std::vector<int64_t> moduliQGenerationReducedError(int logFirstMod,
77+
int logDefaultScale,
78+
int numLevel,
79+
int ringDim) {
80+
auto cyclOrder = ringDim * 2;
81+
std::vector<int64_t> moduliQ(numLevel);
82+
lbcrypto::NativeInteger q =
83+
lbcrypto::FirstPrime<lbcrypto::NativeInteger>(logDefaultScale, cyclOrder);
84+
moduliQ[numLevel - 1] = q.ConvertToInt();
85+
86+
auto maxPrime{q};
87+
auto minPrime{q};
88+
89+
if (numLevel > 2) {
90+
for (size_t i = numLevel - 2, cnt = 0; i >= 1; --i, ++cnt) {
91+
// Comments from OpenFHE ckksrns-parametergeneration.cpp
92+
/* Scaling factors in FLEXIBLEAUTO are a bit fragile,
93+
* in the sense that once one scaling factor gets far enough from the
94+
* original scaling factor, subsequent level scaling factors quickly
95+
* diverge to either 0 or infinity. To mitigate this problem to a certain
96+
* extend, we have a special prime selection process in place. The goal is
97+
* to maintain the scaling factor of all levels as close to the original
98+
* scale factor of level 0 as possible.
99+
*/
100+
double sf = static_cast<double>(moduliQ[numLevel - 1]);
101+
for (size_t i = numLevel - 2, cnt = 0; i >= 1; --i, ++cnt) {
102+
sf = pow(sf, 2) / static_cast<double>(moduliQ[i + 1]);
103+
NativeInteger sfInt = std::llround(sf);
104+
NativeInteger sfRem = sfInt.Mod(cyclOrder);
105+
bool hasSameMod = true;
106+
if ((cnt % 2) == 0) {
107+
NativeInteger qPrev =
108+
sfInt - NativeInteger(cyclOrder) - sfRem + NativeInteger(1);
109+
while (hasSameMod) {
110+
hasSameMod = false;
111+
qPrev = lbcrypto::PreviousPrime(qPrev, cyclOrder);
112+
for (size_t j = i + 1; j < numLevel; j++) {
113+
if (qPrev == moduliQ[j]) {
114+
hasSameMod = true;
115+
break;
116+
}
117+
}
118+
}
119+
moduliQ[i] = qPrev.ConvertToInt();
120+
} else {
121+
NativeInteger qNext =
122+
sfInt + NativeInteger(cyclOrder) - sfRem + NativeInteger(1);
123+
while (hasSameMod) {
124+
hasSameMod = false;
125+
qNext = lbcrypto::NextPrime(qNext, cyclOrder);
126+
for (size_t j = i + 1; j < numLevel; j++) {
127+
if (qNext == moduliQ[j]) {
128+
hasSameMod = true;
129+
break;
130+
}
131+
}
132+
}
133+
moduliQ[i] = qNext.ConvertToInt();
134+
}
135+
if (moduliQ[i] > maxPrime)
136+
maxPrime = moduliQ[i];
137+
else if (moduliQ[i] < minPrime)
138+
minPrime = moduliQ[i];
139+
}
140+
}
141+
}
142+
143+
if (logFirstMod == logDefaultScale) { // this requires dcrtBits < 60
144+
moduliQ[0] =
145+
lbcrypto::NextPrime<lbcrypto::NativeInteger>(maxPrime, cyclOrder)
146+
.ConvertToInt();
147+
} else {
148+
moduliQ[0] =
149+
lbcrypto::LastPrime<lbcrypto::NativeInteger>(logFirstMod, cyclOrder)
150+
.ConvertToInt();
151+
152+
// find if the value of moduliQ[0] is already in the vector starting with
153+
// moduliQ[1] and if there is, then get another prime for moduliQ[0]
154+
const auto pos = std::find(moduliQ.begin() + 1, moduliQ.end(), moduliQ[0]);
155+
if (pos != moduliQ.end()) {
156+
moduliQ[0] =
157+
lbcrypto::NextPrime<lbcrypto::NativeInteger>(maxPrime, cyclOrder)
158+
.ConvertToInt();
159+
}
160+
}
161+
return moduliQ;
162+
}
163+
164+
// numScaleMod is L
14165
SchemeParam SchemeParam::getConcreteSchemeParam(
15-
std::vector<double> logqi, int logDefaultScale, int slotNumber,
16-
bool usePublicKey, bool encryptionTechniqueExtended) {
166+
int logFirstMod, int logDefaultScale, int numScaleMod, int slotNumber,
167+
bool usePublicKey, bool encryptionTechniqueExtended, bool reducedError) {
17168
// CKKS slot number = ringDim / 2
18-
return SchemeParam(RLWESchemeParam::getConcreteRLWESchemeParam(
19-
std::move(logqi), 2 * slotNumber, usePublicKey,
20-
encryptionTechniqueExtended),
21-
logDefaultScale);
169+
auto minRingDim = 2 * slotNumber;
170+
171+
auto dnum = computeDnum(numScaleMod);
172+
173+
// q0 + (q1 + ... + qL) = firstModBits + scalingModBits * L
174+
double logQ = logFirstMod + logDefaultScale * numScaleMod;
175+
// pi can be large
176+
auto sizePi = 60;
177+
178+
// make P > Q / dnum
179+
auto logP = ceil(logQ / dnum);
180+
auto numPi = ceil(logP / sizePi);
181+
// update logP
182+
logP = numPi * sizePi;
183+
184+
auto logPQ = logQ + logP;
185+
186+
// ringDim will change if newLogPQ is too large
187+
auto ringDim = computeRingDim(logPQ, minRingDim);
188+
bool redo = false;
189+
std::vector<int64_t> qiImpl;
190+
std::vector<int64_t> piImpl;
191+
do {
192+
redo = false;
193+
qiImpl.clear();
194+
piImpl.clear();
195+
196+
if (reducedError) {
197+
qiImpl = moduliQGenerationReducedError(logFirstMod, logDefaultScale,
198+
numScaleMod + 1, ringDim);
199+
} else {
200+
qiImpl = moduliQGenerationOpenFHEFixed(logFirstMod, logDefaultScale,
201+
numScaleMod + 1, ringDim);
202+
}
203+
std::vector<int64_t> existingPrimes = qiImpl;
204+
for (size_t i = 0; i < numPi; ++i) {
205+
auto prime = findPrime(sizePi, ringDim, existingPrimes);
206+
piImpl.push_back(prime);
207+
existingPrimes.push_back(prime);
208+
}
209+
210+
// if generated primes are too large, increase ringDim
211+
double newLogPQ = 0;
212+
for (auto qi : qiImpl) {
213+
newLogPQ += log2(qi);
214+
}
215+
for (auto pi : piImpl) {
216+
newLogPQ += log2(pi);
217+
}
218+
auto newRingDim = computeRingDim(newLogPQ, minRingDim);
219+
if (newRingDim != ringDim) {
220+
ringDim = newRingDim;
221+
redo = true;
222+
}
223+
} while (redo);
224+
225+
std::vector<double> logqi;
226+
std::vector<double> logpi;
227+
for (auto qi : qiImpl) {
228+
logqi.push_back(log2(qi));
229+
}
230+
for (auto pi : piImpl) {
231+
logpi.push_back(log2(pi));
232+
}
233+
234+
return SchemeParam(
235+
RLWESchemeParam(ringDim, numScaleMod + 1, logqi, qiImpl, dnum, logpi,
236+
piImpl, usePublicKey, encryptionTechniqueExtended),
237+
logDefaultScale);
22238
}
23239

24240
SchemeParam SchemeParam::getSchemeParamFromAttr(SchemeParamAttr attr) {

‎lib/Parameters/CKKS/Params.h‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,10 +26,9 @@ class SchemeParam : public RLWESchemeParam {
2626
int64_t getLogDefaultScale() const { return logDefaultScale; }
2727
void print(llvm::raw_ostream& os) const override;
2828

29-
static SchemeParam getConcreteSchemeParam(std::vector<double> logqi,
30-
int logDefaultScale, int slotNumber,
31-
bool usePublicKey,
32-
bool encryptionTechniqueExtended);
29+
static SchemeParam getConcreteSchemeParam(
30+
int logFirstMod, int logDefaultScale, int numScaleMod, int slotNumber,
31+
bool usePublicKey, bool encryptionTechniqueExtended, bool reducedError);
3332

3433
static SchemeParam getSchemeParamFromAttr(SchemeParamAttr attr);
3534
};

‎lib/Parameters/RLWEParams.h‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,14 @@ class RLWELocalParam {
131131
int getDimension() const { return dimension; }
132132
};
133133

134+
//===----------------------------------------------------------------------===//
135+
// Helper functions
136+
//===----------------------------------------------------------------------===//
137+
138+
int computeDnum(int level);
139+
int64_t findPrime(int qi, int ringDim,
140+
const std::vector<int64_t>& existingPrimes);
141+
134142
} // namespace heir
135143
} // namespace mlir
136144

‎lib/Transforms/GenerateParam/GenerateParam.td‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -144,9 +144,19 @@ def GenerateParamCKKS : Pass<"generate-param-ckks"> {
144144
the ciphertext level/dimension. These ops and attributes can be added by
145145
a pass like `--secret-insert-mgmt-<scheme>` and `--annotate-mgmt`.
146146

147-
User can provide custom scheme parameters by annotating bgv::SchemeParamAttr
147+
User can provide custom scheme parameters by annotating ckks::SchemeParamAttr
148148
at the module level.
149149

150+
There are two prime selection implementations available:
151+
1. Fixed (from OpenFHE FIXED*)
152+
2. Reduced Error (from https://eprint.iacr.org/2020/1118, OpenFHE FLEXIBLE*)
153+
154+
There is a toggle called `reduced-error` that can choose between them.
155+
The default one is "Fixed".
156+
157+
Reduced Error implementation works better with level-specific scaling
158+
factor.
159+
150160
(* example filepath=tests/Transforms/generate_param_ckks/doctest.mlir *)
151161
}];
152162

@@ -170,6 +180,9 @@ def GenerateParamCKKS : Pass<"generate-param-ckks"> {
170180
Option<"inputRange", "input-range", "int",
171181
/*default=*/"1", "The range of the plaintexts for input ciphertexts "
172182
"for the CKKS scheme; default to [-1, 1]. For other ranges like [-D, D], use D.">,
183+
Option<"reducedError", "reduced-error", "bool",
184+
/*default=*/"false", "If true, uses the prime selection logic in Reduced Error paper "
185+
"(https://eprint.iacr.org/2020/1118).">,
173186
];
174187
}
175188

‎lib/Transforms/GenerateParam/GenerateParamCKKS.cpp‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -122,13 +122,9 @@ struct GenerateParamCKKS : impl::GenerateParamCKKSBase<GenerateParamCKKS> {
122122
encryptionTechniqueExtended = true;
123123
}
124124

125-
// generate scheme parameters
126-
std::vector<double> logPrimes(maxLevel.value_or(0) + 1, scalingModBits);
127-
logPrimes[0] = firstModBits;
128-
129125
auto schemeParam = ckks::SchemeParam::getConcreteSchemeParam(
130-
logPrimes, scalingModBits, slotNumber, usePublicKey,
131-
encryptionTechniqueExtended);
126+
firstModBits, scalingModBits, maxLevel.value_or(0), slotNumber,
127+
usePublicKey, encryptionTechniqueExtended, reducedError);
132128

133129
LLVM_DEBUG(llvm::dbgs() << "Scheme Param:\n" << schemeParam << "\n");
134130

‎tests/Transforms/generate_param_ckks/doctest.mlir‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
// RUN: heir-opt --generate-param-ckks %s | FileCheck %s
22

3-
// CHECK: {ckks.schemeParam = #ckks.scheme_param<logN = 13, Q = [36028797019389953], P = [36028797019488257], logDefaultScale = 45>}
3+
// CHECK: {ckks.schemeParam = #ckks.scheme_param<logN = 13, Q = [36028797018652673], P = [1152921504606994433], logDefaultScale = 45>}
44
module {
55
func.func @add(%arg0: !secret.secret<f16> {mgmt.mgmt = #mgmt.mgmt<level = 0>}) -> (!secret.secret<f16> {mgmt.mgmt = #mgmt.mgmt<level = 0>}) {
66
%0 = secret.generic(%arg0: !secret.secret<f16> {mgmt.mgmt = #mgmt.mgmt<level = 0>}) attrs = {arg0 = {mgmt.mgmt = #mgmt.mgmt<level = 0>}} {

0 commit comments

Comments
 (0)