feat: Add Circuit Bootstrapping

sp301415 · Jan 16, 2025 · b7ae6d7 · b7ae6d7
1 parent 5b1695d
commit b7ae6d7
Show file tree

Hide file tree

Showing 13 changed files with 470 additions and 50 deletions.
diff --git a/README.md b/README.md
@@ -11,7 +11,7 @@
 > TFHE-go is still under heavy development. There may be backward-incompatible changes at any time.
 
 **TFHE-go** is a Go implementation of TFHE[[CGGI16](https://eprint.iacr.org/2016/870)] and Multi-Key TFHE[[KMS22](https://eprint.iacr.org/2022/1460)] scheme. It provides:
-- Support for binary and integer TFHE and its multi-key variant
+- Support for binary and integer TFHE and its multi-key variant, as well as advanced algorithms such as BFV-style evaluation, PBSManyLUT[[CLOT21](https://eprint.iacr.org/2021/729)], Circuit Bootstrapping[[WHS+24](https://eprint.iacr.org/2024/1318)] and LMKCDEY(FHEW)[[LMK+22](https://eprint.iacr.org/2022/198)].
 - Pure Go implementation, along with SIMD-accelerated Go Assembly on amd64 platforms
 - Comparable performance to state-of-the-art C++/Rust libraries
 - Readable code and user-friendly API using modern Go features like generics
@@ -143,10 +143,10 @@ fmt.Println(dec.DecryptLWEBool(ctOut)) // false
 ## Benchmarks
 All benchmarks were measured on a machine equipped with Intel Xeon Platinum 8268 CPU @ 2.90GHz and 384GB of RAM.
 
-|Operation|TFHE-go|TFHE-rs (v0.8.0)|
+|Operation|TFHE-go|TFHE-rs (v0.11.0)|
 |---------|-------|-------|
-|Gate Bootstrapping|9.38ms|15.70ms|
-|Programmable Bootstrapping (6 bits)|21.52ms|105.05ms|
+|Gate Bootstrapping|9.38ms|14.17ms|
+|Programmable Bootstrapping (6 bits)|21.52ms|107.96ms|
 
 You can use the standard go test tool to reproduce benchmarks:
 ```
@@ -190,3 +190,6 @@ TFHE-go logo is designed by [@mlgng2010](https://www.instagram.com/mlgng2010/),
 - New Secret Keys for Enhanced Performance in (T)FHE (https://eprint.iacr.org/2023/979)
 - TFHE Public-Key Encryption Revisited (https://eprint.iacr.org/2023/603)
 - Towards Practical Multi-key TFHE: Parallelizable, Key-Compatible, Quasi-linear Complexity (https://eprint.iacr.org/2022/1460)
+- Improved Programmable Bootstrapping with Larger Precision and Efficient Arithmetic Circuits for TFHE (https://eprint.iacr.org/2021/729)
+- FHEW-like Leveled Homomorphic Evaluation: Refined Workflow and Polished Building Blocks (https://eprint.iacr.org/2024/1318)
+- Efficient FHEW Bootstrapping with Small Evaluation Keys, and Applications to Threshold Homomorphic Encryption (https://eprint.iacr.org/2022/198)
diff --git a/mktfhe/params.go b/mktfhe/params.go
@@ -115,17 +115,19 @@ func (p ParametersLiteral[T]) WithBootstrapOrder(bootstrapOrder tfhe.BootstrapOr
 // Unless you are a cryptographic expert, DO NOT set parameters by yourself;
 // always use the default parameters provided.
 func (p ParametersLiteral[T]) Compile() Parameters[T] {
+	singleKeyParameters := p.SingleKeyParametersLiteral.Compile()
+
 	switch {
 	case p.PartyCount <= 0:
 		panic("PartyCount smaller than zero")
-	case p.SingleKeyParametersLiteral.GLWERank != 1:
+	case singleKeyParameters.GLWERank() != 1:
 		panic("Multi-Key TFHE only supports GLWE dimension 1")
-	case p.SingleKeyParametersLiteral.LookUpTableSize != 0 && p.SingleKeyParametersLiteral.LookUpTableSize != p.SingleKeyParametersLiteral.PolyDegree:
+	case singleKeyParameters.LookUpTableSize() != singleKeyParameters.PolyDegree():
 		panic("Multi-Key TFHE only supports LookUpTableSize equal to PolyDegree")
 	}
 
 	return Parameters[T]{
-		singleKeyParameters: p.SingleKeyParametersLiteral.Compile(),
+		singleKeyParameters: singleKeyParameters,
 
 		partyCount: p.PartyCount,
 

diff --git a/tfhe/bootstrap_lut.go b/tfhe/bootstrap_lut.go
@@ -80,9 +80,9 @@ func (e *Evaluator[T]) GenLookUpTableCustomAssign(f func(int) int, messageModulu
 
 // GenLookUpTableFullCustom generates a lookup table based on function f using custom messageModulus and scale.
 // Output of f is encoded as-is.
-func (e *Evaluator[T]) GenLookUpTableFullCustom(f func(int) T, messageModulus, scale T) LookUpTable[T] {
+func (e *Evaluator[T]) GenLookUpTableFullCustom(f func(int) T, messageModulus T) LookUpTable[T] {
 	lutOut := NewLookUpTable(e.Parameters)
-	e.GenLookUpTableFullAssign(f, lutOut)
+	e.GenLookUpTableFullCustomAssign(f, messageModulus, lutOut)
 	return lutOut
 }
 

diff --git a/xtfhe/bfv_evaluator.go b/xtfhe/bfv_evaluator.go
@@ -36,7 +36,7 @@ type bfvEvaluationBuffer[T tfhe.TorusInt] struct {
 	ctTensorFourier tfhe.FourierGLWECiphertext[T]
 	// ctPermute is the permuted ciphertext for BFV automorphism.
 	ctPermute tfhe.GLWECiphertext[T]
-	// ctPack is the permuted ciphertext for RingPack.
+	// ctPack is the permuted ciphertext for LWEToGLWECiphertext.
 	ctPack tfhe.GLWECiphertext[T]
 }
 

diff --git a/xtfhe/bfv_keygen.go b/xtfhe/bfv_keygen.go
@@ -5,6 +5,15 @@ import (
 	"github.com/sp301415/tfhe-go/tfhe"
 )
 
+// BFVEvaluationKey is a keyswitching key for BFV type operations.
+// It holds relinearization and automorphism keys.
+type BFVEvaluationKey[T tfhe.TorusInt] struct {
+	// RelinKey is a relinearization key.
+	RelinKey tfhe.GLWEKeySwitchKey[T]
+	// GaloisKeys is a map of automorphism keys.
+	GaloisKeys map[int]tfhe.GLWEKeySwitchKey[T]
+}
+
 // BFVKeyGenerator generates keyswitching keys for BFV type operations.
 //
 // BFVKeyGenerator is not safe for concurrent use.
@@ -17,50 +26,36 @@ type BFVKeyGenerator[T tfhe.TorusInt] struct {
 
 	// Parameters is a parameter set for this BFVKeyGenerator.
 	Parameters tfhe.Parameters[T]
-
-	// KeySwitchParams is a gadget parameter set for keyswitching keys in BFV type operations.
-	KeySwitchParams tfhe.GadgetParameters[T]
 }
 
 // NewBFVKeyGenerator creates a new BFVKeyGenerator.
-func NewBFVKeyGenerator[T tfhe.TorusInt](params tfhe.Parameters[T], keySwitchParams tfhe.GadgetParameters[T], sk tfhe.SecretKey[T]) *BFVKeyGenerator[T] {
+func NewBFVKeyGenerator[T tfhe.TorusInt](params tfhe.Parameters[T], sk tfhe.SecretKey[T]) *BFVKeyGenerator[T] {
 	return &BFVKeyGenerator[T]{
-		BaseEncryptor:   tfhe.NewEncryptorWithKey(params, sk),
-		PolyEvaluator:   poly.NewEvaluator[T](params.PolyDegree()),
-		Parameters:      params,
-		KeySwitchParams: keySwitchParams,
+		BaseEncryptor: tfhe.NewEncryptorWithKey(params, sk),
+		PolyEvaluator: poly.NewEvaluator[T](params.PolyDegree()),
+		Parameters:    params,
 	}
 }
 
 // ShallowCopy creates a shallow copy of this BFVKeyGenerator.
 func (kg *BFVKeyGenerator[T]) ShallowCopy() *BFVKeyGenerator[T] {
 	return &BFVKeyGenerator[T]{
-		BaseEncryptor:   kg.BaseEncryptor.ShallowCopy(),
-		PolyEvaluator:   kg.PolyEvaluator.ShallowCopy(),
-		Parameters:      kg.Parameters,
-		KeySwitchParams: kg.KeySwitchParams,
+		BaseEncryptor: kg.BaseEncryptor.ShallowCopy(),
+		PolyEvaluator: kg.PolyEvaluator.ShallowCopy(),
+		Parameters:    kg.Parameters,
 	}
 }
 
-// BFVEvaluationKey is a keyswitching key for BFV type operations.
-// It holds relinearization and automorphism keys.
-type BFVEvaluationKey[T tfhe.TorusInt] struct {
-	// RelinKey is a relinearization key.
-	RelinKey tfhe.GLWEKeySwitchKey[T]
-	// GaloisKeys is a map of automorphism keys.
-	GaloisKeys map[int]tfhe.GLWEKeySwitchKey[T]
-}
-
 // GenEvaluationKey generates an evaluation key for BFV type operations.
-func (kg *BFVKeyGenerator[T]) GenEvaluationKey(idx []int) BFVEvaluationKey[T] {
+func (kg *BFVKeyGenerator[T]) GenEvaluationKey(idx []int, kskParams tfhe.GadgetParameters[T]) BFVEvaluationKey[T] {
 	return BFVEvaluationKey[T]{
-		RelinKey:   kg.GenRelinKey(),
-		GaloisKeys: kg.GenGaloisKeys(idx),
+		RelinKey:   kg.GenRelinKey(kskParams),
+		GaloisKeys: kg.GenGaloisKeys(idx, kskParams),
 	}
 }
 
 // GenRelinKey generates a relinearization key for BFV multiplication.
-func (kg *BFVKeyGenerator[T]) GenRelinKey() tfhe.GLWEKeySwitchKey[T] {
+func (kg *BFVKeyGenerator[T]) GenRelinKey(kskParams tfhe.GadgetParameters[T]) tfhe.GLWEKeySwitchKey[T] {
 	skOut := tfhe.NewGLWESecretKeyCustom[T](kg.Parameters.GLWERank()*(kg.Parameters.GLWERank()+1)/2, kg.Parameters.PolyDegree())
 	fskOut := tfhe.NewFourierGLWESecretKeyCustom[T](kg.Parameters.GLWERank()*(kg.Parameters.GLWERank()+1)/2, kg.Parameters.PolyDegree())
 
@@ -76,51 +71,51 @@ func (kg *BFVKeyGenerator[T]) GenRelinKey() tfhe.GLWEKeySwitchKey[T] {
 		kg.BaseEncryptor.PolyEvaluator.ToPolyAssignUnsafe(fskOut.Value[i], skOut.Value[i])
 	}
 
-	return kg.BaseEncryptor.GenGLWEKeySwitchKey(skOut, kg.KeySwitchParams)
+	return kg.BaseEncryptor.GenGLWEKeySwitchKey(skOut, kskParams)
 }
 
 // GenGaloisKeys generate galois keys for BFV automorphism.
-func (kg *BFVKeyGenerator[T]) GenGaloisKeys(idx []int) map[int]tfhe.GLWEKeySwitchKey[T] {
+func (kg *BFVKeyGenerator[T]) GenGaloisKeys(idx []int, kskParams tfhe.GadgetParameters[T]) map[int]tfhe.GLWEKeySwitchKey[T] {
 	galKeys := make(map[int]tfhe.GLWEKeySwitchKey[T], len(idx))
 	skOut := tfhe.NewGLWESecretKey(kg.Parameters)
 
 	for _, d := range idx {
 		for i := 0; i < kg.Parameters.GLWERank(); i++ {
 			kg.BaseEncryptor.PolyEvaluator.PermutePolyAssign(kg.BaseEncryptor.SecretKey.GLWEKey.Value[i], d, skOut.Value[i])
 		}
-		galKeys[d] = kg.BaseEncryptor.GenGLWEKeySwitchKey(skOut, kg.KeySwitchParams)
+		galKeys[d] = kg.BaseEncryptor.GenGLWEKeySwitchKey(skOut, kskParams)
 	}
 	return galKeys
 }
 
 // GenGaloisKeysAssign generates automorphism keys for BFV automorphism and assigns them to the given map.
 // If a key for a given automorphism degree already exists in the map, it will be overwritten.
-func (kg *BFVKeyGenerator[T]) GenGaloisKeysAssign(idx []int, galKeysOut map[int]tfhe.GLWEKeySwitchKey[T]) {
+func (kg *BFVKeyGenerator[T]) GenGaloisKeysAssign(idx []int, kskParams tfhe.GadgetParameters[T], galKeysOut map[int]tfhe.GLWEKeySwitchKey[T]) {
 	skOut := tfhe.NewGLWESecretKey(kg.Parameters)
 
 	for _, d := range idx {
 		for i := 0; i < kg.Parameters.GLWERank(); i++ {
 			kg.BaseEncryptor.PolyEvaluator.PermutePolyAssign(kg.BaseEncryptor.SecretKey.GLWEKey.Value[i], d, skOut.Value[i])
 		}
-		galKeysOut[d] = kg.BaseEncryptor.GenGLWEKeySwitchKey(skOut, kg.KeySwitchParams)
+		galKeysOut[d] = kg.BaseEncryptor.GenGLWEKeySwitchKey(skOut, kskParams)
 	}
 }
 
 // GenGaloisKeysForLWEToGLWECiphertext generates automorphism keys for BFV automorphism for LWE to GLWE packing.
-func (kg *BFVKeyGenerator[T]) GenGaloisKeysForLWEToGLWECiphertext() map[int]tfhe.GLWEKeySwitchKey[T] {
+func (kg *BFVKeyGenerator[T]) GenGaloisKeysForLWEToGLWECiphertext(kskParams tfhe.GadgetParameters[T]) map[int]tfhe.GLWEKeySwitchKey[T] {
 	auts := make([]int, kg.Parameters.LogPolyDegree())
 	for i := range auts {
 		auts[i] = 1<<(kg.Parameters.LogPolyDegree()-i) + 1
 	}
-	return kg.GenGaloisKeys(auts)
+	return kg.GenGaloisKeys(auts, kskParams)
 }
 
 // GenGaloisKeysForLWEToGLWECiphertextAssign generates automorphism keys for BFV automorphism for LWE to GLWE packing and assigns them to the given map.
 // If a key for a given automorphism degree already exists in the map, it will be overwritten.
-func (kg *BFVKeyGenerator[T]) GenGaloisKeysForLWEToGLWECiphertextAssign(galKeysOut map[int]tfhe.GLWEKeySwitchKey[T]) {
+func (kg *BFVKeyGenerator[T]) GenGaloisKeysForLWEToGLWECiphertextAssign(kskParams tfhe.GadgetParameters[T], galKeysOut map[int]tfhe.GLWEKeySwitchKey[T]) {
 	auts := make([]int, kg.Parameters.LogPolyDegree())
 	for i := range auts {
 		auts[i] = 1<<(kg.Parameters.LogPolyDegree()-i) + 1
 	}
-	kg.GenGaloisKeysAssign(auts, galKeysOut)
+	kg.GenGaloisKeysAssign(auts, kskParams, galKeysOut)
 }
diff --git a/xtfhe/bfv_test.go b/xtfhe/bfv_test.go
@@ -12,10 +12,10 @@ import (
 var (
 	bfvParams = tfhe.ParamsUint3.Compile()
 	bfvEnc    = tfhe.NewEncryptor(bfvParams)
-	bfvKeyGen = xtfhe.NewBFVKeyGenerator(bfvParams, xtfhe.ParamsBFVKeySwitchLogN11.Compile(), bfvEnc.SecretKey)
+	bfvKeyGen = xtfhe.NewBFVKeyGenerator(bfvParams, bfvEnc.SecretKey)
 	bfvEval   = xtfhe.NewBFVEvaluator(bfvParams, xtfhe.BFVEvaluationKey[uint64]{
-		RelinKey:   bfvKeyGen.GenRelinKey(),
-		GaloisKeys: bfvKeyGen.GenGaloisKeysForLWEToGLWECiphertext(),
+		RelinKey:   bfvKeyGen.GenRelinKey(xtfhe.ParamsBFVKeySwitchLogN11.Compile()),
+		GaloisKeys: bfvKeyGen.GenGaloisKeysForLWEToGLWECiphertext(xtfhe.ParamsBFVKeySwitchLogN11.Compile()),
 	})
 )