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.
87 lines
1.9 KiB
Go
87 lines
1.9 KiB
Go
//go:build cgo && gpu
|
|
|
|
package gpu
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/luxfi/lattice/v7/ring"
|
|
)
|
|
|
|
// TestDebug isolates the N=4096 divergence with minimal input.
|
|
func TestDebug_N4096_OneHot(t *testing.T) {
|
|
r, err := ring.NewRing(4096, []uint64{mtQ})
|
|
if err != nil {
|
|
t.Fatalf("NewRing: %v", err)
|
|
}
|
|
s := r.SubRings[0]
|
|
|
|
ctx, err := NewMontgomeryNTTContext(s)
|
|
if err != nil {
|
|
t.Fatalf("NewMontgomeryNTTContext: %v", err)
|
|
}
|
|
defer ctx.Close()
|
|
|
|
for _, hotIdx := range []int{0, 1, 2047, 2048, 4095} {
|
|
t.Run("", func(t *testing.T) {
|
|
input := make([]uint64, 4096)
|
|
input[hotIdx] = 1
|
|
goOut := make([]uint64, 4096)
|
|
gpuOut := make([]uint64, 4096)
|
|
copy(goOut, input)
|
|
copy(gpuOut, input)
|
|
s.NTT(goOut, goOut)
|
|
if err := ctx.Forward(gpuOut, 1); err != nil {
|
|
t.Fatalf("Forward: %v", err)
|
|
}
|
|
diff := 0
|
|
firstDiff := -1
|
|
for i := 0; i < 4096; i++ {
|
|
if goOut[i] != gpuOut[i] {
|
|
if firstDiff < 0 {
|
|
firstDiff = i
|
|
}
|
|
diff++
|
|
}
|
|
}
|
|
if diff > 0 {
|
|
t.Errorf("hot=%d differ at %d coefficients (first at i=%d): go=%016x gpu=%016x",
|
|
hotIdx, diff, firstDiff, goOut[firstDiff], gpuOut[firstDiff])
|
|
} else {
|
|
t.Logf("hot=%d byte-equal", hotIdx)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// Also test N=2048 to confirm sanity.
|
|
func TestDebug_N2048_OneHot(t *testing.T) {
|
|
r, err := ring.NewRing(2048, []uint64{mtQ})
|
|
if err != nil {
|
|
t.Fatalf("NewRing: %v", err)
|
|
}
|
|
s := r.SubRings[0]
|
|
ctx, err := NewMontgomeryNTTContext(s)
|
|
if err != nil {
|
|
t.Fatalf("NewMontgomeryNTTContext: %v", err)
|
|
}
|
|
defer ctx.Close()
|
|
|
|
input := make([]uint64, 2048)
|
|
input[0] = 1
|
|
goOut := make([]uint64, 2048)
|
|
gpuOut := make([]uint64, 2048)
|
|
copy(goOut, input)
|
|
copy(gpuOut, input)
|
|
s.NTT(goOut, goOut)
|
|
if err := ctx.Forward(gpuOut, 1); err != nil {
|
|
t.Fatalf("Forward: %v", err)
|
|
}
|
|
for i := 0; i < 2048; i++ {
|
|
if goOut[i] != gpuOut[i] {
|
|
t.Errorf("differ at i=%d: go=%016x gpu=%016x", i, goOut[i], gpuOut[i])
|
|
return
|
|
}
|
|
}
|
|
}
|