I am working on a project with OpenFHE that requires a large amount of rotation keys so I am looking for ways to reduce my memory usage by switching context of my program in and out of memory. I am using the the following steps:
- Generate Context with application public and private keys and save in disk.
- Generate rotation keys for different sections of the program.
- Clear Context from memory to remove all keys and data not required
- Load Context into memory
- Load just the rotation keys needed for each section. After computing the section, I can clear these keys and context from memory as well. Load the context and keys for the next section and so forth.
For some weird reason, when I clear the context from memory and load it back from disk, I can not use it.
I get the following error
pke/cryptocontext.h:l.344:ValidateKey<std::shared_ptr<lbcrypto::PublicKeyImpl<lbcrypto::DCRTPolyImpl<bigintdyn::mubintvec<bigintdyn::ubint<long unsigned int> > > > > >(): Key was not generated with the same crypto context
Aborted (core dumped)
This suggest that my public key was not generated with the same context which is not true since I do not regenerate the context anywhere in the program.
My implementation looks something like this:
CryptoContext<DCRTPoly> context; // public variable and shared across multiple controllers.
int main(int argc, char *argv[]) {
-----------------
FHEController fheController(context); // creating my fhe controller
fheController.generate_context(15, 16, 12); // generating fhe context with public and private keys and storing it.
-----------------
fheController.generate_bootstrapping_and_rotation_keys("section1_keys.bin");
fheController.clear_bootstrapping_and_rotation_keys();
-----------------
fheController.clear_context(num_slots);
fheController.load_context(true);
-----------------
Ctext enData = fheController.encrypt_packedData(data);
}
From the code snippet above, if I remove the following lines, I won’t have any issues but they are needed for scalability and for the system to work.
fheController.clear_context(num_slots);
fheController.load_context(true);
I believe my issue should be coming from either how I generate, serialization and store the context or how I read the context back into my program. Below are the functions to do all of this;
void FHEController::generate_context(bool serialize, int numSlots, int ringDim, int mlevelBootstrap) {
CCParams<CryptoContextCKKSRNS> parameters;
num_slots = 1 << numSlots;
int dcrtBits = 47;
int firstMod = 52;
auto secretKeyDist = SPARSE_TERNARY;
parameters.SetSecretKeyDist(SPARSE_TERNARY);
parameters.SetSecurityLevel(lbcrypto::HEStd_NotSet);
parameters.SetNumLargeDigits(3);
parameters.SetRingDim(1 << ringDim);
parameters.SetBatchSize(num_slots);
level_budget = {4, 4};
ScalingTechnique rescaleTech = FLEXIBLEAUTO;
parameters.SetScalingModSize(dcrtBits);
parameters.SetFirstModSize(firstMod);
parameters.SetScalingTechnique(rescaleTech);
uint32_t levelsAvailableAfterBootstrap = mlevelBootstrap;
circuit_depth = levelsAvailableAfterBootstrap + FHECKKSRNS::GetBootstrapDepth(level_budget, secretKeyDist);
parameters.SetMultiplicativeDepth(circuit_depth);
context = GenCryptoContext(parameters);
context->Enable(PKE);
context->Enable(KEYSWITCH);
context->Enable(LEVELEDSHE);
context->Enable(ADVANCEDSHE);
context->Enable(FHE);
keyPair = context->KeyGen();
context->EvalMultKeyGen(keyPair.secretKey);
context->EvalSumKeyGen(keyPair.secretKey);
std::vector<uint32_t> bsgsDim = {0, 0};
numSlots = context->GetRingDimension() / 2;
context->EvalBootstrapSetup(level_budget, bsgsDim, numSlots);
context->EvalBootstrapKeyGen(keyPair.secretKey, numSlots);
serialize_context();
}
void FHEController::serialize_context(){
if (fs::exists(sskeys_folder) && fs::is_directory(sskeys_folder)) {
for (const auto& entry : fs::directory_iterator(sskeys_folder)) {
if (fs::is_regular_file(entry)) {
cout << "Deleting file: " << entry.path().filename() << "\n";
fs::remove(entry);
}
}
cout << "All files deleted from " << sskeys_folder << endl;;
}
else {
if (!fs::create_directory(sskeys_folder)) {context->EvalMultKeyGen(keyPair.secretKey);
std::cerr << "Failed to create directory: " << sskeys_folder << endl;
return;
}
std::cout << sskeys_folder << " created to store keys " << endl;
}
ofstream multKeyFile(sskeys_folder + "/mult-keys.txt", ios::out | ios::binary);
if (multKeyFile.is_open()) {
if (!context->SerializeEvalMultKey(multKeyFile, SerType::BINARY)) {
cerr << "Error writing eval mult keys" << std::endl;
exit(1);
}
cout << "Relinearization Keys have been serialized" << std::endl;
multKeyFile.close();
}
else {
cerr << "Error serializing EvalMult keys in \"" << sskeys_folder + "/mult-keys.txt" << "\"" << endl;
exit(1);
}
if (!Serial::SerializeToFile(sskeys_folder + "/crypto-context.txt", context, SerType::BINARY)) {
cerr << "Error writing serialization of the crypto context to crypto-context.txt" << endl;
} else {
cout << "Crypto Context have been serialized" << std::endl;
}
if (!Serial::SerializeToFile(sskeys_folder + "/public-key.txt", keyPair.publicKey, SerType::BINARY)) {
cerr << "Error writing serialization of public key to public-key.txt" << endl;
} else {
cout << "Public Key has been serialized" << std::endl;
}
if (!Serial::SerializeToFile(sskeys_folder + "/secret-key.txt", keyPair.secretKey, SerType::BINARY)) {
cerr << "Error writing serialization of public key to secret-key.txt" << endl;
} else {
cout << "Secret Key has been serialized" << std::endl;
}
}
void FHEController::load_context(bool verbose) {
context->ClearEvalMultKeys();
context->ClearEvalAutomorphismKeys();
CryptoContextFactory<lbcrypto::DCRTPoly>::ReleaseAllContexts();
if (verbose) cout << "Reading serialized context..." << endl;
if (!Serial::DeserializeFromFile(sskeys_folder + "/crypto-context.txt", context, SerType::BINARY)) {
cerr << "I cannot read serialized data from: " << sskeys_folder + "/crypto-context.txt" << endl;
exit(1);
}
PublicKey<DCRTPoly> clientPublicKey;
if (!Serial::DeserializeFromFile(sskeys_folder + "/public-key.txt", clientPublicKey, SerType::BINARY)) {
cerr << "I cannot read serialized data from public-key.txt" << endl;
exit(1);
}
PrivateKey<DCRTPoly> serverSecretKey;
if (!Serial::DeserializeFromFile(sskeys_folder + "/secret-key.txt", serverSecretKey, SerType::BINARY)) {
cerr << "I cannot read serialized data from public-key.txt" << endl;
exit(1);
}
keyPair.publicKey = clientPublicKey;
keyPair.secretKey = serverSecretKey;
std::ifstream multKeyIStream(sskeys_folder + "/mult-keys.txt", ios::in | ios::binary);
if (!multKeyIStream.is_open()) {
cerr << "Cannot read serialization from " << "mult-keys.txt" << endl;
exit(1);
}
if (!context->DeserializeEvalMultKey(multKeyIStream, SerType::BINARY)) {
cerr << "Could not deserialize eval mult key file" << endl;
exit(1);
}
uint32_t approxBootstrapDepth = 4 + 4;
uint32_t levelsUsedBeforeBootstrap = 12;
circuit_depth = levelsUsedBeforeBootstrap + 2 + FHECKKSRNS::GetBootstrapDepth(approxBootstrapDepth, level_budget, SPARSE_TERNARY);
if (verbose) cout << "Circuit depth: " << circuit_depth << ", available multiplications: " << levelsUsedBeforeBootstrap - 2 << endl;
}
Ctext FHEController::encrypt_packedImage(vector<double> packedVector, int level, int num_slots) {
Ptext plaintext = context->MakeCKKSPackedPlaintext(packedVector, 1, level, nullptr, num_slots);
plaintext->SetLength(num_slots);
auto encryptImage = context->Encrypt(plaintext, keyPair.publicKey);
return encryptImage;
}
I have been stocked on this issue for days and I am out of ideas. Any feedback, pointers, ideas etc will be greatly appreciated.