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.