mirror of
https://github.com/luxfi/lattice.git
synced 2026-09-24 00:25:44 +00:00
feat(types): PolyDomain + NTTContext + ReductionBudget for FHE-GPU dispatch
LP-137-FHE-TYPING. Type-system foundation that prevents silent FHE correctness failures at the kernel boundary — the architectural fix for the #121 class of bug (Go Ring.NTT in Montgomery, Metal NTT in Barrett standard form, dispatching with identical signatures). New package github.com/luxfi/lattice/v7/types: - PolyDomain (uint8 enum): Standard, Montgomery, NTTStandard, NTTMontgomery. - NTTContext (48-byte struct): modulus, MontR, MontR2, QInv, N, ModulusID, TwiddleOffset, InputDomain, RootDomain, OutputDomain. Validate() rejects domain mismatches at dispatch. - ReductionMode (uint8 enum) + ReductionBudget (24-byte struct): track lazy reduction op budgets. SafeBoundFor enforces 64/63/62/61-bit modulus caps for Strict/Lazy2/Lazy4/Lazy8. - Sentinel errors: ErrDomainMismatch, ErrModulusZero, ErrNNotPowerOfTwo, ErrModulusTooLargeForLazy. Layout is byte-stable across Go and C++ for cgo / shared-memory dispatch. Verified by unsafe.Sizeof + offset assertions and a cross-language byte image test (TestNTTContext_CrossLangByteImage) matching the C++ side in luxcpp/lattice/test/types_test.cpp. 20 Go tests pass.
This commit is contained in:
@@ -0,0 +1,59 @@
|
||||
// Copyright (c) 2026, Lux Industries Inc.
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
package types
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"testing"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
// TestNTTContext_CrossLangByteImage asserts that the byte image of a
|
||||
// canonical NTTContext is identical between Go and C++.
|
||||
//
|
||||
// The matching C++ test is at luxcpp/lattice/test/types_test.cpp
|
||||
// (test_ntt_context_memcmp). If either side changes the layout, both this
|
||||
// Go test and the C++ memcmp test MUST be updated together.
|
||||
func TestNTTContext_CrossLangByteImage(t *testing.T) {
|
||||
// Same field values as the C++ test_ntt_context_memcmp case.
|
||||
ctx := NTTContext{
|
||||
Modulus: 0x0123456789ABCDEF,
|
||||
MontR: 0x1111111111111111,
|
||||
MontR2: 0x2222222222222222,
|
||||
QInv: 0x3333333333333333,
|
||||
N: 0x00010000,
|
||||
ModulusID: 0x00000007,
|
||||
TwiddleOffset: 0x000000AB,
|
||||
InputDomain: PolyDomainMontgomery,
|
||||
RootDomain: PolyDomainMontgomery,
|
||||
OutputDomain: PolyDomainNTTMontgomery,
|
||||
}
|
||||
|
||||
// Expected little-endian byte image, byte-identical to the C++ side's
|
||||
// `want` array.
|
||||
want := []byte{
|
||||
// modulus 0x0123456789ABCDEF
|
||||
0xEF, 0xCD, 0xAB, 0x89, 0x67, 0x45, 0x23, 0x01,
|
||||
// mont_R 0x1111111111111111
|
||||
0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11,
|
||||
// mont_R2 0x2222222222222222
|
||||
0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22,
|
||||
// q_inv 0x3333333333333333
|
||||
0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33,
|
||||
// N 0x00010000
|
||||
0x00, 0x00, 0x01, 0x00,
|
||||
// modulus_id 0x00000007
|
||||
0x07, 0x00, 0x00, 0x00,
|
||||
// twiddle_offset 0x000000AB
|
||||
0xAB, 0x00, 0x00, 0x00,
|
||||
// input_domain, root_domain, output_domain, _pad
|
||||
0x01, 0x01, 0x03, 0x00,
|
||||
}
|
||||
|
||||
// Reinterpret ctx as raw bytes via unsafe.Pointer.
|
||||
got := unsafe.Slice((*byte)(unsafe.Pointer(&ctx)), unsafe.Sizeof(ctx))
|
||||
if !bytes.Equal(got, want) {
|
||||
t.Fatalf("NTTContext byte image mismatch with C++ side\n got: %x\n want: %x", got, want)
|
||||
}
|
||||
}
|
||||
+109
@@ -0,0 +1,109 @@
|
||||
// Package types defines the FHE-GPU domain typing system that prevents
|
||||
// silent correctness failures in NTT/Montgomery dispatch.
|
||||
//
|
||||
// Reference: LP-137-FHE-TYPING.md.
|
||||
//
|
||||
// Core invariant: every polynomial buffer carries an explicit PolyDomain tag,
|
||||
// and every NTT/arithmetic kernel declares the (Input, Root, Output) tuple
|
||||
// it accepts. Dispatch with mismatched tags is rejected at the type-system
|
||||
// boundary, not silently corrupted at the math layer.
|
||||
//
|
||||
// Copyright (c) 2026, Lux Industries Inc.
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
package types
|
||||
|
||||
// PolyDomain identifies which arithmetic domain a polynomial buffer lives in.
|
||||
//
|
||||
// The four domains form a 2x2 grid:
|
||||
//
|
||||
// | coefficient form | NTT (evaluation) form
|
||||
// ------------------+------------------+----------------------
|
||||
// standard form | Standard | NTTStandard
|
||||
// Montgomery form | Montgomery | NTTMontgomery
|
||||
//
|
||||
// "standard form" means values in [0, q). "Montgomery form" means values
|
||||
// represented as a*R mod q, where R = 2^64 mod q, used by Montgomery reduction
|
||||
// (MRed). NTT/non-NTT denotes whether the buffer is in coefficient form or
|
||||
// in pointwise (evaluation) form after a forward NTT.
|
||||
//
|
||||
// Domain mismatches are correctness bugs, not performance bugs. A polynomial
|
||||
// in Montgomery form cannot be multiplied pointwise with one in standard form;
|
||||
// the result is silently wrong.
|
||||
type PolyDomain uint8
|
||||
|
||||
const (
|
||||
// PolyDomainStandard: coefficient-form, values in [0, q).
|
||||
PolyDomainStandard PolyDomain = 0
|
||||
|
||||
// PolyDomainMontgomery: coefficient-form, Montgomery-encoded (a*R mod q).
|
||||
PolyDomainMontgomery PolyDomain = 1
|
||||
|
||||
// PolyDomainNTTStandard: NTT/evaluation-form, values in [0, q).
|
||||
PolyDomainNTTStandard PolyDomain = 2
|
||||
|
||||
// PolyDomainNTTMontgomery: NTT/evaluation-form, Montgomery-encoded.
|
||||
// This is the form produced by Ring.NTT in luxfi/lattice and consumed by
|
||||
// Montgomery-domain pointwise multiplication.
|
||||
PolyDomainNTTMontgomery PolyDomain = 3
|
||||
)
|
||||
|
||||
// String returns a stable, human-readable name for the domain.
|
||||
// The string values are part of the wire/log contract; do not rename.
|
||||
func (d PolyDomain) String() string {
|
||||
switch d {
|
||||
case PolyDomainStandard:
|
||||
return "Standard"
|
||||
case PolyDomainMontgomery:
|
||||
return "Montgomery"
|
||||
case PolyDomainNTTStandard:
|
||||
return "NTTStandard"
|
||||
case PolyDomainNTTMontgomery:
|
||||
return "NTTMontgomery"
|
||||
default:
|
||||
return "Unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// IsNTT reports whether the domain is in NTT/evaluation form.
|
||||
func (d PolyDomain) IsNTT() bool {
|
||||
return d == PolyDomainNTTStandard || d == PolyDomainNTTMontgomery
|
||||
}
|
||||
|
||||
// IsMontgomery reports whether the domain uses Montgomery encoding.
|
||||
func (d PolyDomain) IsMontgomery() bool {
|
||||
return d == PolyDomainMontgomery || d == PolyDomainNTTMontgomery
|
||||
}
|
||||
|
||||
// AfterForwardNTT returns the domain a Standard or Montgomery buffer enters
|
||||
// after a forward NTT. It is the type-system inverse of AfterInverseNTT.
|
||||
//
|
||||
// Standard -> NTTStandard
|
||||
// Montgomery -> NTTMontgomery
|
||||
// NTT* -> PolyDomain(0xFF) // already NTT, caller error
|
||||
func (d PolyDomain) AfterForwardNTT() PolyDomain {
|
||||
switch d {
|
||||
case PolyDomainStandard:
|
||||
return PolyDomainNTTStandard
|
||||
case PolyDomainMontgomery:
|
||||
return PolyDomainNTTMontgomery
|
||||
default:
|
||||
return PolyDomain(0xFF)
|
||||
}
|
||||
}
|
||||
|
||||
// AfterInverseNTT returns the domain a NTTStandard or NTTMontgomery buffer
|
||||
// enters after an inverse NTT.
|
||||
//
|
||||
// NTTStandard -> Standard
|
||||
// NTTMontgomery -> Montgomery
|
||||
// non-NTT -> PolyDomain(0xFF) // not in NTT form, caller error
|
||||
func (d PolyDomain) AfterInverseNTT() PolyDomain {
|
||||
switch d {
|
||||
case PolyDomainNTTStandard:
|
||||
return PolyDomainStandard
|
||||
case PolyDomainNTTMontgomery:
|
||||
return PolyDomainMontgomery
|
||||
default:
|
||||
return PolyDomain(0xFF)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,211 @@
|
||||
// Copyright (c) 2026, Lux Industries Inc.
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
package types
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"testing"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
// TestPolyDomain_StableValues asserts the on-the-wire byte values for each
|
||||
// PolyDomain variant. Changing these breaks ciphertext header digests and
|
||||
// kernel-tag dispatch tables. DO NOT change.
|
||||
func TestPolyDomain_StableValues(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
got PolyDomain
|
||||
want uint8
|
||||
}{
|
||||
{"Standard", PolyDomainStandard, 0},
|
||||
{"Montgomery", PolyDomainMontgomery, 1},
|
||||
{"NTTStandard", PolyDomainNTTStandard, 2},
|
||||
{"NTTMontgomery", PolyDomainNTTMontgomery, 3},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if uint8(c.got) != c.want {
|
||||
t.Errorf("%s = %d, want %d", c.name, c.got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestPolyDomain_StableStrings asserts the human-readable name returned by
|
||||
// String(). These strings are part of the log/wire contract.
|
||||
func TestPolyDomain_StableStrings(t *testing.T) {
|
||||
cases := []struct {
|
||||
d PolyDomain
|
||||
want string
|
||||
}{
|
||||
{PolyDomainStandard, "Standard"},
|
||||
{PolyDomainMontgomery, "Montgomery"},
|
||||
{PolyDomainNTTStandard, "NTTStandard"},
|
||||
{PolyDomainNTTMontgomery, "NTTMontgomery"},
|
||||
{PolyDomain(0xFF), "Unknown"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := c.d.String(); got != c.want {
|
||||
t.Errorf("PolyDomain(%d).String() = %q, want %q", c.d, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPolyDomain_NTTAndMontgomeryFlags(t *testing.T) {
|
||||
cases := []struct {
|
||||
d PolyDomain
|
||||
wantNTT bool
|
||||
wantMontgom bool
|
||||
}{
|
||||
{PolyDomainStandard, false, false},
|
||||
{PolyDomainMontgomery, false, true},
|
||||
{PolyDomainNTTStandard, true, false},
|
||||
{PolyDomainNTTMontgomery, true, true},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := c.d.IsNTT(); got != c.wantNTT {
|
||||
t.Errorf("%s.IsNTT() = %v, want %v", c.d, got, c.wantNTT)
|
||||
}
|
||||
if got := c.d.IsMontgomery(); got != c.wantMontgom {
|
||||
t.Errorf("%s.IsMontgomery() = %v, want %v", c.d, got, c.wantMontgom)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPolyDomain_ForwardInverseRoundTrip(t *testing.T) {
|
||||
cases := []PolyDomain{PolyDomainStandard, PolyDomainMontgomery}
|
||||
for _, in := range cases {
|
||||
ntt := in.AfterForwardNTT()
|
||||
if ntt == PolyDomain(0xFF) {
|
||||
t.Fatalf("AfterForwardNTT(%s) returned invalid", in)
|
||||
}
|
||||
back := ntt.AfterInverseNTT()
|
||||
if back != in {
|
||||
t.Errorf("round trip %s -> %s -> %s", in, ntt, back)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPolyDomain_ForwardOnAlreadyNTT(t *testing.T) {
|
||||
if got := PolyDomainNTTStandard.AfterForwardNTT(); got != PolyDomain(0xFF) {
|
||||
t.Errorf("forward on NTTStandard should be invalid, got %s", got)
|
||||
}
|
||||
if got := PolyDomainNTTMontgomery.AfterForwardNTT(); got != PolyDomain(0xFF) {
|
||||
t.Errorf("forward on NTTMontgomery should be invalid, got %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNTTContext_ValidatesModulusZero(t *testing.T) {
|
||||
ctx := &NTTContext{N: 16}
|
||||
if err := ctx.Validate(); !errors.Is(err, ErrModulusZero) {
|
||||
t.Errorf("modulus=0 -> err=%v, want ErrModulusZero", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNTTContext_ValidatesNPowerOfTwo(t *testing.T) {
|
||||
bad := []uint32{0, 1, 3, 5, 7, 100, 1000}
|
||||
for _, n := range bad {
|
||||
ctx := &NTTContext{Modulus: 0xFFFFFFFFFFFFFFC5, N: n}
|
||||
if err := ctx.Validate(); !errors.Is(err, ErrNNotPowerOfTwo) {
|
||||
t.Errorf("N=%d -> err=%v, want ErrNNotPowerOfTwo", n, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNTTContext_ValidatesRootDomainMatchesInput(t *testing.T) {
|
||||
ctx := &NTTContext{
|
||||
Modulus: 0xFFFFFFFFFFFFFFC5,
|
||||
N: 16,
|
||||
InputDomain: PolyDomainStandard,
|
||||
RootDomain: PolyDomainMontgomery, // mismatch
|
||||
OutputDomain: PolyDomainNTTStandard,
|
||||
}
|
||||
if err := ctx.Validate(); !errors.Is(err, ErrDomainMismatch) {
|
||||
t.Errorf("root/input mismatch -> err=%v, want ErrDomainMismatch", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNTTContext_ValidatesOutputDomainConsistency(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
input PolyDomain
|
||||
output PolyDomain
|
||||
ok bool
|
||||
}{
|
||||
{"std fwd", PolyDomainStandard, PolyDomainNTTStandard, true},
|
||||
{"std fwd wrong out", PolyDomainStandard, PolyDomainNTTMontgomery, false},
|
||||
{"mont fwd", PolyDomainMontgomery, PolyDomainNTTMontgomery, true},
|
||||
{"mont fwd wrong out", PolyDomainMontgomery, PolyDomainNTTStandard, false},
|
||||
{"std inv", PolyDomainNTTStandard, PolyDomainStandard, true},
|
||||
{"std inv wrong out", PolyDomainNTTStandard, PolyDomainMontgomery, false},
|
||||
{"mont inv", PolyDomainNTTMontgomery, PolyDomainMontgomery, true},
|
||||
{"mont inv wrong out", PolyDomainNTTMontgomery, PolyDomainStandard, false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
ctx := &NTTContext{
|
||||
Modulus: 0xFFFFFFFFFFFFFFC5,
|
||||
N: 16,
|
||||
InputDomain: c.input,
|
||||
RootDomain: c.input,
|
||||
OutputDomain: c.output,
|
||||
}
|
||||
err := ctx.Validate()
|
||||
if c.ok && err != nil {
|
||||
t.Errorf("%s: expected ok, got %v", c.name, err)
|
||||
}
|
||||
if !c.ok && !errors.Is(err, ErrDomainMismatch) {
|
||||
t.Errorf("%s: expected ErrDomainMismatch, got %v", c.name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNTTContext_ForwardInverseClassification(t *testing.T) {
|
||||
fwd := &NTTContext{
|
||||
Modulus: 0xFFFFFFFFFFFFFFC5, N: 16,
|
||||
InputDomain: PolyDomainMontgomery, RootDomain: PolyDomainMontgomery,
|
||||
OutputDomain: PolyDomainNTTMontgomery,
|
||||
}
|
||||
if !fwd.IsForward() || fwd.IsInverse() {
|
||||
t.Errorf("expected forward, got fwd=%v inv=%v", fwd.IsForward(), fwd.IsInverse())
|
||||
}
|
||||
inv := &NTTContext{
|
||||
Modulus: 0xFFFFFFFFFFFFFFC5, N: 16,
|
||||
InputDomain: PolyDomainNTTMontgomery, RootDomain: PolyDomainNTTMontgomery,
|
||||
OutputDomain: PolyDomainMontgomery,
|
||||
}
|
||||
if inv.IsForward() || !inv.IsInverse() {
|
||||
t.Errorf("expected inverse, got fwd=%v inv=%v", inv.IsForward(), inv.IsInverse())
|
||||
}
|
||||
}
|
||||
|
||||
// TestNTTContext_Layout pins the byte layout of NTTContext. Failing this
|
||||
// test means the Go and C++ sides will silently disagree on field
|
||||
// offsets via cgo. STOP and fix the layout, do not paper over.
|
||||
func TestNTTContext_Layout(t *testing.T) {
|
||||
const wantSize = 48
|
||||
if got := unsafe.Sizeof(NTTContext{}); got != wantSize {
|
||||
t.Fatalf("sizeof(NTTContext) = %d, want %d", got, wantSize)
|
||||
}
|
||||
var ctx NTTContext
|
||||
base := uintptr(unsafe.Pointer(&ctx))
|
||||
cases := []struct {
|
||||
name string
|
||||
off uintptr
|
||||
want uintptr
|
||||
}{
|
||||
{"Modulus", uintptr(unsafe.Pointer(&ctx.Modulus)) - base, 0},
|
||||
{"MontR", uintptr(unsafe.Pointer(&ctx.MontR)) - base, 8},
|
||||
{"MontR2", uintptr(unsafe.Pointer(&ctx.MontR2)) - base, 16},
|
||||
{"QInv", uintptr(unsafe.Pointer(&ctx.QInv)) - base, 24},
|
||||
{"N", uintptr(unsafe.Pointer(&ctx.N)) - base, 32},
|
||||
{"ModulusID", uintptr(unsafe.Pointer(&ctx.ModulusID)) - base, 36},
|
||||
{"TwiddleOffset", uintptr(unsafe.Pointer(&ctx.TwiddleOffset)) - base, 40},
|
||||
{"InputDomain", uintptr(unsafe.Pointer(&ctx.InputDomain)) - base, 44},
|
||||
{"RootDomain", uintptr(unsafe.Pointer(&ctx.RootDomain)) - base, 45},
|
||||
{"OutputDomain", uintptr(unsafe.Pointer(&ctx.OutputDomain)) - base, 46},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if c.off != c.want {
|
||||
t.Errorf("offset of %s = %d, want %d", c.name, c.off, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
// Copyright (c) 2026, Lux Industries Inc.
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
package types
|
||||
|
||||
import "errors"
|
||||
|
||||
// Sentinel errors for type-system rejections.
|
||||
//
|
||||
// These errors are returned by Validate() methods at the boundary between
|
||||
// caller and kernel dispatch. They MUST be checked with errors.Is, not by
|
||||
// string comparison.
|
||||
var (
|
||||
// ErrDomainMismatch is returned when a kernel is invoked with input,
|
||||
// root, and output domain tags that are not internally consistent.
|
||||
ErrDomainMismatch = errors.New("lattice/types: domain mismatch")
|
||||
|
||||
// ErrModulusZero is returned when an NTTContext or ReductionBudget is
|
||||
// constructed with modulus == 0.
|
||||
ErrModulusZero = errors.New("lattice/types: modulus must be non-zero")
|
||||
|
||||
// ErrNNotPowerOfTwo is returned when an NTTContext is constructed with
|
||||
// a polynomial degree N that is not a power of two.
|
||||
ErrNNotPowerOfTwo = errors.New("lattice/types: N must be a power of two")
|
||||
|
||||
// ErrModulusTooLargeForLazy is returned when ReductionBudget is asked
|
||||
// for a Lazy mode that the modulus does not fit (e.g. Lazy8 requires
|
||||
// modulus < 2^29).
|
||||
ErrModulusTooLargeForLazy = errors.New("lattice/types: modulus too large for requested lazy mode")
|
||||
)
|
||||
@@ -0,0 +1,135 @@
|
||||
// Copyright (c) 2026, Lux Industries Inc.
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
package types
|
||||
|
||||
import "math/bits"
|
||||
|
||||
// NTTContext is the typed handle that every NTT/arithmetic kernel must accept
|
||||
// in place of an opaque pointer + a pile of scalars.
|
||||
//
|
||||
// Layout is byte-stable across Go and C++ for cgo / shared-memory dispatch.
|
||||
// See lux/lattice/types/domain_layout_test.go and
|
||||
// luxcpp/lattice/include/lux/lattice/types/domain.hpp for the matching
|
||||
// static_assert. The struct has no padding (all fields fall on natural
|
||||
// alignment for an 8-byte aligned struct).
|
||||
//
|
||||
// Field ordering MUST NOT change without updating the C++ mirror and the
|
||||
// layout tests on both sides.
|
||||
type NTTContext struct {
|
||||
// Modulus is the prime q used by this NTT context.
|
||||
// 8 bytes, offset 0.
|
||||
Modulus uint64
|
||||
|
||||
// MontR is R mod q where R = 2^64.
|
||||
// 8 bytes, offset 8.
|
||||
MontR uint64
|
||||
|
||||
// MontR2 is R^2 mod q. Used to convert to/from Montgomery form
|
||||
// without a 128-bit divide.
|
||||
// 8 bytes, offset 16.
|
||||
MontR2 uint64
|
||||
|
||||
// QInv is q^{-1} mod 2^64 (the constant required for Montgomery
|
||||
// reduction MRed). The lattice/v7 ring package calls this
|
||||
// "mredconstant".
|
||||
// 8 bytes, offset 24.
|
||||
QInv uint64
|
||||
|
||||
// N is the polynomial ring degree (must be a power of two).
|
||||
// 4 bytes, offset 32.
|
||||
N uint32
|
||||
|
||||
// ModulusID is an opaque identifier into a modulus chain (RNS basis).
|
||||
// Kernels use it to look up twiddle tables, BRedConstants, etc.
|
||||
// 4 bytes, offset 36.
|
||||
ModulusID uint32
|
||||
|
||||
// TwiddleOffset is the offset into a flat twiddle-factor table where
|
||||
// this context's roots begin. Allows multiple contexts to share a
|
||||
// single contiguous twiddle buffer on GPU.
|
||||
// 4 bytes, offset 40.
|
||||
TwiddleOffset uint32
|
||||
|
||||
// InputDomain is the domain the kernel expects on its input buffer.
|
||||
// 1 byte, offset 44.
|
||||
InputDomain PolyDomain
|
||||
|
||||
// RootDomain is the domain the twiddle (root) table lives in.
|
||||
// 1 byte, offset 45.
|
||||
RootDomain PolyDomain
|
||||
|
||||
// OutputDomain is the domain the kernel will leave the output buffer
|
||||
// in after dispatch.
|
||||
// 1 byte, offset 46.
|
||||
OutputDomain PolyDomain
|
||||
|
||||
// _pad ensures total size is 48 bytes (multiple of 8).
|
||||
// 1 byte, offset 47.
|
||||
_pad uint8
|
||||
}
|
||||
|
||||
// Validate enforces type-system invariants at the kernel boundary.
|
||||
//
|
||||
// A valid NTTContext satisfies:
|
||||
// 1. Modulus != 0
|
||||
// 2. N is a power of two and N >= 2
|
||||
// 3. RootDomain matches InputDomain (twiddles and operand share encoding)
|
||||
// 4. OutputDomain is consistent with the input/root pair under the
|
||||
// forward-NTT transition rule (Montgomery in -> NTTMontgomery out;
|
||||
// Standard in -> NTTStandard out; NTT in -> coefficient out)
|
||||
//
|
||||
// Returns nil on success, or an ErrDomainMismatch / ErrModulusZero /
|
||||
// ErrNNotPowerOfTwo error otherwise. Callers MUST check this before
|
||||
// dispatching to a GPU kernel.
|
||||
func (ctx *NTTContext) Validate() error {
|
||||
if ctx.Modulus == 0 {
|
||||
return ErrModulusZero
|
||||
}
|
||||
if ctx.N < 2 || bits.OnesCount32(ctx.N) != 1 {
|
||||
return ErrNNotPowerOfTwo
|
||||
}
|
||||
|
||||
// Twiddles must live in the same encoding as the input. A standard-form
|
||||
// input multiplied against Montgomery-form roots produces silent
|
||||
// corruption; this is the #121 class of bug.
|
||||
if ctx.RootDomain != ctx.InputDomain {
|
||||
return ErrDomainMismatch
|
||||
}
|
||||
|
||||
// Output must be consistent with the direction implied by Input.
|
||||
switch ctx.InputDomain {
|
||||
case PolyDomainStandard:
|
||||
if ctx.OutputDomain != PolyDomainNTTStandard {
|
||||
return ErrDomainMismatch
|
||||
}
|
||||
case PolyDomainMontgomery:
|
||||
if ctx.OutputDomain != PolyDomainNTTMontgomery {
|
||||
return ErrDomainMismatch
|
||||
}
|
||||
case PolyDomainNTTStandard:
|
||||
if ctx.OutputDomain != PolyDomainStandard {
|
||||
return ErrDomainMismatch
|
||||
}
|
||||
case PolyDomainNTTMontgomery:
|
||||
if ctx.OutputDomain != PolyDomainMontgomery {
|
||||
return ErrDomainMismatch
|
||||
}
|
||||
default:
|
||||
return ErrDomainMismatch
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// IsForward reports whether this context describes a forward NTT
|
||||
// (coefficient -> evaluation form).
|
||||
func (ctx *NTTContext) IsForward() bool {
|
||||
return !ctx.InputDomain.IsNTT() && ctx.OutputDomain.IsNTT()
|
||||
}
|
||||
|
||||
// IsInverse reports whether this context describes an inverse NTT
|
||||
// (evaluation -> coefficient form).
|
||||
func (ctx *NTTContext) IsInverse() bool {
|
||||
return ctx.InputDomain.IsNTT() && !ctx.OutputDomain.IsNTT()
|
||||
}
|
||||
@@ -0,0 +1,175 @@
|
||||
// Copyright (c) 2026, Lux Industries Inc.
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
package types
|
||||
|
||||
import "math/bits"
|
||||
|
||||
// ReductionMode selects how aggressively a kernel chain may defer modular
|
||||
// reduction. Lazy reduction is a correctness-sensitive performance
|
||||
// optimisation: it leaves intermediate values in [0, k*q) for some
|
||||
// small k > 1, and skips the conditional subtract that would normalise
|
||||
// each result back to [0, q).
|
||||
//
|
||||
// Reference: LP-137-FHE-TYPING.md §5 (lazy reduction modes).
|
||||
type ReductionMode uint8
|
||||
|
||||
const (
|
||||
// ReductionStrictEveryOp normalises after every modular operation.
|
||||
// Result of each op is in [0, q). Slowest, safest.
|
||||
ReductionStrictEveryOp ReductionMode = 0
|
||||
|
||||
// ReductionLazy2 allows result range [0, 2q). 1 deferred subtract
|
||||
// allowed before normalisation. Requires q < 2^63.
|
||||
ReductionLazy2 ReductionMode = 1
|
||||
|
||||
// ReductionLazy4 allows result range [0, 4q). Up to 2 deferred adds
|
||||
// in [0, 2q) chained. Requires q < 2^62.
|
||||
ReductionLazy4 ReductionMode = 2
|
||||
|
||||
// ReductionLazy8 allows result range [0, 8q). Up to 3 deferred adds.
|
||||
// Requires q < 2^61.
|
||||
ReductionLazy8 ReductionMode = 3
|
||||
)
|
||||
|
||||
// String returns the stable name for the reduction mode.
|
||||
func (m ReductionMode) String() string {
|
||||
switch m {
|
||||
case ReductionStrictEveryOp:
|
||||
return "Strict"
|
||||
case ReductionLazy2:
|
||||
return "Lazy2"
|
||||
case ReductionLazy4:
|
||||
return "Lazy4"
|
||||
case ReductionLazy8:
|
||||
return "Lazy8"
|
||||
default:
|
||||
return "Unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// ReductionBudget tracks how many lazy ops a kernel chain has consumed and
|
||||
// when it must insert a normalisation step.
|
||||
//
|
||||
// Layout is byte-stable across Go and C++.
|
||||
type ReductionBudget struct {
|
||||
// Modulus is q. Used to recompute MaxOpsBeforeOverflow if Mode changes.
|
||||
// 8 bytes, offset 0.
|
||||
Modulus uint64
|
||||
|
||||
// OpsSinceReduce is the number of lazy ops accumulated since the last
|
||||
// explicit reduction.
|
||||
// 4 bytes, offset 8.
|
||||
OpsSinceReduce uint32
|
||||
|
||||
// MaxOpsBeforeOverflow is the cap derived from (Mode, Modulus) by
|
||||
// SafeBoundFor. The kernel must insert a reduction once
|
||||
// OpsSinceReduce reaches this value.
|
||||
// 4 bytes, offset 12.
|
||||
MaxOpsBeforeOverflow uint32
|
||||
|
||||
// Mode is the lazy mode in effect.
|
||||
// 1 byte, offset 16.
|
||||
Mode ReductionMode
|
||||
|
||||
// _pad pads the struct to 24 bytes (multiple of 8).
|
||||
// 7 bytes, offset 17.
|
||||
_pad [7]uint8
|
||||
}
|
||||
|
||||
// NewReductionBudget constructs a budget for (mode, modulus). Returns
|
||||
// ErrModulusZero if modulus == 0, or ErrModulusTooLargeForLazy if the
|
||||
// modulus exceeds the lazy mode's safe bound.
|
||||
func NewReductionBudget(mode ReductionMode, modulus uint64) (ReductionBudget, error) {
|
||||
if modulus == 0 {
|
||||
return ReductionBudget{}, ErrModulusZero
|
||||
}
|
||||
if !lazyModeFits(mode, modulus) {
|
||||
return ReductionBudget{}, ErrModulusTooLargeForLazy
|
||||
}
|
||||
return ReductionBudget{
|
||||
Modulus: modulus,
|
||||
OpsSinceReduce: 0,
|
||||
MaxOpsBeforeOverflow: SafeBoundFor(mode, modulus),
|
||||
Mode: mode,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// RemainingOps returns the number of additional lazy ops the kernel may
|
||||
// perform before NeedsReduce returns true. Saturating at 0.
|
||||
func (rb *ReductionBudget) RemainingOps() uint32 {
|
||||
if rb.OpsSinceReduce >= rb.MaxOpsBeforeOverflow {
|
||||
return 0
|
||||
}
|
||||
return rb.MaxOpsBeforeOverflow - rb.OpsSinceReduce
|
||||
}
|
||||
|
||||
// NeedsReduce reports whether the next op must be preceded by a reduction.
|
||||
func (rb *ReductionBudget) NeedsReduce() bool {
|
||||
return rb.OpsSinceReduce >= rb.MaxOpsBeforeOverflow
|
||||
}
|
||||
|
||||
// Charge increments OpsSinceReduce by n. Caller invokes this each time
|
||||
// it performs n lazy ops.
|
||||
func (rb *ReductionBudget) Charge(n uint32) {
|
||||
rb.OpsSinceReduce += n
|
||||
}
|
||||
|
||||
// Reset zeroes OpsSinceReduce. Caller invokes this immediately after a
|
||||
// normalisation pass.
|
||||
func (rb *ReductionBudget) Reset() {
|
||||
rb.OpsSinceReduce = 0
|
||||
}
|
||||
|
||||
// SafeBoundFor returns the maximum number of lazy ops permissible before
|
||||
// normalisation, given the lazy mode and modulus.
|
||||
//
|
||||
// The bounds correspond to the bit budget of a 64-bit limb:
|
||||
//
|
||||
// Strict : 1 (always reduce)
|
||||
// Lazy2 : 1 mult / 2 adds for q < 2^63 (range [0, 2q))
|
||||
// Lazy4 : 2 mults for q < 2^62 (range [0, 4q))
|
||||
// Lazy8 : 3 mults for q < 2^61 (range [0, 8q))
|
||||
//
|
||||
// For moduli that do not fit the requested mode (q >= 2^k for the relevant
|
||||
// k), SafeBoundFor returns 0 to force the caller down the strict path.
|
||||
func SafeBoundFor(mode ReductionMode, modulus uint64) uint32 {
|
||||
if modulus == 0 {
|
||||
return 0
|
||||
}
|
||||
if !lazyModeFits(mode, modulus) {
|
||||
return 0
|
||||
}
|
||||
switch mode {
|
||||
case ReductionStrictEveryOp:
|
||||
return 1
|
||||
case ReductionLazy2:
|
||||
return 2
|
||||
case ReductionLazy4:
|
||||
return 4
|
||||
case ReductionLazy8:
|
||||
return 8
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
// lazyModeFits reports whether the modulus is small enough for the lazy
|
||||
// mode without uint64 overflow. The bit thresholds are intentionally
|
||||
// conservative so that intermediate sums of k operands each in [0, k*q)
|
||||
// stay strictly below 2^64.
|
||||
func lazyModeFits(mode ReductionMode, modulus uint64) bool {
|
||||
bitlen := uint(bits.Len64(modulus))
|
||||
switch mode {
|
||||
case ReductionStrictEveryOp:
|
||||
return bitlen <= 64
|
||||
case ReductionLazy2:
|
||||
return bitlen <= 63
|
||||
case ReductionLazy4:
|
||||
return bitlen <= 62
|
||||
case ReductionLazy8:
|
||||
return bitlen <= 61
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,165 @@
|
||||
// Copyright (c) 2026, Lux Industries Inc.
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
package types
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"testing"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
func TestReductionMode_String(t *testing.T) {
|
||||
cases := []struct {
|
||||
m ReductionMode
|
||||
want string
|
||||
}{
|
||||
{ReductionStrictEveryOp, "Strict"},
|
||||
{ReductionLazy2, "Lazy2"},
|
||||
{ReductionLazy4, "Lazy4"},
|
||||
{ReductionLazy8, "Lazy8"},
|
||||
{ReductionMode(0xFF), "Unknown"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := c.m.String(); got != c.want {
|
||||
t.Errorf("ReductionMode(%d) = %q, want %q", c.m, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSafeBoundFor_StrictAlwaysOne(t *testing.T) {
|
||||
moduli := []uint64{
|
||||
0x10001, // 17 bits
|
||||
0x7FFFFFFFFFFFFFC5, // 63 bits, prime
|
||||
0xFFFFFFFFFFFFFFC5, // 64 bits, prime
|
||||
}
|
||||
for _, q := range moduli {
|
||||
if got := SafeBoundFor(ReductionStrictEveryOp, q); got != 1 {
|
||||
t.Errorf("SafeBoundFor(Strict, q=0x%x) = %d, want 1", q, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSafeBoundFor_LazyByModulusBitlen(t *testing.T) {
|
||||
cases := []struct {
|
||||
mode ReductionMode
|
||||
modulus uint64
|
||||
want uint32
|
||||
}{
|
||||
// 30-bit modulus fits all lazy modes
|
||||
{ReductionLazy2, (1 << 30) - 35, 2},
|
||||
{ReductionLazy4, (1 << 30) - 35, 4},
|
||||
{ReductionLazy8, (1 << 30) - 35, 8},
|
||||
|
||||
// 61-bit modulus is the upper bound for Lazy8
|
||||
{ReductionLazy8, (1 << 61) - 1, 8},
|
||||
|
||||
// 62-bit modulus exceeds Lazy8 budget
|
||||
{ReductionLazy8, (1 << 62) - 1, 0},
|
||||
|
||||
// 62-bit fits Lazy4
|
||||
{ReductionLazy4, (1 << 62) - 1, 4},
|
||||
|
||||
// 63-bit exceeds Lazy4
|
||||
{ReductionLazy4, (1 << 63) - 1, 0},
|
||||
|
||||
// 63-bit fits Lazy2
|
||||
{ReductionLazy2, (1 << 63) - 1, 2},
|
||||
|
||||
// 64-bit exceeds Lazy2
|
||||
{ReductionLazy2, ^uint64(0), 0},
|
||||
|
||||
// modulus 0 always returns 0
|
||||
{ReductionLazy4, 0, 0},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := SafeBoundFor(c.mode, c.modulus); got != c.want {
|
||||
t.Errorf("SafeBoundFor(%s, q=0x%x) = %d, want %d",
|
||||
c.mode, c.modulus, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewReductionBudget_RejectsZeroModulus(t *testing.T) {
|
||||
if _, err := NewReductionBudget(ReductionLazy4, 0); !errors.Is(err, ErrModulusZero) {
|
||||
t.Errorf("expected ErrModulusZero, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewReductionBudget_RejectsTooLargeModulus(t *testing.T) {
|
||||
// 64-bit modulus, request Lazy2 (caps at 63-bit) => reject.
|
||||
_, err := NewReductionBudget(ReductionLazy2, ^uint64(0))
|
||||
if !errors.Is(err, ErrModulusTooLargeForLazy) {
|
||||
t.Errorf("expected ErrModulusTooLargeForLazy, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReductionBudget_ChargeAndReset(t *testing.T) {
|
||||
rb, err := NewReductionBudget(ReductionLazy4, (1 << 30) - 35)
|
||||
if err != nil {
|
||||
t.Fatalf("NewReductionBudget: %v", err)
|
||||
}
|
||||
if rb.NeedsReduce() {
|
||||
t.Errorf("fresh budget should not require reduction")
|
||||
}
|
||||
if got := rb.RemainingOps(); got != 4 {
|
||||
t.Errorf("RemainingOps = %d, want 4", got)
|
||||
}
|
||||
rb.Charge(2)
|
||||
if got := rb.RemainingOps(); got != 2 {
|
||||
t.Errorf("after charge 2, RemainingOps = %d, want 2", got)
|
||||
}
|
||||
rb.Charge(2)
|
||||
if !rb.NeedsReduce() {
|
||||
t.Errorf("after charge 4 with cap 4, should need reduce")
|
||||
}
|
||||
if got := rb.RemainingOps(); got != 0 {
|
||||
t.Errorf("RemainingOps saturated should be 0, got %d", got)
|
||||
}
|
||||
rb.Reset()
|
||||
if rb.NeedsReduce() {
|
||||
t.Errorf("after reset, should not need reduce")
|
||||
}
|
||||
if got := rb.RemainingOps(); got != 4 {
|
||||
t.Errorf("after reset, RemainingOps = %d, want 4", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReductionBudget_RemainingOpsSaturates(t *testing.T) {
|
||||
rb, err := NewReductionBudget(ReductionLazy2, 1<<30)
|
||||
if err != nil {
|
||||
t.Fatalf("NewReductionBudget: %v", err)
|
||||
}
|
||||
rb.Charge(1000) // way over cap
|
||||
if got := rb.RemainingOps(); got != 0 {
|
||||
t.Errorf("over-charged budget RemainingOps = %d, want 0", got)
|
||||
}
|
||||
if !rb.NeedsReduce() {
|
||||
t.Errorf("over-charged budget should need reduce")
|
||||
}
|
||||
}
|
||||
|
||||
// TestReductionBudget_Layout pins layout for cgo compatibility.
|
||||
func TestReductionBudget_Layout(t *testing.T) {
|
||||
const wantSize = 24
|
||||
if got := unsafe.Sizeof(ReductionBudget{}); got != wantSize {
|
||||
t.Fatalf("sizeof(ReductionBudget) = %d, want %d", got, wantSize)
|
||||
}
|
||||
var rb ReductionBudget
|
||||
base := uintptr(unsafe.Pointer(&rb))
|
||||
cases := []struct {
|
||||
name string
|
||||
off uintptr
|
||||
want uintptr
|
||||
}{
|
||||
{"Modulus", uintptr(unsafe.Pointer(&rb.Modulus)) - base, 0},
|
||||
{"OpsSinceReduce", uintptr(unsafe.Pointer(&rb.OpsSinceReduce)) - base, 8},
|
||||
{"MaxOpsBeforeOverflow", uintptr(unsafe.Pointer(&rb.MaxOpsBeforeOverflow)) - base, 12},
|
||||
{"Mode", uintptr(unsafe.Pointer(&rb.Mode)) - base, 16},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if c.off != c.want {
|
||||
t.Errorf("offset of %s = %d, want %d", c.name, c.off, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user