fhe: wire batch NTT/INTT to lattice/v7/gpu with byte-equality safety gate

Adds NTTEngine.NTTBatch / INTTBatch that dispatch to luxfi/lattice/v7/gpu (Metal on darwin, CUDA on linux/NVIDIA) when: 1. CGo is on and a GPU device is reachable (lattice/gpu.GPUAvailable()) 2. The engine's (N, Q) has passed a deterministic byte-equality probe against the CPU subring NTT oracle (cached per (N, Q)) Otherwise falls back to per-polynomial CPU subring NTT. The probe ensures a GPU context with different convention encoding (Montgomery vs standard, bit-reversed vs natural) is rejected before it can corrupt ciphertexts. No `gpu` build tag. CGo is the only compile-time gate. Tests: ntt_gpu_test.go::TestNTTBatch_ByteEqualToInPlace (probe enforces byte equality), TestINTTBatch_RoundTrip (INTT(NTT(x)) == x).
2026-07-26 23:16:08 +00:00 · 2026-05-24 14:17:44 -07:00
parent e469f435b8
commit 90cc46fc41
3 changed files with 231 additions and 0 deletions
@@ -574,3 +574,20 @@ All code uses permissive licenses:
 | `core/kms/core/` | KMS core logic |
 | `go/tfhe/cmd/` | Go FHE server |
 | `ml/torus-ml/src/` | ML with FHE |
+
+## NTT (Go side, `ntt_simd.go` + `ntt_gpu.go`)
+
+`fhe.NTTEngine` wraps a cached `subring.SubRing` keyed by `(N, Q)`. Two API tiers:
+
+| Method | Path |
+|--------|------|
+| `NTTInPlace(coeffs)` / `INTTInPlace(coeffs)` | Single polynomial, CPU subring. |
+| `NTTBatch(polys) / INTTBatch(polys)` | Batch dispatch — `luxfi/lattice/v7/gpu` (Metal/CUDA) when `gpu.GPUAvailable()` AND the engine's `(N,Q)` passes byte-equality probe; otherwise per-poly CPU. |
+| `GPUNTTAvailable()` | Reports whether the GPU path serves this engine. |
+
+CGo is the only build-tag axis. The byte-equality probe runs once per `(N, Q)`
+and caches the result — a GPU context that disagrees with the CPU oracle is
+rejected (never silently dispatched).
+
+Tests: `ntt_gpu_test.go::TestNTTBatch_ByteEqualToInPlace`,
+`TestINTTBatch_RoundTrip`.
@@ -0,0 +1,110 @@
+// Copyright (c) 2025-2026, Lux Industries Inc
+// SPDX-License-Identifier: BSD-3-Clause
+
+package fhe
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/luxfi/lattice/v7/gpu"
+)
+
+// Batch NTT/INTT dispatch — routes to lattice/v7/gpu (Metal/CUDA) when a
+// GPU device is reachable AND the engine's (N, Q) probe matches the CPU
+// subring NTT byte-for-byte. Otherwise per-polynomial subring fallback.
+// CGO_ENABLED gates lattice/gpu compilation; no other build tag exists.
+
+var (
+	gpuCtxMu    sync.Mutex
+	gpuCtxCache = map[engineKey]*gpu.NTTContext{}
+)
+
+func fheGPUCtx(e *NTTEngine) *gpu.NTTContext {
+	if !gpu.GPUAvailable() {
+		return nil
+	}
+	gpuCtxMu.Lock()
+	defer gpuCtxMu.Unlock()
+	k := engineKey{N: e.N, Q: e.Q}
+	if ctx, ok := gpuCtxCache[k]; ok {
+		return ctx
+	}
+	ctx, err := gpu.NewNTTContext(e.N, e.Q)
+	if err != nil {
+		gpuCtxCache[k] = nil
+		return nil
+	}
+	if !probeByteEqual(e, ctx) {
+		ctx.Close()
+		gpuCtxCache[k] = nil
+		return nil
+	}
+	gpuCtxCache[k] = ctx
+	return ctx
+}
+
+// probeByteEqual rejects a GPU context whose forward NTT differs from
+// the CPU subring NTT on a deterministic input. Different Montgomery /
+// bit-reversed / twiddle conventions would corrupt ciphertexts silently.
+func probeByteEqual(e *NTTEngine, gpuCtx *gpu.NTTContext) bool {
+	N := int(e.N)
+	probe := make([]uint64, N)
+	for i := range probe {
+		probe[i] = uint64(i+1) % e.Q
+	}
+	cpu := make([]uint64, N)
+	copy(cpu, probe)
+	e.sr.NTT(cpu, cpu)
+	gpuOut, err := gpuCtx.NTT([][]uint64{probe})
+	if err != nil || len(gpuOut) != 1 || len(gpuOut[0]) != N {
+		return false
+	}
+	for i := 0; i < N; i++ {
+		if cpu[i] != gpuOut[0][i] {
+			return false
+		}
+	}
+	return true
+}
+
+// NTTBatch performs forward NTT on every polynomial in polys. Output is
+// byte-equal to repeated NTTInPlace calls regardless of dispatch path.
+func (e *NTTEngine) NTTBatch(polys [][]uint64) ([][]uint64, error) {
+	if ctx := fheGPUCtx(e); ctx != nil {
+		return ctx.NTT(polys)
+	}
+	N := int(e.N)
+	out := make([][]uint64, len(polys))
+	for i, p := range polys {
+		if len(p) != N {
+			return nil, fmt.Errorf("fhe.NTTBatch: poly %d size %d != %d", i, len(p), N)
+		}
+		out[i] = make([]uint64, N)
+		copy(out[i], p)
+		e.sr.NTT(out[i], out[i])
+	}
+	return out, nil
+}
+
+// INTTBatch is the inverse of NTTBatch.
+func (e *NTTEngine) INTTBatch(polys [][]uint64) ([][]uint64, error) {
+	if ctx := fheGPUCtx(e); ctx != nil {
+		return ctx.INTT(polys)
+	}
+	N := int(e.N)
+	out := make([][]uint64, len(polys))
+	for i, p := range polys {
+		if len(p) != N {
+			return nil, fmt.Errorf("fhe.INTTBatch: poly %d size %d != %d", i, len(p), N)
+		}
+		out[i] = make([]uint64, N)
+		copy(out[i], p)
+		e.sr.INTT(out[i], out[i])
+	}
+	return out, nil
+}
+
+// GPUNTTAvailable reports whether the GPU path serves this engine's (N, Q).
+// First call may run the byte-equality probe; subsequent calls are cached.
+func (e *NTTEngine) GPUNTTAvailable() bool { return fheGPUCtx(e) != nil }
@@ -0,0 +1,104 @@
+// Copyright (c) 2025-2026, Lux Industries Inc
+// SPDX-License-Identifier: BSD-3-Clause
+
+package fhe
+
+import (
+	"testing"
+)
+
+// TestNTTBatch_ByteEqualToInPlace pins the byte-equality contract: the
+// batch path (whether routed to GPU or CPU subring fallback) MUST produce
+// identical output to repeated NTTInPlace calls on the same inputs.
+func TestNTTBatch_ByteEqualToInPlace(t *testing.T) {
+	const (
+		N = uint32(1024)
+		Q = uint64(0x3FFFFFFFFED001) // 54-bit NTT-friendly prime (PN11QP54)
+	)
+	eng, err := NewNTTEngine(N, Q)
+	if err != nil {
+		t.Skipf("subring init failed for N=%d Q=%x: %v", N, Q, err)
+	}
+
+	batchSize := 8
+	polys := make([][]uint64, batchSize)
+	for i := range polys {
+		p := make([]uint64, N)
+		for j := range p {
+			p[j] = uint64(i*1000+j+1) % Q
+		}
+		polys[i] = p
+	}
+
+	// Reference: per-poly in-place.
+	want := make([][]uint64, batchSize)
+	for i, p := range polys {
+		c := make([]uint64, N)
+		copy(c, p)
+		eng.NTTInPlace(c)
+		want[i] = c
+	}
+
+	got, err := eng.NTTBatch(polys)
+	if err != nil {
+		t.Fatalf("NTTBatch error: %v", err)
+	}
+	if len(got) != batchSize {
+		t.Fatalf("NTTBatch len=%d want=%d", len(got), batchSize)
+	}
+	for i := 0; i < batchSize; i++ {
+		for j := uint32(0); j < N; j++ {
+			if got[i][j] != want[i][j] {
+				t.Fatalf("NTTBatch byte mismatch at [%d][%d]: got=%d want=%d (gpu_available=%v)",
+					i, j, got[i][j], want[i][j], eng.GPUNTTAvailable())
+			}
+		}
+	}
+}
+
+// TestINTTBatch_RoundTrip pins INTT(NTT(x)) == x for the batch path.
+func TestINTTBatch_RoundTrip(t *testing.T) {
+	const (
+		N = uint32(1024)
+		Q = uint64(0x3FFFFFFFFED001)
+	)
+	eng, err := NewNTTEngine(N, Q)
+	if err != nil {
+		t.Skipf("subring init failed: %v", err)
+	}
+
+	batchSize := 4
+	original := make([][]uint64, batchSize)
+	for i := range original {
+		p := make([]uint64, N)
+		for j := range p {
+			p[j] = uint64(j*7+i+1) % Q
+		}
+		original[i] = p
+	}
+
+	// Deep-copy because NTT/INTTBatch shouldn't mutate input but contracts vary.
+	inputCopy := make([][]uint64, batchSize)
+	for i, p := range original {
+		inputCopy[i] = make([]uint64, N)
+		copy(inputCopy[i], p)
+	}
+
+	forward, err := eng.NTTBatch(inputCopy)
+	if err != nil {
+		t.Fatalf("NTTBatch: %v", err)
+	}
+	back, err := eng.INTTBatch(forward)
+	if err != nil {
+		t.Fatalf("INTTBatch: %v", err)
+	}
+
+	for i := 0; i < batchSize; i++ {
+		for j := uint32(0); j < N; j++ {
+			if back[i][j] != original[i][j] {
+				t.Fatalf("round-trip mismatch at [%d][%d]: got=%d want=%d (gpu_available=%v)",
+					i, j, back[i][j], original[i][j], eng.GPUNTTAvailable())
+			}
+		}
+	}
+}