diff --git a/LLM.md b/LLM.md index 41a736f..8a96b3d 100644 --- a/LLM.md +++ b/LLM.md @@ -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`. diff --git a/ntt_gpu.go b/ntt_gpu.go new file mode 100644 index 0000000..fd7aa36 --- /dev/null +++ b/ntt_gpu.go @@ -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 } diff --git a/ntt_gpu_test.go b/ntt_gpu_test.go new file mode 100644 index 0000000..b256163 --- /dev/null +++ b/ntt_gpu_test.go @@ -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()) + } + } + } +}