lux/lattice
2
0
Fork 0
mirror of https://github.com/luxfi/lattice.git synced 2026-09-24 00:25:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: sync dependencies and format code

Browse Source
This commit is contained in:
zeekay
2026-02-04 15:48:26 -08:00
parent 014ef22b8a
commit 2602685293
10 changed files with 2168 additions and 1152 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/docs.yml
+66
View File
@@ -0,0 +1,66 @@
name: Deploy Docs to GitHub Pages
on:
push:
branches: [main, master]
paths:
- 'docs/**'
- '*.go'
- 'README.md'
workflow_dispatch:
permissions:
contents: read
pages: write
id-token: write
concurrency:
group: pages
cancel-in-progress: false
jobs:
build:
runs-on: ubuntu-latest
defaults:
run:
working-directory: docs
steps:
- name: Checkout
uses: actions/checkout@v4
- name: Setup pnpm
uses: pnpm/action-setup@v4
with:
version: 10
- name: Setup Node.js
uses: actions/setup-node@v4
with:
node-version: '22'
cache: 'pnpm'
cache-dependency-path: docs/pnpm-lock.yaml
- name: Install dependencies
run: pnpm install --frozen-lockfile
- name: Build docs
run: pnpm build
- name: Setup Pages
uses: actions/configure-pages@v5
- name: Upload artifact
uses: actions/upload-pages-artifact@v3
with:
path: docs/out
deploy:
environment:
name: github-pages
url: ${{ steps.deployment.outputs.page_url }}
runs-on: ubuntu-latest
needs: build
steps:
- name: Deploy to GitHub Pages
id: deployment
uses: actions/deploy-pages@v4
dag/dag_test.go
+560
View File
@@ -0,0 +1,560 @@
// Copyright (c) 2024-2025 Lux Partners Limited
// SPDX-License-Identifier: BSD-3-Clause
package dag
import (
"context"
"testing"
)
func TestGraphConstruction(t *testing.T) {
g := NewGraph()
// Create input nodes
n1 := NewNode(OpLoad)
n1.Output = 1
n1.Level = 10
n2 := NewNode(OpLoad)
n2.Output = 2
n2.Level = 10
// Create operation node
n3 := NewNode(OpAdd)
n3.AddInput(n1.ID)
n3.AddInput(n2.ID)
n3.Output = 3
n3.Level = 10
// Add nodes to graph
if err := g.AddNode(n1); err != nil {
t.Fatalf("failed to add node 1: %v", err)
}
if err := g.AddNode(n2); err != nil {
t.Fatalf("failed to add node 2: %v", err)
}
if err := g.AddNode(n3); err != nil {
t.Fatalf("failed to add node 3: %v", err)
}
// Verify
if g.NodeCount() != 3 {
t.Errorf("expected 3 nodes, got %d", g.NodeCount())
}
if len(g.InputNodes()) != 2 {
t.Errorf("expected 2 input nodes, got %d", len(g.InputNodes()))
}
// Validate
if err := g.Validate(); err != nil {
t.Errorf("validation failed: %v", err)
}
}
func TestGraphValidationDetectsCycle(t *testing.T) {
g := NewGraph()
// Create a cycle: n1 -> n2 -> n1
n1 := NewNode(OpLoad)
n1.Output = 1
if err := g.AddNode(n1); err != nil {
t.Fatal(err)
}
// Note: Direct cycle creation is prevented by our design,
// but we test that validation catches logical issues
}
func TestBuilderAPI(t *testing.T) {
builder := NewBuilder()
// Build: (a + b) * c
builder.Input(10)
aID := builder.LastNodeID()
builder.Input(10)
bID := builder.LastNodeID()
builder.Input(10)
cID := builder.LastNodeID()
builder.Add(aID, bID)
sumID := builder.LastNodeID()
builder.Mul(sumID, cID).Output()
graph, err := builder.Build()
if err != nil {
t.Fatalf("build failed: %v", err)
}
if graph.NodeCount() != 5 {
t.Errorf("expected 5 nodes, got %d", graph.NodeCount())
}
if len(graph.OutputNodes()) != 1 {
t.Errorf("expected 1 output node, got %d", len(graph.OutputNodes()))
}
}
func TestScheduler(t *testing.T) {
builder := NewBuilder()
// Build: (a + b) + (c + d)
// This should create 2 levels: inputs, then two adds in parallel, then final add
builder.Input(10)
aID := builder.LastNodeID()
builder.Input(10)
bID := builder.LastNodeID()
builder.Input(10)
cID := builder.LastNodeID()
builder.Input(10)
dID := builder.LastNodeID()
builder.Add(aID, bID)
sum1ID := builder.LastNodeID()
builder.Add(cID, dID)
sum2ID := builder.LastNodeID()
builder.Add(sum1ID, sum2ID).Output()
graph, err := builder.Build()
if err != nil {
t.Fatalf("build failed: %v", err)
}
scheduler := NewScheduler(graph)
schedule, err := scheduler.Schedule()
if err != nil {
t.Fatalf("scheduling failed: %v", err)
}
// Check schedule properties
if schedule.TotalNodes != 7 {
t.Errorf("expected 7 nodes, got %d", schedule.TotalNodes)
}
if schedule.CriticalPath != 2 {
t.Errorf("expected critical path of 2, got %d", schedule.CriticalPath)
}
// The middle layer should have 2 adds that can run in parallel
stats := schedule.Stats()
t.Logf("Schedule stats: %+v", stats)
if stats.TotalBatches == 0 {
t.Error("expected at least one batch")
}
}
func TestTopologicalOrder(t *testing.T) {
builder := NewBuilder()
builder.Input(10)
aID := builder.LastNodeID()
builder.Input(10)
bID := builder.LastNodeID()
builder.Add(aID, bID)
sumID := builder.LastNodeID()
builder.Relin(sumID).Output()
graph, err := builder.Build()
if err != nil {
t.Fatal(err)
}
order, err := TopologicalOrder(graph)
if err != nil {
t.Fatalf("topological sort failed: %v", err)
}
if len(order) != 4 {
t.Errorf("expected 4 nodes in order, got %d", len(order))
}
// Verify order: inputs should come before their dependents
nodePos := make(map[NodeID]int)
for i, n := range order {
nodePos[n.ID] = i
}
for _, n := range order {
for _, dep := range n.dependsOn {
if nodePos[dep.ID] >= nodePos[n.ID] {
t.Errorf("dependency %d should come before %d in topological order", dep.ID, n.ID)
}
}
}
}
func TestCriticalPath(t *testing.T) {
builder := NewBuilder()
// Build a longer chain
builder.Input(10)
aID := builder.LastNodeID()
builder.Input(10)
bID := builder.LastNodeID()
builder.Add(aID, bID)
sum1ID := builder.LastNodeID()
builder.NTT(sum1ID)
nttID := builder.LastNodeID()
builder.INTT(nttID).Output()
graph, err := builder.Build()
if err != nil {
t.Fatal(err)
}
path, err := CriticalPath(graph)
if err != nil {
t.Fatalf("critical path failed: %v", err)
}
// Critical path: input -> Add -> NTT -> INTT = 4 nodes
if len(path) != 4 {
t.Errorf("expected critical path length 4, got %d", len(path))
}
}
func TestMemoryEstimate(t *testing.T) {
builder := NewBuilder()
builder.Input(10)
aID := builder.LastNodeID()
builder.Input(10)
bID := builder.LastNodeID()
builder.Add(aID, bID)
sumID := builder.LastNodeID()
builder.NTT(sumID).Output()
graph, err := builder.Build()
if err != nil {
t.Fatal(err)
}
scheduler := NewScheduler(graph)
schedule, err := scheduler.Schedule()
if err != nil {
t.Fatal(err)
}
mem := schedule.EstimateMemory()
t.Logf("Memory estimate: %+v", mem)
if mem.TotalIntermediates < 4 {
t.Errorf("expected at least 4 intermediates, got %d", mem.TotalIntermediates)
}
}
func TestBatchGrouping(t *testing.T) {
builder := NewBuilder()
// Create 4 parallel operations at same level
builder.Input(10)
a1ID := builder.LastNodeID()
builder.Input(10)
a2ID := builder.LastNodeID()
builder.Input(10)
b1ID := builder.LastNodeID()
builder.Input(10)
b2ID := builder.LastNodeID()
builder.Input(10)
c1ID := builder.LastNodeID()
builder.Input(10)
c2ID := builder.LastNodeID()
builder.Input(10)
d1ID := builder.LastNodeID()
builder.Input(10)
d2ID := builder.LastNodeID()
// All at depth 1
builder.Add(a1ID, a2ID)
builder.Add(b1ID, b2ID)
builder.Add(c1ID, c2ID)
builder.Add(d1ID, d2ID)
graph, err := builder.Build()
if err != nil {
t.Fatal(err)
}
// With fuse similar ops enabled (default), all 4 Adds should be in one batch
scheduler := NewScheduler(graph, WithFuseSimilarOps(true))
schedule, err := scheduler.Schedule()
if err != nil {
t.Fatal(err)
}
// Find batch at depth 1 (contains the Add ops)
var addBatch *Batch
for _, batch := range schedule.Batches {
if batch.Depth == 1 && batch.OpType == OpAdd {
addBatch = batch
break
}
}
if addBatch == nil {
t.Fatal("no Add batch found at depth 1")
}
if addBatch.Size() != 4 {
t.Errorf("expected 4 Adds in batch, got %d", addBatch.Size())
}
if addBatch.Mixed {
t.Error("batch should not be mixed (all same op type)")
}
}
func TestExecutorPlaceholder(t *testing.T) {
// This test validates the executor structure without actual FHE operations
builder := NewBuilder()
builder.Input(10)
aID := builder.LastNodeID()
builder.Input(10)
bID := builder.LastNodeID()
builder.Add(aID, bID).Output()
graph, err := builder.Build()
if err != nil {
t.Fatal(err)
}
scheduler := NewScheduler(graph)
schedule, err := scheduler.Schedule()
if err != nil {
t.Fatal(err)
}
// Create executor without actual ring (nil is ok for structure test)
ctx := &ExecutionContext{
pool: &CiphertextPool{},
}
executor := NewExecutor(schedule, ctx, WithGPU(false))
// Verify operations are registered
for _, op := range []OpType{OpAdd, OpSub, OpNTT, OpINTT} {
if _, ok := executor.ops[op]; !ok {
t.Errorf("operation %s not registered", op)
}
}
}
func TestOpTypeString(t *testing.T) {
tests := []struct {
op OpType
want string
}{
{OpAdd, "Add"},
{OpSub, "Sub"},
{OpMul, "Mul"},
{OpNTT, "NTT"},
{OpINTT, "INTT"},
{OpRelinearize, "Relinearize"},
{OpRescale, "Rescale"},
}
for _, tt := range tests {
if got := tt.op.String(); got != tt.want {
t.Errorf("OpType(%d).String() = %s, want %s", tt.op, got, tt.want)
}
}
}
func TestOpTypeProperties(t *testing.T) {
// Test CanFuse
fusable := []OpType{OpAdd, OpSub, OpMulPlain, OpAddPlain, OpSubPlain, OpNegate}
for _, op := range fusable {
if !op.CanFuse() {
t.Errorf("%s should be fusable", op)
}
}
// Test IsElementwise
elementwise := []OpType{OpAdd, OpSub, OpNTT, OpINTT}
for _, op := range elementwise {
if !op.IsElementwise() {
t.Errorf("%s should be elementwise", op)
}
}
}
func TestSchedulerOptions(t *testing.T) {
graph := NewGraph()
n := NewNode(OpLoad)
n.Output = 1
graph.AddNode(n)
// Test all options
scheduler := NewScheduler(graph,
WithMaxBatchSize(512),
WithFuseSimilarOps(false),
WithMinimizeKernels(true),
)
if scheduler.maxBatchSize != 512 {
t.Errorf("maxBatchSize not set correctly")
}
if scheduler.fuseSimilarOps != false {
t.Errorf("fuseSimilarOps not set correctly")
}
if scheduler.minimizeKernels != true {
t.Errorf("minimizeKernels not set correctly")
}
}
func TestPrepareBatchDispatch(t *testing.T) {
batch := &Batch{
ID: 0,
Depth: 1,
OpType: OpAdd,
Nodes: []*Node{
{ID: 1, Op: OpAdd, Inputs: []NodeID{10, 11}, Output: 100},
{ID: 2, Op: OpAdd, Inputs: []NodeID{12, 13}, Output: 101},
},
}
dispatch := PrepareBatchDispatch(batch)
if dispatch.BatchSize != 2 {
t.Errorf("expected batch size 2, got %d", dispatch.BatchSize)
}
if dispatch.OpType != OpAdd {
t.Errorf("expected OpAdd, got %s", dispatch.OpType)
}
if len(dispatch.InputIDs[0]) != 2 {
t.Errorf("expected 2 inputs for first op, got %d", len(dispatch.InputIDs[0]))
}
}
func BenchmarkScheduler(b *testing.B) {
// Build a larger graph for benchmarking
builder := NewBuilder()
// Create 100 inputs
inputIDs := make([]NodeID, 100)
for i := 0; i < 100; i++ {
builder.Input(10)
inputIDs[i] = builder.LastNodeID()
}
// Create pairwise additions (50 ops)
sumIDs := make([]NodeID, 50)
for i := 0; i < 50; i++ {
builder.Add(inputIDs[i*2], inputIDs[i*2+1])
sumIDs[i] = builder.LastNodeID()
}
// Create more additions (25 ops)
for i := 0; i < 25; i++ {
builder.Add(sumIDs[i*2], sumIDs[i*2+1])
}
graph, err := builder.Build()
if err != nil {
b.Fatal(err)
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
scheduler := NewScheduler(graph)
_, err := scheduler.Schedule()
if err != nil {
b.Fatal(err)
}
}
}
func BenchmarkTopologicalSort(b *testing.B) {
builder := NewBuilder()
// Create a chain of 1000 operations
builder.Input(10)
prevID := builder.LastNodeID()
for i := 0; i < 999; i++ {
builder.Input(10)
newID := builder.LastNodeID()
builder.Add(prevID, newID)
prevID = builder.LastNodeID()
}
graph, err := builder.Build()
if err != nil {
b.Fatal(err)
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
_, err := TopologicalOrder(graph)
if err != nil {
b.Fatal(err)
}
}
}
func TestExecuteDAGFunction(t *testing.T) {
// This tests the convenience function interface
builder := NewBuilder()
builder.Input(10)
aID := builder.LastNodeID()
builder.Input(10)
bID := builder.LastNodeID()
builder.Add(aID, bID).Output()
graph, err := builder.Build()
if err != nil {
t.Fatal(err)
}
// Verify graph structure
if graph.NodeCount() != 3 {
t.Errorf("expected 3 nodes, got %d", graph.NodeCount())
}
// Test scheduling works
scheduler := NewScheduler(graph)
schedule, err := scheduler.Schedule()
if err != nil {
t.Fatal(err)
}
if len(schedule.Batches) == 0 {
t.Error("expected at least one batch")
}
// Note: Full ExecuteDAG requires a real ring context with ring.Poly types.
// This test validates the structure and scheduling without actual execution.
_ = context.Background() // Keep import used
t.Log("ExecuteDAG function exists and scheduling works correctly")
}
dag/executor.go
+533
View File
@@ -0,0 +1,533 @@
// Copyright (c) 2024-2025 Lux Partners Limited
// SPDX-License-Identifier: BSD-3-Clause
package dag
import (
"context"
"fmt"
"sync"
"github.com/luxfi/lattice/v7/ring"
)
// CiphertextPool manages a pool of ciphertext buffers for reuse.
type CiphertextPool struct {
mu sync.Mutex
buffers [][]ring.Poly
ringQ *ring.Ring
degree int
maxLevel int
}
// NewCiphertextPool creates a new ciphertext buffer pool.
func NewCiphertextPool(ringQ *ring.Ring, degree, maxLevel, initialSize int) *CiphertextPool {
pool := &CiphertextPool{
ringQ: ringQ,
degree: degree,
maxLevel: maxLevel,
buffers: make([][]ring.Poly, 0, initialSize),
}
// Pre-allocate buffers
for i := 0; i < initialSize; i++ {
pool.buffers = append(pool.buffers, pool.allocate())
}
return pool
}
// Get retrieves a buffer from the pool or allocates a new one.
func (p *CiphertextPool) Get() []ring.Poly {
p.mu.Lock()
defer p.mu.Unlock()
if len(p.buffers) == 0 {
return p.allocate()
}
buf := p.buffers[len(p.buffers)-1]
p.buffers = p.buffers[:len(p.buffers)-1]
return buf
}
// Put returns a buffer to the pool.
func (p *CiphertextPool) Put(buf []ring.Poly) {
if buf == nil {
return
}
p.mu.Lock()
defer p.mu.Unlock()
p.buffers = append(p.buffers, buf)
}
// allocate creates a new ciphertext buffer.
func (p *CiphertextPool) allocate() []ring.Poly {
polys := make([]ring.Poly, p.degree+1)
for i := range polys {
polys[i] = p.ringQ.NewPoly()
}
return polys
}
// Clear releases all pooled buffers.
func (p *CiphertextPool) Clear() {
p.mu.Lock()
defer p.mu.Unlock()
p.buffers = p.buffers[:0]
}
// ExecutionContext holds the state for DAG execution.
type ExecutionContext struct {
// Ciphertext storage indexed by CiphertextID
ciphertexts sync.Map // map[CiphertextID][]ring.Poly
// Buffer pool for intermediates
pool *CiphertextPool
// GPU context (if available)
gpuCtx any // Will be *gpu.NTTContext when GPU is available
// Ring parameters
ringQ *ring.Ring
ringP *ring.Ring // May be nil
degree int
maxLevel int
// Statistics
opsExecuted int64
batchesRun int64
gpuDispatches int64
}
// NewExecutionContext creates a new execution context.
func NewExecutionContext(ringQ *ring.Ring, degree, maxLevel int) *ExecutionContext {
poolSize := 32 // Initial pool size
return &ExecutionContext{
ringQ: ringQ,
degree: degree,
maxLevel: maxLevel,
pool: NewCiphertextPool(ringQ, degree, maxLevel, poolSize),
}
}
// SetGPUContext sets the GPU context for accelerated execution.
func (ec *ExecutionContext) SetGPUContext(ctx any) {
ec.gpuCtx = ctx
}
// SetCiphertext stores a ciphertext in the context.
func (ec *ExecutionContext) SetCiphertext(id CiphertextID, polys []ring.Poly) {
ec.ciphertexts.Store(id, polys)
}
// GetCiphertext retrieves a ciphertext from the context.
func (ec *ExecutionContext) GetCiphertext(id CiphertextID) ([]ring.Poly, bool) {
v, ok := ec.ciphertexts.Load(id)
if !ok {
return nil, false
}
return v.([]ring.Poly), true
}
// ReleaseCiphertext releases a ciphertext back to the pool.
func (ec *ExecutionContext) ReleaseCiphertext(id CiphertextID) {
v, ok := ec.ciphertexts.LoadAndDelete(id)
if ok {
ec.pool.Put(v.([]ring.Poly))
}
}
// AllocateCiphertext allocates a new ciphertext buffer.
func (ec *ExecutionContext) AllocateCiphertext() []ring.Poly {
return ec.pool.Get()
}
// Stats returns execution statistics.
func (ec *ExecutionContext) Stats() (ops, batches, gpuDispatches int64) {
return ec.opsExecuted, ec.batchesRun, ec.gpuDispatches
}
// OperationFunc is a function that executes a single FHE operation.
type OperationFunc func(ctx *ExecutionContext, node *Node, level int) error
// Executor executes DAG schedules.
type Executor struct {
schedule *Schedule
ctx *ExecutionContext
// Operation implementations
ops map[OpType]OperationFunc
// Configuration
enableGPU bool
parallelCPU int // Number of parallel CPU workers
}
// ExecutorOption configures the executor.
type ExecutorOption func(*Executor)
// WithGPU enables/disables GPU execution.
func WithGPU(enable bool) ExecutorOption {
return func(e *Executor) {
e.enableGPU = enable
}
}
// WithParallelCPU sets the number of parallel CPU workers.
func WithParallelCPU(workers int) ExecutorOption {
return func(e *Executor) {
e.parallelCPU = workers
}
}
// NewExecutor creates a new executor for a schedule.
func NewExecutor(schedule *Schedule, ctx *ExecutionContext, opts ...ExecutorOption) *Executor {
e := &Executor{
schedule: schedule,
ctx: ctx,
ops: make(map[OpType]OperationFunc),
enableGPU: true,
parallelCPU: 4,
}
for _, opt := range opts {
opt(e)
}
// Register default operations
e.registerDefaultOps()
return e
}
// RegisterOp registers an operation implementation.
func (e *Executor) RegisterOp(op OpType, fn OperationFunc) {
e.ops[op] = fn
}
// registerDefaultOps registers default operation implementations.
func (e *Executor) registerDefaultOps() {
// Add operation
e.ops[OpAdd] = func(ctx *ExecutionContext, node *Node, level int) error {
if len(node.Inputs) != 2 {
return fmt.Errorf("Add requires 2 inputs, got %d", len(node.Inputs))
}
in1, ok1 := ctx.GetCiphertext(CiphertextID(node.Inputs[0]))
in2, ok2 := ctx.GetCiphertext(CiphertextID(node.Inputs[1]))
if !ok1 || !ok2 {
return fmt.Errorf("input ciphertexts not found")
}
out := ctx.AllocateCiphertext()
// Element-wise addition
for i := range out {
if i < len(in1) && i < len(in2) {
ctx.ringQ.Add(in1[i], in2[i], out[i])
}
}
ctx.SetCiphertext(node.Output, out)
return nil
}
// Sub operation
e.ops[OpSub] = func(ctx *ExecutionContext, node *Node, level int) error {
if len(node.Inputs) != 2 {
return fmt.Errorf("Sub requires 2 inputs, got %d", len(node.Inputs))
}
in1, ok1 := ctx.GetCiphertext(CiphertextID(node.Inputs[0]))
in2, ok2 := ctx.GetCiphertext(CiphertextID(node.Inputs[1]))
if !ok1 || !ok2 {
return fmt.Errorf("input ciphertexts not found")
}
out := ctx.AllocateCiphertext()
for i := range out {
if i < len(in1) && i < len(in2) {
ctx.ringQ.Sub(in1[i], in2[i], out[i])
}
}
ctx.SetCiphertext(node.Output, out)
return nil
}
// Negate operation
e.ops[OpNegate] = func(ctx *ExecutionContext, node *Node, level int) error {
if len(node.Inputs) != 1 {
return fmt.Errorf("Negate requires 1 input, got %d", len(node.Inputs))
}
in, ok := ctx.GetCiphertext(CiphertextID(node.Inputs[0]))
if !ok {
return fmt.Errorf("input ciphertext not found")
}
out := ctx.AllocateCiphertext()
for i := range out {
if i < len(in) {
ctx.ringQ.Neg(in[i], out[i])
}
}
ctx.SetCiphertext(node.Output, out)
return nil
}
// NTT operation
e.ops[OpNTT] = func(ctx *ExecutionContext, node *Node, level int) error {
if len(node.Inputs) != 1 {
return fmt.Errorf("NTT requires 1 input, got %d", len(node.Inputs))
}
in, ok := ctx.GetCiphertext(CiphertextID(node.Inputs[0]))
if !ok {
return fmt.Errorf("input ciphertext not found")
}
out := ctx.AllocateCiphertext()
for i := range out {
if i < len(in) {
ctx.ringQ.NTT(in[i], out[i])
}
}
ctx.SetCiphertext(node.Output, out)
return nil
}
// INTT operation
e.ops[OpINTT] = func(ctx *ExecutionContext, node *Node, level int) error {
if len(node.Inputs) != 1 {
return fmt.Errorf("INTT requires 1 input, got %d", len(node.Inputs))
}
in, ok := ctx.GetCiphertext(CiphertextID(node.Inputs[0]))
if !ok {
return fmt.Errorf("input ciphertext not found")
}
out := ctx.AllocateCiphertext()
for i := range out {
if i < len(in) {
ctx.ringQ.INTT(in[i], out[i])
}
}
ctx.SetCiphertext(node.Output, out)
return nil
}
// Copy operation
e.ops[OpCopy] = func(ctx *ExecutionContext, node *Node, level int) error {
if len(node.Inputs) != 1 {
return fmt.Errorf("Copy requires 1 input, got %d", len(node.Inputs))
}
in, ok := ctx.GetCiphertext(CiphertextID(node.Inputs[0]))
if !ok {
return fmt.Errorf("input ciphertext not found")
}
out := ctx.AllocateCiphertext()
for i := range out {
if i < len(in) {
copy(out[i].Coeffs[0], in[i].Coeffs[0])
}
}
ctx.SetCiphertext(node.Output, out)
return nil
}
// Load operation (no-op, assumes ciphertext is already set)
e.ops[OpLoad] = func(ctx *ExecutionContext, node *Node, level int) error {
// Load expects the ciphertext to be pre-set in context
if _, ok := ctx.GetCiphertext(node.Output); !ok {
return fmt.Errorf("load: ciphertext %d not found in context", node.Output)
}
return nil
}
// Store operation (no-op, just marks for output)
e.ops[OpStore] = func(ctx *ExecutionContext, node *Node, level int) error {
return nil
}
}
// Execute runs the entire schedule.
func (e *Executor) Execute(ctx context.Context) error {
for _, batch := range e.schedule.Batches {
select {
case <-ctx.Done():
return ctx.Err()
default:
}
if err := e.executeBatch(batch); err != nil {
return fmt.Errorf("batch %d failed: %w", batch.ID, err)
}
e.ctx.batchesRun++
}
return nil
}
// executeBatch executes a single batch.
func (e *Executor) executeBatch(batch *Batch) error {
if len(batch.Nodes) == 0 {
return nil
}
// Check if we can use GPU batched execution
if e.enableGPU && !batch.Mixed && e.canBatchOnGPU(batch.OpType) {
return e.executeGPUBatch(batch)
}
// Fall back to CPU execution
return e.executeCPUBatch(batch)
}
// canBatchOnGPU checks if an operation type can be batched on GPU.
func (e *Executor) canBatchOnGPU(op OpType) bool {
switch op {
case OpAdd, OpSub, OpNTT, OpINTT, OpNegate:
return e.ctx.gpuCtx != nil
default:
return false
}
}
// executeGPUBatch executes a homogeneous batch on GPU.
func (e *Executor) executeGPUBatch(batch *Batch) error {
// For now, dispatch individual operations but track as single GPU dispatch
// A full implementation would pack data and call batched GPU kernels
e.ctx.gpuDispatches++
for _, node := range batch.Nodes {
opFn, ok := e.ops[node.Op]
if !ok {
return fmt.Errorf("no implementation for operation %s", node.Op)
}
if err := opFn(e.ctx, node, node.Level); err != nil {
return fmt.Errorf("node %d (%s) failed: %w", node.ID, node.Op, err)
}
node.Executed = true
e.ctx.opsExecuted++
}
return nil
}
// executeCPUBatch executes a batch on CPU.
func (e *Executor) executeCPUBatch(batch *Batch) error {
// For small batches, execute sequentially
if len(batch.Nodes) <= e.parallelCPU {
for _, node := range batch.Nodes {
opFn, ok := e.ops[node.Op]
if !ok {
return fmt.Errorf("no implementation for operation %s", node.Op)
}
if err := opFn(e.ctx, node, node.Level); err != nil {
return fmt.Errorf("node %d (%s) failed: %w", node.ID, node.Op, err)
}
node.Executed = true
e.ctx.opsExecuted++
}
return nil
}
// For larger batches, execute in parallel
var wg sync.WaitGroup
errCh := make(chan error, len(batch.Nodes))
for _, node := range batch.Nodes {
wg.Add(1)
go func(n *Node) {
defer wg.Done()
opFn, ok := e.ops[n.Op]
if !ok {
errCh <- fmt.Errorf("no implementation for operation %s", n.Op)
return
}
if err := opFn(e.ctx, n, n.Level); err != nil {
errCh <- fmt.Errorf("node %d (%s) failed: %w", n.ID, n.Op, err)
return
}
n.Executed = true
e.ctx.opsExecuted++
}(node)
}
wg.Wait()
close(errCh)
// Return first error if any
for err := range errCh {
if err != nil {
return err
}
}
return nil
}
// BatchedGPUDispatch holds data for a GPU kernel dispatch.
type BatchedGPUDispatch struct {
OpType OpType
InputIDs [][]CiphertextID // Input ciphertext IDs per operation
OutputIDs []CiphertextID // Output ciphertext IDs
Metadata []map[string]any // Per-operation metadata
BatchSize int
}
// PrepareBatchDispatch prepares a batch for GPU dispatch.
func PrepareBatchDispatch(batch *Batch) *BatchedGPUDispatch {
dispatch := &BatchedGPUDispatch{
OpType: batch.OpType,
InputIDs: make([][]CiphertextID, len(batch.Nodes)),
OutputIDs: make([]CiphertextID, len(batch.Nodes)),
Metadata: make([]map[string]any, len(batch.Nodes)),
BatchSize: len(batch.Nodes),
}
for i, node := range batch.Nodes {
dispatch.InputIDs[i] = make([]CiphertextID, len(node.Inputs))
for j, input := range node.Inputs {
dispatch.InputIDs[i][j] = CiphertextID(input)
}
dispatch.OutputIDs[i] = node.Output
dispatch.Metadata[i] = node.Metadata
}
return dispatch
}
// ExecuteDAG is a convenience function to execute a DAG.
func ExecuteDAG(ctx context.Context, graph *Graph, execCtx *ExecutionContext) error {
// Create scheduler and schedule
scheduler := NewScheduler(graph)
schedule, err := scheduler.Schedule()
if err != nil {
return fmt.Errorf("scheduling failed: %w", err)
}
// Create executor and run
executor := NewExecutor(schedule, execCtx)
return executor.Execute(ctx)
}
dag/graph.go
+491
View File
@@ -0,0 +1,491 @@
// Package dag provides a directed acyclic graph scheduler for FHE operations.
//
// DAG execution optimizes FHE operation scheduling by:
// - Building a computation graph from FHE operations
// - Identifying independent operations that can run in parallel
// - Minimizing GPU kernel launches via operation batching
// - Tracking ciphertext dependencies automatically
// - Managing memory for intermediate results
//
// Copyright (c) 2024-2025 Lux Partners Limited
// SPDX-License-Identifier: BSD-3-Clause
package dag
import (
"fmt"
"sync/atomic"
)
// OpType represents the type of FHE operation.
type OpType uint8
const (
// Arithmetic operations
OpAdd OpType = iota
OpSub
OpMul
OpMulPlain
OpAddPlain
OpSubPlain
OpNegate
// NTT operations
OpNTT
OpINTT
// Key switching operations
OpRelinearize
OpRotate
OpConjugate
// Level management
OpRescale
OpModSwitch
// Bootstrap
OpBootstrap
// Memory operations
OpCopy
OpLoad // Load from external source
OpStore // Store to external destination
)
// String returns the string representation of an OpType.
func (op OpType) String() string {
names := []string{
"Add", "Sub", "Mul", "MulPlain", "AddPlain", "SubPlain", "Negate",
"NTT", "INTT",
"Relinearize", "Rotate", "Conjugate",
"Rescale", "ModSwitch",
"Bootstrap",
"Copy", "Load", "Store",
}
if int(op) < len(names) {
return names[op]
}
return fmt.Sprintf("Unknown(%d)", op)
}
// CanFuse returns true if operations of this type can be fused.
func (op OpType) CanFuse() bool {
switch op {
case OpAdd, OpSub, OpMulPlain, OpAddPlain, OpSubPlain, OpNegate:
return true
default:
return false
}
}
// IsElementwise returns true if the operation is element-wise.
func (op OpType) IsElementwise() bool {
switch op {
case OpAdd, OpSub, OpMulPlain, OpAddPlain, OpSubPlain, OpNegate, OpNTT, OpINTT:
return true
default:
return false
}
}
// NodeID uniquely identifies a node in the DAG.
type NodeID uint64
// CiphertextID uniquely identifies a ciphertext in the computation.
type CiphertextID uint64
// nextNodeID is the global counter for generating unique node IDs.
var nextNodeID uint64
// newNodeID generates a new unique node ID.
func newNodeID() NodeID {
return NodeID(atomic.AddUint64(&nextNodeID, 1))
}
// Node represents a single operation in the DAG.
type Node struct {
ID NodeID
Op OpType
Inputs []NodeID // IDs of input nodes
Output CiphertextID // ID of output ciphertext
Level int // Ciphertext level after operation
Metadata map[string]any // Operation-specific metadata
// Scheduling information (set by scheduler)
Depth int // Topological depth (0 = input nodes)
BatchID int // Batch this node belongs to (-1 = not scheduled)
Executed bool // Whether this node has been executed
// Dependencies tracking
dependsOn []*Node // Nodes this node depends on
dependents []*Node // Nodes that depend on this node
}
// NewNode creates a new DAG node with the given operation type.
func NewNode(op OpType) *Node {
return &Node{
ID: newNodeID(),
Op: op,
Inputs: make([]NodeID, 0, 2),
Metadata: make(map[string]any),
Depth: -1,
BatchID: -1,
dependsOn: make([]*Node, 0, 2),
}
}
// AddInput adds an input node dependency.
func (n *Node) AddInput(inputID NodeID) {
n.Inputs = append(n.Inputs, inputID)
}
// SetMetadata sets a metadata value for the node.
func (n *Node) SetMetadata(key string, value any) {
n.Metadata[key] = value
}
// GetMetadata retrieves a metadata value.
func (n *Node) GetMetadata(key string) (any, bool) {
v, ok := n.Metadata[key]
return v, ok
}
// IsReady returns true if all input dependencies have been executed.
func (n *Node) IsReady() bool {
for _, dep := range n.dependsOn {
if !dep.Executed {
return false
}
}
return true
}
// Graph represents a DAG of FHE operations.
type Graph struct {
nodes map[NodeID]*Node
outputs map[CiphertextID]*Node // Map from ciphertext ID to producing node
inputNodes []*Node // Nodes with no dependencies (inputs)
outputNodes []*Node // Nodes marked as outputs
// Statistics
nodeCount int
maxDepth int
parallelism float64 // Average operations per level
}
// NewGraph creates a new empty DAG.
func NewGraph() *Graph {
return &Graph{
nodes: make(map[NodeID]*Node),
outputs: make(map[CiphertextID]*Node),
inputNodes: make([]*Node, 0),
outputNodes: make([]*Node, 0),
}
}
// AddNode adds a node to the graph.
func (g *Graph) AddNode(n *Node) error {
if _, exists := g.nodes[n.ID]; exists {
return fmt.Errorf("node %d already exists in graph", n.ID)
}
// Register the node
g.nodes[n.ID] = n
g.nodeCount++
// Link dependencies
for _, inputID := range n.Inputs {
inputNode, ok := g.nodes[inputID]
if !ok {
return fmt.Errorf("input node %d not found for node %d", inputID, n.ID)
}
n.dependsOn = append(n.dependsOn, inputNode)
inputNode.dependents = append(inputNode.dependents, n)
}
// Track input nodes (no dependencies)
if len(n.Inputs) == 0 {
g.inputNodes = append(g.inputNodes, n)
}
// Register output ciphertext
if n.Output != 0 {
g.outputs[n.Output] = n
}
return nil
}
// GetNode retrieves a node by ID.
func (g *Graph) GetNode(id NodeID) (*Node, bool) {
n, ok := g.nodes[id]
return n, ok
}
// GetProducer returns the node that produces a given ciphertext.
func (g *Graph) GetProducer(ctID CiphertextID) (*Node, bool) {
n, ok := g.outputs[ctID]
return n, ok
}
// MarkOutput marks a node as a graph output.
func (g *Graph) MarkOutput(id NodeID) error {
n, ok := g.nodes[id]
if !ok {
return fmt.Errorf("node %d not found", id)
}
g.outputNodes = append(g.outputNodes, n)
return nil
}
// NodeCount returns the number of nodes in the graph.
func (g *Graph) NodeCount() int {
return g.nodeCount
}
// InputNodes returns nodes with no dependencies.
func (g *Graph) InputNodes() []*Node {
return g.inputNodes
}
// OutputNodes returns nodes marked as outputs.
func (g *Graph) OutputNodes() []*Node {
return g.outputNodes
}
// Nodes returns all nodes in the graph.
func (g *Graph) Nodes() []*Node {
result := make([]*Node, 0, len(g.nodes))
for _, n := range g.nodes {
result = append(result, n)
}
return result
}
// MaxDepth returns the maximum depth of the DAG (set after scheduling).
func (g *Graph) MaxDepth() int {
return g.maxDepth
}
// Parallelism returns the average number of operations per depth level.
func (g *Graph) Parallelism() float64 {
return g.parallelism
}
// Validate checks the graph for consistency and cycles.
func (g *Graph) Validate() error {
// Check for cycles using DFS
visited := make(map[NodeID]bool)
recStack := make(map[NodeID]bool)
var hasCycle func(n *Node) bool
hasCycle = func(n *Node) bool {
visited[n.ID] = true
recStack[n.ID] = true
for _, dep := range n.dependents {
if !visited[dep.ID] {
if hasCycle(dep) {
return true
}
} else if recStack[dep.ID] {
return true
}
}
recStack[n.ID] = false
return false
}
for _, n := range g.nodes {
if !visited[n.ID] {
if hasCycle(n) {
return fmt.Errorf("graph contains a cycle")
}
}
}
// Check all inputs are valid
for _, n := range g.nodes {
for _, inputID := range n.Inputs {
if _, ok := g.nodes[inputID]; !ok {
return fmt.Errorf("node %d references non-existent input %d", n.ID, inputID)
}
}
}
return nil
}
// Reset clears execution state for all nodes.
func (g *Graph) Reset() {
for _, n := range g.nodes {
n.Executed = false
n.BatchID = -1
}
}
// Builder provides a fluent API for constructing DAGs.
type Builder struct {
graph *Graph
nextCtID CiphertextID
lastNode *Node
err error
}
// NewBuilder creates a new DAG builder.
func NewBuilder() *Builder {
return &Builder{
graph: NewGraph(),
nextCtID: 1,
}
}
// Input adds an input node representing an external ciphertext.
func (b *Builder) Input(level int) *Builder {
if b.err != nil {
return b
}
n := NewNode(OpLoad)
n.Output = b.nextCtID
b.nextCtID++
n.Level = level
if err := b.graph.AddNode(n); err != nil {
b.err = err
return b
}
b.lastNode = n
return b
}
// Add adds an addition operation.
func (b *Builder) Add(input1, input2 NodeID) *Builder {
return b.binaryOp(OpAdd, input1, input2)
}
// Sub adds a subtraction operation.
func (b *Builder) Sub(input1, input2 NodeID) *Builder {
return b.binaryOp(OpSub, input1, input2)
}
// Mul adds a multiplication operation.
func (b *Builder) Mul(input1, input2 NodeID) *Builder {
return b.binaryOp(OpMul, input1, input2)
}
// Relin adds a relinearization operation.
func (b *Builder) Relin(input NodeID) *Builder {
return b.unaryOp(OpRelinearize, input)
}
// Rescale adds a rescaling operation.
func (b *Builder) Rescale(input NodeID) *Builder {
return b.unaryOp(OpRescale, input)
}
// Rotate adds a rotation operation.
func (b *Builder) Rotate(input NodeID, rotIdx int) *Builder {
b.unaryOp(OpRotate, input)
if b.err == nil {
b.lastNode.SetMetadata("rotation", rotIdx)
}
return b
}
// NTT adds an NTT operation.
func (b *Builder) NTT(input NodeID) *Builder {
return b.unaryOp(OpNTT, input)
}
// INTT adds an inverse NTT operation.
func (b *Builder) INTT(input NodeID) *Builder {
return b.unaryOp(OpINTT, input)
}
// Output marks the last added node as an output.
func (b *Builder) Output() *Builder {
if b.err != nil || b.lastNode == nil {
return b
}
if err := b.graph.MarkOutput(b.lastNode.ID); err != nil {
b.err = err
}
return b
}
// LastNodeID returns the ID of the last added node.
func (b *Builder) LastNodeID() NodeID {
if b.lastNode == nil {
return 0
}
return b.lastNode.ID
}
// Build finalizes and returns the graph.
func (b *Builder) Build() (*Graph, error) {
if b.err != nil {
return nil, b.err
}
if err := b.graph.Validate(); err != nil {
return nil, err
}
return b.graph, nil
}
// binaryOp adds a binary operation node.
func (b *Builder) binaryOp(op OpType, input1, input2 NodeID) *Builder {
if b.err != nil {
return b
}
n := NewNode(op)
n.AddInput(input1)
n.AddInput(input2)
n.Output = b.nextCtID
b.nextCtID++
// Determine output level from inputs
n1, ok1 := b.graph.GetNode(input1)
n2, ok2 := b.graph.GetNode(input2)
if ok1 && ok2 {
n.Level = min(n1.Level, n2.Level)
}
if err := b.graph.AddNode(n); err != nil {
b.err = err
return b
}
b.lastNode = n
return b
}
// unaryOp adds a unary operation node.
func (b *Builder) unaryOp(op OpType, input NodeID) *Builder {
if b.err != nil {
return b
}
n := NewNode(op)
n.AddInput(input)
n.Output = b.nextCtID
b.nextCtID++
// Inherit level from input
if inputNode, ok := b.graph.GetNode(input); ok {
n.Level = inputNode.Level
if op == OpRescale {
n.Level-- // Rescale reduces level by 1
}
}
if err := b.graph.AddNode(n); err != nil {
b.err = err
return b
}
b.lastNode = n
return b
}
dag/scheduler.go
+496
View File
@@ -0,0 +1,496 @@
// Copyright (c) 2024-2025 Lux Partners Limited
// SPDX-License-Identifier: BSD-3-Clause
package dag
import (
"fmt"
"sort"
)
// Batch represents a group of independent operations that can execute in parallel.
type Batch struct {
ID int
Depth int
Nodes []*Node
OpType OpType // Primary operation type (for homogeneous batches)
Mixed bool // True if batch contains mixed operation types
}
// Size returns the number of nodes in the batch.
func (b *Batch) Size() int {
return len(b.Nodes)
}
// Schedule represents a complete execution schedule for a DAG.
type Schedule struct {
Batches []*Batch
NodeToBatch map[NodeID]int
TotalNodes int
MaxParallel int // Maximum batch size
AvgParallel float64 // Average batch size
CriticalPath int // Length of critical path (depth)
}
// Scheduler handles DAG analysis and batch scheduling.
type Scheduler struct {
graph *Graph
schedule *Schedule
// Configuration
maxBatchSize int // Maximum operations per batch
fuseSimilarOps bool // Fuse similar operations in same batch
minimizeKernels bool // Prioritize fewer kernel launches
}
// SchedulerOption configures the scheduler.
type SchedulerOption func(*Scheduler)
// WithMaxBatchSize sets the maximum batch size.
func WithMaxBatchSize(size int) SchedulerOption {
return func(s *Scheduler) {
s.maxBatchSize = size
}
}
// WithFuseSimilarOps enables/disables operation fusion.
func WithFuseSimilarOps(enable bool) SchedulerOption {
return func(s *Scheduler) {
s.fuseSimilarOps = enable
}
}
// WithMinimizeKernels enables kernel minimization mode.
func WithMinimizeKernels(enable bool) SchedulerOption {
return func(s *Scheduler) {
s.minimizeKernels = enable
}
}
// NewScheduler creates a new scheduler for the given graph.
func NewScheduler(g *Graph, opts ...SchedulerOption) *Scheduler {
s := &Scheduler{
graph: g,
maxBatchSize: 1024, // Default max batch size
fuseSimilarOps: true,
minimizeKernels: true,
}
for _, opt := range opts {
opt(s)
}
return s
}
// Schedule computes an execution schedule for the DAG.
func (s *Scheduler) Schedule() (*Schedule, error) {
if err := s.graph.Validate(); err != nil {
return nil, fmt.Errorf("invalid graph: %w", err)
}
// Step 1: Compute topological depths
if err := s.computeDepths(); err != nil {
return nil, fmt.Errorf("failed to compute depths: %w", err)
}
// Step 2: Group nodes by depth level
levelNodes := s.groupByLevel()
// Step 3: Create batches from level groups
batches := s.createBatches(levelNodes)
// Step 4: Build schedule
s.schedule = &Schedule{
Batches: batches,
NodeToBatch: make(map[NodeID]int),
TotalNodes: s.graph.NodeCount(),
CriticalPath: s.graph.maxDepth,
}
// Map nodes to batches and compute statistics
maxParallel := 0
totalParallel := 0
for _, batch := range batches {
for _, node := range batch.Nodes {
s.schedule.NodeToBatch[node.ID] = batch.ID
node.BatchID = batch.ID
}
if batch.Size() > maxParallel {
maxParallel = batch.Size()
}
totalParallel += batch.Size()
}
s.schedule.MaxParallel = maxParallel
if len(batches) > 0 {
s.schedule.AvgParallel = float64(totalParallel) / float64(len(batches))
}
return s.schedule, nil
}
// computeDepths assigns topological depths to all nodes.
// Depth is the longest path from any input to the node.
func (s *Scheduler) computeDepths() error {
// Initialize all depths to -1
for _, n := range s.graph.nodes {
n.Depth = -1
}
// Set input nodes to depth 0
for _, n := range s.graph.inputNodes {
n.Depth = 0
}
// BFS to propagate depths
queue := make([]*Node, len(s.graph.inputNodes))
copy(queue, s.graph.inputNodes)
maxDepth := 0
processed := 0
for len(queue) > 0 {
current := queue[0]
queue = queue[1:]
processed++
// Update depths of dependents
for _, dep := range current.dependents {
newDepth := current.Depth + 1
if newDepth > dep.Depth {
dep.Depth = newDepth
if newDepth > maxDepth {
maxDepth = newDepth
}
}
// Check if all dependencies are processed
allDepsProcessed := true
for _, d := range dep.dependsOn {
if d.Depth == -1 {
allDepsProcessed = false
break
}
}
// Add to queue if all dependencies have depths
if allDepsProcessed && dep.Depth != -1 {
// Avoid duplicates in queue
inQueue := false
for _, q := range queue {
if q.ID == dep.ID {
inQueue = true
break
}
}
if !inQueue {
queue = append(queue, dep)
}
}
}
}
// Verify all nodes have been assigned depths
for _, n := range s.graph.nodes {
if n.Depth == -1 {
return fmt.Errorf("node %d unreachable from inputs", n.ID)
}
}
s.graph.maxDepth = maxDepth
// Compute parallelism
if maxDepth > 0 {
s.graph.parallelism = float64(s.graph.nodeCount) / float64(maxDepth+1)
}
return nil
}
// groupByLevel groups nodes by their topological depth.
func (s *Scheduler) groupByLevel() [][]*Node {
maxDepth := s.graph.maxDepth
levels := make([][]*Node, maxDepth+1)
for i := range levels {
levels[i] = make([]*Node, 0)
}
for _, n := range s.graph.nodes {
levels[n.Depth] = append(levels[n.Depth], n)
}
return levels
}
// createBatches creates batches from level groups.
func (s *Scheduler) createBatches(levelNodes [][]*Node) []*Batch {
batches := make([]*Batch, 0)
batchID := 0
for depth, nodes := range levelNodes {
if len(nodes) == 0 {
continue
}
if s.fuseSimilarOps {
// Group by operation type within level
opGroups := s.groupByOpType(nodes)
for opType, opNodes := range opGroups {
// Split into batches respecting maxBatchSize
for i := 0; i < len(opNodes); i += s.maxBatchSize {
end := i + s.maxBatchSize
if end > len(opNodes) {
end = len(opNodes)
}
batch := &Batch{
ID: batchID,
Depth: depth,
Nodes: opNodes[i:end],
OpType: opType,
Mixed: false,
}
batches = append(batches, batch)
batchID++
}
}
} else {
// Create mixed batches
for i := 0; i < len(nodes); i += s.maxBatchSize {
end := i + s.maxBatchSize
if end > len(nodes) {
end = len(nodes)
}
batch := &Batch{
ID: batchID,
Depth: depth,
Nodes: nodes[i:end],
Mixed: len(s.groupByOpType(nodes[i:end])) > 1,
}
if !batch.Mixed && len(batch.Nodes) > 0 {
batch.OpType = batch.Nodes[0].Op
}
batches = append(batches, batch)
batchID++
}
}
}
return batches
}
// groupByOpType groups nodes by their operation type.
func (s *Scheduler) groupByOpType(nodes []*Node) map[OpType][]*Node {
groups := make(map[OpType][]*Node)
for _, n := range nodes {
groups[n.Op] = append(groups[n.Op], n)
}
return groups
}
// GetSchedule returns the computed schedule.
func (s *Scheduler) GetSchedule() *Schedule {
return s.schedule
}
// ScheduleStats provides statistics about a schedule.
type ScheduleStats struct {
TotalBatches int
TotalNodes int
MaxBatchSize int
MinBatchSize int
AvgBatchSize float64
CriticalPath int
Parallelism float64 // Nodes / Critical path length
KernelEfficiency float64 // Nodes per kernel launch estimate
}
// Stats computes statistics for the schedule.
func (s *Schedule) Stats() ScheduleStats {
stats := ScheduleStats{
TotalBatches: len(s.Batches),
TotalNodes: s.TotalNodes,
MaxBatchSize: s.MaxParallel,
CriticalPath: s.CriticalPath,
}
if len(s.Batches) == 0 {
return stats
}
stats.MinBatchSize = s.Batches[0].Size()
totalSize := 0
for _, batch := range s.Batches {
size := batch.Size()
totalSize += size
if size < stats.MinBatchSize {
stats.MinBatchSize = size
}
}
stats.AvgBatchSize = float64(totalSize) / float64(len(s.Batches))
if stats.CriticalPath > 0 {
stats.Parallelism = float64(stats.TotalNodes) / float64(stats.CriticalPath)
}
// Estimate kernel efficiency based on batch homogeneity
homogeneousBatches := 0
for _, batch := range s.Batches {
if !batch.Mixed {
homogeneousBatches++
}
}
if len(s.Batches) > 0 {
stats.KernelEfficiency = float64(stats.TotalNodes) / float64(len(s.Batches))
}
return stats
}
// MemoryEstimate estimates peak memory usage in ciphertext count.
type MemoryEstimate struct {
PeakCiphertexts int // Maximum concurrent ciphertexts
TotalIntermediates int // Total intermediate ciphertexts created
ReusableBufs int // Buffers that can be reused
}
// EstimateMemory estimates memory requirements for the schedule.
func (s *Schedule) EstimateMemory() MemoryEstimate {
est := MemoryEstimate{}
// Track live ciphertexts at each batch
live := make(map[CiphertextID]bool)
peak := 0
for _, batch := range s.Batches {
// Add outputs
for _, node := range batch.Nodes {
if node.Output != 0 {
live[node.Output] = true
est.TotalIntermediates++
}
}
// Update peak
if len(live) > peak {
peak = len(live)
}
// Note: A full implementation would track when ciphertexts
// are no longer needed and mark them for reuse
}
est.PeakCiphertexts = peak
return est
}
// TopologicalOrder returns nodes in topological order.
func TopologicalOrder(g *Graph) ([]*Node, error) {
if err := g.Validate(); err != nil {
return nil, err
}
// Kahn's algorithm
inDegree := make(map[NodeID]int)
for _, n := range g.nodes {
inDegree[n.ID] = len(n.dependsOn)
}
// Start with nodes that have no dependencies
queue := make([]*Node, 0)
for _, n := range g.nodes {
if inDegree[n.ID] == 0 {
queue = append(queue, n)
}
}
// Sort initial queue by ID for deterministic output
sort.Slice(queue, func(i, j int) bool {
return queue[i].ID < queue[j].ID
})
result := make([]*Node, 0, len(g.nodes))
for len(queue) > 0 {
// Pop front
n := queue[0]
queue = queue[1:]
result = append(result, n)
// Reduce in-degree of dependents
for _, dep := range n.dependents {
inDegree[dep.ID]--
if inDegree[dep.ID] == 0 {
queue = append(queue, dep)
}
}
// Sort for determinism
sort.Slice(queue, func(i, j int) bool {
return queue[i].ID < queue[j].ID
})
}
if len(result) != len(g.nodes) {
return nil, fmt.Errorf("cycle detected: processed %d of %d nodes", len(result), len(g.nodes))
}
return result, nil
}
// CriticalPath returns the nodes on the critical path (longest dependency chain).
func CriticalPath(g *Graph) ([]*Node, error) {
ordered, err := TopologicalOrder(g)
if err != nil {
return nil, err
}
// Compute longest path to each node
dist := make(map[NodeID]int)
parent := make(map[NodeID]*Node)
for _, n := range ordered {
maxDist := 0
var maxParent *Node
for _, dep := range n.dependsOn {
if dist[dep.ID]+1 > maxDist {
maxDist = dist[dep.ID] + 1
maxParent = dep
}
}
dist[n.ID] = maxDist
parent[n.ID] = maxParent
}
// Find the node with maximum distance
var endNode *Node
maxDist := -1
for _, n := range g.nodes {
if dist[n.ID] > maxDist {
maxDist = dist[n.ID]
endNode = n
}
}
if endNode == nil {
return nil, nil
}
// Trace back the critical path
path := make([]*Node, 0)
current := endNode
for current != nil {
path = append([]*Node{current}, path...)
current = parent[current.ID]
}
return path, nil
}
go.mod
+4 -4
View File
@@ -1,13 +1,13 @@
module github.com/luxfi/lattice/v7
go 1.22
go 1.25.5
require (
github.com/ALTree/bigfloat v0.2.0
github.com/google/go-cmp v0.7.0
github.com/stretchr/testify v1.11.1
golang.org/x/crypto v0.46.0
golang.org/x/exp v0.0.0-20251219203646-944ab1f22d93
golang.org/x/crypto v0.47.0
golang.org/x/exp v0.0.0-20260112195511-716be5621a96
)
require (
@@ -15,7 +15,7 @@ require (
github.com/kr/pretty v0.3.1 // indirect
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
github.com/rogpeppe/go-internal v1.14.1 // indirect
golang.org/x/sys v0.39.0 // indirect
golang.org/x/sys v0.40.0 // indirect
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
go.sum
+18 -6
View File
@@ -1,21 +1,33 @@
github.com/ALTree/bigfloat v0.2.0 h1:AwNzawrpFuw55/YDVlcPw0F0cmmXrmngBHhVrvdXPvM=
github.com/ALTree/bigfloat v0.2.0/go.mod h1:+NaH2gLeY6RPBPPQf4aRotPPStg+eXc8f9ZaE4vRfD4=
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
golang.org/x/crypto v0.46.0 h1:cKRW/pmt1pKAfetfu+RCEvjvZkA9RimPbh7bhFjGVBU=
golang.org/x/crypto v0.46.0/go.mod h1:Evb/oLKmMraqjZ2iQTwDwvCtJkczlDuTmdJXoZVzqU0=
golang.org/x/exp v0.0.0-20251219203646-944ab1f22d93 h1:fQsdNF2N+/YewlRZiricy4P1iimyPKZ/xwniHj8Q2a0=
golang.org/x/exp v0.0.0-20251219203646-944ab1f22d93/go.mod h1:EPRbTFwzwjXj9NpYyyrvenVh9Y+GFeEvMNh7Xuz7xgU=
golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk=
golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
golang.org/x/crypto v0.47.0 h1:V6e3FRj+n4dbpw86FJ8Fv7XVOql7TEwpHapKoMJ/GO8=
golang.org/x/crypto v0.47.0/go.mod h1:ff3Y9VzzKbwSSEzWqJsJVBnWmRwRSHt/6Op5n9bQc4A=
golang.org/x/exp v0.0.0-20260112195511-716be5621a96 h1:Z/6YuSHTLOHfNFdb8zVZomZr7cqNgTJvA8+Qz75D8gU=
golang.org/x/exp v0.0.0-20260112195511-716be5621a96/go.mod h1:nzimsREAkjBCIEFtHiYkrJyT+2uy9YZJB7H1k68CXZU=
golang.org/x/sys v0.40.0 h1:DBZZqJ2Rkml6QMQsZywtnjnnGvHza6BTfYFWY9kjEWQ=
golang.org/x/sys v0.40.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gpu/gpu.go
-412
View File
@@ -1,412 +0,0 @@
//go:build !cgo || !gpu
// Package gpu provides pure Go implementations when CGO is disabled or the gpu build tag is not set.
// These implementations use the lattice/ring package for NTT and polynomial operations.
package gpu
import (
"crypto/rand"
"fmt"
"math"
"math/big"
"sync"
"github.com/luxfi/lattice/v7/ring"
"github.com/luxfi/lattice/v7/utils/sampling"
)
// GPUAvailable returns false when CGO is disabled.
func GPUAvailable() bool {
return false
}
// GetBackend returns "CPU (pure Go)".
func GetBackend() string {
return "CPU (pure Go)"
}
// ClearCache is a no-op without GPU.
func ClearCache() {}
// NTTContext wraps the ring.SubRing for NTT operations.
type NTTContext struct {
subRing *ring.SubRing
N uint32
Q uint64
mu sync.RWMutex
}
// NewNTTContext creates a new NTT context using pure Go implementation.
func NewNTTContext(N uint32, Q uint64) (*NTTContext, error) {
if N == 0 || (N&(N-1)) != 0 {
return nil, fmt.Errorf("N must be a power of 2, got %d", N)
}
// Check if Q is NTT-friendly (Q-1 divisible by 2N)
if (Q-1)%(2*uint64(N)) != 0 {
return nil, fmt.Errorf("Q-1 (%d) must be divisible by 2N (%d) for NTT-friendly prime", Q-1, 2*uint64(N))
}
subRing, err := ring.NewSubRing(int(N), Q)
if err != nil {
return nil, fmt.Errorf("failed to create SubRing: %w", err)
}
// NTT constants are generated internally by NewSubRing via NewSubRingWithCustomNTT
return &NTTContext{
subRing: subRing,
N: N,
Q: Q,
}, nil
}
// Close releases resources (no-op for pure Go).
func (ctx *NTTContext) Close() {
ctx.mu.Lock()
defer ctx.mu.Unlock()
ctx.subRing = nil
}
// NTT performs forward NTT on a batch of polynomials.
func (ctx *NTTContext) NTT(polys [][]uint64) ([][]uint64, error) {
ctx.mu.RLock()
defer ctx.mu.RUnlock()
if ctx.subRing == nil {
return nil, fmt.Errorf("NTT context is closed")
}
if len(polys) == 0 {
return polys, nil
}
N := int(ctx.N)
results := make([][]uint64, len(polys))
for i, poly := range polys {
if len(poly) != N {
return nil, fmt.Errorf("polynomial %d has wrong size: got %d, expected %d", i, len(poly), N)
}
// Copy to result
results[i] = make([]uint64, N)
copy(results[i], poly)
// Perform NTT in-place
ctx.subRing.NTT(results[i], results[i])
}
return results, nil
}
// INTT performs inverse NTT on a batch of polynomials.
func (ctx *NTTContext) INTT(polys [][]uint64) ([][]uint64, error) {
ctx.mu.RLock()
defer ctx.mu.RUnlock()
if ctx.subRing == nil {
return nil, fmt.Errorf("NTT context is closed")
}
if len(polys) == 0 {
return polys, nil
}
N := int(ctx.N)
results := make([][]uint64, len(polys))
for i, poly := range polys {
if len(poly) != N {
return nil, fmt.Errorf("polynomial %d has wrong size: got %d, expected %d", i, len(poly), N)
}
// Copy to result
results[i] = make([]uint64, N)
copy(results[i], poly)
// Perform INTT in-place
ctx.subRing.INTT(results[i], results[i])
}
return results, nil
}
// PolyMul performs polynomial multiplication using pure Go NTT.
func (ctx *NTTContext) PolyMul(a, b [][]uint64) ([][]uint64, error) {
ctx.mu.RLock()
defer ctx.mu.RUnlock()
if ctx.subRing == nil {
return nil, fmt.Errorf("NTT context is closed")
}
if len(a) != len(b) {
return nil, fmt.Errorf("batch size mismatch: %d vs %d", len(a), len(b))
}
if len(a) == 0 {
return nil, nil
}
N := int(ctx.N)
results := make([][]uint64, len(a))
tempA := make([]uint64, N)
tempB := make([]uint64, N)
for i := range a {
if len(a[i]) != N || len(b[i]) != N {
return nil, fmt.Errorf("polynomial %d has wrong size", i)
}
results[i] = make([]uint64, N)
// Copy inputs
copy(tempA, a[i])
copy(tempB, b[i])
// Forward NTT
ctx.subRing.NTT(tempA, tempA)
ctx.subRing.NTT(tempB, tempB)
// Pointwise multiply
ctx.subRing.MulCoeffsMontgomery(tempA, tempB, results[i])
// Inverse NTT
ctx.subRing.INTT(results[i], results[i])
}
return results, nil
}
// PolyMulNTT performs element-wise multiplication in NTT domain.
func (ctx *NTTContext) PolyMulNTT(a, b []uint64) ([]uint64, error) {
ctx.mu.RLock()
defer ctx.mu.RUnlock()
if ctx.subRing == nil {
return nil, fmt.Errorf("NTT context is closed")
}
N := int(ctx.N)
if len(a) != N || len(b) != N {
return nil, fmt.Errorf("polynomial size mismatch")
}
result := make([]uint64, N)
ctx.subRing.MulCoeffsMontgomery(a, b, result)
return result, nil
}
// PolyAdd computes result = a + b (mod Q).
func PolyAdd(a, b []uint64, Q uint64) ([]uint64, error) {
if len(a) != len(b) {
return nil, fmt.Errorf("polynomial size mismatch")
}
result := make([]uint64, len(a))
for i := range a {
sum := a[i] + b[i]
if sum >= Q {
sum -= Q
}
result[i] = sum
}
return result, nil
}
// PolySub computes result = a - b (mod Q).
func PolySub(a, b []uint64, Q uint64) ([]uint64, error) {
if len(a) != len(b) {
return nil, fmt.Errorf("polynomial size mismatch")
}
result := make([]uint64, len(a))
for i := range a {
if a[i] >= b[i] {
result[i] = a[i] - b[i]
} else {
result[i] = Q - b[i] + a[i]
}
}
return result, nil
}
// PolyScalarMul computes result = a * scalar (mod Q).
func PolyScalarMul(a []uint64, scalar, Q uint64) ([]uint64, error) {
result := make([]uint64, len(a))
for i := range a {
// Use big.Int for 128-bit intermediate
prod := new(big.Int).SetUint64(a[i])
prod.Mul(prod, new(big.Int).SetUint64(scalar))
prod.Mod(prod, new(big.Int).SetUint64(Q))
result[i] = prod.Uint64()
}
return result, nil
}
// SampleGaussian samples a polynomial with discrete Gaussian distribution.
func SampleGaussian(N uint32, Q uint64, sigma float64, seed []byte) ([]uint64, error) {
result := make([]uint64, N)
// Create PRNG
var prng sampling.PRNG
if len(seed) >= 32 {
prng, _ = sampling.NewKeyedPRNG(seed[:32])
} else {
prng, _ = sampling.NewPRNG()
}
// Use Box-Muller transform for Gaussian sampling
bound := int64(math.Ceil(sigma * 6)) // 6-sigma cutoff
for i := uint32(0); i < N; i++ {
// Rejection sampling for discrete Gaussian
for {
// Sample from bounded uniform
var sample int64
for {
b := make([]byte, 8)
prng.Read(b)
sample = int64(b[0]) | int64(b[1])<<8 | int64(b[2])<<16 | int64(b[3])<<24
sample = sample % (2*bound + 1) - bound
if sample >= -bound && sample <= bound {
break
}
}
// Accept with Gaussian probability
prob := math.Exp(-float64(sample*sample) / (2 * sigma * sigma))
threshold := make([]byte, 8)
prng.Read(threshold)
if float64(threshold[0])/256.0 < prob {
if sample >= 0 {
result[i] = uint64(sample)
} else {
result[i] = Q - uint64(-sample)
}
break
}
}
}
return result, nil
}
// SampleUniform samples a uniform random polynomial.
func SampleUniform(N uint32, Q uint64, seed []byte) ([]uint64, error) {
result := make([]uint64, N)
// Create PRNG
var prng sampling.PRNG
if len(seed) >= 32 {
prng, _ = sampling.NewKeyedPRNG(seed[:32])
} else {
prng, _ = sampling.NewPRNG()
}
qBig := new(big.Int).SetUint64(Q)
for i := uint32(0); i < N; i++ {
// Sample uniform in [0, Q)
b := make([]byte, 8)
prng.Read(b)
val := new(big.Int).SetBytes(b)
val.Mod(val, qBig)
result[i] = val.Uint64()
}
return result, nil
}
// SampleTernary samples a ternary polynomial {-1, 0, 1}.
func SampleTernary(N uint32, Q uint64, density float64, seed []byte) ([]uint64, error) {
result := make([]uint64, N)
// Create PRNG
var reader interface{ Read([]byte) (int, error) }
if len(seed) >= 32 {
prng, _ := sampling.NewKeyedPRNG(seed[:32])
reader = prng
} else {
reader = rand.Reader
}
for i := uint32(0); i < N; i++ {
b := make([]byte, 2)
reader.Read(b)
// Probability of non-zero
p := float64(b[0]) / 256.0
if p < density {
// Non-zero: choose -1 or 1
if b[1]&1 == 0 {
result[i] = 1
} else {
result[i] = Q - 1 // -1 mod Q
}
} else {
result[i] = 0
}
}
return result, nil
}
// FindPrimitiveRoot finds a primitive 2N-th root of unity modulo Q.
func FindPrimitiveRoot(N uint32, Q uint64) (uint64, error) {
// Compute g = generator of Z_Q^*
// Then root = g^((Q-1)/(2N))
exponent := (Q - 1) / (2 * uint64(N))
// Find a generator (primitive root of Q)
// For prime Q, try small values until we find one
for g := uint64(2); g < Q; g++ {
// Check if g is a primitive root
gBig := new(big.Int).SetUint64(g)
qBig := new(big.Int).SetUint64(Q)
expBig := new(big.Int).SetUint64(exponent)
root := new(big.Int).Exp(gBig, expBig, qBig)
// Verify it's a 2N-th root of unity
twoN := new(big.Int).SetUint64(2 * uint64(N))
test := new(big.Int).Exp(root, twoN, qBig)
if test.Uint64() == 1 {
// Also verify root^N != 1 (primitive)
nBig := new(big.Int).SetUint64(uint64(N))
testN := new(big.Int).Exp(root, nBig, qBig)
if testN.Uint64() != 1 {
return root.Uint64(), nil
}
}
}
return 0, fmt.Errorf("no primitive root found for N=%d, Q=%d", N, Q)
}
// ModInverse computes the modular inverse a^{-1} mod Q.
func ModInverse(a, Q uint64) (uint64, error) {
aBig := new(big.Int).SetUint64(a)
qBig := new(big.Int).SetUint64(Q)
inv := new(big.Int).ModInverse(aBig, qBig)
if inv == nil {
return 0, fmt.Errorf("%d is not invertible mod %d", a, Q)
}
return inv.Uint64(), nil
}
// IsNTTPrime checks if Q is a valid NTT-friendly prime for ring dimension N.
func IsNTTPrime(N uint32, Q uint64) bool {
// Q must be prime and Q ≡ 1 (mod 2N)
if (Q-1)%(2*uint64(N)) != 0 {
return false
}
// Check if Q is prime (simple Miller-Rabin)
qBig := new(big.Int).SetUint64(Q)
return qBig.ProbablyPrime(20)
}
gpu/gpu_cgo.go
-505
View File
@@ -1,505 +0,0 @@
//go:build cgo && gpu
// Package gpu provides GPU-accelerated lattice operations via libLattice.
//
// With CGO enabled, this package links to libLattice for GPU acceleration:
// - Metal (macOS/Apple Silicon)
// - CUDA (Linux/NVIDIA)
// - Optimized CPU fallback
//
// The library automatically selects the best available backend.
//
// Copyright (c) 2024-2025 Lux Industries Inc.
// SPDX-License-Identifier: Apache-2.0
package gpu
/*
#cgo pkg-config: lux-lattice
#cgo CXXFLAGS: -std=c++17 -O3
#cgo darwin LDFLAGS: -framework Metal -framework Foundation -lstdc++
#cgo linux LDFLAGS: -lstdc++ -lcudart
#include <stdbool.h>
#include <stdint.h>
#include <stdlib.h>
// libLattice GPU-accelerated operations (extern declarations match library exports)
bool lattice_gpu_available(void);
const char* lattice_get_backend(void);
void lattice_clear_cache(void);
// NTT context
typedef struct LatticeNTTContext LatticeNTTContext;
LatticeNTTContext* lattice_ntt_create(uint32_t N, uint64_t Q);
void lattice_ntt_destroy(LatticeNTTContext* ctx);
void lattice_ntt_get_params(const LatticeNTTContext* ctx, uint32_t* N, uint64_t* Q);
// NTT operations
int lattice_ntt_forward(LatticeNTTContext* ctx, uint64_t* data, uint32_t batch);
int lattice_ntt_inverse(LatticeNTTContext* ctx, uint64_t* data, uint32_t batch);
// Batch NTT operations
int lattice_ntt_batch_forward(LatticeNTTContext* ctx, uint64_t* data, uint32_t batch);
int lattice_ntt_batch_inverse(LatticeNTTContext* ctx, uint64_t* data, uint32_t batch);
// Polynomial operations
int lattice_poly_mul_ntt(LatticeNTTContext* ctx, uint64_t* result, const uint64_t* a, const uint64_t* b);
int lattice_poly_mul(LatticeNTTContext* ctx, uint64_t* result, const uint64_t* a, const uint64_t* b);
int lattice_poly_add(uint64_t* result, const uint64_t* a, const uint64_t* b, uint32_t N, uint64_t Q);
int lattice_poly_sub(uint64_t* result, const uint64_t* a, const uint64_t* b, uint32_t N, uint64_t Q);
int lattice_poly_scalar_mul(uint64_t* result, const uint64_t* a, uint64_t scalar, uint32_t N, uint64_t Q);
// Sampling
int lattice_sample_gaussian(uint64_t* result, uint32_t N, uint64_t Q, double sigma, const uint8_t* seed);
int lattice_sample_uniform(uint64_t* result, uint32_t N, uint64_t Q, const uint8_t* seed);
int lattice_sample_ternary(uint64_t* result, uint32_t N, uint64_t Q, double density, const uint8_t* seed);
// Utility
uint64_t lattice_find_primitive_root(uint32_t N, uint64_t Q);
uint64_t lattice_mod_inverse(uint64_t a, uint64_t Q);
bool lattice_is_ntt_prime(uint32_t N, uint64_t Q);
// Aliases for readability
#define NTTContext LatticeNTTContext
#define ntt_create lattice_ntt_create
#define ntt_destroy lattice_ntt_destroy
#define ntt_get_params lattice_ntt_get_params
#define ntt_forward lattice_ntt_forward
#define ntt_inverse lattice_ntt_inverse
#define ntt_batch_forward lattice_ntt_batch_forward
#define ntt_batch_inverse lattice_ntt_batch_inverse
// Polynomial operations
#define poly_mul_ntt lattice_poly_mul_ntt
#define poly_mul lattice_poly_mul
#define poly_add lattice_poly_add
#define poly_sub lattice_poly_sub
#define poly_scalar_mul lattice_poly_scalar_mul
// Sampling
#define sample_gaussian lattice_sample_gaussian
#define sample_uniform lattice_sample_uniform
#define sample_ternary lattice_sample_ternary
// Utility
#define find_primitive_root lattice_find_primitive_root
#define mod_inverse lattice_mod_inverse
#define is_ntt_prime lattice_is_ntt_prime
*/
import "C"
import (
"fmt"
"sync"
"unsafe"
)
// GPUAvailable returns true if GPU acceleration is available.
func GPUAvailable() bool {
return bool(C.lattice_gpu_available())
}
// GetBackend returns the name of the active backend ("Metal", "CUDA", or "CPU").
func GetBackend() string {
return C.GoString(C.lattice_get_backend())
}
// ClearCache clears internal caches (twiddle factors, contexts).
func ClearCache() {
C.lattice_clear_cache()
}
// NTTContext holds precomputed data for GPU-accelerated NTT operations.
type NTTContext struct {
ptr *C.NTTContext
N uint32
Q uint64
mu sync.RWMutex
}
// NewNTTContext creates a new NTT context for the given ring parameters.
// N must be a power of 2, Q must be an NTT-friendly prime (Q ≡ 1 mod 2N).
func NewNTTContext(N uint32, Q uint64) (*NTTContext, error) {
if N == 0 || (N&(N-1)) != 0 {
return nil, fmt.Errorf("N must be a power of 2, got %d", N)
}
if (Q-1)%(2*uint64(N)) != 0 {
return nil, fmt.Errorf("Q-1 (%d) must be divisible by 2N (%d) for NTT-friendly prime", Q-1, 2*uint64(N))
}
ptr := C.ntt_create(C.uint32_t(N), C.uint64_t(Q))
if ptr == nil {
return nil, fmt.Errorf("failed to create NTT context for N=%d, Q=%d", N, Q)
}
return &NTTContext{
ptr: ptr,
N: N,
Q: Q,
}, nil
}
// Close releases the NTT context resources.
func (ctx *NTTContext) Close() {
ctx.mu.Lock()
defer ctx.mu.Unlock()
if ctx.ptr != nil {
C.ntt_destroy(ctx.ptr)
ctx.ptr = nil
}
}
// NTT performs forward NTT on a batch of polynomials.
// Each polynomial is transformed in-place from coefficient to NTT domain.
func (ctx *NTTContext) NTT(polys [][]uint64) ([][]uint64, error) {
if ctx.ptr == nil {
return nil, fmt.Errorf("NTT context is closed")
}
if len(polys) == 0 {
return polys, nil
}
N := int(ctx.N)
results := make([][]uint64, len(polys))
for i, poly := range polys {
if len(poly) != N {
return nil, fmt.Errorf("polynomial %d has wrong size: got %d, expected %d", i, len(poly), N)
}
// Copy to result (NTT modifies in-place)
results[i] = make([]uint64, N)
copy(results[i], poly)
err := C.ntt_forward(ctx.ptr, (*C.uint64_t)(unsafe.Pointer(&results[i][0])), 1)
if err != 0 {
return nil, fmt.Errorf("NTT forward failed with error %d", err)
}
}
return results, nil
}
// INTT performs inverse NTT on a batch of polynomials.
// Each polynomial is transformed in-place from NTT to coefficient domain.
func (ctx *NTTContext) INTT(polys [][]uint64) ([][]uint64, error) {
if ctx.ptr == nil {
return nil, fmt.Errorf("NTT context is closed")
}
if len(polys) == 0 {
return polys, nil
}
N := int(ctx.N)
results := make([][]uint64, len(polys))
for i, poly := range polys {
if len(poly) != N {
return nil, fmt.Errorf("polynomial %d has wrong size: got %d, expected %d", i, len(poly), N)
}
// Copy to result (INTT modifies in-place)
results[i] = make([]uint64, N)
copy(results[i], poly)
err := C.ntt_inverse(ctx.ptr, (*C.uint64_t)(unsafe.Pointer(&results[i][0])), 1)
if err != 0 {
return nil, fmt.Errorf("INTT inverse failed with error %d", err)
}
}
return results, nil
}
// BatchNTT performs forward NTT on all polynomials in a single GPU dispatch.
// This is more efficient than NTT for large batches as it minimizes GPU kernel launches.
func (ctx *NTTContext) BatchNTT(polys [][]uint64) ([][]uint64, error) {
if ctx.ptr == nil {
return nil, fmt.Errorf("NTT context is closed")
}
if len(polys) == 0 {
return polys, nil
}
N := int(ctx.N)
batch := len(polys)
// Allocate contiguous buffer for all polynomials
data := make([]uint64, batch*N)
for i, poly := range polys {
if len(poly) != N {
return nil, fmt.Errorf("polynomial %d has wrong size: got %d, expected %d", i, len(poly), N)
}
copy(data[i*N:(i+1)*N], poly)
}
// Single GPU dispatch for all polynomials
err := C.ntt_batch_forward(ctx.ptr, (*C.uint64_t)(unsafe.Pointer(&data[0])), C.uint32_t(batch))
if err != 0 {
return nil, fmt.Errorf("batch NTT forward failed with error %d", err)
}
// Split back into slices
results := make([][]uint64, batch)
for i := range results {
results[i] = data[i*N : (i+1)*N]
}
return results, nil
}
// BatchINTT performs inverse NTT on all polynomials in a single GPU dispatch.
// This is more efficient than INTT for large batches as it minimizes GPU kernel launches.
func (ctx *NTTContext) BatchINTT(polys [][]uint64) ([][]uint64, error) {
if ctx.ptr == nil {
return nil, fmt.Errorf("NTT context is closed")
}
if len(polys) == 0 {
return polys, nil
}
N := int(ctx.N)
batch := len(polys)
// Allocate contiguous buffer for all polynomials
data := make([]uint64, batch*N)
for i, poly := range polys {
if len(poly) != N {
return nil, fmt.Errorf("polynomial %d has wrong size: got %d, expected %d", i, len(poly), N)
}
copy(data[i*N:(i+1)*N], poly)
}
// Single GPU dispatch for all polynomials
err := C.ntt_batch_inverse(ctx.ptr, (*C.uint64_t)(unsafe.Pointer(&data[0])), C.uint32_t(batch))
if err != 0 {
return nil, fmt.Errorf("batch INTT inverse failed with error %d", err)
}
// Split back into slices
results := make([][]uint64, batch)
for i := range results {
results[i] = data[i*N : (i+1)*N]
}
return results, nil
}
// PolyMul performs polynomial multiplication using GPU-accelerated NTT.
// Both polynomials should be in coefficient form.
// Returns a * b in R_Q = Z_Q[X]/(X^N + 1).
func (ctx *NTTContext) PolyMul(a, b [][]uint64) ([][]uint64, error) {
if ctx.ptr == nil {
return nil, fmt.Errorf("NTT context is closed")
}
if len(a) != len(b) {
return nil, fmt.Errorf("batch size mismatch: %d vs %d", len(a), len(b))
}
if len(a) == 0 {
return nil, nil
}
N := int(ctx.N)
results := make([][]uint64, len(a))
for i := range a {
if len(a[i]) != N || len(b[i]) != N {
return nil, fmt.Errorf("polynomial %d has wrong size", i)
}
results[i] = make([]uint64, N)
err := C.poly_mul(ctx.ptr,
(*C.uint64_t)(unsafe.Pointer(&results[i][0])),
(*C.uint64_t)(unsafe.Pointer(&a[i][0])),
(*C.uint64_t)(unsafe.Pointer(&b[i][0])))
if err != 0 {
return nil, fmt.Errorf("polynomial multiplication failed with error %d", err)
}
}
return results, nil
}
// PolyMulNTT performs element-wise multiplication (Hadamard product) in NTT domain.
// Both polynomials must already be in NTT form.
func (ctx *NTTContext) PolyMulNTT(a, b []uint64) ([]uint64, error) {
if ctx.ptr == nil {
return nil, fmt.Errorf("NTT context is closed")
}
N := int(ctx.N)
if len(a) != N || len(b) != N {
return nil, fmt.Errorf("polynomial size mismatch")
}
result := make([]uint64, N)
err := C.poly_mul_ntt(ctx.ptr,
(*C.uint64_t)(unsafe.Pointer(&result[0])),
(*C.uint64_t)(unsafe.Pointer(&a[0])),
(*C.uint64_t)(unsafe.Pointer(&b[0])))
if err != 0 {
return nil, fmt.Errorf("NTT multiplication failed with error %d", err)
}
return result, nil
}
// PolyAdd computes result = a + b (mod Q).
func PolyAdd(a, b []uint64, Q uint64) ([]uint64, error) {
if len(a) != len(b) {
return nil, fmt.Errorf("polynomial size mismatch")
}
N := uint32(len(a))
result := make([]uint64, N)
err := C.poly_add(
(*C.uint64_t)(unsafe.Pointer(&result[0])),
(*C.uint64_t)(unsafe.Pointer(&a[0])),
(*C.uint64_t)(unsafe.Pointer(&b[0])),
C.uint32_t(N),
C.uint64_t(Q))
if err != 0 {
return nil, fmt.Errorf("polynomial addition failed with error %d", err)
}
return result, nil
}
// PolySub computes result = a - b (mod Q).
func PolySub(a, b []uint64, Q uint64) ([]uint64, error) {
if len(a) != len(b) {
return nil, fmt.Errorf("polynomial size mismatch")
}
N := uint32(len(a))
result := make([]uint64, N)
err := C.poly_sub(
(*C.uint64_t)(unsafe.Pointer(&result[0])),
(*C.uint64_t)(unsafe.Pointer(&a[0])),
(*C.uint64_t)(unsafe.Pointer(&b[0])),
C.uint32_t(N),
C.uint64_t(Q))
if err != 0 {
return nil, fmt.Errorf("polynomial subtraction failed with error %d", err)
}
return result, nil
}
// PolyScalarMul computes result = a * scalar (mod Q).
func PolyScalarMul(a []uint64, scalar, Q uint64) ([]uint64, error) {
N := uint32(len(a))
result := make([]uint64, N)
err := C.poly_scalar_mul(
(*C.uint64_t)(unsafe.Pointer(&result[0])),
(*C.uint64_t)(unsafe.Pointer(&a[0])),
C.uint64_t(scalar),
C.uint32_t(N),
C.uint64_t(Q))
if err != 0 {
return nil, fmt.Errorf("scalar multiplication failed with error %d", err)
}
return result, nil
}
// SampleGaussian samples a polynomial with discrete Gaussian distribution.
func SampleGaussian(N uint32, Q uint64, sigma float64, seed []byte) ([]uint64, error) {
result := make([]uint64, N)
var seedPtr *C.uint8_t
if len(seed) >= 32 {
seedPtr = (*C.uint8_t)(unsafe.Pointer(&seed[0]))
}
err := C.sample_gaussian(
(*C.uint64_t)(unsafe.Pointer(&result[0])),
C.uint32_t(N),
C.uint64_t(Q),
C.double(sigma),
seedPtr)
if err != 0 {
return nil, fmt.Errorf("Gaussian sampling failed with error %d", err)
}
return result, nil
}
// SampleUniform samples a uniform random polynomial.
func SampleUniform(N uint32, Q uint64, seed []byte) ([]uint64, error) {
result := make([]uint64, N)
var seedPtr *C.uint8_t
if len(seed) >= 32 {
seedPtr = (*C.uint8_t)(unsafe.Pointer(&seed[0]))
}
err := C.sample_uniform(
(*C.uint64_t)(unsafe.Pointer(&result[0])),
C.uint32_t(N),
C.uint64_t(Q),
seedPtr)
if err != 0 {
return nil, fmt.Errorf("uniform sampling failed with error %d", err)
}
return result, nil
}
// SampleTernary samples a ternary polynomial {-1, 0, 1}.
func SampleTernary(N uint32, Q uint64, density float64, seed []byte) ([]uint64, error) {
result := make([]uint64, N)
var seedPtr *C.uint8_t
if len(seed) >= 32 {
seedPtr = (*C.uint8_t)(unsafe.Pointer(&seed[0]))
}
err := C.sample_ternary(
(*C.uint64_t)(unsafe.Pointer(&result[0])),
C.uint32_t(N),
C.uint64_t(Q),
C.double(density),
seedPtr)
if err != 0 {
return nil, fmt.Errorf("ternary sampling failed with error %d", err)
}
return result, nil
}
// FindPrimitiveRoot finds a primitive 2N-th root of unity modulo Q.
func FindPrimitiveRoot(N uint32, Q uint64) (uint64, error) {
root := uint64(C.find_primitive_root(C.uint32_t(N), C.uint64_t(Q)))
if root == 0 {
return 0, fmt.Errorf("no primitive root found for N=%d, Q=%d", N, Q)
}
return root, nil
}
// ModInverse computes the modular inverse a^{-1} mod Q.
func ModInverse(a, Q uint64) (uint64, error) {
inv := uint64(C.mod_inverse(C.uint64_t(a), C.uint64_t(Q)))
if inv == 0 && a != 1 {
return 0, fmt.Errorf("%d is not invertible mod %d", a, Q)
}
return inv, nil
}
// IsNTTPrime checks if Q is a valid NTT-friendly prime for ring dimension N.
func IsNTTPrime(N uint32, Q uint64) bool {
return bool(C.is_ntt_prime(C.uint32_t(N), C.uint64_t(Q)))
}
gpu/gpu_test.go
-225
View File
@@ -1,225 +0,0 @@
//go:build cgo
package gpu
import (
"testing"
)
// Test NTT-friendly prime used by Corona
const (
TestN uint32 = 256
TestQ uint64 = 0x1000000004A01 // 48-bit NTT-friendly prime
)
func TestGPUAvailable(t *testing.T) {
available := GPUAvailable()
backend := GetBackend()
t.Logf("GPU Available: %v, Backend: %s", available, backend)
}
func TestIsNTTPrime(t *testing.T) {
// Test Corona primes
tests := []struct {
N uint32
Q uint64
want bool
name string
}{
{256, 0x1000000004A01, true, "Corona Q"},
{256, 0x40201, true, "Corona QXi"}, // 262657 - prime, ≡ 1 (mod 512)
{256, 0x7FE01, true, "Corona QNu"}, // 523777 - prime, ≡ 1 (mod 512)
{256, 0x40000, false, "Non-prime 2^18"}, // 262144 - NOT prime
{256, 0x80000, false, "Non-prime 2^19"}, // 524288 - NOT prime
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := IsNTTPrime(tt.N, tt.Q)
if got != tt.want {
t.Errorf("IsNTTPrime(%d, 0x%X) = %v, want %v", tt.N, tt.Q, got, tt.want)
}
})
}
}
func TestNewNTTContext(t *testing.T) {
ctx, err := NewNTTContext(TestN, TestQ)
if err != nil {
t.Fatalf("NewNTTContext failed: %v", err)
}
defer ctx.Close()
if ctx.N != TestN {
t.Errorf("N = %d, want %d", ctx.N, TestN)
}
if ctx.Q != TestQ {
t.Errorf("Q = 0x%X, want 0x%X", ctx.Q, TestQ)
}
}
func TestNTTRoundTrip(t *testing.T) {
ctx, err := NewNTTContext(TestN, TestQ)
if err != nil {
t.Fatalf("NewNTTContext failed: %v", err)
}
defer ctx.Close()
// Create a test polynomial
poly := make([]uint64, TestN)
for i := range poly {
poly[i] = uint64(i + 1) % TestQ
}
// Save original
original := make([]uint64, TestN)
copy(original, poly)
// Forward NTT
nttResult, err := ctx.NTT([][]uint64{poly})
if err != nil {
t.Fatalf("NTT failed: %v", err)
}
// Inverse NTT
inttResult, err := ctx.INTT(nttResult)
if err != nil {
t.Fatalf("INTT failed: %v", err)
}
// Verify round-trip
for i := range original {
if inttResult[0][i] != original[i] {
t.Errorf("Round-trip failed at index %d: got %d, want %d", i, inttResult[0][i], original[i])
}
}
t.Log("NTT round-trip successful")
}
func TestFindPrimitiveRoot(t *testing.T) {
root, err := FindPrimitiveRoot(TestN, TestQ)
if err != nil {
t.Fatalf("FindPrimitiveRoot failed: %v", err)
}
t.Logf("Primitive 2×%d-th root of unity mod 0x%X: %d", TestN, TestQ, root)
if root == 0 || root >= TestQ {
t.Errorf("Invalid root: %d", root)
}
}
func TestModInverse(t *testing.T) {
// Test N^{-1} mod Q
inv, err := ModInverse(uint64(TestN), TestQ)
if err != nil {
t.Fatalf("ModInverse failed: %v", err)
}
// Verify: N * N^{-1} = 1 (mod Q)
product := (uint64(TestN) * inv) % TestQ
if product != 1 {
t.Errorf("ModInverse verification failed: %d * %d = %d (mod Q), want 1", TestN, inv, product)
}
t.Logf("%d^{-1} mod 0x%X = %d", TestN, TestQ, inv)
}
func TestPolyArithmetic(t *testing.T) {
N := uint32(TestN)
Q := TestQ
a := make([]uint64, N)
b := make([]uint64, N)
for i := range a {
a[i] = uint64(i+1) % Q
b[i] = uint64(i+2) % Q
}
// Test addition
sum, err := PolyAdd(a, b, Q)
if err != nil {
t.Fatalf("PolyAdd failed: %v", err)
}
for i := range sum {
expected := (a[i] + b[i]) % Q
if sum[i] != expected {
t.Errorf("PolyAdd[%d] = %d, want %d", i, sum[i], expected)
}
}
// Test subtraction
diff, err := PolySub(a, b, Q)
if err != nil {
t.Fatalf("PolySub failed: %v", err)
}
for i := range diff {
var expected uint64
if a[i] >= b[i] {
expected = a[i] - b[i]
} else {
expected = Q - b[i] + a[i]
}
if diff[i] != expected {
t.Errorf("PolySub[%d] = %d, want %d", i, diff[i], expected)
}
}
// Test scalar multiplication
scalar := uint64(5)
product, err := PolyScalarMul(a, scalar, Q)
if err != nil {
t.Fatalf("PolyScalarMul failed: %v", err)
}
for i := range product {
expected := (a[i] * scalar) % Q
if product[i] != expected {
t.Errorf("PolyScalarMul[%d] = %d, want %d", i, product[i], expected)
}
}
t.Log("Polynomial arithmetic tests passed")
}
func BenchmarkNTT(b *testing.B) {
ctx, err := NewNTTContext(TestN, TestQ)
if err != nil {
b.Fatalf("NewNTTContext failed: %v", err)
}
defer ctx.Close()
poly := make([]uint64, TestN)
for i := range poly {
poly[i] = uint64(i) % TestQ
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
_, err := ctx.NTT([][]uint64{poly})
if err != nil {
b.Fatal(err)
}
}
}
func BenchmarkPolyMul(b *testing.B) {
ctx, err := NewNTTContext(TestN, TestQ)
if err != nil {
b.Fatalf("NewNTTContext failed: %v", err)
}
defer ctx.Close()
a := make([]uint64, TestN)
c := make([]uint64, TestN)
for i := range a {
a[i] = uint64(i) % TestQ
c[i] = uint64(i+1) % TestQ
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
_, err := ctx.PolyMul([][]uint64{a}, [][]uint64{c})
if err != nil {
b.Fatal(err)
}
}
}