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.
102 lines
2.5 KiB
Go
102 lines
2.5 KiB
Go
//go:build cgo && gpu
|
||
|
||
package gpu
|
||
|
||
import (
|
||
"testing"
|
||
|
||
"github.com/luxfi/lattice/v7/ring"
|
||
)
|
||
|
||
// BenchmarkMontGPU_Forward measures the Montgomery-form Metal NTT in
|
||
// single-poly mode. Compare against BenchmarkRingNTT_Forward (pure Go
|
||
// in the ring package) for the per-N crossover analysis.
|
||
func BenchmarkMontGPU_Forward(b *testing.B) {
|
||
for _, N := range []int{1024, 2048, 4096, 8192, 16384} {
|
||
b.Run("", func(b *testing.B) {
|
||
r, err := ring.NewRing(N, []uint64{mtQ})
|
||
if err != nil {
|
||
b.Fatalf("NewRing: %v", err)
|
||
}
|
||
s := r.SubRings[0]
|
||
ctx, err := NewMontgomeryNTTContext(s)
|
||
if err != nil {
|
||
b.Fatalf("NewMontgomeryNTTContext: %v", err)
|
||
}
|
||
defer ctx.Close()
|
||
|
||
data := genVec(0xCAFE, N, mtQ)
|
||
work := make([]uint64, N)
|
||
b.ResetTimer()
|
||
for i := 0; i < b.N; i++ {
|
||
copy(work, data)
|
||
if err := ctx.Forward(work, 1); err != nil {
|
||
b.Fatalf("Forward: %v", err)
|
||
}
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// BenchmarkMontGPU_BatchForward measures the per-poly amortised cost
|
||
// when batching B polynomials in a single GPU dispatch. This is the
|
||
// "14× speedup" claim in #88.
|
||
func BenchmarkMontGPU_BatchForward(b *testing.B) {
|
||
configs := []struct {
|
||
N int
|
||
B int
|
||
}{
|
||
{1024, 16}, {1024, 64},
|
||
{2048, 16}, {2048, 64},
|
||
{4096, 16}, {4096, 64}, {4096, 128},
|
||
{8192, 16}, {8192, 64}, {8192, 128},
|
||
}
|
||
for _, cfg := range configs {
|
||
b.Run("", func(b *testing.B) {
|
||
r, err := ring.NewRing(cfg.N, []uint64{mtQ})
|
||
if err != nil {
|
||
b.Fatalf("NewRing: %v", err)
|
||
}
|
||
s := r.SubRings[0]
|
||
ctx, err := NewMontgomeryNTTContext(s)
|
||
if err != nil {
|
||
b.Fatalf("NewMontgomeryNTTContext: %v", err)
|
||
}
|
||
defer ctx.Close()
|
||
|
||
seed := genVec(0xCAFE, cfg.N, mtQ)
|
||
work := make([]uint64, cfg.N*cfg.B)
|
||
for j := 0; j < cfg.B; j++ {
|
||
copy(work[j*cfg.N:(j+1)*cfg.N], seed)
|
||
}
|
||
b.ResetTimer()
|
||
for i := 0; i < b.N; i++ {
|
||
if err := ctx.Forward(work, uint32(cfg.B)); err != nil {
|
||
b.Fatalf("Forward batched: %v", err)
|
||
}
|
||
}
|
||
b.ReportMetric(float64(b.N)*float64(cfg.B)/b.Elapsed().Seconds(), "ntts/sec")
|
||
})
|
||
}
|
||
}
|
||
|
||
// BenchmarkRingNTT_Forward is the pure-Go baseline.
|
||
func BenchmarkRingNTT_Forward(b *testing.B) {
|
||
for _, N := range []int{1024, 2048, 4096, 8192, 16384} {
|
||
b.Run("", func(b *testing.B) {
|
||
r, err := ring.NewRing(N, []uint64{mtQ})
|
||
if err != nil {
|
||
b.Fatalf("NewRing: %v", err)
|
||
}
|
||
s := r.SubRings[0]
|
||
data := genVec(0xCAFE, N, mtQ)
|
||
work := make([]uint64, N)
|
||
b.ResetTimer()
|
||
for i := 0; i < b.N; i++ {
|
||
copy(work, data)
|
||
s.NTT(work, work)
|
||
}
|
||
})
|
||
}
|
||
}
|