From ee42b96e778facf3a9029c20c6325ca3b795f296 Mon Sep 17 00:00:00 2001 From: lehugueni Date: Thu, 16 Jan 2025 11:07:19 +0100 Subject: [PATCH] refactor: keep only BuffFromUintPool and cleanup --- core/rlwe/evaluator.go | 56 ++++-------------------------- ring/basis_extension.go | 18 ++-------- ring/ring.go | 11 ++++++ ring/ringqp/ring.go | 12 +++++++ schemes/ckks/encoder.go | 21 +++-------- utils/structs/concurrent_buffer.go | 49 +++++++------------------- 6 files changed, 49 insertions(+), 118 deletions(-) diff --git a/core/rlwe/evaluator.go b/core/rlwe/evaluator.go index 0bf0a6c2..b04f3c1a 100644 --- a/core/rlwe/evaluator.go +++ b/core/rlwe/evaluator.go @@ -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()) diff --git a/ring/basis_extension.go b/ring/basis_extension.go index 84bb145b..8d553cb1 100644 --- a/ring/basis_extension.go +++ b/ring/basis_extension.go @@ -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 } diff --git a/ring/ring.go b/ring/ring.go index dca13b2c..b51cf5ee 100644 --- a/ring/ring.go +++ b/ring/ring.go @@ -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) { diff --git a/ring/ringqp/ring.go b/ring/ringqp/ring.go index bac41c79..af22b1f6 100644 --- a/ring/ringqp/ring.go +++ b/ring/ringqp/ring.go @@ -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) + }, + ) +} diff --git a/schemes/ckks/encoder.go b/schemes/ckks/encoder.go index 2483c516..0f677721 100644 --- a/schemes/ckks/encoder.go +++ b/schemes/ckks/encoder.go @@ -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++ { diff --git a/utils/structs/concurrent_buffer.go b/utils/structs/concurrent_buffer.go index 37cc46a1..7f4c25ed 100644 --- a/utils/structs/concurrent_buffer.go +++ b/utils/structs/concurrent_buffer.go @@ -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: - } -}