mirror of
https://github.com/luxfi/chains.git
synced 2026-09-24 00:26:00 +00:00
mpcvm/fhe: one device round trip, three kernels
BatchNTTForward, BatchNTTInverse and BatchPolyMul were three copies of one loop — snapshot the session, allocate operand tensors, run a kernel, sync, read back, copy out, and on any failure at any step close what is open and do the operation on the CPU instead. ~45 statements each, differing in how many polynomials an element reads, which kernel runs, and what the CPU does instead. Three copies of a resource lifecycle are three places to forget a Close, and a fourth kernel would have been a fourth copy. Now one offload plus one element, and three wrappers that name their three differences. 101 lines go, no behaviour with them. The floors are deliberately NOT unified: forward NTT reaches the device at 64 polynomials and ring degree 8192, the other two at 4 and any degree, and the reasoning written on the forward path applies to all three. That reads as drift, but changing a performance threshold needs a device to measure on and this host has none, so the numbers are carried through as parameters and the disagreement is stated where the next reader will see it. Also drops the count == 0 guard, which count < minBatch already subsumes for every floor; reverting it changed no test result, which is what redundant means.
This commit is contained in:
@@ -0,0 +1,94 @@
|
||||
// Copyright (C) 2019-2026, Lux Industries, Inc. All rights reserved.
|
||||
// See the file LICENSE for licensing terms.
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/luxfi/chains/mpcvm"
|
||||
"github.com/luxfi/vm/rpc/runtime"
|
||||
)
|
||||
|
||||
// The `version` subcommand prints the id the binary must be INSTALLED as, not
|
||||
// just the software version.
|
||||
//
|
||||
// The two facts fail differently and an operator needs both. A wrong software
|
||||
// version starts a chain that misbehaves; a wrong FILENAME starts nothing at
|
||||
// all — the node's registry resolves a CreateChainTx's vmID by looking for a
|
||||
// file with that name, so a correctly built plugin under the wrong name is
|
||||
// invisible, and the symptom is a chain that never appears with no error
|
||||
// anywhere. Printing the id here puts it where an operator already looks.
|
||||
func TestVersionNamesTheIdTheBinaryMustBeInstalledAs(t *testing.T) {
|
||||
var out, errOut bytes.Buffer
|
||||
if code := run(context.Background(), []string{"version"}, &out, &errOut); code != 0 {
|
||||
t.Fatalf("version exited %d, want 0; stderr=%q", code, errOut.String())
|
||||
}
|
||||
got := out.String()
|
||||
if !strings.Contains(got, mpcvm.Version.String()) {
|
||||
t.Errorf("version output %q does not name the software version %s", got, mpcvm.Version)
|
||||
}
|
||||
if !strings.Contains(got, mpcvm.VMID.String()) {
|
||||
t.Errorf("version output %q does not name the vmID %s the plugin must be installed as",
|
||||
got, mpcvm.VMID)
|
||||
}
|
||||
if errOut.Len() != 0 {
|
||||
t.Errorf("version wrote to stderr: %q", errOut.String())
|
||||
}
|
||||
}
|
||||
|
||||
// Anything that is not `version` serves, and serving needs a node. Started by
|
||||
// hand there is no engine address in the environment, so the plugin says why it
|
||||
// is not serving and exits non-zero.
|
||||
//
|
||||
// The exit code is the load-bearing half. A plugin that fails to serve and
|
||||
// exits 0 looks to the node's runtime like a plugin that ran and finished, so
|
||||
// the chain never starts and nothing reports a failure — the same silent
|
||||
// absence a misnamed binary produces.
|
||||
func TestAPluginWithNoNodeToServeSaysSoAndFails(t *testing.T) {
|
||||
t.Setenv(runtime.EngineAddressKey, "")
|
||||
|
||||
var out, errOut bytes.Buffer
|
||||
if code := run(context.Background(), nil, &out, &errOut); code != 1 {
|
||||
t.Fatalf("run with no engine address exited %d, want 1", code)
|
||||
}
|
||||
if !strings.Contains(errOut.String(), "mpcvm plugin:") {
|
||||
t.Errorf("stderr %q does not name the plugin as the source of the failure", errOut.String())
|
||||
}
|
||||
if !strings.Contains(errOut.String(), runtime.EngineAddressKey) {
|
||||
t.Errorf("stderr %q does not name the missing environment variable %q",
|
||||
errOut.String(), runtime.EngineAddressKey)
|
||||
}
|
||||
}
|
||||
|
||||
// An empty argument list and an argument that is not `version` take the same
|
||||
// path: there is one subcommand, and everything else is "serve".
|
||||
func TestOnlyVersionIsASubcommand(t *testing.T) {
|
||||
t.Setenv(runtime.EngineAddressKey, "")
|
||||
|
||||
for _, args := range [][]string{nil, {}, {"serve"}, {"Version"}, {"--version"}} {
|
||||
var out, errOut bytes.Buffer
|
||||
if code := run(context.Background(), args, &out, &errOut); code != 1 {
|
||||
t.Errorf("run(%q) exited %d, want 1 (it must have tried to serve)", args, code)
|
||||
}
|
||||
if out.Len() != 0 {
|
||||
t.Errorf("run(%q) wrote %q to stdout; only `version` reports on stdout", args, out.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// fail is the one place a refusal is reported, so the message names the plugin
|
||||
// and the code is non-zero whatever went wrong.
|
||||
func TestEveryRefusalIsAttributedAndNonZero(t *testing.T) {
|
||||
var w bytes.Buffer
|
||||
if code := fail(&w, errors.New("disk fell off")); code != 1 {
|
||||
t.Fatalf("fail() = %d, want 1", code)
|
||||
}
|
||||
if got := w.String(); got != "mpcvm plugin: disk fell off\n" {
|
||||
t.Fatalf("fail() wrote %q", got)
|
||||
}
|
||||
}
|
||||
+127
-228
@@ -129,21 +129,86 @@ func (g *FHEAccelerator) Stats() FHEStats {
|
||||
return *g.stats
|
||||
}
|
||||
|
||||
// BatchNTTForward performs forward NTT on multiple polynomials.
|
||||
// This is the primary use case for GPU acceleration - batch operations.
|
||||
func (g *FHEAccelerator) BatchNTTForward(r *ring.Ring, polys []ring.Poly) error {
|
||||
if len(polys) == 0 {
|
||||
return nil
|
||||
// =============================================================================
|
||||
// One device round trip, three kernels
|
||||
// =============================================================================
|
||||
//
|
||||
// The three batch entry points below differ in exactly three things: how many
|
||||
// polynomials one element reads, which lattice kernel runs, and what the CPU
|
||||
// does instead. Everything else — the size floor, the session snapshot, the
|
||||
// tensor lifecycle, and the rule that ANY device failure falls back to the CPU
|
||||
// rather than failing the batch — used to be written out three times, ~45
|
||||
// statements each. Three copies of a resource lifecycle are three places to
|
||||
// forget a Close, and a kernel added later is a fourth.
|
||||
//
|
||||
// The floors are NOT unified, deliberately. Forward NTT goes to the device only
|
||||
// at 64 polynomials and ring degree 8192; inverse NTT and multiply go at 4, at
|
||||
// any degree. The reasoning written on the forward path — that below some size
|
||||
// the host/device transfer costs more than the kernel saves — applies to all
|
||||
// three, so the other two look like drift rather than a decision. Changing a
|
||||
// performance threshold needs a device to measure on, and this host has none,
|
||||
// so the numbers are carried through as parameters and the disagreement is
|
||||
// stated rather than guessed at.
|
||||
|
||||
// kernel is one lattice operation on operands already resident on the device.
|
||||
type kernel func(ops accel.LatticeOps, in []*accel.UntypedTensor, out *accel.UntypedTensor, q uint32) error
|
||||
|
||||
// element moves one batch element's operands onto the device, runs k, and
|
||||
// brings the result back. Every tensor it opens it closes, on every path,
|
||||
// including the ones it never got to use.
|
||||
func element(s *accel.Session, n int, q uint32, operands [][]uint64, k kernel) ([]uint64, error) {
|
||||
var open []*accel.Tensor[uint64]
|
||||
defer func() {
|
||||
for _, t := range open {
|
||||
t.Close()
|
||||
}
|
||||
}()
|
||||
|
||||
in := make([]*accel.UntypedTensor, 0, len(operands))
|
||||
for _, src := range operands {
|
||||
t, err := accel.NewTensorWithData[uint64](s, []int{n}, src[:n])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
open = append(open, t)
|
||||
in = append(in, t.Untyped())
|
||||
}
|
||||
|
||||
// Get parameters from ring
|
||||
N := r.N()
|
||||
out, err := accel.NewTensor[uint64](s, []int{n})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
open = append(open, out)
|
||||
|
||||
// GPU only beneficial for large batches (>64 polys with N>=8192)
|
||||
// Otherwise CPU is faster due to data transfer overhead
|
||||
if !g.enabled || len(polys) < 64 || N < 8192 {
|
||||
for i := range polys {
|
||||
r.NTT(polys[i], polys[i])
|
||||
if err := k(s.Lattice(), in, out.Untyped(), q); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.Sync(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out.ToSlice()
|
||||
}
|
||||
|
||||
// offload runs k over a batch of count elements, writing each result into the
|
||||
// slice dst names, and does cpu(i) instead for any element the device will not
|
||||
// take — a short coefficient slice, a failed allocation, a kernel error, a
|
||||
// failed readback. A device that cannot do the work is a slower batch, never a
|
||||
// failed one.
|
||||
func (g *FHEAccelerator) offload(
|
||||
r *ring.Ring,
|
||||
count, minBatch, minDegree int,
|
||||
cpu func(i int),
|
||||
operands func(i int) [][]uint64,
|
||||
dst func(i int) []uint64,
|
||||
k kernel,
|
||||
) error {
|
||||
// An empty batch needs no special case: count < minBatch is already true of
|
||||
// it for every floor, so it takes the same path as a batch too small to be
|
||||
// worth a device, and loops zero times.
|
||||
n := r.N()
|
||||
if !g.enabled || count < minBatch || n < minDegree {
|
||||
for i := 0; i < count; i++ {
|
||||
cpu(i)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -151,10 +216,9 @@ func (g *FHEAccelerator) BatchNTTForward(r *ring.Ring, polys []ring.Poly) error
|
||||
g.mu.RLock()
|
||||
session := g.session
|
||||
g.mu.RUnlock()
|
||||
|
||||
if session == nil {
|
||||
for i := range polys {
|
||||
r.NTT(polys[i], polys[i])
|
||||
for i := 0; i < count; i++ {
|
||||
cpu(i)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -162,145 +226,63 @@ func (g *FHEAccelerator) BatchNTTForward(r *ring.Ring, polys []ring.Poly) error
|
||||
if len(r.ModuliChain()) == 0 {
|
||||
return fmt.Errorf("ring has no moduli")
|
||||
}
|
||||
Q := r.ModuliChain()[0]
|
||||
q := uint32(r.ModuliChain()[0])
|
||||
|
||||
// Process each polynomial using GPU
|
||||
for i := range polys {
|
||||
if len(polys[i].Coeffs) == 0 || len(polys[i].Coeffs[0]) < N {
|
||||
r.NTT(polys[i], polys[i])
|
||||
for i := 0; i < count; i++ {
|
||||
src := operands(i)
|
||||
if src == nil {
|
||||
cpu(i)
|
||||
continue
|
||||
}
|
||||
|
||||
inputTensor, err := accel.NewTensorWithData[uint64](session, []int{N}, polys[i].Coeffs[0][:N])
|
||||
result, err := element(session, n, q, src, k)
|
||||
if err != nil {
|
||||
r.NTT(polys[i], polys[i])
|
||||
cpu(i)
|
||||
continue
|
||||
}
|
||||
|
||||
outputTensor, err := accel.NewTensor[uint64](session, []int{N})
|
||||
if err != nil {
|
||||
inputTensor.Close()
|
||||
r.NTT(polys[i], polys[i])
|
||||
continue
|
||||
}
|
||||
|
||||
err = session.Lattice().PolynomialNTT(inputTensor.Untyped(), outputTensor.Untyped(), uint32(Q))
|
||||
if err != nil {
|
||||
inputTensor.Close()
|
||||
outputTensor.Close()
|
||||
r.NTT(polys[i], polys[i])
|
||||
continue
|
||||
}
|
||||
|
||||
if err := session.Sync(); err != nil {
|
||||
inputTensor.Close()
|
||||
outputTensor.Close()
|
||||
r.NTT(polys[i], polys[i])
|
||||
continue
|
||||
}
|
||||
|
||||
result, err := outputTensor.ToSlice()
|
||||
if err != nil {
|
||||
inputTensor.Close()
|
||||
outputTensor.Close()
|
||||
r.NTT(polys[i], polys[i])
|
||||
continue
|
||||
}
|
||||
|
||||
copy(polys[i].Coeffs[0], result)
|
||||
inputTensor.Close()
|
||||
outputTensor.Close()
|
||||
copy(dst(i), result)
|
||||
}
|
||||
|
||||
g.mu.Lock()
|
||||
g.stats.BatchCalls++
|
||||
g.mu.Unlock()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// wide returns the coefficient rows of polys[i] if the element is wide enough
|
||||
// to move, and nil if it is not — which offload reads as "do this one on the
|
||||
// CPU".
|
||||
func wide(n int, polys ...ring.Poly) [][]uint64 {
|
||||
out := make([][]uint64, 0, len(polys))
|
||||
for _, p := range polys {
|
||||
if len(p.Coeffs) == 0 || len(p.Coeffs[0]) < n {
|
||||
return nil
|
||||
}
|
||||
out = append(out, p.Coeffs[0])
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// BatchNTTForward performs forward NTT on multiple polynomials.
|
||||
// This is the primary use case for GPU acceleration - batch operations.
|
||||
func (g *FHEAccelerator) BatchNTTForward(r *ring.Ring, polys []ring.Poly) error {
|
||||
return g.offload(r, len(polys), 64, 8192,
|
||||
func(i int) { r.NTT(polys[i], polys[i]) },
|
||||
func(i int) [][]uint64 { return wide(r.N(), polys[i]) },
|
||||
func(i int) []uint64 { return polys[i].Coeffs[0] },
|
||||
func(ops accel.LatticeOps, in []*accel.UntypedTensor, out *accel.UntypedTensor, q uint32) error {
|
||||
return ops.PolynomialNTT(in[0], out, q)
|
||||
})
|
||||
}
|
||||
|
||||
// BatchNTTInverse performs inverse NTT on multiple polynomials.
|
||||
func (g *FHEAccelerator) BatchNTTInverse(r *ring.Ring, polys []ring.Poly) error {
|
||||
if len(polys) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
if !g.enabled || len(polys) < 4 {
|
||||
for i := range polys {
|
||||
r.INTT(polys[i], polys[i])
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
g.mu.RLock()
|
||||
session := g.session
|
||||
g.mu.RUnlock()
|
||||
|
||||
if session == nil {
|
||||
for i := range polys {
|
||||
r.INTT(polys[i], polys[i])
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
N := r.N()
|
||||
if len(r.ModuliChain()) == 0 {
|
||||
return fmt.Errorf("ring has no moduli")
|
||||
}
|
||||
Q := r.ModuliChain()[0]
|
||||
|
||||
for i := range polys {
|
||||
if len(polys[i].Coeffs) == 0 || len(polys[i].Coeffs[0]) < N {
|
||||
r.INTT(polys[i], polys[i])
|
||||
continue
|
||||
}
|
||||
|
||||
inputTensor, err := accel.NewTensorWithData[uint64](session, []int{N}, polys[i].Coeffs[0][:N])
|
||||
if err != nil {
|
||||
r.INTT(polys[i], polys[i])
|
||||
continue
|
||||
}
|
||||
|
||||
outputTensor, err := accel.NewTensor[uint64](session, []int{N})
|
||||
if err != nil {
|
||||
inputTensor.Close()
|
||||
r.INTT(polys[i], polys[i])
|
||||
continue
|
||||
}
|
||||
|
||||
err = session.Lattice().PolynomialINTT(inputTensor.Untyped(), outputTensor.Untyped(), uint32(Q))
|
||||
if err != nil {
|
||||
inputTensor.Close()
|
||||
outputTensor.Close()
|
||||
r.INTT(polys[i], polys[i])
|
||||
continue
|
||||
}
|
||||
|
||||
if err := session.Sync(); err != nil {
|
||||
inputTensor.Close()
|
||||
outputTensor.Close()
|
||||
r.INTT(polys[i], polys[i])
|
||||
continue
|
||||
}
|
||||
|
||||
result, err := outputTensor.ToSlice()
|
||||
if err != nil {
|
||||
inputTensor.Close()
|
||||
outputTensor.Close()
|
||||
r.INTT(polys[i], polys[i])
|
||||
continue
|
||||
}
|
||||
|
||||
copy(polys[i].Coeffs[0], result)
|
||||
inputTensor.Close()
|
||||
outputTensor.Close()
|
||||
}
|
||||
|
||||
g.mu.Lock()
|
||||
g.stats.BatchCalls++
|
||||
g.mu.Unlock()
|
||||
|
||||
return nil
|
||||
return g.offload(r, len(polys), 4, 0,
|
||||
func(i int) { r.INTT(polys[i], polys[i]) },
|
||||
func(i int) [][]uint64 { return wide(r.N(), polys[i]) },
|
||||
func(i int) []uint64 { return polys[i].Coeffs[0] },
|
||||
func(ops accel.LatticeOps, in []*accel.UntypedTensor, out *accel.UntypedTensor, q uint32) error {
|
||||
return ops.PolynomialINTT(in[0], out, q)
|
||||
})
|
||||
}
|
||||
|
||||
// BatchPolyMul performs polynomial multiplication on batches using GPU.
|
||||
@@ -308,96 +290,13 @@ func (g *FHEAccelerator) BatchPolyMul(r *ring.Ring, a, b, out []ring.Poly) error
|
||||
if len(a) != len(b) || len(a) != len(out) {
|
||||
return fmt.Errorf("batch size mismatch")
|
||||
}
|
||||
|
||||
if !g.enabled || len(a) < 4 {
|
||||
// CPU fallback
|
||||
for i := range a {
|
||||
r.MulCoeffsBarrett(a[i], b[i], out[i])
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
g.mu.RLock()
|
||||
session := g.session
|
||||
g.mu.RUnlock()
|
||||
|
||||
if session == nil {
|
||||
for i := range a {
|
||||
r.MulCoeffsBarrett(a[i], b[i], out[i])
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
N := r.N()
|
||||
if len(r.ModuliChain()) == 0 {
|
||||
return fmt.Errorf("ring has no moduli")
|
||||
}
|
||||
Q := r.ModuliChain()[0]
|
||||
|
||||
for i := range a {
|
||||
if len(a[i].Coeffs) == 0 || len(b[i].Coeffs) == 0 || len(out[i].Coeffs) == 0 {
|
||||
r.MulCoeffsBarrett(a[i], b[i], out[i])
|
||||
continue
|
||||
}
|
||||
|
||||
aTensor, err := accel.NewTensorWithData[uint64](session, []int{N}, a[i].Coeffs[0][:N])
|
||||
if err != nil {
|
||||
r.MulCoeffsBarrett(a[i], b[i], out[i])
|
||||
continue
|
||||
}
|
||||
|
||||
bTensor, err := accel.NewTensorWithData[uint64](session, []int{N}, b[i].Coeffs[0][:N])
|
||||
if err != nil {
|
||||
aTensor.Close()
|
||||
r.MulCoeffsBarrett(a[i], b[i], out[i])
|
||||
continue
|
||||
}
|
||||
|
||||
outTensor, err := accel.NewTensor[uint64](session, []int{N})
|
||||
if err != nil {
|
||||
aTensor.Close()
|
||||
bTensor.Close()
|
||||
r.MulCoeffsBarrett(a[i], b[i], out[i])
|
||||
continue
|
||||
}
|
||||
|
||||
err = session.Lattice().PolynomialMul(aTensor.Untyped(), bTensor.Untyped(), outTensor.Untyped(), uint32(Q))
|
||||
if err != nil {
|
||||
aTensor.Close()
|
||||
bTensor.Close()
|
||||
outTensor.Close()
|
||||
r.MulCoeffsBarrett(a[i], b[i], out[i])
|
||||
continue
|
||||
}
|
||||
|
||||
if err := session.Sync(); err != nil {
|
||||
aTensor.Close()
|
||||
bTensor.Close()
|
||||
outTensor.Close()
|
||||
r.MulCoeffsBarrett(a[i], b[i], out[i])
|
||||
continue
|
||||
}
|
||||
|
||||
result, err := outTensor.ToSlice()
|
||||
if err != nil {
|
||||
aTensor.Close()
|
||||
bTensor.Close()
|
||||
outTensor.Close()
|
||||
r.MulCoeffsBarrett(a[i], b[i], out[i])
|
||||
continue
|
||||
}
|
||||
|
||||
copy(out[i].Coeffs[0], result)
|
||||
aTensor.Close()
|
||||
bTensor.Close()
|
||||
outTensor.Close()
|
||||
}
|
||||
|
||||
g.mu.Lock()
|
||||
g.stats.PolyMulCalls++
|
||||
g.mu.Unlock()
|
||||
|
||||
return nil
|
||||
return g.offload(r, len(a), 4, 0,
|
||||
func(i int) { r.MulCoeffsBarrett(a[i], b[i], out[i]) },
|
||||
func(i int) [][]uint64 { return wide(r.N(), a[i], b[i], out[i]) },
|
||||
func(i int) []uint64 { return out[i].Coeffs[0] },
|
||||
func(ops accel.LatticeOps, in []*accel.UntypedTensor, o *accel.UntypedTensor, q uint32) error {
|
||||
return ops.PolynomialMul(in[0], in[1], o, q)
|
||||
})
|
||||
}
|
||||
|
||||
// ClearCache clears any cached state.
|
||||
|
||||
Reference in New Issue
Block a user