Skip to content

Add Engine abstraction to enable multiple backends #26

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ryanorendorff wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add Engine abstraction to enable multiple backends #26

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -225,6 +225,8 @@ target_sources(fideslib PUBLIC

target_include_directories(fideslib
PRIVATE nvtx/c/include
# api/ is the include root so headers use api-relative paths (e.g. "engine/cuda/CudaEngine.hpp").
PUBLIC $<BUILD_INTERFACE:${API_DIR}>
)

target_link_libraries(fideslib
Expand Down
64 changes: 36 additions & 28 deletions api/CCParams.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,84 +4,83 @@
#include <openfhe.h>

#include <cassert>
#include <iostream>
#include <string>

namespace fideslib {

CCParams<CryptoContextCKKSRNS>::CCParams() {
lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS> params;
this->cpu = std::make_any<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>>(std::move(params));
this->host = std::make_any<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>>(std::move(params));
}

// ---- CKKS Parameters ----

void CCParams<CryptoContextCKKSRNS>::SetMultiplicativeDepth(uint32_t depth) {
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(cpu);
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(host);
params.SetMultiplicativeDepth(depth);
}

void CCParams<CryptoContextCKKSRNS>::SetScalingModSize(uint32_t size) {
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(cpu);
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(host);
params.SetScalingModSize(size);
}

void CCParams<CryptoContextCKKSRNS>::SetBatchSize(uint32_t size) {
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(cpu);
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(host);
params.SetBatchSize(size);
}

void CCParams<CryptoContextCKKSRNS>::SetRingDim(uint32_t dim) {
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(cpu);
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(host);
params.SetRingDim(dim);
}

void CCParams<CryptoContextCKKSRNS>::SetScalingTechnique(ScalingTechnique tech) {
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(cpu);
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(host);
auto scale_openfhe = static_cast<lbcrypto::ScalingTechnique>(tech);
assert((int)scale_openfhe == (int)tech);
params.SetScalingTechnique(scale_openfhe);
}

void CCParams<CryptoContextCKKSRNS>::SetNumLargeDigits(uint32_t numDigits) {
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(cpu);
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(host);
params.SetNumLargeDigits(numDigits);
}

void CCParams<CryptoContextCKKSRNS>::SetFirstModSize(uint32_t size) {
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(cpu);
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(host);
params.SetFirstModSize(size);
}

void CCParams<CryptoContextCKKSRNS>::SetDigitSize(uint32_t size) {
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(cpu);
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(host);
params.SetDigitSize(size);
}

void CCParams<CryptoContextCKKSRNS>::SetKeySwitchTechnique(KeySwitchTechnique tech) {
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(cpu);
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(host);
auto ks_openfhe = static_cast<lbcrypto::KeySwitchTechnique>(tech);
assert((int)ks_openfhe == (int)tech);
params.SetKeySwitchTechnique(ks_openfhe);
}

void CCParams<CryptoContextCKKSRNS>::SetSecretKeyDist(SecretKeyDist dist) {
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(cpu);

if (this->devices.empty()) {
if (dist == SecretKeyDist::SPARSE_TERNARY) {
params.SetSecretKeyDist(lbcrypto::SPARSE_TERNARY);
}
else {
params.SetSecretKeyDist(lbcrypto::UNIFORM_TERNARY);
}
}
else {
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(host);

// Record the requested distribution as-is. The CUDA backend's restriction against sparse-ternary keys
// is enforced at context construction (GenCryptoContext), so an unsupported choice fails loudly
// there rather than being silently downgraded here.
if (dist == SecretKeyDist::SPARSE_TERNARY) {
params.SetSecretKeyDist(lbcrypto::SPARSE_TERNARY);
} else {
params.SetSecretKeyDist(lbcrypto::UNIFORM_TERNARY);
}
keyDist = dist;
}

void CCParams<CryptoContextCKKSRNS>::SetSecurityLevel(SecurityLevel level) {
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(cpu);
auto& params = std::any_cast<lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(host);
auto sl_openfhe = static_cast<lbcrypto::SecurityLevel>(level);
assert((int)sl_openfhe == (int)level);
params.SetSecurityLevel(sl_openfhe);
Expand All @@ -90,25 +89,34 @@ void CCParams<CryptoContextCKKSRNS>::SetSecurityLevel(SecurityLevel level) {
// ---- Getters ----

SecretKeyDist CCParams<CryptoContextCKKSRNS>::GetSecretKeyDist() const {
auto& params = std::any_cast<const lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(cpu);
auto& params = std::any_cast<const lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(host);
auto skd_openfhe = params.GetSecretKeyDist();
return static_cast<SecretKeyDist>(skd_openfhe);
}

uint32_t CCParams<CryptoContextCKKSRNS>::GetMultiplicativeDepth() const {
auto& params = std::any_cast<const lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(cpu);
auto& params = std::any_cast<const lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(host);
return params.GetMultiplicativeDepth();
}

uint32_t CCParams<CryptoContextCKKSRNS>::GetBatchSize() const {
auto& params = std::any_cast<const lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(cpu);
auto& params = std::any_cast<const lbcrypto::CCParams<lbcrypto::CryptoContextCKKSRNS>&>(host);
return params.GetBatchSize();
}

// ---- Device Parameters ----
// ---- Backend Parameters ----

void CCParams<CryptoContextCKKSRNS>::SetDevices(std::vector<int>&& devices) {
this->devices = std::move(devices);
void CCParams<CryptoContextCKKSRNS>::SetBackend(Backend backend) {
if (!IsBackendAvailable(backend)) {
// Build a descriptive name for the unavailable backend.
const char* name = "unknown";
switch (backend) {
case Backend::CPU: name = "CPU"; break;
case Backend::CUDA: name = "CUDA"; break;
}
OPENFHE_THROW(std::string(name) + " backend requested but not available (FIDESlib was not built with support for this backend)");
}
this->backend = backend;
}

void CCParams<CryptoContextCKKSRNS>::SetPlaintextAutoload(bool autoload) {
Expand Down
17 changes: 10 additions & 7 deletions api/CCParams.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
#include <vector>

#include "Definitions.hpp"
#include "engine/Backend.hpp"

namespace fideslib {

Expand Down Expand Up @@ -44,9 +45,11 @@ template <> class CCParams<CryptoContextCKKSRNS> {
void SetSecretKeyDist(SecretKeyDist dist);
void SetSecurityLevel(SecurityLevel level);

// ---- Device Parameters ----
// ---- Backend Parameters ----

void SetDevices(std::vector<int>&& devices);
/// @brief Select the execution backend (default CPU). Multi-GPU device selection is done
/// later through the facade method CryptoContextImpl::SetCudaDevices, before LoadContext.
void SetBackend(Backend backend);
void SetPlaintextAutoload(bool autoload);
void SetCiphertextAutoload(bool autoload);

Expand All @@ -57,11 +60,11 @@ template <> class CCParams<CryptoContextCKKSRNS> {

// ---- Internal State ----

std::any cpu;
std::vector<int> devices = { };
SecretKeyDist keyDist = UNIFORM_TERNARY;
bool plaintextAutoload = false;
bool ciphertextAutoload = true;
std::any host;
Backend backend = Backend::CPU;
SecretKeyDist keyDist = UNIFORM_TERNARY;
bool plaintextAutoload = false;
bool ciphertextAutoload = true;
};

} // namespace fideslib
Expand Down
139 changes: 55 additions & 84 deletions api/Ciphertext.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,19 +1,11 @@
#include "Ciphertext.hpp"
#include "CKKS/Ciphertext.cuh"
#include "Definitions.hpp"

#include <iostream>
#include <openfhe.h>

namespace fideslib {

CiphertextImpl<DCRTPoly>::~CiphertextImpl() {
if (this->loaded && this->gpu != 0 && this->parent_context) {
this->parent_context->EvictDeviceCiphertext(this->gpu);
this->gpu = 0;
}
}

CiphertextImpl<DCRTPoly>::CiphertextImpl(const CryptoContext<DCRTPoly>&& context) : parent_context(context) {
if (!context) {
OPENFHE_THROW("Cannot create Ciphertext with null CryptoContext");
Expand All @@ -23,22 +15,17 @@ CiphertextImpl<DCRTPoly>::CiphertextImpl(const CryptoContext<DCRTPoly>&& context
// ---- Copy ----

CiphertextImpl<DCRTPoly>::CiphertextImpl(const CiphertextImpl<DCRTPoly>& other) {
// Share CPU ciphertext and detach only on first CPU mutation.
auto const& other_cpu = std::any_cast<const lbcrypto::Ciphertext<lbcrypto::DCRTPoly>&>(other.cpu);
this->cpu = std::make_any<lbcrypto::Ciphertext<lbcrypto::DCRTPoly>>(other_cpu);
this->need_lazy_copy = true;

// Copy underlying GPU ciphertext if loaded.
this->loaded = other.loaded;
if (this->loaded) {
this->gpu = other.parent_context->CopyDeviceCiphertext(other);
} else {
this->gpu = 0;
}
// Share host ciphertext and detach only on first host mutation.
auto const& other_host = std::any_cast<const lbcrypto::Ciphertext<lbcrypto::DCRTPoly>&>(other.host);
this->host = std::make_any<lbcrypto::Ciphertext<lbcrypto::DCRTPoly>>(other_host);
this->need_lazy_copy = true;

// Copy parent context.
this->parent_context = other.parent_context;
this->original_level = other.original_level;

// Cloning any backend-resident payload is the engine's decision, not ours: it returns an empty
// slot when `other` is not resident (always so on the CPU backend).
this->device = other.parent_context->CloneCiphertextBackend(other);
}

CiphertextImpl<DCRTPoly>::CiphertextImpl(const Ciphertext<DCRTPoly>& other) : CiphertextImpl<DCRTPoly>(static_cast<const CiphertextImpl<DCRTPoly>&>(other)) {
Expand All @@ -51,94 +38,78 @@ Ciphertext<DCRTPoly> CiphertextImpl<DCRTPoly>::Clone() const {
return clone;
}

// ---- Getters ----
// ---- Getters / setters ----
//
// These delegate unconditionally to the backend; the engine decides whether to read the host value or
// the device-resident copy. The host primitives below hold the host computation the engines reuse.

size_t CiphertextImpl<DCRTPoly>::GetLevel() const {

if (!this->loaded) {
// Fall back to CPU.
auto& ct = std::any_cast<const lbcrypto::Ciphertext<lbcrypto::DCRTPoly>&>(this->cpu);
return ct->GetLevel();
}

// GPU path. Depth is reversed in FIDESlib, must do depth = maxDepth - depth
auto ct_gpu = std::static_pointer_cast<FIDESlib::CKKS::Ciphertext>(this->parent_context->GetDeviceCiphertext(this->gpu));
auto maxDepth = this->parent_context->multiplicative_depth;
return maxDepth - ct_gpu->getLevel();
return this->parent_context->CiphertextLevel(*this);
}

size_t CiphertextImpl<DCRTPoly>::GetNoiseScaleDeg() const {
return this->parent_context->CiphertextNoiseScaleDeg(*this);
}

if (!this->loaded) {
// Fall back to CPU.
auto& ct = std::any_cast<const lbcrypto::Ciphertext<lbcrypto::DCRTPoly>&>(this->cpu);
return ct->GetNoiseScaleDeg();
}
void CiphertextImpl<DCRTPoly>::SetSlots(size_t slots) {
this->parent_context->SetCiphertextSlots(*this, slots);
}

// GPU path.
auto ct_gpu = std::static_pointer_cast<FIDESlib::CKKS::Ciphertext>(this->parent_context->GetDeviceCiphertext(this->gpu));
return ct_gpu->NoiseLevel;
void CiphertextImpl<DCRTPoly>::SetLevel(size_t level) {
this->parent_context->SetCiphertextLevel(*this, level);
}

// ---- Setters ----
// ---- Host primitives (invoked by the engine backends; no backend logic of their own) ----

void CiphertextImpl<DCRTPoly>::SetSlots(size_t slots) {
size_t CiphertextImpl<DCRTPoly>::GetLevelHost() const {
auto& ct = std::any_cast<const lbcrypto::Ciphertext<lbcrypto::DCRTPoly>&>(this->host);
return ct->GetLevel();
}

if (!this->loaded) {
// Fall back to CPU.
this->EnsureLazyCPUCopy();
auto& ct = std::any_cast<lbcrypto::Ciphertext<lbcrypto::DCRTPoly>&>(this->cpu);
ct->SetSlots(slots);
return;
}
// GPU path.
auto ct_gpu = std::static_pointer_cast<FIDESlib::CKKS::Ciphertext>(this->parent_context->GetDeviceCiphertext(this->gpu));
ct_gpu->slots = static_cast<int>(slots);
size_t CiphertextImpl<DCRTPoly>::GetNoiseScaleDegHost() const {
auto& ct = std::any_cast<const lbcrypto::Ciphertext<lbcrypto::DCRTPoly>&>(this->host);
return ct->GetNoiseScaleDeg();
}

void CiphertextImpl<DCRTPoly>::SetLevel(size_t level) {
void CiphertextImpl<DCRTPoly>::SetSlotsHost(size_t slots) {
this->EnsureLazyHostCopy();
auto& ct = std::any_cast<lbcrypto::Ciphertext<lbcrypto::DCRTPoly>&>(this->host);
ct->SetSlots(slots);
}

if (!this->loaded) {
// Fall back to CPU.
this->EnsureLazyCPUCopy();
auto& ct = std::any_cast<lbcrypto::Ciphertext<lbcrypto::DCRTPoly>&>(this->cpu);
void CiphertextImpl<DCRTPoly>::SetLevelHost(size_t level) {
this->EnsureLazyHostCopy();
auto& ct = std::any_cast<lbcrypto::Ciphertext<lbcrypto::DCRTPoly>&>(this->host);

size_t currentTowers = ct->GetElements()[0].GetNumOfElements();
size_t currentLevel = ct->GetLevel();
size_t currentTowers = ct->GetElements()[0].GetNumOfElements();
size_t currentLevel = ct->GetLevel();

size_t totalPrimes = currentTowers + currentLevel;
size_t targetTowers = totalPrimes - level;
size_t totalPrimes = currentTowers + currentLevel;
size_t targetTowers = totalPrimes - level;

if (currentTowers > targetTowers) {
// Need to drop towers
size_t towersToDrop = currentTowers - targetTowers;
if (currentTowers > targetTowers) {
// Need to drop towers
size_t towersToDrop = currentTowers - targetTowers;

auto& elements = ct->GetElements();
for (auto& elem : elements) {
elem.DropLastElements(towersToDrop);
}
auto& elements = ct->GetElements();
for (auto& elem : elements) {
elem.DropLastElements(towersToDrop);
}

ct->SetLevel(level);

return;
}

// GPU path.
auto ct_gpu = std::static_pointer_cast<FIDESlib::CKKS::Ciphertext>(this->parent_context->GetDeviceCiphertext(this->gpu));
auto maxDepth = this->parent_context->multiplicative_depth;
ct_gpu->dropToLevel(maxDepth - level);
ct->SetLevel(level);
}

void CiphertextImpl<DCRTPoly>::EnsureLazyCPUCopy() {
if (!this->need_lazy_copy) {
void CiphertextImpl<DCRTPoly>::EnsureLazyHostCopy() {
auto const& ct_host = std::any_cast<const lbcrypto::Ciphertext<lbcrypto::DCRTPoly>&>(this->host);
// Detach before in-place mutation if still lazily shared or the underlying
// OpenFHE ciphertext has other owners (e.g. a Clone): copy-on-write.
if (!this->need_lazy_copy && ct_host.use_count() <= 1) {
return;
}

auto const& ct_cpu = std::any_cast<const lbcrypto::Ciphertext<lbcrypto::DCRTPoly>&>(this->cpu);
lbcrypto::Ciphertext<lbcrypto::DCRTPoly> cpu_copy = std::make_shared<lbcrypto::CiphertextImpl<lbcrypto::DCRTPoly>>(*ct_cpu);
this->cpu = std::make_any<lbcrypto::Ciphertext<lbcrypto::DCRTPoly>>(std::move(cpu_copy));
this->need_lazy_copy = false;
lbcrypto::Ciphertext<lbcrypto::DCRTPoly> host_copy = std::make_shared<lbcrypto::CiphertextImpl<lbcrypto::DCRTPoly>>(*ct_host);
this->host = std::make_any<lbcrypto::Ciphertext<lbcrypto::DCRTPoly>>(std::move(host_copy));
this->need_lazy_copy = false;
}

// ---- Operators ----
Expand Down
Loading