mirrors/lattigo
1
0
Fork 0
mirror of https://github.com/tuneinsight/lattigo.git synced 2025-09-13 03:27:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: keep only BuffFromUintPool and cleanup

Browse Source
This commit is contained in:
lehugueni
2025-01-16 11:07:19 +01:00
parent 77c62e6421
commit ee42b96e77
6 changed files with 49 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
core/rlwe/evaluator.go
View File

@@ -29,33 +29,15 @@ type EvaluatorBuffers struct {
BuffCtPool structs.BufferPool[*Ciphertext]
}
func newBuffer[T any](f func() T) structs.BufferPool[T] {
// Uncomment to try with free lists instead of sync pool:
// nbItemsInPool := 10
// return structs.NewFreeList(nbItemsInPool, f)
return structs.NewSyncPool(f)
}
func NewEvaluatorBuffersWithUintPool(params Parameters) *EvaluatorBuffers {
// NewEvaluatorBuffers creates the buffers that are used to recycle large obejcts instead of instantiating new ones.
// Under the hood, all buffers use the same sync.Pool of *[]uint64.
func NewEvaluatorBuffers(params Parameters) *EvaluatorBuffers {
buff := new(EvaluatorBuffers)
ringQP := params.RingQP()
ringQ := params.ringQ
buff.BuffQPPool = structs.NewBuffFromUintPool(
func() *ringqp.Poly {
return ringQP.NewPolyQPFromUintPool()
},
func(poly *ringqp.Poly) {
ringQP.RecyclePolyQPFromUintPool(poly)
},
)
buff.BuffQPool = structs.NewBuffFromUintPool(
func() *ring.Poly {
return ringQ.NewPolyFromUintPool()
},
func(poly *ring.Poly) {
ringQ.RecyclePolyInUintPool(poly)
},
)
buff.BuffQPPool = ringQP.NewBuffFromUintPool()
buff.BuffQPool = ringQ.NewBuffFromUintPool()
buff.BuffCtPool = structs.NewBuffFromUintPool(
func() *Ciphertext {
return NewCiphertextFromUintPool(params, 2, params.MaxLevel())
@@ -68,38 +50,14 @@ func NewEvaluatorBuffersWithUintPool(params Parameters) *EvaluatorBuffers {
return buff
}
func NewEvaluatorBuffers(params Parameters) *EvaluatorBuffers {
buff := new(EvaluatorBuffers)
ringQP := params.RingQP()
buff.BuffQPPool = newBuffer(func() *ringqp.Poly {
poly := ringQP.NewPoly()
return &poly
})
buff.BuffQPool = newBuffer(func() *ring.Poly {
poly := params.RingQ().NewPoly()
return &poly
})
buff.BuffCtPool = newBuffer(func() *Ciphertext {
return NewCiphertext(params, 2, params.MaxLevel())
})
buff.BuffBitPool = newBuffer(func() *[]uint64 {
buff := make([]uint64, params.RingQ().N())
return &buff
})
return buff
}
// NewEvaluator creates a new [Evaluator].
func NewEvaluator(params ParameterProvider, evk EvaluationKeySet) (eval *Evaluator) {
eval = new(Evaluator)
p := params.GetRLWEParameters()
eval.params = *p
// All buffer use the same sync.Pool of *[]uint64
eval.EvaluatorBuffers = NewEvaluatorBuffersWithUintPool(eval.params)
// Uncomment following line to have one sync.Pool per buffer type
// eval.EvaluatorBuffers = NewEvaluatorBuffers(eval.params)
eval.EvaluatorBuffers = NewEvaluatorBuffers(eval.params)
if p.RingP() != nil {
eval.BasisExtender = ring.NewBasisExtender(p.RingQ(), p.RingP())

18
ring/basis_extension.go
View File

@@ -73,22 +73,8 @@ func NewBasisExtender(ringQ, ringP *Ring) (be *BasisExtender) {
be.modDownConstantsPtoQ = genmodDownConstants(ringQ, ringP)
be.modDownConstantsQtoP = genmodDownConstants(ringP, ringQ)
be.buffQPool = structs.NewBuffFromUintPool(
func() *Poly {
return ringQ.NewPolyFromUintPool()
},
func(poly *Poly) {
ringQ.RecyclePolyInUintPool(poly)
},
)
be.buffPPool = structs.NewBuffFromUintPool(
func() *Poly {
return ringP.NewPolyFromUintPool()
},
func(poly *Poly) {
ringP.RecyclePolyInUintPool(poly)
},
)
be.buffQPool = ringQ.NewBuffFromUintPool()
be.buffPPool = ringP.NewBuffFromUintPool()
return
}

11
ring/ring.go
View File

@@ -394,6 +394,17 @@ func (r Ring) RecyclePolyInUintPool(pol *Poly) {
RecyclePolyInUintPool(r.bufferPool, pol)
}
func (r Ring) NewBuffFromUintPool() *structs.BuffFromUintPool[*Poly] {
return structs.NewBuffFromUintPool(
func() *Poly {
return r.NewPolyFromUintPool()
},
func(poly *Poly) {
r.RecyclePolyInUintPool(poly)
},
)
}
// NewMonomialXi returns a polynomial X^{i}.
func (r Ring) NewMonomialXi(i int) (p Poly) {

12
ring/ringqp/ring.go
View File

@@ -7,6 +7,7 @@ import (
"github.com/tuneinsight/lattigo/v6/ring"
"github.com/tuneinsight/lattigo/v6/utils/bignum"
"github.com/tuneinsight/lattigo/v6/utils/structs"
)
// Ring is a structure that implements the operation in the ring R_QP.
@@ -236,3 +237,14 @@ func (r Ring) RecyclePolyQPFromUintPool(poly *Poly) {
r.RingQ.RecyclePolyInUintPool(&poly.Q)
r.RingP.RecyclePolyInUintPool(&poly.P)
}
func (r Ring) NewBuffFromUintPool() *structs.BuffFromUintPool[*Poly] {
return structs.NewBuffFromUintPool(
func() *Poly {
return r.NewPolyQPFromUintPool()
},
func(poly *Poly) {
r.RecyclePolyQPFromUintPool(poly)
},
)
}

21
schemes/ckks/encoder.go
View File

@@ -67,7 +67,8 @@ type Encoder struct {
rotGroup []int
roots interface{}
// buffCmplx interface{}
// Pools used to recycle large objects.
BuffPolyPool structs.BufferPool[*ring.Poly]
BuffBigIntPool structs.BufferPool[*[]*big.Int]
BuffComplexPool structs.BufferPool[Complex]
@@ -112,15 +113,7 @@ func NewEncoder(parameters Parameters, precision ...uint) (ecd *Encoder) {
})
ringQ := parameters.RingQ()
ecd.BuffPolyPool = structs.NewBuffFromUintPool(
func() *ring.Poly {
return ringQ.NewPolyFromUintPool()
},
func(poly *ring.Poly) {
ringQ.RecyclePolyInUintPool(poly)
},
)
ecd.BuffPolyPool = ringQ.NewBuffFromUintPool()
if prec <= 53 {
@@ -133,14 +126,8 @@ func NewEncoder(parameters Parameters, precision ...uint) (ecd *Encoder) {
} else {
// tmp := make([]*bignum.Complex, ecd.m>>2)
//
// for i := 0; i < ecd.m>>2; i++ {
// tmp[i] = &bignum.Complex{bignum.NewFloat(0, prec), bignum.NewFloat(0, prec)}
// }
ecd.roots = GetRootsBigComplex(ecd.m, prec)
// ecd.buffCmplx = tmp
ecd.BuffComplexPool = structs.NewSyncPool(func() Complex {
buff := make([]*bignum.Complex, ecd.m>>2)
for i := 0; i < ecd.m>>2; i++ {

49
utils/structs/concurrent_buffer.go
View File

@@ -2,15 +2,19 @@ package structs
import "sync"
// BufferPool is an interface for all pools of buffers.
type BufferPool[T any] interface {
Get() T
Put(T)
}
// SyncPool is a wrapper around [sync.Pool] (it avoids doing type conversion after Get()).
type SyncPool[T any] struct {
pool *sync.Pool
}
// NewSyncPool creates a new SyncPool.
// The input function f is the function that is used to create new objects if none is available in the pool.
func NewSyncPool[T any](f func() T) *SyncPool[T] {
pool := &sync.Pool{
New: func() any {
@@ -20,19 +24,26 @@ func NewSyncPool[T any](f func() T) *SyncPool[T] {
return &SyncPool[T]{pool: pool}
}
// Get returns a new object of type T from the pool.
func (spool *SyncPool[T]) Get() T {
return spool.pool.Get().(T)
}
// Put returns the buff to the pool.
func (spool *SyncPool[T]) Put(buff T) {
spool.pool.Put(buff)
}
// BuffFromUintPool represents a pool of objects built on []uint64 backing arrays.
// It implements the [BufferPool] interface.
type BuffFromUintPool[T any] struct {
createObject func() T
recycleObject func(T)
}
// NewBuffFromUintPool returns a new BuffFromUintPool structure.
// The create (resp. recycle) function are meant to use an underlying
// pool of []uint64 to build (resp. recycle) an object of type T.
func NewBuffFromUintPool[T any](create func() T, recycle func(T)) *BuffFromUintPool[T] {
return &BuffFromUintPool[T]{
createObject: create,
@@ -40,46 +51,12 @@ func NewBuffFromUintPool[T any](create func() T, recycle func(T)) *BuffFromUintP
}
}
// Get returns a new object of type T built from a []uint64 backing array obtained from a pool.
func (bu *BuffFromUintPool[T]) Get() T {
return bu.createObject()
}
// Put recycle an object of type T. I.e. it returns the []uint64 backing arrays of obj to their pool.
func (bu *BuffFromUintPool[T]) Put(obj T) {
bu.recycleObject(obj)
}
type FreeList[T any] struct {
pool chan T
newObject func() T
capacity int
}
func NewFreeList[T any](capacity int, f func() T) *FreeList[T] {
pool := make(chan T, capacity)
for i := 0; i < capacity; i++ {
pool <- f()
}
return &FreeList[T]{
pool: pool,
newObject: f,
capacity: capacity,
}
}
func (fl *FreeList[T]) Get() T {
var obj T
select {
case obj = <-fl.pool:
default:
obj = fl.newObject()
}
return obj
}
func (fl *FreeList[T]) Put(obj T) {
select {
case fl.pool <- obj:
default:
}
}