Skip to content

Commit

Permalink
Make PrivKSK aligned
Browse files Browse the repository at this point in the history
  • Loading branch information
@nindanaoto
nindanaoto committed Jul 15, 2024
1 parent f554442 commit d5c7f63
Show file tree
Hide file tree
Showing 7 changed files with 78 additions and 26 deletions.
22 changes: 11 additions & 11 deletions include/circuitbootstrapping.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ void CircuitBootstrapping(TRGSW<typename privksP::targetP> &trgsw,
const TLWE<typename bkP::domainP> &tlwe,
const EvalKey &ek)
{
std::array<TLWE<typename bkP::targetP>, privksP::targetP::l> temp;
alignas(64) std::array<TLWE<typename bkP::targetP>, privksP::targetP::l> temp;
GateBootstrappingManyLUT<bkP, privksP::targetP::l>(
temp, tlwe, ek.getbkfft<bkP>(), CBtestvector<privksP>());
for (int i = 0; i < privksP::targetP::l; i++) {
Expand Down Expand Up @@ -55,7 +55,7 @@ void CircuitBootstrappingFFT(TRGSWFFT<typename privksP::targetP> &trgswfft,
const TLWE<typename brP::domainP> &tlwe,
const EvalKey &ek)
{
TRGSW<typename privksP::targetP> trgsw;
alignas(64) TRGSW<typename privksP::targetP> trgsw;
CircuitBootstrapping<brP, privksP>(trgsw, tlwe, ek);
for (int i = 0; i < (privksP::targetP::k + 1) * privksP::targetP::l; i++)
for (int j = 0; j < privksP::targetP::k + 1; j++)
Expand All @@ -67,7 +67,7 @@ void CircuitBootstrappingFFT(TRGSWFFT<typename privksP::targetP> &trgswfft,
const TLWE<typename iksP::domainP> &tlwe,
const EvalKey &ek)
{
TRGSW<typename privksP::targetP> trgsw;
alignas(64) TRGSW<typename privksP::targetP> trgsw;
CircuitBootstrapping<iksP, bkP, privksP>(trgsw, tlwe, ek);
for (int i = 0; i < (privksP::targetP::k + 1) * privksP::targetP::l; i++)
for (int j = 0; j < privksP::targetP::k + 1; j++)
Expand All @@ -79,17 +79,17 @@ void CircuitBootstrappingSub(TRGSW<typename privksP::targetP> &trgsw,
const TLWE<typename iksP::domainP> &tlwe,
const EvalKey &ek)
{
TLWE<typename bkP::domainP> tlwelvl0;
alignas(64) TLWE<typename bkP::domainP> tlwelvl0;
IdentityKeySwitch<iksP>(tlwelvl0, tlwe, ek.getiksk<iksP>());
std::array<TLWE<typename bkP::targetP>, privksP::targetP::l> temp;
alignas(64) std::array<TLWE<typename bkP::targetP>, privksP::targetP::l> temp;
GateBootstrappingManyLUT<bkP, privksP::targetP::l>(
temp, tlwelvl0, ek.getbkfft<bkP>(), CBtestvector<privksP>());
for (int i = 0; i < privksP::targetP::l; i++) {
temp[i][privksP::domainP::k * privksP::domainP::n] +=
1ULL << (numeric_limits<typename privksP::domainP::T>::digits -
(i + 1) * privksP::targetP::Bgbit - 1);
for (int k = 0; k < privksP::targetP::k + 1; k++) {
TLWE<typename privksP::targetP> subsettlwe;
alignas(64) TLWE<typename privksP::targetP> subsettlwe;
SubsetIdentityKeySwitch<privksP>(subsettlwe, temp[i],
ek.getsubiksk<privksP>());
SubsetPrivKeySwitch<privksP>(
Expand All @@ -105,7 +105,7 @@ void CircuitBootstrappingSubFFT(TRGSWFFT<typename privksP::targetP> &trgswfft,
const TLWE<typename iksP::domainP> &tlwe,
const EvalKey &ek)
{
TRGSW<typename privksP::targetP> trgsw;
alignas(64) TRGSW<typename privksP::targetP> trgsw;
CircuitBootstrappingSub<iksP, bkP, privksP>(trgsw, tlwe, ek);
for (int i = 0; i < (privksP::targetP::k + 1) * privksP::targetP::l; i++)
for (int j = 0; j < privksP::targetP::k + 1; j++)
Expand All @@ -117,7 +117,7 @@ void CircuitBootstrappingFFTInv(
TRGSWFFT<typename privksP::targetP> &invtrgswfft,
const TLWE<typename brP::domainP> &tlwe, const EvalKey &ek)
{
TLWE<typename brP::domainP> invtlwe;
alignas(64) TLWE<typename brP::domainP> invtlwe;
// HomNot
for (int i = 0; i <= brP::domainP::k * brP::domainP::n; i++)
invtlwe[i] = -tlwe[i];
Expand All @@ -129,7 +129,7 @@ void CircuitBootstrappingFFTInv(
TRGSWFFT<typename privksP::targetP> &invtrgswfft,
const TLWE<typename iksP::domainP> &tlwe, const EvalKey &ek)
{
TLWE<typename iksP::domainP> invtlwe;
alignas(64) TLWE<typename iksP::domainP> invtlwe;
// HomNot
for (int i = 0; i <= iksP::domainP::k * iksP::domainP::n; i++)
invtlwe[i] = -tlwe[i];
Expand All @@ -145,7 +145,7 @@ void CircuitBootstrappingFFTwithInv(
constexpr array<typename privksP::targetP::T, privksP::targetP::l> h =
hgen<typename privksP::targetP>();

TRGSW<typename privksP::targetP> trgsw;
alignas(64) TRGSW<typename privksP::targetP> trgsw;
CircuitBootstrapping<brP, privksP>(trgsw, tlwe, ek);
for (int i = 0; i < (privksP::targetP::k + 1) * privksP::targetP::l; i++)
for (int j = 0; j < privksP::targetP::k + 1; j++) {
Expand All @@ -171,7 +171,7 @@ void CircuitBootstrappingFFTwithInv(
constexpr array<typename privksP::targetP::T, privksP::targetP::l> h =
hgen<typename privksP::targetP>();

TRGSW<typename privksP::targetP> trgsw;
alignas(64) TRGSW<typename privksP::targetP> trgsw;
CircuitBootstrapping<iksP, bkP, privksP>(trgsw, tlwe, ek);
for (int i = 0; i < (privksP::targetP::k + 1) * privksP::targetP::l; i++)
for (int j = 0; j < privksP::targetP::k + 1; j++) {
Expand Down
9 changes: 3 additions & 6 deletions include/cloudkey.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -576,18 +576,15 @@ struct EvalKey {
const SecretKey& sk)
{
if constexpr (std::is_same_v<P, lvl11param>) {
privksklvl11[key] = std::make_unique_for_overwrite<
PrivateKeySwitchingKey<lvl11param>>();
privksklvl11[key] = std::unique_ptr<PrivateKeySwitchingKey<lvl11param>>(new (std::align_val_t(64)) PrivateKeySwitchingKey<lvl11param>());
privkskgen<lvl11param>(*privksklvl11[key], func, sk);
}
else if constexpr (std::is_same_v<P, lvl21param>) {
privksklvl21[key] = std::make_unique_for_overwrite<
PrivateKeySwitchingKey<lvl21param>>();
privksklvl21[key] = std::unique_ptr<PrivateKeySwitchingKey<lvl21param>>(new (std::align_val_t(64)) PrivateKeySwitchingKey<lvl21param>());
privkskgen<lvl21param>(*privksklvl21[key], func, sk);
}
else if constexpr (std::is_same_v<P, lvl22param>) {
privksklvl22[key] = std::make_unique_for_overwrite<
PrivateKeySwitchingKey<lvl22param>>();
privksklvl22[key] = std::unique_ptr<PrivateKeySwitchingKey<lvl22param>>(new (std::align_val_t(64)) PrivateKeySwitchingKey<lvl22param>());
privkskgen<lvl22param>(*privksklvl22[key], func, sk);
}
else
Expand Down
2 changes: 1 addition & 1 deletion include/detwfa.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ void CMUXFFTwithPolynomialMulByXaiMinusOne(
const BootstrappingKeyElementFFT<bkP> &cs, const int a)
{
if constexpr (bkP::domainP::key_value_diff == 1) {
TRLWE<typename bkP::targetP> temp;
alignas(64) TRLWE<typename bkP::targetP> temp;
for (int k = 0; k < bkP::targetP::k + 1; k++)
PolynomialMulByXaiMinusOne<typename bkP::targetP>(temp[k], acc[k],
a);
Expand Down
4 changes: 2 additions & 2 deletions include/keyswitch.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -264,8 +264,8 @@ void PrivKeySwitch(TRLWE<typename P::targetP> &res,
mask;

if (aij != 0) {
for (int p = 0; p < P::targetP::n; p++)
for (int k = 0; k < P::targetP::k + 1; k++)
for (int k = 0; k < P::targetP::k + 1; k++)
for (int p = 0; p < P::targetP::n; p++)
res[k][p] -= privksk[i][j][aij - 1][k][p];
}
}
Expand Down
55 changes: 55 additions & 0 deletions include/utils.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#include <functional>
#include <limits>
#include <random>
#include <cstdlib>

namespace TFHEpp {
#ifdef USE_RANDEN
Expand All @@ -36,6 +37,60 @@ concept hasqbit = requires
T::qbit;
};


// https://github.com/zhourrr/aligned-memory-allocator/blob/main/aligned_allocator.h
// A minimal implementation of an allocator for C++ Standard Library, which
// allocates aligned memory (specified by the alignment argument).
// Note:
// A minimal custom allocator is preferred because C++ allocator_traits class
// provides default implementation for you. Take a look at Microsoft's
// documentation about Allocators and allocator class.
template <typename T, std::size_t alignment> class AlignedAllocator {
public:
using value_type = T;

public:
// According to Microsoft's documentation, default constructor is not required
// by C++ Standard Library.
AlignedAllocator() noexcept {};

template <typename U> AlignedAllocator(const AlignedAllocator<U, alignment>& other) noexcept {};

template <typename U>
inline bool operator==(const AlignedAllocator<U, alignment>& other) const noexcept {
return true;
}

template <typename U>
inline bool operator!=(const AlignedAllocator<U, alignment>& other) const noexcept {
return false;
}

template <typename U> struct rebind {
using other = AlignedAllocator<U, alignment>;
};

// STL containers call this function to allocate uninitialized memory block to
// store (no more than n) elements of type T (value_type).
inline value_type* allocate(const std::size_t n) const {
auto size = n;
/*
If you wish, for some strange reason, that the size of allocated buffer is
also aligned to alignment, uncomment the following statement.
Note: this increases the size of underlying memory, but STL containers
still treat it as a memory block of size n, i.e., STL containers will not
put more than n elements into the returned memory.
*/
// size = (n + alignment - 1) / alignment * alignment;
return static_cast<value_type *>(std::aligned_alloc(alignment, sizeof(T) * size));
};

// STL containers call this function to free a memory block beginning at a
// specified position.
inline void deallocate(value_type* const ptr, std::size_t n) const noexcept { std::free(ptr); }
};

// Double to Torus(32bit fixed-point number)
inline uint16_t dtot16(double d)
{
Expand Down
6 changes: 3 additions & 3 deletions test/circuitbootstrapping.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,9 +36,9 @@ int main()
for (int j = 0; j < privksP::targetP::n; j++)
pmu[i][j] = pa[i][j] ? privksP::targetP::μ : -privksP::targetP::μ;
for (int i = 0; i < num_test; i++) pones[i] = true;
std::vector<TFHEpp::TRLWE<typename privksP::targetP>> ca(num_test);
std::vector<TFHEpp::TLWE<typename iksP::domainP>> cones(num_test);
std::vector<TFHEpp::TRGSWFFT<typename privksP::targetP>> bootedTGSW(
alignas(64) std::vector<TFHEpp::TRLWE<typename privksP::targetP>> ca(num_test);
alignas(64) std::vector<TFHEpp::TLWE<typename iksP::domainP>> cones(num_test);
std::vector<TFHEpp::TRGSWFFT<typename privksP::targetP>,TFHEpp::AlignedAllocator<TFHEpp::TRGSWFFT<typename privksP::targetP>,64>> bootedTGSW(
num_test);

for (int i = 0; i < num_test; i++)
Expand Down
6 changes: 3 additions & 3 deletions test/nested_cmux.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,16 +49,16 @@ int main()

SecretKey skey;

std::vector<TRGSWLvl1FFT> guard;
alignas(64) std::vector<TRGSWLvl1FFT> guard;
TFHEpp::Polynomial<TFHEpp::lvl1param> plainpoly = {};
plainpoly[0] = 1;
for (size_t i = 0; i < N; i++)
guard.push_back(
TFHEpp::trgswfftSymEncrypt<Lvl1>(plainpoly, skey.key.lvl1));

TRLWELvl1 c1 = trivial_TRLWELvl1(uint2weight(1)),
alignas(64) TRLWELvl1 c1 = trivial_TRLWELvl1(uint2weight(1)),
c0 = trivial_TRLWELvl1(uint2weight(0));
TRLWELvl1 res = c1;
alignas(64) TRLWELvl1 res = c1;
dump_histgram_of_phase_of_TRLWELvl1(
std::cout, TFHEpp::trlwePhase<TFHEpp::lvl1param>(res, skey.key.lvl1));
for (size_t i = 0; i < N; i++) {
Expand Down

0 comments on commit d5c7f63

Please sign in to comment.