Added InvSbox

Author: nindanaoto

CMakeLists.txt

@@ -165,6 +165,7 @@ add_subdirectory(src)
 
 if(ENABLE_TEST)
   add_subdirectory(test)
+  add_subdirectory(thirdparties/AES)
 endif()
 if(ENABLE_BENCHMARK)
   add_subdirectory(benchmark)

include/aes.hpp

@@ -1,7 +1,41 @@
 #pragma once
+#include <AES.h>
 //Transciphering by AES
 // Based on Hippogryph
 namespace TFHEpp
 {
 
+template <class P>
+inline Polynomial<P> AESInvSboxPoly(const uint8_t upperindex){
+    Polynomial<P> poly;
+    constexpr uint segment = P::n/16;
+    for(int i = 0; i < 16; i++)
+        for(int j = 0; j < segment/2; j++){
+            poly[i*segment+2*j] = (inv_sbox[upperindex][i]&0xF)*(1ULL<<(std::numeric_limits<typename P::T>::digits-5));
+            poly[i*segment+2*j+1] = (inv_sbox[upperindex][i]>>4)*(1ULL<<(std::numeric_limits<typename P::T>::digits-5));
+        }
+    return poly;
+}
+
+template <class iksP, class brP>
+void AESInvSbox(std::array<TLWE<typename brP::targetP>, 2> &res,
+             const std::array<TLWE<typename iksP::domainP>, 2> &tlwe,
+             const EvalKey &ek)
+{
+    std::array<std::array<TLWE<typename brP::targetP>,2>,16> midtlwes;
+    TLWE<typename iksP::targetP> shifted;
+    IdentityKeySwitch<iksP>(shifted, tlwe[0], ek.getiksk<iksP>());
+    shifted[iksP::targetP::k*iksP::targetP::n] += 1ULL<<(std::numeric_limits<typename iksP::targetP::T>::digits-6);
+    for (int i = 0; i < 16; i++)
+        GateBootstrappingManyLUT<brP, 2>(midtlwes[i], shifted, ek.getbkfft<brP>(), AESInvSboxPoly<typename brP::targetP>(i));
+    IdentityKeySwitch<iksP>(shifted, tlwe[1], ek.getiksk<iksP>());
+    shifted[iksP::targetP::k*iksP::targetP::n] += 1ULL<<(std::numeric_limits<typename iksP::targetP::T>::digits-6);
+    for(int i = 0; i < 2; i++){
+        TRLWE<typename brP::targetP> trlwe;
+        std::array<TLWE<typename iksP::domainP>, 16> tabletlwe;
+        for(int j = 0; j < 16; j++) tabletlwe[j] = midtlwes[j][i];
+        TLWE2TablePacking<typename brP::targetP, 16>(trlwe, tabletlwe, ek.getahk<typename brP::targetP>());
+        GateBootstrappingTLWE2TLWEFFT<brP>(res[i], shifted, ek.getbkfft<brP>(), trlwe);
+    }
+}
 }

include/gatebootstrapping.hpp

@@ -119,7 +119,7 @@ void BlindRotate(TRLWE<typename P::targetP> &res,
                 << bitwidth;
         if (ā == 0) continue;
         // Do not use CMUXFFT to avoid unnecessary copy.
-        CMUXFFTwithPolynomialMulByXaiMinusOne<typename P::targetP>(res,
+        CMUXFFTwithPolynomialMulByXaiMinusOne<P>(res,
                                                                    bkfft[i], ā);
     }
 #endif
@@ -196,6 +196,17 @@ void GateBootstrappingTLWE2TLWEFFT(
     SampleExtractIndex<typename P::targetP>(res, acc, 0);
 }
 
+template <class P>
+void GateBootstrappingTLWE2TLWEFFT(
+    TLWE<typename P::targetP> &res, const TLWE<typename P::domainP> &tlwe,
+    const BootstrappingKeyFFT<P> &bkfft,
+    const TRLWE<typename P::targetP> &testvector)
+{
+    alignas(64) TRLWE<typename P::targetP> acc;
+    BlindRotate<P>(acc, tlwe, bkfft, testvector);
+    SampleExtractIndex<typename P::targetP>(res, acc, 0);
+}
+
 template <class P>
 void GateBootstrappingTLWE2TLWENTT(
     TLWE<typename P::targetP> &res, const TLWE<typename P::domainP> &tlwe,

include/params/128bit.hpp

@@ -102,7 +102,7 @@ struct lvl2param {
     static constexpr std::uint32_t Bgₐ = 1 << Bgₐbit;
     static constexpr ErrorDistribution errordist =
         ErrorDistribution::ModularGaussian;
-    static const inline double α = std::pow(2.0, -47);  // fresh noise
+    static const inline double α = std::pow(2.0, -51);  // fresh noise
     using T = uint64_t;                                 // Torus representation
     static constexpr std::make_signed_t<T> μ = 1LL << 61;
     static constexpr uint32_t plain_modulus = 8;

include/tfhe++.hpp

@@ -28,4 +28,7 @@
 #include "externs/circuitbootstrapping.hpp"
 #include "externs/gate.hpp"
 #include "externs/gatebootstrapping.hpp"
-#endif
+#endif
+
+// Application
+#include "aes.hpp"

include/tlwe.hpp

@@ -104,7 +104,7 @@ bool tlweSymDecrypt(const TLWE<P> &c, const Key<P> &key)
     return res;
 }
 
-template <class P, const uint plain_modulus>
+template <class P, uint plain_modulus = P::plain_modulus>
 typename P::T tlweSymIntDecrypt(const TLWE<P> &c, const Key<P> &key)
 {
     constexpr double Δ =
@@ -117,16 +117,10 @@ typename P::T tlweSymIntDecrypt(const TLWE<P> &c, const Key<P> &key)
     return res >= plain_modulus / 2 ? res - plain_modulus : res;
 }
 
-template <class P>
-typename P::T tlweSymIntDecrypt(const TLWE<P> &c, const Key<P> &key)
-{
-    return tlweSymIntDecrypt<P, P::plain_modulus>(c, key);
-}
-
-template <class P>
+template <class P, uint plain_modulus=P::plain_modulus>
 typename P::T tlweSymIntDecrypt(const TLWE<P> &c, const SecretKey &sk)
 {
-    return tlweSymIntDecrypt<P, P::plain_modulus>(c, sk.key.get<P>());
+    return tlweSymIntDecrypt<P, plain_modulus>(c, sk.key.get<P>());
 }
 
 template <class P, std::make_signed_t<typename P::T> μ>

src/CMakeLists.txt

@@ -18,6 +18,7 @@ target_include_directories(
     ${PROJECT_SOURCE_DIR}/thirdparties/spqliox_aarch64/xbyak_aarch64/xbyak_aarch64
     ${PROJECT_SOURCE_DIR}/thirdparties/hexl/hexl/hexl/include
     ${PROJECT_SOURCE_DIR}/thirdparties/randen
+    ${PROJECT_SOURCE_DIR}/thirdparties/AES/src
     ${PROJECT_SOURCE_DIR}/thirdparties/HLS_arbitrary_Precision_Types/include
     ${PROJECT_SOURCE_DIR}/thirdparties/cereal/include)
 

test/CMakeLists.txt

@@ -9,9 +9,9 @@ foreach(test_source ${test_sources})
     add_executable(${test_name} ${test_source})
     if(USE_PERF)
         add_compile_definitions(USE_PERF)
-        target_link_libraries(${test_name} tfhe++  profiler)
+        target_link_libraries(${test_name} tfhe++ AES profiler)
     else()
-        target_link_libraries(${test_name} tfhe++)
+        target_link_libraries(${test_name} tfhe++ AES)
     endif()
 endforeach(test_source ${test_sources})
 

test/aesinvsbox.cpp

@@ -0,0 +1,47 @@
+#include <cassert>
+#include <chrono>
+#include <iostream>
+#include <memory>
+#include <random>
+#include <tfhe++.hpp>
+
+int main(){
+    using brP = TFHEpp::lvlh2param;
+    using iksP = TFHEpp::lvl2hparam;
+    std::random_device seed_gen;
+    std::default_random_engine engine(seed_gen());
+    constexpr uint32_t plain_modulus = 1<<(4+1);
+    std::unique_ptr<TFHEpp::SecretKey> sk(new TFHEpp::SecretKey());
+    constexpr uint num_test = 1U<<8;
+    // constexpr uint num_test = 1U<<4;
+    std::vector<std::array<TFHEpp::TLWE<typename iksP::domainP>,2>> cin(num_test);
+
+    for (int i = 0; i < num_test; i++){
+        cin[i][0] = TFHEpp::tlweSymIntEncrypt<typename iksP::domainP, plain_modulus>(i&0xF, *sk);
+        cin[i][1] = TFHEpp::tlweSymIntEncrypt<typename iksP::domainP, plain_modulus>((i>>4), *sk);
+    }
+    std::vector<std::array<TFHEpp::TLWE<typename brP::targetP>,2>> cres(num_test);
+    TFHEpp::EvalKey ek;
+    ek.emplacebkfft<brP>(*sk);
+    ek.emplaceiksk<iksP>(*sk);
+    ek.emplaceahk<typename brP::targetP>(*sk);
+
+    std::chrono::system_clock::time_point start, end;
+    start = std::chrono::system_clock::now();
+    for (int test = 0; test < num_test; test++) {
+        std::cout<<"test: " << test << std::endl;
+        TFHEpp::AESInvSbox<iksP,brP>(cres[test], cin[test], ek);
+    }
+
+    end = std::chrono::system_clock::now();
+    double elapsed =
+        std::chrono::duration_cast<std::chrono::milliseconds>(end - start)
+            .count();
+    std::cout << elapsed / num_test << "ms" << std::endl;
+    for (int i = 0; i < num_test; i++) {
+        const uint8_t pres = (TFHEpp::tlweSymIntDecrypt<typename brP::targetP, plain_modulus>(cres[i][1], *sk)<<4)+TFHEpp::tlweSymIntDecrypt<typename brP::targetP, plain_modulus>(cres[i][0], *sk);
+        // std::cout << "test: " << i << " pres: " << (int)pres << " expected: " << (int)inv_sbox[i>>4][i&0xF] << std::endl;
+        assert(pres == inv_sbox[i>>4][i&0xF]);
+    }
+    std::cout << "Passed" << std::endl;
+}