mirror of
https://github.com/luxfi/lattice.git
synced 2026-09-24 00:25:44 +00:00
zeekay
5d4473dfb8
Go-side surface for the Montgomery Metal NTT shipped in
luxcpp/lattice. New constructors:
- gpu.NewMontgomeryNTTContext(s *ring.SubRing): binds a SubRing to a
GPU-resident NTT context, capturing RootsForward/RootsBackward/
NInv/MRedConstant. Forward/Backward output is byte-equal to
ring.SubRing.NTT/INTT for the same inputs.
- gpu.RegisterSubRing / LookupSubRing / UnregisterSubRing: per-SubRing
registry so the Ring.NTT fast-path can route calls to GPU.
- gpu.SetNTTThreshold / NTTThreshold: tunable per-N crossover.
ring.SubRing now consults pluggable NTT dispatchers via
ring.SetGPUDispatchers (called from gpu/init). Default behaviour is
unchanged -- dispatcher returns false unless a SubRing was explicitly
registered AND the threshold is non-zero. Pure-Go path remains the
default.
Tests:
- TestMontgomery_ByteEqual_Forward: 1024 random vectors x N in
{1024, 2048, 4096} byte-equal to ring.SubRing.NTT
- TestMontgomery_ByteEqual_Inverse: same for INTT
- TestMontgomery_RoundTrip: at N up to 8192
- TestBatchNTT_NoSegfault: previously-segfaulting (N=4096, B=128) now
runs and produces byte-equal output to per-poly Go pure-Go.
Bench (M1 Max, benchtime=10x):
- N=4096 B=128 batched: 12.6 us per NTT (4.57x faster than
Go pure-Go)
- N=4096 B=64 batched: 18.2 us per NTT (3.18x faster)
- Single-poly Metal remains slower than Go at all measured N
(8-100x slower); single-poly dispatch is opt-in via threshold.
Also fixes a pre-existing bug in gpu.go pure-Go NewNTTContext:
ring.NewSubRing alone does not generate NTT constants; switched to
ring.NewRing which calls generateNTTConstants under the hood.
301 lines
12 KiB
Go
301 lines
12 KiB
Go
package ring
|
|
|
|
// Add evaluates p3 = p1 + p2 (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) Add(p1, p2, p3 []uint64) {
|
|
addvec(p1, p2, p3, s.Modulus)
|
|
}
|
|
|
|
// AddLazy evaluates p3 = p1 + p2.
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) AddLazy(p1, p2, p3 []uint64) {
|
|
addlazyvec(p1, p2, p3)
|
|
}
|
|
|
|
// Sub evaluates p3 = p1 - p2 (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) Sub(p1, p2, p3 []uint64) {
|
|
subvec(p1, p2, p3, s.Modulus)
|
|
}
|
|
|
|
// SubLazy evaluates p3 = p1 - p2.
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) SubLazy(p1, p2, p3 []uint64) {
|
|
sublazyvec(p1, p2, p3, s.Modulus)
|
|
}
|
|
|
|
// Neg evaluates p2 = -p1 (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) Neg(p1, p2 []uint64) {
|
|
negvec(p1, p2, s.Modulus)
|
|
}
|
|
|
|
// Reduce evaluates p2 = p1 (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) Reduce(p1, p2 []uint64) {
|
|
reducevec(p1, p2, s.Modulus, s.BRedConstant)
|
|
}
|
|
|
|
// ReduceLazy evaluates p2 = p1 (mod modulus) with p2 in range [0, 2*modulus-1].
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) ReduceLazy(p1, p2 []uint64) {
|
|
reducelazyvec(p1, p2, s.Modulus, s.BRedConstant)
|
|
}
|
|
|
|
// MulCoeffsLazy evaluates p3 = p1*p2.
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulCoeffsLazy(p1, p2, p3 []uint64) {
|
|
mulcoeffslazyvec(p1, p2, p3)
|
|
}
|
|
|
|
// MulCoeffsLazyThenAddLazy evaluates p3 = p3 + p1*p2.
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulCoeffsLazyThenAddLazy(p1, p2, p3 []uint64) {
|
|
mulcoeffslazythenaddlazyvec(p1, p2, p3)
|
|
}
|
|
|
|
// MulCoeffsBarrett evaluates p3 = p1*p2 (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulCoeffsBarrett(p1, p2, p3 []uint64) {
|
|
mulcoeffsbarrettvec(p1, p2, p3, s.Modulus, s.BRedConstant)
|
|
}
|
|
|
|
// MulCoeffsBarrettLazy evaluates p3 = p1*p2 (mod modulus) with p3 in [0, 2*modulus-1].
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulCoeffsBarrettLazy(p1, p2, p3 []uint64) {
|
|
mulcoeffsbarrettlazyvec(p1, p2, p3, s.Modulus, s.BRedConstant)
|
|
}
|
|
|
|
// MulCoeffsBarrettThenAdd evaluates p3 = p3 + (p1*p2) (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulCoeffsBarrettThenAdd(p1, p2, p3 []uint64) {
|
|
mulcoeffsthenaddvec(p1, p2, p3, s.Modulus, s.BRedConstant)
|
|
}
|
|
|
|
// MulCoeffsBarrettThenAddLazy evaluates p3 = p3 + p1*p2 (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulCoeffsBarrettThenAddLazy(p1, p2, p3 []uint64) {
|
|
mulcoeffsbarrettthenaddlazyvec(p1, p2, p3, s.Modulus, s.BRedConstant)
|
|
}
|
|
|
|
// MulCoeffsMontgomery evaluates p3 = p1*p2 (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulCoeffsMontgomery(p1, p2, p3 []uint64) {
|
|
mulcoeffsmontgomeryvec(p1, p2, p3, s.Modulus, s.MRedConstant)
|
|
}
|
|
|
|
// MulCoeffsMontgomeryLazy evaluates p3 = p1*p2 (mod modulus) with p3 in range [0, 2*modulus-1].
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulCoeffsMontgomeryLazy(p1, p2, p3 []uint64) {
|
|
mulcoeffsmontgomerylazyvec(p1, p2, p3, s.Modulus, s.MRedConstant)
|
|
}
|
|
|
|
// MulCoeffsMontgomeryThenAdd evaluates p3 = p3 + (p1*p2) (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulCoeffsMontgomeryThenAdd(p1, p2, p3 []uint64) {
|
|
mulcoeffsmontgomerythenaddvec(p1, p2, p3, s.Modulus, s.MRedConstant)
|
|
}
|
|
|
|
// MulCoeffsMontgomeryThenAddLazy evaluates p3 = p3 + (p1*p2 (mod modulus)).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulCoeffsMontgomeryThenAddLazy(p1, p2, p3 []uint64) {
|
|
mulcoeffsmontgomerythenaddlazyvec(p1, p2, p3, s.Modulus, s.MRedConstant)
|
|
}
|
|
|
|
// MulCoeffsMontgomeryLazyThenAddLazy evaluates p3 = p3 + p1*p2 (mod modulus) with p3 in range [0, 3modulus-2].
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulCoeffsMontgomeryLazyThenAddLazy(p1, p2, p3 []uint64) {
|
|
mulcoeffsmontgomerylazythenaddlazyvec(p1, p2, p3, s.Modulus, s.MRedConstant)
|
|
}
|
|
|
|
// MulCoeffsMontgomeryThenSub evaluates p3 = p3 - p1*p2 (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulCoeffsMontgomeryThenSub(p1, p2, p3 []uint64) {
|
|
mulcoeffsmontgomerythensubvec(p1, p2, p3, s.Modulus, s.MRedConstant)
|
|
}
|
|
|
|
// MulCoeffsMontgomeryThenSubLazy evaluates p3 = p3 - p1*p2 (mod modulus) with p3 in range [0, 2*modulus-2].
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulCoeffsMontgomeryThenSubLazy(p1, p2, p3 []uint64) {
|
|
mulcoeffsmontgomerythensublazyvec(p1, p2, p3, s.Modulus, s.MRedConstant)
|
|
}
|
|
|
|
// MulCoeffsMontgomeryLazyThenSubLazy evaluates p3 = p3 - p1*p2 (mod modulus) with p3 in range [0, 3*modulus-2].
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulCoeffsMontgomeryLazyThenSubLazy(p1, p2, p3 []uint64) {
|
|
mulcoeffsmontgomerylazythensublazyvec(p1, p2, p3, s.Modulus, s.MRedConstant)
|
|
}
|
|
|
|
// MulCoeffsMontgomeryLazyThenNeg evaluates p3 = - p1*p2 (mod modulus) with p3 in range [0, 2*modulus-2].
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulCoeffsMontgomeryLazyThenNeg(p1, p2, p3 []uint64) {
|
|
mulcoeffsmontgomerylazythenNegvec(p1, p2, p3, s.Modulus, s.MRedConstant)
|
|
}
|
|
|
|
// AddLazyThenMulScalarMontgomery evaluates p3 = (p1+p2)*scalarMont (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) AddLazyThenMulScalarMontgomery(p1, p2 []uint64, scalarMont uint64, p3 []uint64) {
|
|
addlazythenmulscalarmontgomeryvec(p1, p2, scalarMont, p3, s.Modulus, s.MRedConstant)
|
|
}
|
|
|
|
// AddScalarLazyThenMulScalarMontgomery evaluates p3 = (scalarMont0+p2)*scalarMont1 (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) AddScalarLazyThenMulScalarMontgomery(p1 []uint64, scalar0, scalarMont1 uint64, p2 []uint64) {
|
|
addscalarlazythenmulscalarmontgomeryvec(p1, scalar0, scalarMont1, p2, s.Modulus, s.MRedConstant)
|
|
}
|
|
|
|
// AddScalar evaluates p2 = p1 + scalar (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) AddScalar(p1 []uint64, scalar uint64, p2 []uint64) {
|
|
addscalarvec(p1, scalar, p2, s.Modulus)
|
|
}
|
|
|
|
// AddScalarLazy evaluates p2 = p1 + scalar.
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) AddScalarLazy(p1 []uint64, scalar uint64, p2 []uint64) {
|
|
addscalarlazyvec(p1, scalar, p2)
|
|
}
|
|
|
|
// AddScalarLazyThenNegTwoModulusLazy evaluates p2 = 2*modulus - p1 + scalar.
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) AddScalarLazyThenNegTwoModulusLazy(p1 []uint64, scalar uint64, p2 []uint64) {
|
|
addscalarlazythenNegTwoModuluslazyvec(p1, scalar, p2, s.Modulus)
|
|
}
|
|
|
|
// SubScalar evaluates p2 = p1 - scalar (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) SubScalar(p1 []uint64, scalar uint64, p2 []uint64) {
|
|
subscalarvec(p1, scalar, p2, s.Modulus)
|
|
}
|
|
|
|
// MulScalarMontgomery evaluates p2 = p1*scalarMont (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulScalarMontgomery(p1 []uint64, scalarMont uint64, p2 []uint64) {
|
|
mulscalarmontgomeryvec(p1, scalarMont, p2, s.Modulus, s.MRedConstant)
|
|
}
|
|
|
|
// MulScalarMontgomeryLazy evaluates p2 = p1*scalarMont (mod modulus) with p2 in range [0, 2*modulus-1].
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulScalarMontgomeryLazy(p1 []uint64, scalarMont uint64, p2 []uint64) {
|
|
mulscalarmontgomerylazyvec(p1, scalarMont, p2, s.Modulus, s.MRedConstant)
|
|
}
|
|
|
|
// MulScalarMontgomeryThenAdd evaluates p2 = p2 + p1*scalarMont (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulScalarMontgomeryThenAdd(p1 []uint64, scalarMont uint64, p2 []uint64) {
|
|
mulscalarmontgomerythenaddvec(p1, scalarMont, p2, s.Modulus, s.MRedConstant)
|
|
}
|
|
|
|
// MulScalarMontgomeryThenAddScalar evaluates p2 = scalar + p1*scalarMont (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MulScalarMontgomeryThenAddScalar(p1 []uint64, scalar0, scalarMont1 uint64, p2 []uint64) {
|
|
mulscalarmontgomerythenaddscalarvec(p1, scalar0, scalarMont1, p2, s.Modulus, s.MRedConstant)
|
|
}
|
|
|
|
// SubThenMulScalarMontgomeryTwoModulus evaluates p3 = (p1 + twomodulus - p2) * scalarMont (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) SubThenMulScalarMontgomeryTwoModulus(p1, p2 []uint64, scalarMont uint64, p3 []uint64) {
|
|
subthenmulscalarmontgomeryTwoModulusvec(p1, p2, scalarMont, p3, s.Modulus, s.MRedConstant)
|
|
}
|
|
|
|
// nttGPUDispatcher is the optional GPU dispatch hook. The gpu package
|
|
// installs a non-nil dispatcher (build-tagged on cgo + gpu) that returns
|
|
// true when it has handled the call. The pure-Go path runs only when
|
|
// the dispatcher returns false (no GPU bound for this SubRing, threshold
|
|
// not reached, GPU unavailable, etc.). Pure-Go behaviour is unchanged
|
|
// when the dispatcher is left nil.
|
|
type subRingNTTDispatcher func(s *SubRing, src, dst []uint64) bool
|
|
|
|
var (
|
|
nttForwardDispatcher subRingNTTDispatcher
|
|
nttBackwardDispatcher subRingNTTDispatcher
|
|
)
|
|
|
|
// SetGPUDispatchers installs the forward/backward GPU dispatchers. The
|
|
// gpu package is the only intended caller; passing nil reverts to the
|
|
// pure-Go path. Safe to call from package init.
|
|
func SetGPUDispatchers(forward, backward subRingNTTDispatcher) {
|
|
nttForwardDispatcher = forward
|
|
nttBackwardDispatcher = backward
|
|
}
|
|
|
|
// NTT evaluates p2 = NTT(p1).
|
|
func (s *SubRing) NTT(p1, p2 []uint64) {
|
|
if nttForwardDispatcher != nil && nttForwardDispatcher(s, p1, p2) {
|
|
return
|
|
}
|
|
s.ntt.Forward(p1, p2)
|
|
}
|
|
|
|
// NTTLazy evaluates p2 = NTT(p1) with p2 in [0, 6*modulus-2].
|
|
func (s *SubRing) NTTLazy(p1, p2 []uint64) {
|
|
s.ntt.ForwardLazy(p1, p2)
|
|
}
|
|
|
|
// INTT evaluates p2 = INTT(p1).
|
|
func (s *SubRing) INTT(p1, p2 []uint64) {
|
|
if nttBackwardDispatcher != nil && nttBackwardDispatcher(s, p1, p2) {
|
|
return
|
|
}
|
|
s.ntt.Backward(p1, p2)
|
|
}
|
|
|
|
// INTTLazy evaluates p2 = INTT(p1) with p2 in [0, 2*modulus-1].
|
|
func (s *SubRing) INTTLazy(p1, p2 []uint64) {
|
|
s.ntt.BackwardLazy(p1, p2)
|
|
}
|
|
|
|
// MForm evaluates p2 = p1 * 2^64 (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MForm(p1, p2 []uint64) {
|
|
mformvec(p1, p2, s.Modulus, s.BRedConstant)
|
|
}
|
|
|
|
// MFormLazy evaluates p2 = p1 * 2^64 (mod modulus) with p2 in the range [0, 2*modulus-1].
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) MFormLazy(p1, p2 []uint64) {
|
|
mformlazyvec(p1, p2, s.Modulus, s.BRedConstant)
|
|
}
|
|
|
|
// IMForm evaluates p2 = p1 * (2^64)^-1 (mod modulus).
|
|
// Iteration is done with respect to len(p1).
|
|
// All input must have a size which is a multiple of 8.
|
|
func (s *SubRing) IMForm(p1, p2 []uint64) {
|
|
imformvec(p1, p2, s.Modulus, s.MRedConstant)
|
|
}
|