mirror of
https://github.com/luxfi/lattice.git
synced 2026-09-24 00:25:44 +00:00
feat(blindrot): BatchEvaluate parallel API for N independent blind rotations
Fans N independent (*Evaluator).Evaluate calls across a goroutine pool of GOMAXPROCS workers. Each worker owns a fresh NewEvaluator clone because the receiver carries mutable scratch buffers (poolMod2N, accumulator) and the package exposes no ShallowCopy. The BRK key set is shared read-only. Per-iteration output is byte-equal to a serial Evaluate call. Measured 4.61x at N=16 on M1 Max P-cores with blindrot test params (LogN=10/9, Q=0x7fff801). Geometry matches the 5.84x ceiling documented in fhe/policy/PERFORMANCE.md section 4 at full bootstrap params; the gap is the relative weight of per-worker NewEvaluator allocation against the cheaper test bootstrap. Tests: - TestBatchEvaluate_ByteEqualSerial: N=16 byte-equal vs serial - TestBatchEvaluate_Empty: nil in -> nil out - TestBatchEvaluate_Single: N=1 equivalent to direct Evaluate - TestBatchEvaluate_LengthMismatch: rejects mismatched len(testPolys) - TestBatchEvaluate_RaceClean: passes -race
This commit is contained in:
@@ -0,0 +1,104 @@
|
||||
package blindrot
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"runtime"
|
||||
"sync"
|
||||
|
||||
"github.com/luxfi/lattice/v7/core/rlwe"
|
||||
"github.com/luxfi/lattice/v7/ring"
|
||||
)
|
||||
|
||||
// errMismatchedTestPolyLen is returned for every input slot when len(testPolys)
|
||||
// does not match len(cts). The whole batch is rejected — failing per-slot here
|
||||
// would mask a programmer error.
|
||||
var errMismatchedTestPolyLen = errors.New("blindrot: BatchEvaluate: len(testPolys) != len(cts)")
|
||||
|
||||
// BatchEvaluate runs N independent blind-rotation evaluations in parallel.
|
||||
//
|
||||
// Each entry in cts is paired with the corresponding entry in testPolys; the
|
||||
// per-iteration call is byte-equivalent to a serial Evaluate(ct, tp, BRK).
|
||||
//
|
||||
// Concurrency model: the receiver Evaluator carries mutable scratch buffers
|
||||
// (poolMod2N, accumulator) and an embedded *rgsw.Evaluator with its own buffer
|
||||
// pool. Those structures are NOT safe for concurrent use, and the package
|
||||
// exposes no ShallowCopy. BatchEvaluate therefore allocates one fresh
|
||||
// Evaluator per worker goroutine via NewEvaluator(paramsBR, paramsLWE) and
|
||||
// fans the work out across runtime.GOMAXPROCS workers. The receiver Evaluator
|
||||
// is read-only here (only paramsBR and paramsLWE are inspected for cloning).
|
||||
//
|
||||
// The BlindRotationEvaluationKeySet is treated as read-only and shared across
|
||||
// goroutines: GetBlindRotationKey / GetEvaluationKeySet must return the same
|
||||
// keys for the same indices (the in-tree MemBlindRotationEvaluationKeySet
|
||||
// satisfies this — it is a fixed slice of pointers).
|
||||
//
|
||||
// Per PERFORMANCE.md §4 (luxfi/fhe), the M1 Max P-core ceiling is ~5.84× at
|
||||
// N≥16; the goroutine fan-out saturates the 8 P-cores. The Metal batch path
|
||||
// is excluded by design until the metal_batch_bootstrap kernel ships in
|
||||
// luxcpp/lattice/src/metal/ (PERFORMANCE.md §6 G3). When it lands, this
|
||||
// function gains a fast path that dispatches the whole batch to one Metal
|
||||
// command buffer instead of fanning out goroutines.
|
||||
func (eval *Evaluator) BatchEvaluate(
|
||||
cts []*rlwe.Ciphertext,
|
||||
testPolys []map[int]*ring.Poly,
|
||||
BRK BlindRotationEvaluationKeySet,
|
||||
) ([]map[int]*rlwe.Ciphertext, []error) {
|
||||
|
||||
n := len(cts)
|
||||
if n == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
if len(testPolys) != n {
|
||||
errs := make([]error, n)
|
||||
for i := range errs {
|
||||
errs[i] = errMismatchedTestPolyLen
|
||||
}
|
||||
return make([]map[int]*rlwe.Ciphertext, n), errs
|
||||
}
|
||||
|
||||
out := make([]map[int]*rlwe.Ciphertext, n)
|
||||
errs := make([]error, n)
|
||||
|
||||
workers := runtime.GOMAXPROCS(0)
|
||||
if workers < 1 {
|
||||
workers = 1
|
||||
}
|
||||
if workers > n {
|
||||
workers = n
|
||||
}
|
||||
|
||||
// Single-worker fast path — fall through to a plain serial loop on the
|
||||
// receiver. Avoids the cost of spinning up a worker pool for n=1.
|
||||
if workers == 1 {
|
||||
for i := range cts {
|
||||
out[i], errs[i] = eval.Evaluate(cts[i], testPolys[i], BRK)
|
||||
}
|
||||
return out, errs
|
||||
}
|
||||
|
||||
// Fan-out: each worker owns one fresh Evaluator. Re-creating an Evaluator
|
||||
// is cheap relative to a single Evaluate call (one ring poly pair + one
|
||||
// ciphertext + the discrete-log map; ~N entries). The BRK is shared.
|
||||
jobs := make(chan int, n)
|
||||
for i := 0; i < n; i++ {
|
||||
jobs <- i
|
||||
}
|
||||
close(jobs)
|
||||
|
||||
paramsBR := eval.paramsBR
|
||||
paramsLWE := eval.paramsLWE
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(workers)
|
||||
for w := 0; w < workers; w++ {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
worker := NewEvaluator(paramsBR, paramsLWE)
|
||||
for i := range jobs {
|
||||
out[i], errs[i] = worker.Evaluate(cts[i], testPolys[i], BRK)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
return out, errs
|
||||
}
|
||||
@@ -0,0 +1,242 @@
|
||||
package blindrot
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"testing"
|
||||
|
||||
"github.com/luxfi/lattice/v7/core/rlwe"
|
||||
"github.com/luxfi/lattice/v7/ring"
|
||||
"github.com/luxfi/lattice/v7/utils"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// batchHarness builds the parameter set and shared inputs used by every
|
||||
// batch test. Mirrors testBlindRotation in blindrot_test.go.
|
||||
type batchHarness struct {
|
||||
paramsBR rlwe.Parameters
|
||||
paramsLWE rlwe.Parameters
|
||||
eval *Evaluator
|
||||
BRK MemBlindRotationEvaluationKeySet
|
||||
ctLWE *rlwe.Ciphertext
|
||||
testPoly ring.Poly
|
||||
}
|
||||
|
||||
func newBatchHarness(t testing.TB) *batchHarness {
|
||||
paramsBR, err := rlwe.NewParametersFromLiteral(rlwe.ParametersLiteral{
|
||||
LogN: 10,
|
||||
Q: []uint64{0x7fff801},
|
||||
NTTFlag: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
paramsLWE, err := rlwe.NewParametersFromLiteral(rlwe.ParametersLiteral{
|
||||
LogN: 9,
|
||||
Q: []uint64{0x3001},
|
||||
NTTFlag: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
evkParams := rlwe.EvaluationKeyParameters{BaseTwoDecomposition: utils.Pointy(7)}
|
||||
|
||||
scaleLWE := float64(paramsLWE.Q()[0]) / 4.0
|
||||
scaleBR := float64(paramsBR.Q()[0]) / 4.0
|
||||
testPoly := InitTestPolynomial(sign, rlwe.NewScale(scaleBR), paramsBR.RingQ(), -1, 1)
|
||||
|
||||
skLWE := rlwe.NewKeyGenerator(paramsLWE).GenSecretKeyNew()
|
||||
encryptorLWE := rlwe.NewEncryptor(paramsLWE, skLWE)
|
||||
|
||||
ptLWE := rlwe.NewPlaintext(paramsLWE, paramsLWE.MaxLevel())
|
||||
slots := 16
|
||||
values := make([]float64, slots)
|
||||
for i := 0; i < slots; i++ {
|
||||
values[i] = -1 + float64(2*i)/float64(slots)
|
||||
}
|
||||
for i := range values {
|
||||
if values[i] < 0 {
|
||||
ptLWE.Value.Coeffs[0][i] = paramsLWE.Q()[0] - uint64(-values[i]*scaleLWE)
|
||||
} else {
|
||||
ptLWE.Value.Coeffs[0][i] = uint64(values[i] * scaleLWE)
|
||||
}
|
||||
}
|
||||
if ptLWE.IsNTT {
|
||||
paramsLWE.RingQ().NTT(ptLWE.Value, ptLWE.Value)
|
||||
}
|
||||
ctLWE := rlwe.NewCiphertext(paramsLWE, 1, paramsLWE.MaxLevel())
|
||||
require.NoError(t, encryptorLWE.Encrypt(ptLWE, ctLWE))
|
||||
|
||||
skBR := rlwe.NewKeyGenerator(paramsBR).GenSecretKeyNew()
|
||||
BRK := GenEvaluationKeyNew(paramsBR, skBR, paramsLWE, skLWE, evkParams)
|
||||
|
||||
return &batchHarness{
|
||||
paramsBR: paramsBR,
|
||||
paramsLWE: paramsLWE,
|
||||
eval: NewEvaluator(paramsBR, paramsLWE),
|
||||
BRK: BRK,
|
||||
ctLWE: ctLWE,
|
||||
testPoly: testPoly,
|
||||
}
|
||||
}
|
||||
|
||||
// makeBatch returns N copies of the same ciphertext and test-poly map. Each
|
||||
// Evaluate call against (ctLWE, testPolyMap, BRK) is deterministic for fixed
|
||||
// inputs, so N copies must produce N byte-equal outputs both serially and in
|
||||
// batch.
|
||||
func (h *batchHarness) makeBatch(n int) ([]*rlwe.Ciphertext, []map[int]*ring.Poly) {
|
||||
cts := make([]*rlwe.Ciphertext, n)
|
||||
polys := make([]map[int]*ring.Poly, n)
|
||||
tp := h.testPoly
|
||||
for i := 0; i < n; i++ {
|
||||
// CopyNew gives each batch slot an independent input so a buggy
|
||||
// implementation that aliases through the shared receiver buffers
|
||||
// shows up as a byte mismatch.
|
||||
cts[i] = h.ctLWE.CopyNew()
|
||||
m := make(map[int]*ring.Poly, 16)
|
||||
for s := 0; s < 16; s++ {
|
||||
m[s] = &tp
|
||||
}
|
||||
polys[i] = m
|
||||
}
|
||||
return cts, polys
|
||||
}
|
||||
|
||||
func ctBytesEqual(a, b *rlwe.Ciphertext) bool {
|
||||
for i := range a.Value {
|
||||
if i >= len(b.Value) {
|
||||
return false
|
||||
}
|
||||
for j := range a.Value[i].Coeffs {
|
||||
if !bytes.Equal(uint64SliceBytes(a.Value[i].Coeffs[j]), uint64SliceBytes(b.Value[i].Coeffs[j])) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
return a.IsNTT == b.IsNTT
|
||||
}
|
||||
|
||||
func uint64SliceBytes(s []uint64) []byte {
|
||||
out := make([]byte, 8*len(s))
|
||||
for i, v := range s {
|
||||
for k := 0; k < 8; k++ {
|
||||
out[8*i+k] = byte(v >> (8 * k))
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// TestBatchEvaluate_ByteEqualSerial verifies that BatchEvaluate(N) produces
|
||||
// the same byte sequence per slot as N serial Evaluate calls on a fresh
|
||||
// Evaluator.
|
||||
func TestBatchEvaluate_ByteEqualSerial(t *testing.T) {
|
||||
h := newBatchHarness(t)
|
||||
const N = 16
|
||||
|
||||
cts, polys := h.makeBatch(N)
|
||||
|
||||
serial := make([]map[int]*rlwe.Ciphertext, N)
|
||||
for i := 0; i < N; i++ {
|
||||
serialEval := NewEvaluator(h.paramsBR, h.paramsLWE)
|
||||
out, err := serialEval.Evaluate(cts[i], polys[i], h.BRK)
|
||||
require.NoError(t, err)
|
||||
serial[i] = out
|
||||
}
|
||||
|
||||
cts2, polys2 := h.makeBatch(N)
|
||||
batchEval := NewEvaluator(h.paramsBR, h.paramsLWE)
|
||||
parallel, errs := batchEval.BatchEvaluate(cts2, polys2, h.BRK)
|
||||
|
||||
require.Len(t, parallel, N)
|
||||
require.Len(t, errs, N)
|
||||
for i := 0; i < N; i++ {
|
||||
require.NoErrorf(t, errs[i], "batch slot %d", i)
|
||||
require.Equalf(t, len(serial[i]), len(parallel[i]), "slot %d map size", i)
|
||||
for k, sv := range serial[i] {
|
||||
pv, ok := parallel[i][k]
|
||||
require.Truef(t, ok, "slot %d missing key %d", i, k)
|
||||
require.Truef(t, ctBytesEqual(sv, pv), "slot %d key %d byte mismatch", i, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchEvaluate_Empty(t *testing.T) {
|
||||
h := newBatchHarness(t)
|
||||
out, errs := h.eval.BatchEvaluate(nil, nil, h.BRK)
|
||||
require.Nil(t, out)
|
||||
require.Nil(t, errs)
|
||||
}
|
||||
|
||||
func TestBatchEvaluate_Single(t *testing.T) {
|
||||
h := newBatchHarness(t)
|
||||
cts, polys := h.makeBatch(1)
|
||||
|
||||
cts2, polys2 := h.makeBatch(1)
|
||||
serialEval := NewEvaluator(h.paramsBR, h.paramsLWE)
|
||||
expected, err := serialEval.Evaluate(cts2[0], polys2[0], h.BRK)
|
||||
require.NoError(t, err)
|
||||
|
||||
got, errs := h.eval.BatchEvaluate(cts, polys, h.BRK)
|
||||
require.Len(t, got, 1)
|
||||
require.Len(t, errs, 1)
|
||||
require.NoError(t, errs[0])
|
||||
require.Equal(t, len(expected), len(got[0]))
|
||||
for k, sv := range expected {
|
||||
pv, ok := got[0][k]
|
||||
require.Truef(t, ok, "missing key %d", k)
|
||||
require.Truef(t, ctBytesEqual(sv, pv), "key %d byte mismatch", k)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchEvaluate_LengthMismatch(t *testing.T) {
|
||||
h := newBatchHarness(t)
|
||||
cts, _ := h.makeBatch(4)
|
||||
out, errs := h.eval.BatchEvaluate(cts, nil, h.BRK)
|
||||
require.Len(t, out, 4)
|
||||
require.Len(t, errs, 4)
|
||||
for i := range errs {
|
||||
require.ErrorIs(t, errs[i], errMismatchedTestPolyLen)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBatchEvaluate_RaceClean — runs concurrent BatchEvaluate dispatches.
|
||||
// Together with `go test -race` this surfaces any unintended sharing of
|
||||
// receiver buffers across worker goroutines.
|
||||
func TestBatchEvaluate_RaceClean(t *testing.T) {
|
||||
h := newBatchHarness(t)
|
||||
const N = 8
|
||||
cts, polys := h.makeBatch(N)
|
||||
out, errs := h.eval.BatchEvaluate(cts, polys, h.BRK)
|
||||
require.Len(t, out, N)
|
||||
for i, e := range errs {
|
||||
require.NoErrorf(t, e, "slot %d", i)
|
||||
require.NotNil(t, out[i])
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkBatchEvaluate_Serial_N16(b *testing.B) {
|
||||
h := newBatchHarness(b)
|
||||
cts, polys := h.makeBatch(16)
|
||||
|
||||
b.ResetTimer()
|
||||
for n := 0; n < b.N; n++ {
|
||||
eval := NewEvaluator(h.paramsBR, h.paramsLWE)
|
||||
for i := 0; i < len(cts); i++ {
|
||||
if _, err := eval.Evaluate(cts[i], polys[i], h.BRK); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkBatchEvaluate_Parallel_N16(b *testing.B) {
|
||||
h := newBatchHarness(b)
|
||||
cts, polys := h.makeBatch(16)
|
||||
|
||||
b.ResetTimer()
|
||||
for n := 0; n < b.N; n++ {
|
||||
_, errs := h.eval.BatchEvaluate(cts, polys, h.BRK)
|
||||
for _, e := range errs {
|
||||
if e != nil {
|
||||
b.Fatal(e)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user