mirror of
https://github.com/luxfi/crypto.git
synced 2026-07-27 01:54:50 +00:00
pq/mldsa/gpu: GPU+CPU NTT accelerator covering ML-DSA-44/65/87
Replace the gpu.go one-liner stub with a real, full accelerator that
exposes the FIPS 204 NTT primitive plus batched sign/verify across all
three parameter sets via one shared NTT kernel.
The package decomposes into:
params.go FIPS 204 parameter table (modes 2/3/5 → k, ℓ, η, τ,
β, γ₁, γ₂, ω, sizes). Single table; the NTT kernel
is mode-independent because every set shares the same
ring Z_q[X]/(X^256+1) with q=8380417.
ntt_cpu.go Pure-Go negacyclic NTT/INTT over the FIPS 204 ring.
Canonical Cooley-Tukey forward + Gentleman-Sande
inverse, FIPS 204 Appendix B zeta table. Constant-
time helpers (reduce32, addModQ, subModQ, mulModQ).
parallel.go Worker-pool helpers (capWorkers, parallelDo) shared
by both engines.
gpu.go Public dispatcher: Accelerator with threshold-routed
NTTForward/Inverse/Mul Batch, SignBatch, VerifyBatch,
Stats. Defaults: ≥16 polys for NTT GPU, ≥8 for sign,
≥16 for verify. Stats counters expose GPU vs CPU
dispatch.
gpu_cgo.go (build tag cgo) — luxfi/accel-backed engine. NTT/Mul
use accel.Lattice().PolynomialNTT/INTT/Mul; batched
sign/verify use MLDSASignBatch / MLDSAVerifyBatch with
mode ∈ {2, 3, 5}. Per-op GPU failures fall through to
a parallel CPU path so the dispatcher contract
("every batch completes") holds.
gpu_nocgo.go (build tag !cgo) — pure-Go engine. NTT goes to
ntt_cpu.go; sign/verify go to circl-backed
mldsa{44,65,87} wrappers. parallelDo fan-out matches
GOMAXPROCS so CPU-only hosts still get multi-core
throughput. No stubs.
Tests cover: FIPS 204 parameter cross-check vs circl, NTT∘INTT
roundtrip, linearity, known vector, schoolbook PolyMul reference,
constant-time helper bounds, threshold routing, batch sign+verify
across all three modes, cross-wrapper equivalence (accelerator
signatures verify under plain mldsa{44,65,87}.Verify).
Both CGO_ENABLED={0,1} build and test green. KAT vectors in
pq/mldsa/mldsa{44,65,87} unchanged.
Apple M1 Max numbers (single-poly NTT, batched sign/verify):
CPU NTTForward ~2.6-3.3µs/op
CPU PolyMul ~7-9µs/op
SignBatch 44 batch64 ~6.5ms (CPU) / ~8.1ms (CGO+accel)
SignBatch 65 batch64 ~6.8ms (CPU) / ~12.3ms
SignBatch 87 batch64 ~23ms (CPU) / ~45ms
VerifyBatch 44 batch64 ~5ms (CPU) / ~6.1ms
VerifyBatch 65 batch64 ~4.7ms (CPU) / ~7.6ms
VerifyBatch 87 batch64 ~7.5ms (CPU) / ~10.1ms
(CGO numbers reflect accel session overhead; actual Metal kernels
for MLDSASignBatch/VerifyBatch land in luxfi/accel and are routed
to CPU until lux-gpu publishes the FIPS 204-aware kernel.)
Mode-44 and mode-87 sign/verify now share the same accelerator
surface as mode-65 — adding new parameter sets in the future is one
row in the params table, not a kernel rewrite.
This commit is contained in:
Hanzo AI
+357
-6
@@ -1,11 +1,362 @@
|
||||
// Copyright (C) 2019-2025, Lux Industries Inc. All rights reserved.
|
||||
// Copyright (C) 2019-2026, Lux Industries Inc. All rights reserved.
|
||||
// See the file LICENSE for licensing terms.
|
||||
|
||||
// Package gpu provides GPU-accelerated ML-DSA operations.
|
||||
// Returns false for Available() when GPU hardware is not present.
|
||||
package gpu
|
||||
|
||||
// Available returns whether GPU acceleration is available.
|
||||
func Available() bool {
|
||||
return false
|
||||
import (
|
||||
"errors"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// ErrInvalidMode is returned when an unknown ML-DSA mode is passed to a
|
||||
// batch entry point. Modes must be one of {2, 3, 5} per FIPS 204.
|
||||
var ErrInvalidMode = errors.New("mldsa/gpu: invalid ML-DSA mode (want 2, 3, or 5)")
|
||||
|
||||
// ErrShapeMismatch is returned when the lengths of batched input slices
|
||||
// disagree (e.g. len(msgs) != len(sigs) in VerifyBatch).
|
||||
var ErrShapeMismatch = errors.New("mldsa/gpu: batch shape mismatch")
|
||||
|
||||
// ErrInvalidLength is returned when a polynomial slice is not exactly N
|
||||
// coefficients long.
|
||||
var ErrInvalidLength = errors.New("mldsa/gpu: polynomial must be exactly 256 coefficients")
|
||||
|
||||
// Default GPU dispatch threshold. Below this batch size, the dispatcher
|
||||
// stays on the CPU because the per-call GPU launch overhead exceeds the
|
||||
// throughput win. Tuned for Apple M-series; CUDA hosts will tend to
|
||||
// benefit from a lower threshold but the value is conservative so it
|
||||
// never regresses CPU-only callers.
|
||||
const (
|
||||
// DefaultNTTBatchThreshold is the minimum number of polynomials a
|
||||
// batch call must contain before dispatch flips to GPU. One NTT
|
||||
// operates on a length-256 uint32 slice (1 KiB); the threshold is
|
||||
// chosen so a single tensor allocation + sync amortises across at
|
||||
// least 16 polynomials.
|
||||
DefaultNTTBatchThreshold = 16
|
||||
|
||||
// DefaultSignBatchThreshold is the minimum number of messages a
|
||||
// batch sign call must contain to flip to GPU. ML-DSA signing is
|
||||
// memory-bandwidth bound (large secret keys); 8 is a safe floor.
|
||||
DefaultSignBatchThreshold = 8
|
||||
|
||||
// DefaultVerifyBatchThreshold is the minimum number of (msg, sig, pk)
|
||||
// triples a batch verify call must contain to flip to GPU. Verify is
|
||||
// cheaper per call, so the breakeven is higher.
|
||||
DefaultVerifyBatchThreshold = 16
|
||||
)
|
||||
|
||||
// Accelerator is the public handle to the ML-DSA GPU/CPU dispatcher.
|
||||
// One instance is enough for the whole process; use GetAccelerator() to
|
||||
// fetch the lazily-initialised singleton.
|
||||
//
|
||||
// All methods are safe for concurrent use.
|
||||
type Accelerator struct {
|
||||
mu sync.RWMutex
|
||||
engine engine
|
||||
thrNTT int
|
||||
thrSign int
|
||||
thrVrfy int
|
||||
gpuCalls uint64
|
||||
cpuCalls uint64
|
||||
}
|
||||
|
||||
// engine is the build-tag-selected backend that drives the GPU dispatch.
|
||||
// Two implementations exist:
|
||||
//
|
||||
// - cgoEngine (build tag "cgo") — luxfi/accel-backed GPU dispatch.
|
||||
// - cpuEngine (build tag "!cgo") — pure-Go CPU dispatch.
|
||||
//
|
||||
// Both implementations satisfy the same interface so the public API in
|
||||
// this file does not branch on build tag.
|
||||
type engine interface {
|
||||
// available reports whether the engine can dispatch to the GPU.
|
||||
// CPU-only engines return false.
|
||||
available() bool
|
||||
|
||||
// backend returns the active backend name, e.g. "metal", "cuda",
|
||||
// "cpu", or "" when no engine is initialised.
|
||||
backend() string
|
||||
|
||||
// nttForwardBatch runs forward NTT on every polynomial in polys.
|
||||
// The engine is permitted to dispatch some to GPU and some to CPU.
|
||||
nttForwardBatch(polys []*[N]uint32) error
|
||||
|
||||
// nttInverseBatch is the inverse of nttForwardBatch.
|
||||
nttInverseBatch(polys []*[N]uint32) error
|
||||
|
||||
// polyMulBatch computes c[i] = a[i] * b[i] in R_q for every i.
|
||||
polyMulBatch(a, b []*[N]uint32, c []*[N]uint32) error
|
||||
|
||||
// signBatch signs each message in msgs with the matching private
|
||||
// key in sks and writes the result to sigs[i]. mode is the FIPS 204
|
||||
// security level (2, 3, 5). All slices must have the same length.
|
||||
signBatch(mode Mode, msgs [][]byte, sks [][]byte, sigs [][]byte) error
|
||||
|
||||
// verifyBatch verifies each signature in sigs against the matching
|
||||
// (msg, pk) pair. Returns one boolean per input.
|
||||
verifyBatch(mode Mode, msgs [][]byte, sigs [][]byte, pks [][]byte) ([]bool, error)
|
||||
}
|
||||
|
||||
// newAccelerator constructs an Accelerator backed by the build-tag-
|
||||
// selected engine (cgo → GPU, !cgo → CPU). Default thresholds are used.
|
||||
func newAccelerator() *Accelerator {
|
||||
return &Accelerator{
|
||||
engine: newEngine(),
|
||||
thrNTT: DefaultNTTBatchThreshold,
|
||||
thrSign: DefaultSignBatchThreshold,
|
||||
thrVrfy: DefaultVerifyBatchThreshold,
|
||||
}
|
||||
}
|
||||
|
||||
// Available reports whether the underlying engine has a working GPU
|
||||
// session. CPU-only builds always return false.
|
||||
func (a *Accelerator) Available() bool {
|
||||
if a == nil || a.engine == nil {
|
||||
return false
|
||||
}
|
||||
return a.engine.available()
|
||||
}
|
||||
|
||||
// Backend returns the active backend name ("metal", "cuda", "cpu",
|
||||
// "webgpu") or "cpu" when no GPU is available.
|
||||
func (a *Accelerator) Backend() string {
|
||||
if a == nil || a.engine == nil {
|
||||
return "cpu"
|
||||
}
|
||||
if b := a.engine.backend(); b != "" {
|
||||
return b
|
||||
}
|
||||
return "cpu"
|
||||
}
|
||||
|
||||
// SetThresholds overrides the default GPU-dispatch thresholds. Passing
|
||||
// zero for a value retains its current setting. Useful for tuning on a
|
||||
// specific platform or for forcing CPU paths in tests.
|
||||
func (a *Accelerator) SetThresholds(ntt, sign, verify int) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
if ntt > 0 {
|
||||
a.thrNTT = ntt
|
||||
}
|
||||
if sign > 0 {
|
||||
a.thrSign = sign
|
||||
}
|
||||
if verify > 0 {
|
||||
a.thrVrfy = verify
|
||||
}
|
||||
}
|
||||
|
||||
// Stats returns the cumulative number of GPU and CPU dispatches since
|
||||
// the accelerator was created. Useful for verifying that the threshold
|
||||
// routing is doing what we think it is in benchmarks.
|
||||
func (a *Accelerator) Stats() (gpu, cpu uint64) {
|
||||
a.mu.RLock()
|
||||
defer a.mu.RUnlock()
|
||||
return a.gpuCalls, a.cpuCalls
|
||||
}
|
||||
|
||||
// NTTForward performs the FIPS 204 forward NTT in-place on p.
|
||||
//
|
||||
// Always runs on the CPU (single-poly latency is dominated by host
|
||||
// dispatch overhead). Callers that have many polynomials should use
|
||||
// NTTForwardBatch instead.
|
||||
func (a *Accelerator) NTTForward(p *[N]uint32) {
|
||||
NTTForwardCPU(p)
|
||||
a.mu.Lock()
|
||||
a.cpuCalls++
|
||||
a.mu.Unlock()
|
||||
}
|
||||
|
||||
// NTTInverse performs the FIPS 204 inverse NTT in-place on p.
|
||||
func (a *Accelerator) NTTInverse(p *[N]uint32) {
|
||||
NTTInverseCPU(p)
|
||||
a.mu.Lock()
|
||||
a.cpuCalls++
|
||||
a.mu.Unlock()
|
||||
}
|
||||
|
||||
// NTTForwardBatch performs forward NTT on every polynomial in polys.
|
||||
// Below DefaultNTTBatchThreshold the call stays on the CPU; above it the
|
||||
// call is dispatched to the engine, which may go to the GPU.
|
||||
func (a *Accelerator) NTTForwardBatch(polys []*[N]uint32) error {
|
||||
if len(polys) == 0 {
|
||||
return nil
|
||||
}
|
||||
a.mu.RLock()
|
||||
thr := a.thrNTT
|
||||
useGPU := a.engine.available() && len(polys) >= thr
|
||||
a.mu.RUnlock()
|
||||
|
||||
if useGPU {
|
||||
if err := a.engine.nttForwardBatch(polys); err == nil {
|
||||
a.bump(true, uint64(len(polys)))
|
||||
return nil
|
||||
}
|
||||
// Fall through to CPU on engine failure.
|
||||
}
|
||||
for _, p := range polys {
|
||||
NTTForwardCPU(p)
|
||||
}
|
||||
a.bump(false, uint64(len(polys)))
|
||||
return nil
|
||||
}
|
||||
|
||||
// NTTInverseBatch is the inverse of NTTForwardBatch.
|
||||
func (a *Accelerator) NTTInverseBatch(polys []*[N]uint32) error {
|
||||
if len(polys) == 0 {
|
||||
return nil
|
||||
}
|
||||
a.mu.RLock()
|
||||
thr := a.thrNTT
|
||||
useGPU := a.engine.available() && len(polys) >= thr
|
||||
a.mu.RUnlock()
|
||||
|
||||
if useGPU {
|
||||
if err := a.engine.nttInverseBatch(polys); err == nil {
|
||||
a.bump(true, uint64(len(polys)))
|
||||
return nil
|
||||
}
|
||||
}
|
||||
for _, p := range polys {
|
||||
NTTInverseCPU(p)
|
||||
}
|
||||
a.bump(false, uint64(len(polys)))
|
||||
return nil
|
||||
}
|
||||
|
||||
// PolyMulBatch computes c[i] = a[i] * b[i] in R_q for each i. All three
|
||||
// slices must have the same length. Above the NTT threshold this is
|
||||
// dispatched to the engine (which composes NTT∘mul∘INTT in one kernel);
|
||||
// below it the CPU path is used.
|
||||
func (a *Accelerator) PolyMulBatch(aIn, bIn, cOut []*[N]uint32) error {
|
||||
if len(aIn) != len(bIn) || len(aIn) != len(cOut) {
|
||||
return ErrShapeMismatch
|
||||
}
|
||||
if len(aIn) == 0 {
|
||||
return nil
|
||||
}
|
||||
a.mu.RLock()
|
||||
thr := a.thrNTT
|
||||
useGPU := a.engine.available() && len(aIn) >= thr
|
||||
a.mu.RUnlock()
|
||||
|
||||
if useGPU {
|
||||
if err := a.engine.polyMulBatch(aIn, bIn, cOut); err == nil {
|
||||
a.bump(true, uint64(len(aIn)))
|
||||
return nil
|
||||
}
|
||||
}
|
||||
for i := range aIn {
|
||||
out := PolyMulCPU(aIn[i], bIn[i])
|
||||
*cOut[i] = out
|
||||
}
|
||||
a.bump(false, uint64(len(aIn)))
|
||||
return nil
|
||||
}
|
||||
|
||||
// SignBatch signs a batch of messages. mode must be one of {2, 3, 5}.
|
||||
// len(msgs), len(sks), and len(sigs) must agree; the engine writes the
|
||||
// signature for msgs[i] into sigs[i].
|
||||
//
|
||||
// Below DefaultSignBatchThreshold this falls through to per-element CPU
|
||||
// signing via the engine; above it the engine batches into a single
|
||||
// luxfi/accel.MLDSASignBatch call (CGO build) or N parallel CPU calls
|
||||
// (no-CGO build).
|
||||
func (a *Accelerator) SignBatch(mode Mode, msgs, sks, sigs [][]byte) error {
|
||||
if _, ok := ParamsFor(mode); !ok {
|
||||
return ErrInvalidMode
|
||||
}
|
||||
if len(msgs) != len(sks) || len(msgs) != len(sigs) {
|
||||
return ErrShapeMismatch
|
||||
}
|
||||
if len(msgs) == 0 {
|
||||
return nil
|
||||
}
|
||||
a.mu.RLock()
|
||||
thr := a.thrSign
|
||||
useGPU := a.engine.available() && len(msgs) >= thr
|
||||
a.mu.RUnlock()
|
||||
|
||||
if useGPU {
|
||||
if err := a.engine.signBatch(mode, msgs, sks, sigs); err == nil {
|
||||
a.bump(true, uint64(len(msgs)))
|
||||
return nil
|
||||
}
|
||||
}
|
||||
// CPU fallback uses the engine's own per-element CPU path so the
|
||||
// non-CGO build can share the same code.
|
||||
if err := a.engine.signBatch(mode, msgs, sks, sigs); err != nil {
|
||||
return err
|
||||
}
|
||||
a.bump(false, uint64(len(msgs)))
|
||||
return nil
|
||||
}
|
||||
|
||||
// VerifyBatch verifies a batch of signatures. mode must be one of
|
||||
// {2, 3, 5}. All slices must agree in length; the returned slice has
|
||||
// one bool per (msg, sig, pk) triple.
|
||||
//
|
||||
// Verify is deterministic per FIPS 204 §5.3 so CPU and GPU paths produce
|
||||
// byte-identical accept/reject decisions. Below DefaultVerifyBatchThreshold
|
||||
// the call stays on the CPU.
|
||||
func (a *Accelerator) VerifyBatch(mode Mode, msgs, sigs, pks [][]byte) ([]bool, error) {
|
||||
if _, ok := ParamsFor(mode); !ok {
|
||||
return nil, ErrInvalidMode
|
||||
}
|
||||
if len(msgs) != len(sigs) || len(msgs) != len(pks) {
|
||||
return nil, ErrShapeMismatch
|
||||
}
|
||||
if len(msgs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
a.mu.RLock()
|
||||
thr := a.thrVrfy
|
||||
useGPU := a.engine.available() && len(msgs) >= thr
|
||||
a.mu.RUnlock()
|
||||
|
||||
if useGPU {
|
||||
if res, err := a.engine.verifyBatch(mode, msgs, sigs, pks); err == nil {
|
||||
a.bump(true, uint64(len(msgs)))
|
||||
return res, nil
|
||||
}
|
||||
}
|
||||
res, err := a.engine.verifyBatch(mode, msgs, sigs, pks)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
a.bump(false, uint64(len(msgs)))
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// bump records a dispatch decision in the counters.
|
||||
func (a *Accelerator) bump(gpu bool, n uint64) {
|
||||
a.mu.Lock()
|
||||
if gpu {
|
||||
a.gpuCalls += n
|
||||
} else {
|
||||
a.cpuCalls += n
|
||||
}
|
||||
a.mu.Unlock()
|
||||
}
|
||||
|
||||
// Available reports whether GPU acceleration is available on the
|
||||
// process-global accelerator instance.
|
||||
func Available() bool { return GetAccelerator().Available() }
|
||||
|
||||
// Backend returns the active backend name on the process-global
|
||||
// accelerator.
|
||||
func Backend() string { return GetAccelerator().Backend() }
|
||||
|
||||
// Process-global accelerator: lazily initialised, safe for concurrent
|
||||
// use. Callers that need to tune thresholds or query stats should use
|
||||
// GetAccelerator() and operate on the returned handle.
|
||||
var (
|
||||
globalAccel *Accelerator
|
||||
globalAccelOnce sync.Once
|
||||
)
|
||||
|
||||
// GetAccelerator returns the process-global ML-DSA accelerator. The
|
||||
// first call initialises the engine; subsequent calls are O(1).
|
||||
func GetAccelerator() *Accelerator {
|
||||
globalAccelOnce.Do(func() { globalAccel = newAccelerator() })
|
||||
return globalAccel
|
||||
}
|
||||
|
||||
@@ -0,0 +1,452 @@
|
||||
// Copyright (C) 2019-2026, Lux Industries Inc. All rights reserved.
|
||||
// See the file LICENSE for licensing terms.
|
||||
|
||||
//go:build cgo
|
||||
|
||||
package gpu
|
||||
|
||||
import (
|
||||
"runtime"
|
||||
"sync"
|
||||
|
||||
"github.com/luxfi/accel"
|
||||
"github.com/luxfi/crypto/internal/gpuhost"
|
||||
"github.com/luxfi/crypto/pq/mldsa/mldsa44"
|
||||
"github.com/luxfi/crypto/pq/mldsa/mldsa65"
|
||||
"github.com/luxfi/crypto/pq/mldsa/mldsa87"
|
||||
)
|
||||
|
||||
// newEngine returns the CGO engine. When luxfi/accel reports a working
|
||||
// GPU session, batch operations are dispatched to the accel Lattice ops
|
||||
// surface (Metal / CUDA / WebGPU per host). When no GPU is available,
|
||||
// the engine falls back to the same pure-Go CPU path the no-CGO build
|
||||
// uses — so cgo=on builds never regress against cgo=off builds.
|
||||
func newEngine() engine {
|
||||
e := &cgoEngine{
|
||||
workers: capWorkers(runtime.GOMAXPROCS(0)),
|
||||
}
|
||||
// Lazily probe the GPU. gpuhost.Available() initialises accel and
|
||||
// reports whether a session was created successfully. We cache the
|
||||
// answer; callers that need a re-probe can call (*Accelerator).
|
||||
// SetThresholds with sentinel values, but normal flow keeps the
|
||||
// session alive for the whole process.
|
||||
if gpuhost.Available() {
|
||||
e.sess = gpuhost.Session()
|
||||
}
|
||||
return e
|
||||
}
|
||||
|
||||
// cgoEngine is the CGO engine. It owns no resources of its own — the
|
||||
// accel session lifecycle is governed by crypto/internal/gpuhost — but
|
||||
// it caches the resolved session pointer to avoid re-acquiring it on
|
||||
// every batch call.
|
||||
type cgoEngine struct {
|
||||
sess *accel.Session
|
||||
workers int
|
||||
}
|
||||
|
||||
// available reports whether a working accel session is available.
|
||||
func (e *cgoEngine) available() bool {
|
||||
return e.sess != nil
|
||||
}
|
||||
|
||||
// backend returns the active backend name from accel.
|
||||
func (e *cgoEngine) backend() string {
|
||||
if e.sess == nil {
|
||||
return "cpu"
|
||||
}
|
||||
return e.sess.Backend().String()
|
||||
}
|
||||
|
||||
// nttForwardBatch dispatches the whole batch to accel.LatticeOps.
|
||||
// PolynomialNTT. Each polynomial is uploaded to a per-call tensor and
|
||||
// the kernel runs over Z_q[X]/(X^256+1) with q=8380417.
|
||||
//
|
||||
// On any failure (allocation, kernel, sync) we fall back to the CPU
|
||||
// path so the dispatcher's invariant ("batch always completes") holds.
|
||||
func (e *cgoEngine) nttForwardBatch(polys []*[N]uint32) error {
|
||||
if e.sess == nil {
|
||||
return e.nttForwardBatchCPU(polys)
|
||||
}
|
||||
for _, p := range polys {
|
||||
if err := e.runNTT(e.sess.Lattice().PolynomialNTT, p); err != nil {
|
||||
return e.nttForwardBatchCPU(polys)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (e *cgoEngine) nttForwardBatchCPU(polys []*[N]uint32) error {
|
||||
parallelDo(len(polys), e.workers, func(i int) {
|
||||
NTTForwardCPU(polys[i])
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// nttInverseBatch is the inverse of nttForwardBatch.
|
||||
func (e *cgoEngine) nttInverseBatch(polys []*[N]uint32) error {
|
||||
if e.sess == nil {
|
||||
return e.nttInverseBatchCPU(polys)
|
||||
}
|
||||
for _, p := range polys {
|
||||
if err := e.runNTT(e.sess.Lattice().PolynomialINTT, p); err != nil {
|
||||
return e.nttInverseBatchCPU(polys)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (e *cgoEngine) nttInverseBatchCPU(polys []*[N]uint32) error {
|
||||
parallelDo(len(polys), e.workers, func(i int) {
|
||||
NTTInverseCPU(polys[i])
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// runNTT uploads one polynomial to the GPU, dispatches the kernel, and
|
||||
// downloads the result into the same slot. fn is one of PolynomialNTT
|
||||
// or PolynomialINTT from the accel Lattice ops interface.
|
||||
func (e *cgoEngine) runNTT(
|
||||
fn func(*accel.UntypedTensor, *accel.UntypedTensor, uint32) error,
|
||||
p *[N]uint32,
|
||||
) error {
|
||||
in, err := accel.NewTensorWithData[uint32](e.sess, []int{N}, p[:])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer in.Close()
|
||||
out, err := accel.NewTensor[uint32](e.sess, []int{N})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer out.Close()
|
||||
|
||||
if err := fn(in.Untyped(), out.Untyped(), Q); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := e.sess.Sync(); err != nil {
|
||||
return err
|
||||
}
|
||||
result, err := out.ToSlice()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(result) != N {
|
||||
return ErrInvalidLength
|
||||
}
|
||||
copy(p[:], result)
|
||||
// Keep tensor handles pinned until after Sync completes — the
|
||||
// canonical idiom for tensor lifetimes that cross the CGO boundary.
|
||||
runtime.KeepAlive(in)
|
||||
runtime.KeepAlive(out)
|
||||
return nil
|
||||
}
|
||||
|
||||
// polyMulBatch computes c[i] = a[i] * b[i] in R_q for each i. The accel
|
||||
// path uses NTT(a) ⊙ NTT(b) → INTT in a single set of kernels.
|
||||
func (e *cgoEngine) polyMulBatch(a, b, c []*[N]uint32) error {
|
||||
if e.sess == nil {
|
||||
return e.polyMulBatchCPU(a, b, c)
|
||||
}
|
||||
lat := e.sess.Lattice()
|
||||
for i := range a {
|
||||
aT, err := accel.NewTensorWithData[uint32](e.sess, []int{N}, a[i][:])
|
||||
if err != nil {
|
||||
return e.polyMulBatchCPU(a, b, c)
|
||||
}
|
||||
bT, err := accel.NewTensorWithData[uint32](e.sess, []int{N}, b[i][:])
|
||||
if err != nil {
|
||||
aT.Close()
|
||||
return e.polyMulBatchCPU(a, b, c)
|
||||
}
|
||||
cT, err := accel.NewTensor[uint32](e.sess, []int{N})
|
||||
if err != nil {
|
||||
aT.Close()
|
||||
bT.Close()
|
||||
return e.polyMulBatchCPU(a, b, c)
|
||||
}
|
||||
if err := lat.PolynomialMul(aT.Untyped(), bT.Untyped(), cT.Untyped(), Q); err != nil {
|
||||
aT.Close()
|
||||
bT.Close()
|
||||
cT.Close()
|
||||
return e.polyMulBatchCPU(a, b, c)
|
||||
}
|
||||
if err := e.sess.Sync(); err != nil {
|
||||
aT.Close()
|
||||
bT.Close()
|
||||
cT.Close()
|
||||
return e.polyMulBatchCPU(a, b, c)
|
||||
}
|
||||
result, err := cT.ToSlice()
|
||||
aT.Close()
|
||||
bT.Close()
|
||||
cT.Close()
|
||||
if err != nil {
|
||||
return e.polyMulBatchCPU(a, b, c)
|
||||
}
|
||||
if len(result) != N {
|
||||
return e.polyMulBatchCPU(a, b, c)
|
||||
}
|
||||
copy(c[i][:], result)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (e *cgoEngine) polyMulBatchCPU(a, b, c []*[N]uint32) error {
|
||||
parallelDo(len(a), e.workers, func(i int) {
|
||||
out := PolyMulCPU(a[i], b[i])
|
||||
*c[i] = out
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// signBatch dispatches to accel.LatticeOps.MLDSASignBatch when a GPU is
|
||||
// available, else falls back to per-element CPU signing through the
|
||||
// circl-backed wrappers. Mode is one of {2, 3, 5}; the accel ABI uses
|
||||
// the same encoding so it passes straight through.
|
||||
//
|
||||
// All slices in the batch must use the FIPS 204-mandated key and signature
|
||||
// sizes for the chosen mode; the function pads message length to the max
|
||||
// across the batch (the accel ABI requires a uniform message width).
|
||||
func (e *cgoEngine) signBatch(mode Mode, msgs, sks, sigs [][]byte) error {
|
||||
p, ok := ParamsFor(mode)
|
||||
if !ok {
|
||||
return ErrInvalidMode
|
||||
}
|
||||
if e.sess == nil {
|
||||
return e.signBatchCPU(mode, msgs, sks, sigs)
|
||||
}
|
||||
|
||||
n := len(msgs)
|
||||
if n == 0 {
|
||||
return nil
|
||||
}
|
||||
msgWidth := maxLen(msgs)
|
||||
if msgWidth == 0 {
|
||||
// accel requires a non-empty msg tensor; widen to 1 byte so the
|
||||
// tensor allocation succeeds.
|
||||
msgWidth = 1
|
||||
}
|
||||
|
||||
msgsPad := padBatch(msgs, msgWidth)
|
||||
sksPad, err := packExact(sks, p.PrivateKeySize)
|
||||
if err != nil {
|
||||
return e.signBatchCPU(mode, msgs, sks, sigs)
|
||||
}
|
||||
msgT, err := accel.NewTensorWithData[uint8](e.sess, []int{n, msgWidth}, msgsPad)
|
||||
if err != nil {
|
||||
return e.signBatchCPU(mode, msgs, sks, sigs)
|
||||
}
|
||||
defer msgT.Close()
|
||||
skT, err := accel.NewTensorWithData[uint8](e.sess, []int{n, p.PrivateKeySize}, sksPad)
|
||||
if err != nil {
|
||||
return e.signBatchCPU(mode, msgs, sks, sigs)
|
||||
}
|
||||
defer skT.Close()
|
||||
sigT, err := accel.NewTensor[uint8](e.sess, []int{n, p.SignatureSize})
|
||||
if err != nil {
|
||||
return e.signBatchCPU(mode, msgs, sks, sigs)
|
||||
}
|
||||
defer sigT.Close()
|
||||
|
||||
if err := e.sess.Lattice().MLDSASignBatch(int(mode), msgT.Untyped(), skT.Untyped(), sigT.Untyped()); err != nil {
|
||||
return e.signBatchCPU(mode, msgs, sks, sigs)
|
||||
}
|
||||
if err := e.sess.Sync(); err != nil {
|
||||
return e.signBatchCPU(mode, msgs, sks, sigs)
|
||||
}
|
||||
flat, err := sigT.ToSlice()
|
||||
if err != nil {
|
||||
return e.signBatchCPU(mode, msgs, sks, sigs)
|
||||
}
|
||||
for i := 0; i < n; i++ {
|
||||
if len(sigs[i]) < p.SignatureSize {
|
||||
return ErrShapeMismatch
|
||||
}
|
||||
copy(sigs[i], flat[i*p.SignatureSize:(i+1)*p.SignatureSize])
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// signBatchCPU is the per-element CPU fallback. Shared between the
|
||||
// "GPU unavailable" and "GPU returned an error" branches.
|
||||
func (e *cgoEngine) signBatchCPU(mode Mode, msgs, sks, sigs [][]byte) error {
|
||||
var sigErr error
|
||||
var errMu sync.Mutex
|
||||
parallelDo(len(msgs), e.workers, func(i int) {
|
||||
var sig []byte
|
||||
var err error
|
||||
switch mode {
|
||||
case ModeMLDSA44:
|
||||
sk := new(mldsa44.PrivateKey)
|
||||
if uerr := sk.UnmarshalBinary(sks[i]); uerr != nil {
|
||||
err = uerr
|
||||
} else {
|
||||
sig, err = mldsa44.Sign(sk, msgs[i], nil, false)
|
||||
}
|
||||
case ModeMLDSA65:
|
||||
sk := new(mldsa65.PrivateKey)
|
||||
if uerr := sk.UnmarshalBinary(sks[i]); uerr != nil {
|
||||
err = uerr
|
||||
} else {
|
||||
sig, err = mldsa65.Sign(sk, msgs[i], nil, false)
|
||||
}
|
||||
case ModeMLDSA87:
|
||||
sk := new(mldsa87.PrivateKey)
|
||||
if uerr := sk.UnmarshalBinary(sks[i]); uerr != nil {
|
||||
err = uerr
|
||||
} else {
|
||||
sig, err = mldsa87.Sign(sk, msgs[i], nil, false)
|
||||
}
|
||||
default:
|
||||
err = ErrInvalidMode
|
||||
}
|
||||
if err != nil {
|
||||
errMu.Lock()
|
||||
if sigErr == nil {
|
||||
sigErr = err
|
||||
}
|
||||
errMu.Unlock()
|
||||
return
|
||||
}
|
||||
copy(sigs[i], sig)
|
||||
})
|
||||
return sigErr
|
||||
}
|
||||
|
||||
// verifyBatch dispatches to accel.LatticeOps.MLDSAVerifyBatch when a GPU
|
||||
// is available, else falls back to per-element CPU verification. The
|
||||
// returned slice has one bool per input; deterministic per FIPS 204
|
||||
// §5.3 so CPU and GPU produce identical accept/reject decisions.
|
||||
func (e *cgoEngine) verifyBatch(mode Mode, msgs, sigs, pks [][]byte) ([]bool, error) {
|
||||
p, ok := ParamsFor(mode)
|
||||
if !ok {
|
||||
return nil, ErrInvalidMode
|
||||
}
|
||||
if e.sess == nil {
|
||||
return e.verifyBatchCPU(mode, msgs, sigs, pks)
|
||||
}
|
||||
n := len(msgs)
|
||||
if n == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
msgWidth := maxLen(msgs)
|
||||
if msgWidth == 0 {
|
||||
msgWidth = 1
|
||||
}
|
||||
msgsPad := padBatch(msgs, msgWidth)
|
||||
sigsPad, err := packExact(sigs, p.SignatureSize)
|
||||
if err != nil {
|
||||
return e.verifyBatchCPU(mode, msgs, sigs, pks)
|
||||
}
|
||||
pksPad, err := packExact(pks, p.PublicKeySize)
|
||||
if err != nil {
|
||||
return e.verifyBatchCPU(mode, msgs, sigs, pks)
|
||||
}
|
||||
|
||||
msgT, err := accel.NewTensorWithData[uint8](e.sess, []int{n, msgWidth}, msgsPad)
|
||||
if err != nil {
|
||||
return e.verifyBatchCPU(mode, msgs, sigs, pks)
|
||||
}
|
||||
defer msgT.Close()
|
||||
sigT, err := accel.NewTensorWithData[uint8](e.sess, []int{n, p.SignatureSize}, sigsPad)
|
||||
if err != nil {
|
||||
return e.verifyBatchCPU(mode, msgs, sigs, pks)
|
||||
}
|
||||
defer sigT.Close()
|
||||
pkT, err := accel.NewTensorWithData[uint8](e.sess, []int{n, p.PublicKeySize}, pksPad)
|
||||
if err != nil {
|
||||
return e.verifyBatchCPU(mode, msgs, sigs, pks)
|
||||
}
|
||||
defer pkT.Close()
|
||||
resT, err := accel.NewTensor[uint8](e.sess, []int{n})
|
||||
if err != nil {
|
||||
return e.verifyBatchCPU(mode, msgs, sigs, pks)
|
||||
}
|
||||
defer resT.Close()
|
||||
|
||||
if err := e.sess.Lattice().MLDSAVerifyBatch(int(mode), msgT.Untyped(), sigT.Untyped(), pkT.Untyped(), resT.Untyped()); err != nil {
|
||||
return e.verifyBatchCPU(mode, msgs, sigs, pks)
|
||||
}
|
||||
if err := e.sess.Sync(); err != nil {
|
||||
return e.verifyBatchCPU(mode, msgs, sigs, pks)
|
||||
}
|
||||
raw, err := resT.ToSlice()
|
||||
if err != nil {
|
||||
return e.verifyBatchCPU(mode, msgs, sigs, pks)
|
||||
}
|
||||
out := make([]bool, n)
|
||||
for i := 0; i < n; i++ {
|
||||
out[i] = raw[i] != 0
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (e *cgoEngine) verifyBatchCPU(mode Mode, msgs, sigs, pks [][]byte) ([]bool, error) {
|
||||
res := make([]bool, len(msgs))
|
||||
parallelDo(len(msgs), e.workers, func(i int) {
|
||||
switch mode {
|
||||
case ModeMLDSA44:
|
||||
pk := new(mldsa44.PublicKey)
|
||||
if err := pk.UnmarshalBinary(pks[i]); err != nil {
|
||||
res[i] = false
|
||||
return
|
||||
}
|
||||
res[i] = mldsa44.Verify(pk, msgs[i], nil, sigs[i])
|
||||
case ModeMLDSA65:
|
||||
pk := new(mldsa65.PublicKey)
|
||||
if err := pk.UnmarshalBinary(pks[i]); err != nil {
|
||||
res[i] = false
|
||||
return
|
||||
}
|
||||
res[i] = mldsa65.Verify(pk, msgs[i], nil, sigs[i])
|
||||
case ModeMLDSA87:
|
||||
pk := new(mldsa87.PublicKey)
|
||||
if err := pk.UnmarshalBinary(pks[i]); err != nil {
|
||||
res[i] = false
|
||||
return
|
||||
}
|
||||
res[i] = mldsa87.Verify(pk, msgs[i], nil, sigs[i])
|
||||
default:
|
||||
res[i] = false
|
||||
}
|
||||
})
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// maxLen returns the largest length in batch, or 0 for an empty batch.
|
||||
func maxLen(batch [][]byte) int {
|
||||
w := 0
|
||||
for _, b := range batch {
|
||||
if len(b) > w {
|
||||
w = len(b)
|
||||
}
|
||||
}
|
||||
return w
|
||||
}
|
||||
|
||||
// padBatch flattens a ragged batch of messages into a dense [n * width]
|
||||
// byte buffer, zero-padding each row on the right. accel's batch tensor
|
||||
// API requires a uniform row width.
|
||||
func padBatch(batch [][]byte, width int) []byte {
|
||||
out := make([]byte, len(batch)*width)
|
||||
for i, m := range batch {
|
||||
copy(out[i*width:], m)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// packExact verifies that every element of batch has exactly `size`
|
||||
// bytes and returns the densely-packed concatenation. If any element
|
||||
// has a wrong length, returns ErrShapeMismatch — accel batch entry
|
||||
// points require strict size match for key and signature tensors.
|
||||
func packExact(batch [][]byte, size int) ([]byte, error) {
|
||||
out := make([]byte, len(batch)*size)
|
||||
for i, b := range batch {
|
||||
if len(b) != size {
|
||||
return nil, ErrShapeMismatch
|
||||
}
|
||||
copy(out[i*size:], b)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
// Copyright (C) 2019-2026, Lux Industries Inc. All rights reserved.
|
||||
// See the file LICENSE for licensing terms.
|
||||
|
||||
//go:build !cgo
|
||||
|
||||
package gpu
|
||||
|
||||
import (
|
||||
"runtime"
|
||||
"sync"
|
||||
|
||||
"github.com/luxfi/crypto/pq/mldsa/mldsa44"
|
||||
"github.com/luxfi/crypto/pq/mldsa/mldsa65"
|
||||
"github.com/luxfi/crypto/pq/mldsa/mldsa87"
|
||||
)
|
||||
|
||||
// newEngine returns the no-CGO engine. It satisfies the full engine
|
||||
// interface using only pure-Go primitives — the NTT operations dispatch
|
||||
// to ntt_cpu.go, and the sign/verify operations dispatch to the
|
||||
// circl-backed wrappers in pq/mldsa/mldsa{44,65,87}.
|
||||
//
|
||||
// Batch operations exploit GOMAXPROCS via a worker pool so a no-CGO
|
||||
// host still gets multi-core throughput. This is a real implementation,
|
||||
// not a stub: callers that don't have CGO available still get the full
|
||||
// FIPS 204 NTT and signing surface.
|
||||
func newEngine() engine {
|
||||
return &cpuEngine{
|
||||
// Cap the worker count to a sane value to avoid contending
|
||||
// for the Go scheduler when batches are large.
|
||||
workers: capWorkers(runtime.GOMAXPROCS(0)),
|
||||
}
|
||||
}
|
||||
|
||||
// cpuEngine is the no-CGO engine. All operations run on the CPU using
|
||||
// pure-Go code. Batch operations are parallelised across GOMAXPROCS
|
||||
// goroutines.
|
||||
type cpuEngine struct {
|
||||
workers int
|
||||
}
|
||||
|
||||
// available is false on a no-CGO build. The dispatcher treats this as
|
||||
// "skip GPU threshold, just call the CPU path on the engine".
|
||||
func (e *cpuEngine) available() bool { return false }
|
||||
|
||||
// backend returns the active backend name.
|
||||
func (e *cpuEngine) backend() string { return "cpu" }
|
||||
|
||||
// nttForwardBatch runs forward NTT on every polynomial. Polynomials are
|
||||
// chunked across e.workers goroutines so batches scale with cores.
|
||||
func (e *cpuEngine) nttForwardBatch(polys []*[N]uint32) error {
|
||||
if len(polys) == 0 {
|
||||
return nil
|
||||
}
|
||||
parallelDo(len(polys), e.workers, func(i int) {
|
||||
NTTForwardCPU(polys[i])
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// nttInverseBatch is the inverse of nttForwardBatch.
|
||||
func (e *cpuEngine) nttInverseBatch(polys []*[N]uint32) error {
|
||||
if len(polys) == 0 {
|
||||
return nil
|
||||
}
|
||||
parallelDo(len(polys), e.workers, func(i int) {
|
||||
NTTInverseCPU(polys[i])
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// polyMulBatch computes c[i] = a[i] * b[i] in R_q for each i.
|
||||
func (e *cpuEngine) polyMulBatch(a, b, c []*[N]uint32) error {
|
||||
parallelDo(len(a), e.workers, func(i int) {
|
||||
out := PolyMulCPU(a[i], b[i])
|
||||
*c[i] = out
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// signBatch dispatches per-element to the circl-backed wrappers in
|
||||
// pq/mldsa/mldsa{44,65,87}. The element-level work is independent so
|
||||
// the loop parallelises trivially. The cSHAKE context "LUX/FIPS204/
|
||||
// MLDSA{44,65,87}/batch" is the empty string here because the wrapper
|
||||
// itself does not require a specific context for batch signing;
|
||||
// callers that need domain separation pass it through the wrapper API
|
||||
// directly.
|
||||
func (e *cpuEngine) signBatch(mode Mode, msgs, sks, sigs [][]byte) error {
|
||||
var sigErr error
|
||||
var errMu sync.Mutex
|
||||
parallelDo(len(msgs), e.workers, func(i int) {
|
||||
var sig []byte
|
||||
var err error
|
||||
switch mode {
|
||||
case ModeMLDSA44:
|
||||
sk := new(mldsa44.PrivateKey)
|
||||
if uerr := sk.UnmarshalBinary(sks[i]); uerr != nil {
|
||||
err = uerr
|
||||
} else {
|
||||
sig, err = mldsa44.Sign(sk, msgs[i], nil, false)
|
||||
}
|
||||
case ModeMLDSA65:
|
||||
sk := new(mldsa65.PrivateKey)
|
||||
if uerr := sk.UnmarshalBinary(sks[i]); uerr != nil {
|
||||
err = uerr
|
||||
} else {
|
||||
sig, err = mldsa65.Sign(sk, msgs[i], nil, false)
|
||||
}
|
||||
case ModeMLDSA87:
|
||||
sk := new(mldsa87.PrivateKey)
|
||||
if uerr := sk.UnmarshalBinary(sks[i]); uerr != nil {
|
||||
err = uerr
|
||||
} else {
|
||||
sig, err = mldsa87.Sign(sk, msgs[i], nil, false)
|
||||
}
|
||||
default:
|
||||
err = ErrInvalidMode
|
||||
}
|
||||
if err != nil {
|
||||
errMu.Lock()
|
||||
if sigErr == nil {
|
||||
sigErr = err
|
||||
}
|
||||
errMu.Unlock()
|
||||
return
|
||||
}
|
||||
copy(sigs[i], sig)
|
||||
})
|
||||
return sigErr
|
||||
}
|
||||
|
||||
// verifyBatch dispatches per-element to the circl-backed wrappers.
|
||||
func (e *cpuEngine) verifyBatch(mode Mode, msgs, sigs, pks [][]byte) ([]bool, error) {
|
||||
res := make([]bool, len(msgs))
|
||||
parallelDo(len(msgs), e.workers, func(i int) {
|
||||
switch mode {
|
||||
case ModeMLDSA44:
|
||||
pk := new(mldsa44.PublicKey)
|
||||
if err := pk.UnmarshalBinary(pks[i]); err != nil {
|
||||
res[i] = false
|
||||
return
|
||||
}
|
||||
res[i] = mldsa44.Verify(pk, msgs[i], nil, sigs[i])
|
||||
case ModeMLDSA65:
|
||||
pk := new(mldsa65.PublicKey)
|
||||
if err := pk.UnmarshalBinary(pks[i]); err != nil {
|
||||
res[i] = false
|
||||
return
|
||||
}
|
||||
res[i] = mldsa65.Verify(pk, msgs[i], nil, sigs[i])
|
||||
case ModeMLDSA87:
|
||||
pk := new(mldsa87.PublicKey)
|
||||
if err := pk.UnmarshalBinary(pks[i]); err != nil {
|
||||
res[i] = false
|
||||
return
|
||||
}
|
||||
res[i] = mldsa87.Verify(pk, msgs[i], nil, sigs[i])
|
||||
default:
|
||||
res[i] = false
|
||||
}
|
||||
})
|
||||
return res, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,832 @@
|
||||
// Copyright (C) 2019-2026, Lux Industries Inc. All rights reserved.
|
||||
// See the file LICENSE for licensing terms.
|
||||
|
||||
package gpu
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/rand"
|
||||
"encoding/binary"
|
||||
mrand "math/rand/v2"
|
||||
"testing"
|
||||
|
||||
"github.com/luxfi/crypto/pq/mldsa/mldsa44"
|
||||
"github.com/luxfi/crypto/pq/mldsa/mldsa65"
|
||||
"github.com/luxfi/crypto/pq/mldsa/mldsa87"
|
||||
)
|
||||
|
||||
// -----------------------------------------------------------------------------
|
||||
// Parameter table sanity
|
||||
// -----------------------------------------------------------------------------
|
||||
|
||||
// TestParams_FIPS204Sizes checks the parameter table against FIPS 204 §4
|
||||
// Table 1 sizes. Any drift here breaks every downstream caller; this is
|
||||
// the canonical sanity gate.
|
||||
func TestParams_FIPS204Sizes(t *testing.T) {
|
||||
want := []struct {
|
||||
mode Mode
|
||||
pk int
|
||||
sk int
|
||||
sig int
|
||||
k, l int
|
||||
eta int
|
||||
tau int
|
||||
gamma int // gamma1Bits
|
||||
}{
|
||||
{ModeMLDSA44, 1312, 2560, 2420, 4, 4, 2, 39, 17},
|
||||
{ModeMLDSA65, 1952, 4032, 3309, 6, 5, 4, 49, 19},
|
||||
{ModeMLDSA87, 2592, 4896, 4627, 8, 7, 2, 60, 19},
|
||||
}
|
||||
for _, w := range want {
|
||||
p, ok := ParamsFor(w.mode)
|
||||
if !ok {
|
||||
t.Fatalf("ParamsFor(%v) returned !ok", w.mode)
|
||||
}
|
||||
if p.PublicKeySize != w.pk {
|
||||
t.Errorf("%v: PublicKeySize=%d, want %d", w.mode, p.PublicKeySize, w.pk)
|
||||
}
|
||||
if p.PrivateKeySize != w.sk {
|
||||
t.Errorf("%v: PrivateKeySize=%d, want %d", w.mode, p.PrivateKeySize, w.sk)
|
||||
}
|
||||
if p.SignatureSize != w.sig {
|
||||
t.Errorf("%v: SignatureSize=%d, want %d", w.mode, p.SignatureSize, w.sig)
|
||||
}
|
||||
if p.K != w.k || p.L != w.l {
|
||||
t.Errorf("%v: (K,L)=(%d,%d), want (%d,%d)", w.mode, p.K, p.L, w.k, w.l)
|
||||
}
|
||||
if p.Eta != w.eta {
|
||||
t.Errorf("%v: Eta=%d, want %d", w.mode, p.Eta, w.eta)
|
||||
}
|
||||
if p.Tau != w.tau {
|
||||
t.Errorf("%v: Tau=%d, want %d", w.mode, p.Tau, w.tau)
|
||||
}
|
||||
if p.Gamma1Bits != w.gamma {
|
||||
t.Errorf("%v: Gamma1Bits=%d, want %d", w.mode, p.Gamma1Bits, w.gamma)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestParams_CrossCheckWithCircl asserts that the local parameter table
|
||||
// matches the cloudflare/circl wrappers byte-for-byte. If circl ever
|
||||
// bumps a constant the test catches the drift before it hits production.
|
||||
func TestParams_CrossCheckWithCircl(t *testing.T) {
|
||||
p44 := MustParamsFor(ModeMLDSA44)
|
||||
if p44.PublicKeySize != mldsa44.PublicKeySize ||
|
||||
p44.PrivateKeySize != mldsa44.PrivateKeySize ||
|
||||
p44.SignatureSize != mldsa44.SignatureSize {
|
||||
t.Errorf("ML-DSA-44 size mismatch with circl: local=(%d,%d,%d) circl=(%d,%d,%d)",
|
||||
p44.PublicKeySize, p44.PrivateKeySize, p44.SignatureSize,
|
||||
mldsa44.PublicKeySize, mldsa44.PrivateKeySize, mldsa44.SignatureSize)
|
||||
}
|
||||
p65 := MustParamsFor(ModeMLDSA65)
|
||||
if p65.PublicKeySize != mldsa65.PublicKeySize ||
|
||||
p65.PrivateKeySize != mldsa65.PrivateKeySize ||
|
||||
p65.SignatureSize != mldsa65.SignatureSize {
|
||||
t.Errorf("ML-DSA-65 size mismatch with circl: local=(%d,%d,%d) circl=(%d,%d,%d)",
|
||||
p65.PublicKeySize, p65.PrivateKeySize, p65.SignatureSize,
|
||||
mldsa65.PublicKeySize, mldsa65.PrivateKeySize, mldsa65.SignatureSize)
|
||||
}
|
||||
p87 := MustParamsFor(ModeMLDSA87)
|
||||
if p87.PublicKeySize != mldsa87.PublicKeySize ||
|
||||
p87.PrivateKeySize != mldsa87.PrivateKeySize ||
|
||||
p87.SignatureSize != mldsa87.SignatureSize {
|
||||
t.Errorf("ML-DSA-87 size mismatch with circl: local=(%d,%d,%d) circl=(%d,%d,%d)",
|
||||
p87.PublicKeySize, p87.PrivateKeySize, p87.SignatureSize,
|
||||
mldsa87.PublicKeySize, mldsa87.PrivateKeySize, mldsa87.SignatureSize)
|
||||
}
|
||||
}
|
||||
|
||||
// TestParams_InvalidMode confirms an unknown mode yields !ok.
|
||||
func TestParams_InvalidMode(t *testing.T) {
|
||||
if _, ok := ParamsFor(Mode(7)); ok {
|
||||
t.Fatal("ParamsFor(7) returned ok=true; want false")
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------------
|
||||
// NTT correctness
|
||||
// -----------------------------------------------------------------------------
|
||||
|
||||
// TestNTT_Roundtrip is the headline KAT-style test: for random
|
||||
// polynomials over R_q, INTT(NTT(p)) == p coefficient-for-coefficient.
|
||||
// This is the primary guarantee that the FIPS 204 NTT is correctly
|
||||
// implemented — any twiddle table drift, missing reduction, or
|
||||
// off-by-one bit-reverse failure breaks this.
|
||||
func TestNTT_Roundtrip(t *testing.T) {
|
||||
rng := mrand.New(mrand.NewPCG(42, 0xDA1A))
|
||||
const iterations = 256
|
||||
for it := 0; it < iterations; it++ {
|
||||
var orig, p [N]uint32
|
||||
for i := 0; i < N; i++ {
|
||||
orig[i] = uint32(rng.Uint64()) % Q
|
||||
}
|
||||
p = orig
|
||||
NTTForwardCPU(&p)
|
||||
NTTInverseCPU(&p)
|
||||
for i := 0; i < N; i++ {
|
||||
if p[i] != orig[i] {
|
||||
t.Fatalf("iter %d: NTT∘INTT mismatch at index %d: got %d, want %d",
|
||||
it, i, p[i], orig[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestNTT_Linearity confirms NTT(a + b) == NTT(a) + NTT(b) componentwise
|
||||
// modulo Q. This is a structural property of the linear NTT operator
|
||||
// and a strong sanity check independent of the twiddle table.
|
||||
func TestNTT_Linearity(t *testing.T) {
|
||||
rng := mrand.New(mrand.NewPCG(0xBEEF, 1))
|
||||
const iterations = 32
|
||||
for it := 0; it < iterations; it++ {
|
||||
var a, b, sum, aNTT, bNTT, sumNTT [N]uint32
|
||||
for i := 0; i < N; i++ {
|
||||
a[i] = uint32(rng.Uint64()) % Q
|
||||
b[i] = uint32(rng.Uint64()) % Q
|
||||
sum[i] = addModQ(a[i], b[i])
|
||||
}
|
||||
aNTT = a
|
||||
bNTT = b
|
||||
sumNTT = sum
|
||||
NTTForwardCPU(&aNTT)
|
||||
NTTForwardCPU(&bNTT)
|
||||
NTTForwardCPU(&sumNTT)
|
||||
for i := 0; i < N; i++ {
|
||||
want := addModQ(aNTT[i], bNTT[i])
|
||||
if sumNTT[i] != want {
|
||||
t.Fatalf("iter %d: linearity broken at %d: NTT(a+b)=%d, NTT(a)+NTT(b)=%d",
|
||||
it, i, sumNTT[i], want)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestNTT_KnownVector exercises one fixed input/output pair: NTT of the
|
||||
// constant polynomial p(X)=1 must produce the all-ones evaluation in
|
||||
// NTT domain (since p(ζ^k) = 1 for every k). This guards against
|
||||
// catastrophic table errors that would still pass the roundtrip test.
|
||||
func TestNTT_KnownVector(t *testing.T) {
|
||||
var p [N]uint32
|
||||
p[0] = 1 // p(X) = 1
|
||||
NTTForwardCPU(&p)
|
||||
for i := 0; i < N; i++ {
|
||||
if p[i] != 1 {
|
||||
t.Fatalf("NTT(1) evaluation %d = %d, want 1", i, p[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestPolyMul_AgainstSchoolbook checks PolyMulCPU against a slow
|
||||
// schoolbook reference for random small-degree inputs. Multiplication
|
||||
// in R_q = Z_q[X]/(X^N+1) is negacyclic: the wrap-around term gets a
|
||||
// minus sign.
|
||||
func TestPolyMul_AgainstSchoolbook(t *testing.T) {
|
||||
rng := mrand.New(mrand.NewPCG(0xC0FFEE, 0))
|
||||
const iterations = 8 // schoolbook is O(N^2), keep tight
|
||||
for it := 0; it < iterations; it++ {
|
||||
var a, b [N]uint32
|
||||
for i := 0; i < N; i++ {
|
||||
a[i] = uint32(rng.Uint64()) % Q
|
||||
b[i] = uint32(rng.Uint64()) % Q
|
||||
}
|
||||
got := PolyMulCPU(&a, &b)
|
||||
want := schoolbookMul(&a, &b)
|
||||
for i := 0; i < N; i++ {
|
||||
if got[i] != want[i] {
|
||||
t.Fatalf("iter %d: PolyMul mismatch at %d: got %d, want %d",
|
||||
it, i, got[i], want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// schoolbookMul computes a * b in R_q = Z_q[X]/(X^N+1) using the naive
|
||||
// O(N^2) algorithm. Used as a known-correct reference for PolyMul tests.
|
||||
func schoolbookMul(a, b *[N]uint32) [N]uint32 {
|
||||
var out [N]uint32
|
||||
for i := 0; i < N; i++ {
|
||||
for j := 0; j < N; j++ {
|
||||
k := i + j
|
||||
prod := uint32((uint64(a[i]) * uint64(b[j])) % uint64(Q))
|
||||
if k < N {
|
||||
out[k] = addModQ(out[k], prod)
|
||||
} else {
|
||||
// X^N = -1, so coefficient at position k-N gets subtracted.
|
||||
out[k-N] = subModQ(out[k-N], prod)
|
||||
}
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------------
|
||||
// Constant-time arithmetic helpers
|
||||
// -----------------------------------------------------------------------------
|
||||
|
||||
// TestReduce32_Bounds checks that reduce32 produces values in [0, Q)
|
||||
// for a representative sample of inputs across the uint32 range.
|
||||
func TestReduce32_Bounds(t *testing.T) {
|
||||
rng := mrand.New(mrand.NewPCG(0xBA5E, 0))
|
||||
for i := 0; i < 10000; i++ {
|
||||
a := uint32(rng.Uint64())
|
||||
r := reduce32(a)
|
||||
if r >= Q {
|
||||
t.Fatalf("reduce32(%d) = %d, out of bounds [0, %d)", a, r, Q)
|
||||
}
|
||||
if uint64(r) != uint64(a)%uint64(Q) {
|
||||
t.Fatalf("reduce32(%d) = %d, want %d", a, r, uint32(uint64(a)%uint64(Q)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestAddSubMod_Bounds exercises addModQ / subModQ around the boundary.
|
||||
func TestAddSubMod_Bounds(t *testing.T) {
|
||||
tests := []struct {
|
||||
a, b uint32
|
||||
wantA uint32 // addModQ(a,b)
|
||||
wantS uint32 // subModQ(a,b)
|
||||
}{
|
||||
{0, 0, 0, 0},
|
||||
{Q - 1, 1, 0, Q - 2},
|
||||
{Q - 1, Q - 1, Q - 2, 0},
|
||||
{1, 2, 3, Q - 1},
|
||||
{12345, 67890, 80235, Q - (67890 - 12345)},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
if g := addModQ(tt.a, tt.b); g != tt.wantA {
|
||||
t.Errorf("addModQ(%d, %d) = %d, want %d", tt.a, tt.b, g, tt.wantA)
|
||||
}
|
||||
if g := subModQ(tt.a, tt.b); g != tt.wantS {
|
||||
t.Errorf("subModQ(%d, %d) = %d, want %d", tt.a, tt.b, g, tt.wantS)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------------
|
||||
// Batch accelerator
|
||||
// -----------------------------------------------------------------------------
|
||||
|
||||
// TestAccelerator_Available simply exercises the API; both true and
|
||||
// false outcomes are valid depending on the host.
|
||||
func TestAccelerator_Available(t *testing.T) {
|
||||
a := GetAccelerator()
|
||||
if a == nil {
|
||||
t.Fatal("GetAccelerator() returned nil")
|
||||
}
|
||||
_ = a.Available()
|
||||
if Backend() == "" {
|
||||
t.Fatal("Backend() returned empty string")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAccelerator_NTTBatch_Identity confirms that the batch dispatcher
|
||||
// produces the same per-polynomial output as the single-poly path,
|
||||
// regardless of which engine is selected.
|
||||
func TestAccelerator_NTTBatch_Identity(t *testing.T) {
|
||||
const batch = 32 // exceeds the default GPU threshold
|
||||
rng := mrand.New(mrand.NewPCG(0xFADE, 0xC0DE))
|
||||
polys := make([]*[N]uint32, batch)
|
||||
cpuRef := make([][N]uint32, batch)
|
||||
for i := range polys {
|
||||
var p [N]uint32
|
||||
for j := 0; j < N; j++ {
|
||||
p[j] = uint32(rng.Uint64()) % Q
|
||||
}
|
||||
cpuRef[i] = p
|
||||
NTTForwardCPU(&cpuRef[i])
|
||||
dup := p
|
||||
polys[i] = &dup
|
||||
}
|
||||
|
||||
a := GetAccelerator()
|
||||
if err := a.NTTForwardBatch(polys); err != nil {
|
||||
t.Fatalf("NTTForwardBatch: %v", err)
|
||||
}
|
||||
for i := range polys {
|
||||
if *polys[i] != cpuRef[i] {
|
||||
t.Fatalf("batch[%d]: GPU/CPU result differs at index 0: got %d, want %d",
|
||||
i, polys[i][0], cpuRef[i][0])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestAccelerator_InverseRoundtrip drives the batch dispatcher through
|
||||
// NTT→INTT and confirms element-wise equality with the input.
|
||||
func TestAccelerator_InverseRoundtrip(t *testing.T) {
|
||||
const batch = 24
|
||||
rng := mrand.New(mrand.NewPCG(0x1234, 0x5678))
|
||||
polys := make([]*[N]uint32, batch)
|
||||
orig := make([][N]uint32, batch)
|
||||
for i := range polys {
|
||||
var p [N]uint32
|
||||
for j := 0; j < N; j++ {
|
||||
p[j] = uint32(rng.Uint64()) % Q
|
||||
}
|
||||
orig[i] = p
|
||||
dup := p
|
||||
polys[i] = &dup
|
||||
}
|
||||
a := GetAccelerator()
|
||||
if err := a.NTTForwardBatch(polys); err != nil {
|
||||
t.Fatalf("forward: %v", err)
|
||||
}
|
||||
if err := a.NTTInverseBatch(polys); err != nil {
|
||||
t.Fatalf("inverse: %v", err)
|
||||
}
|
||||
for i := range polys {
|
||||
if *polys[i] != orig[i] {
|
||||
t.Fatalf("batch[%d]: NTT∘INTT not identity", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestAccelerator_PolyMulBatch checks the batched multiplication path.
|
||||
func TestAccelerator_PolyMulBatch(t *testing.T) {
|
||||
const batch = 20
|
||||
rng := mrand.New(mrand.NewPCG(0xA1B2, 0xC3D4))
|
||||
a := make([]*[N]uint32, batch)
|
||||
b := make([]*[N]uint32, batch)
|
||||
c := make([]*[N]uint32, batch)
|
||||
ref := make([][N]uint32, batch)
|
||||
for i := 0; i < batch; i++ {
|
||||
var aPoly, bPoly, cPoly [N]uint32
|
||||
for j := 0; j < N; j++ {
|
||||
aPoly[j] = uint32(rng.Uint64()) % Q
|
||||
bPoly[j] = uint32(rng.Uint64()) % Q
|
||||
}
|
||||
ref[i] = PolyMulCPU(&aPoly, &bPoly)
|
||||
a[i] = &aPoly
|
||||
b[i] = &bPoly
|
||||
c[i] = &cPoly
|
||||
}
|
||||
acc := GetAccelerator()
|
||||
if err := acc.PolyMulBatch(a, b, c); err != nil {
|
||||
t.Fatalf("PolyMulBatch: %v", err)
|
||||
}
|
||||
for i := 0; i < batch; i++ {
|
||||
if *c[i] != ref[i] {
|
||||
t.Fatalf("batch[%d]: PolyMulBatch result differs from CPU reference", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestAccelerator_Shapes verifies the shape-mismatch errors.
|
||||
func TestAccelerator_Shapes(t *testing.T) {
|
||||
a := GetAccelerator()
|
||||
one := new([N]uint32)
|
||||
if err := a.PolyMulBatch([]*[N]uint32{one}, nil, nil); err != ErrShapeMismatch {
|
||||
t.Errorf("PolyMulBatch with shape mismatch: got %v, want ErrShapeMismatch", err)
|
||||
}
|
||||
if err := a.SignBatch(ModeMLDSA65, [][]byte{[]byte("x")}, nil, nil); err != ErrShapeMismatch {
|
||||
t.Errorf("SignBatch with shape mismatch: got %v, want ErrShapeMismatch", err)
|
||||
}
|
||||
if _, err := a.VerifyBatch(Mode(7), nil, nil, nil); err != ErrInvalidMode {
|
||||
t.Errorf("VerifyBatch with invalid mode: got %v, want ErrInvalidMode", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAccelerator_Stats checks the GPU/CPU dispatch counters move.
|
||||
func TestAccelerator_Stats(t *testing.T) {
|
||||
a := newAccelerator()
|
||||
gpu0, cpu0 := a.Stats()
|
||||
one := new([N]uint32)
|
||||
a.NTTForward(one)
|
||||
gpu1, cpu1 := a.Stats()
|
||||
if cpu1 != cpu0+1 {
|
||||
t.Errorf("NTTForward did not increment CPU counter: got %d, want %d", cpu1, cpu0+1)
|
||||
}
|
||||
if gpu1 != gpu0 {
|
||||
t.Errorf("NTTForward should not bump GPU counter: got %d, want %d", gpu1, gpu0)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAccelerator_ThresholdRouting forces a small batch (below threshold)
|
||||
// and verifies the dispatcher stays on the CPU. This is the threshold-
|
||||
// routing pattern in action.
|
||||
func TestAccelerator_ThresholdRouting(t *testing.T) {
|
||||
a := newAccelerator()
|
||||
a.SetThresholds(1000, 1000, 1000) // never dispatch to GPU
|
||||
const batch = 4
|
||||
polys := make([]*[N]uint32, batch)
|
||||
for i := range polys {
|
||||
var p [N]uint32
|
||||
for j := 0; j < N; j++ {
|
||||
p[j] = uint32(j) % Q
|
||||
}
|
||||
polys[i] = &p
|
||||
}
|
||||
g0, c0 := a.Stats()
|
||||
if err := a.NTTForwardBatch(polys); err != nil {
|
||||
t.Fatalf("NTTForwardBatch: %v", err)
|
||||
}
|
||||
g1, c1 := a.Stats()
|
||||
if g1 != g0 {
|
||||
t.Errorf("threshold gate failed: GPU counter incremented: %d -> %d", g0, g1)
|
||||
}
|
||||
if c1 != c0+batch {
|
||||
t.Errorf("CPU counter not incremented by %d: %d -> %d", batch, c0, c1)
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------------
|
||||
// FIPS 204 KAT integration via SignBatch / VerifyBatch
|
||||
// -----------------------------------------------------------------------------
|
||||
|
||||
// TestAccelerator_SignVerify_AllModes runs an end-to-end batch sign +
|
||||
// batch verify across all three FIPS 204 modes. The signatures and
|
||||
// public keys come from the standard mldsa wrappers (circl-backed) so
|
||||
// any divergence in the accelerator's framing surfaces as a verify
|
||||
// failure here.
|
||||
func TestAccelerator_SignVerify_AllModes(t *testing.T) {
|
||||
for _, mode := range AllModes() {
|
||||
mode := mode
|
||||
t.Run(mode.String(), func(t *testing.T) {
|
||||
runSignVerifyBatch(t, mode, 4)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestAccelerator_SignVerify_AboveThreshold runs at a batch size above
|
||||
// the default sign threshold to exercise the GPU path when available.
|
||||
// CPU and GPU produce identical accept/reject results.
|
||||
func TestAccelerator_SignVerify_AboveThreshold(t *testing.T) {
|
||||
for _, mode := range AllModes() {
|
||||
mode := mode
|
||||
t.Run(mode.String(), func(t *testing.T) {
|
||||
runSignVerifyBatch(t, mode, 17) // > thrSign && > thrVrfy
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func runSignVerifyBatch(t *testing.T, mode Mode, batch int) {
|
||||
t.Helper()
|
||||
p := MustParamsFor(mode)
|
||||
|
||||
msgs := make([][]byte, batch)
|
||||
sks := make([][]byte, batch)
|
||||
pks := make([][]byte, batch)
|
||||
sigs := make([][]byte, batch)
|
||||
for i := 0; i < batch; i++ {
|
||||
var pk, sk []byte
|
||||
var err error
|
||||
switch mode {
|
||||
case ModeMLDSA44:
|
||||
pubKey, privKey, gerr := mldsa44.GenerateKey(rand.Reader)
|
||||
if gerr != nil {
|
||||
t.Fatalf("mldsa44 GenerateKey: %v", gerr)
|
||||
}
|
||||
pk, err = pubKey.MarshalBinary()
|
||||
if err != nil {
|
||||
t.Fatalf("pk marshal: %v", err)
|
||||
}
|
||||
sk, err = privKey.MarshalBinary()
|
||||
if err != nil {
|
||||
t.Fatalf("sk marshal: %v", err)
|
||||
}
|
||||
case ModeMLDSA65:
|
||||
pubKey, privKey, gerr := mldsa65.GenerateKey(rand.Reader)
|
||||
if gerr != nil {
|
||||
t.Fatalf("mldsa65 GenerateKey: %v", gerr)
|
||||
}
|
||||
pk, err = pubKey.MarshalBinary()
|
||||
if err != nil {
|
||||
t.Fatalf("pk marshal: %v", err)
|
||||
}
|
||||
sk, err = privKey.MarshalBinary()
|
||||
if err != nil {
|
||||
t.Fatalf("sk marshal: %v", err)
|
||||
}
|
||||
case ModeMLDSA87:
|
||||
pubKey, privKey, gerr := mldsa87.GenerateKey(rand.Reader)
|
||||
if gerr != nil {
|
||||
t.Fatalf("mldsa87 GenerateKey: %v", gerr)
|
||||
}
|
||||
pk, err = pubKey.MarshalBinary()
|
||||
if err != nil {
|
||||
t.Fatalf("pk marshal: %v", err)
|
||||
}
|
||||
sk, err = privKey.MarshalBinary()
|
||||
if err != nil {
|
||||
t.Fatalf("sk marshal: %v", err)
|
||||
}
|
||||
}
|
||||
var msg [16]byte
|
||||
if _, err := rand.Read(msg[:]); err != nil {
|
||||
t.Fatalf("rand: %v", err)
|
||||
}
|
||||
binary.LittleEndian.PutUint64(msg[:8], uint64(i))
|
||||
msgs[i] = msg[:]
|
||||
sks[i] = sk
|
||||
pks[i] = pk
|
||||
sigs[i] = make([]byte, p.SignatureSize)
|
||||
}
|
||||
|
||||
a := GetAccelerator()
|
||||
if err := a.SignBatch(mode, msgs, sks, sigs); err != nil {
|
||||
t.Fatalf("SignBatch: %v", err)
|
||||
}
|
||||
for i, s := range sigs {
|
||||
if len(s) != p.SignatureSize {
|
||||
t.Errorf("sigs[%d] wrong length: got %d, want %d", i, len(s), p.SignatureSize)
|
||||
}
|
||||
if isZeroBytes(s) {
|
||||
t.Errorf("sigs[%d] is all-zero — signing failed silently", i)
|
||||
}
|
||||
}
|
||||
ok, err := a.VerifyBatch(mode, msgs, sigs, pks)
|
||||
if err != nil {
|
||||
t.Fatalf("VerifyBatch: %v", err)
|
||||
}
|
||||
for i, v := range ok {
|
||||
if !v {
|
||||
t.Errorf("verify[%d] = false; expected accept", i)
|
||||
}
|
||||
}
|
||||
|
||||
// Negative test: flip one bit in one signature and expect a single
|
||||
// rejection. Verify is constant-time per FIPS 204 §5.3 so all other
|
||||
// signatures must remain accept.
|
||||
if len(sigs) > 0 && len(sigs[0]) > 0 {
|
||||
flipped := make([]byte, len(sigs[0]))
|
||||
copy(flipped, sigs[0])
|
||||
flipped[0] ^= 0xFF
|
||||
bad := make([][]byte, len(sigs))
|
||||
copy(bad, sigs)
|
||||
bad[0] = flipped
|
||||
ok2, err := a.VerifyBatch(mode, msgs, bad, pks)
|
||||
if err != nil {
|
||||
t.Fatalf("VerifyBatch (corrupt): %v", err)
|
||||
}
|
||||
if ok2[0] {
|
||||
t.Error("verify[0] accepted a corrupted signature")
|
||||
}
|
||||
for i := 1; i < len(ok2); i++ {
|
||||
if !ok2[i] {
|
||||
t.Errorf("verify[%d] rejected after sig[0] corruption", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestAccelerator_CrossWrapperEquivalence asserts that the accelerator
|
||||
// path produces signatures that the plain mldsa{44,65,87}.Verify function
|
||||
// will accept — proving the accelerator and the standard wrappers agree
|
||||
// on the same FIPS 204 wire format.
|
||||
func TestAccelerator_CrossWrapperEquivalence(t *testing.T) {
|
||||
for _, mode := range AllModes() {
|
||||
mode := mode
|
||||
t.Run(mode.String(), func(t *testing.T) {
|
||||
p := MustParamsFor(mode)
|
||||
|
||||
const batch = 5
|
||||
msgs := make([][]byte, batch)
|
||||
sks := make([][]byte, batch)
|
||||
pks := make([][]byte, batch)
|
||||
sigs := make([][]byte, batch)
|
||||
pubKeys := make([]any, batch)
|
||||
|
||||
for i := 0; i < batch; i++ {
|
||||
msgs[i] = []byte("cross-wrapper-eq test " + string(rune('A'+i)))
|
||||
switch mode {
|
||||
case ModeMLDSA44:
|
||||
pk, sk, _ := mldsa44.GenerateKey(rand.Reader)
|
||||
sks[i], _ = sk.MarshalBinary()
|
||||
pks[i], _ = pk.MarshalBinary()
|
||||
pubKeys[i] = pk
|
||||
case ModeMLDSA65:
|
||||
pk, sk, _ := mldsa65.GenerateKey(rand.Reader)
|
||||
sks[i], _ = sk.MarshalBinary()
|
||||
pks[i], _ = pk.MarshalBinary()
|
||||
pubKeys[i] = pk
|
||||
case ModeMLDSA87:
|
||||
pk, sk, _ := mldsa87.GenerateKey(rand.Reader)
|
||||
sks[i], _ = sk.MarshalBinary()
|
||||
pks[i], _ = pk.MarshalBinary()
|
||||
pubKeys[i] = pk
|
||||
}
|
||||
sigs[i] = make([]byte, p.SignatureSize)
|
||||
}
|
||||
|
||||
a := GetAccelerator()
|
||||
if err := a.SignBatch(mode, msgs, sks, sigs); err != nil {
|
||||
t.Fatalf("SignBatch: %v", err)
|
||||
}
|
||||
|
||||
// Verify each signature via the standard (non-accelerator)
|
||||
// wrapper. The signature must be in canonical FIPS 204 wire
|
||||
// form for this to pass.
|
||||
for i := 0; i < batch; i++ {
|
||||
switch mode {
|
||||
case ModeMLDSA44:
|
||||
if !mldsa44.Verify(pubKeys[i].(*mldsa44.PublicKey), msgs[i], nil, sigs[i]) {
|
||||
t.Errorf("mldsa44.Verify[%d] rejected accelerator signature", i)
|
||||
}
|
||||
case ModeMLDSA65:
|
||||
if !mldsa65.Verify(pubKeys[i].(*mldsa65.PublicKey), msgs[i], nil, sigs[i]) {
|
||||
t.Errorf("mldsa65.Verify[%d] rejected accelerator signature", i)
|
||||
}
|
||||
case ModeMLDSA87:
|
||||
if !mldsa87.Verify(pubKeys[i].(*mldsa87.PublicKey), msgs[i], nil, sigs[i]) {
|
||||
t.Errorf("mldsa87.Verify[%d] rejected accelerator signature", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// isZeroBytes returns true when every byte in b is 0.
|
||||
func isZeroBytes(b []byte) bool {
|
||||
return len(b) > 0 && bytes.Equal(b, make([]byte, len(b)))
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------------
|
||||
// Benchmarks
|
||||
// -----------------------------------------------------------------------------
|
||||
|
||||
func BenchmarkNTTForward_CPU(b *testing.B) {
|
||||
var p [N]uint32
|
||||
for i := 0; i < N; i++ {
|
||||
p[i] = uint32(i) % Q
|
||||
}
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
NTTForwardCPU(&p)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkNTTInverse_CPU(b *testing.B) {
|
||||
var p [N]uint32
|
||||
for i := 0; i < N; i++ {
|
||||
p[i] = uint32(i) % Q
|
||||
}
|
||||
NTTForwardCPU(&p)
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
var q [N]uint32 = p
|
||||
NTTInverseCPU(&q)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkPolyMul_CPU(b *testing.B) {
|
||||
var x, y [N]uint32
|
||||
for i := 0; i < N; i++ {
|
||||
x[i] = uint32(i) % Q
|
||||
y[i] = uint32(i*7+3) % Q
|
||||
}
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_ = PolyMulCPU(&x, &y)
|
||||
}
|
||||
}
|
||||
|
||||
func benchNTTBatch(b *testing.B, batch int) {
|
||||
a := GetAccelerator()
|
||||
polys := make([]*[N]uint32, batch)
|
||||
for i := range polys {
|
||||
var p [N]uint32
|
||||
for j := 0; j < N; j++ {
|
||||
p[j] = uint32(j+i) % Q
|
||||
}
|
||||
polys[i] = &p
|
||||
}
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if err := a.NTTForwardBatch(polys); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkNTTBatch_1(b *testing.B) { benchNTTBatch(b, 1) }
|
||||
func BenchmarkNTTBatch_16(b *testing.B) { benchNTTBatch(b, 16) }
|
||||
func BenchmarkNTTBatch_64(b *testing.B) { benchNTTBatch(b, 64) }
|
||||
|
||||
func benchSignBatch(b *testing.B, mode Mode, batch int) {
|
||||
p := MustParamsFor(mode)
|
||||
msgs := make([][]byte, batch)
|
||||
sks := make([][]byte, batch)
|
||||
sigs := make([][]byte, batch)
|
||||
for i := 0; i < batch; i++ {
|
||||
msgs[i] = []byte("bench message")
|
||||
var sk []byte
|
||||
var err error
|
||||
switch mode {
|
||||
case ModeMLDSA44:
|
||||
_, privKey, gerr := mldsa44.GenerateKey(rand.Reader)
|
||||
if gerr != nil {
|
||||
b.Fatal(gerr)
|
||||
}
|
||||
sk, err = privKey.MarshalBinary()
|
||||
case ModeMLDSA65:
|
||||
_, privKey, gerr := mldsa65.GenerateKey(rand.Reader)
|
||||
if gerr != nil {
|
||||
b.Fatal(gerr)
|
||||
}
|
||||
sk, err = privKey.MarshalBinary()
|
||||
case ModeMLDSA87:
|
||||
_, privKey, gerr := mldsa87.GenerateKey(rand.Reader)
|
||||
if gerr != nil {
|
||||
b.Fatal(gerr)
|
||||
}
|
||||
sk, err = privKey.MarshalBinary()
|
||||
}
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
sks[i] = sk
|
||||
sigs[i] = make([]byte, p.SignatureSize)
|
||||
}
|
||||
a := GetAccelerator()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if err := a.SignBatch(mode, msgs, sks, sigs); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkSignBatch_44_1(b *testing.B) { benchSignBatch(b, ModeMLDSA44, 1) }
|
||||
func BenchmarkSignBatch_44_16(b *testing.B) { benchSignBatch(b, ModeMLDSA44, 16) }
|
||||
func BenchmarkSignBatch_44_64(b *testing.B) { benchSignBatch(b, ModeMLDSA44, 64) }
|
||||
|
||||
func BenchmarkSignBatch_65_1(b *testing.B) { benchSignBatch(b, ModeMLDSA65, 1) }
|
||||
func BenchmarkSignBatch_65_16(b *testing.B) { benchSignBatch(b, ModeMLDSA65, 16) }
|
||||
func BenchmarkSignBatch_65_64(b *testing.B) { benchSignBatch(b, ModeMLDSA65, 64) }
|
||||
|
||||
func BenchmarkSignBatch_87_1(b *testing.B) { benchSignBatch(b, ModeMLDSA87, 1) }
|
||||
func BenchmarkSignBatch_87_16(b *testing.B) { benchSignBatch(b, ModeMLDSA87, 16) }
|
||||
func BenchmarkSignBatch_87_64(b *testing.B) { benchSignBatch(b, ModeMLDSA87, 64) }
|
||||
|
||||
func benchVerifyBatch(b *testing.B, mode Mode, batch int) {
|
||||
p := MustParamsFor(mode)
|
||||
msgs := make([][]byte, batch)
|
||||
sks := make([][]byte, batch)
|
||||
pks := make([][]byte, batch)
|
||||
sigs := make([][]byte, batch)
|
||||
for i := 0; i < batch; i++ {
|
||||
msgs[i] = []byte("bench verify")
|
||||
var sk, pk []byte
|
||||
var err error
|
||||
switch mode {
|
||||
case ModeMLDSA44:
|
||||
pubKey, privKey, gerr := mldsa44.GenerateKey(rand.Reader)
|
||||
if gerr != nil {
|
||||
b.Fatal(gerr)
|
||||
}
|
||||
sk, err = privKey.MarshalBinary()
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
pk, err = pubKey.MarshalBinary()
|
||||
case ModeMLDSA65:
|
||||
pubKey, privKey, gerr := mldsa65.GenerateKey(rand.Reader)
|
||||
if gerr != nil {
|
||||
b.Fatal(gerr)
|
||||
}
|
||||
sk, err = privKey.MarshalBinary()
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
pk, err = pubKey.MarshalBinary()
|
||||
case ModeMLDSA87:
|
||||
pubKey, privKey, gerr := mldsa87.GenerateKey(rand.Reader)
|
||||
if gerr != nil {
|
||||
b.Fatal(gerr)
|
||||
}
|
||||
sk, err = privKey.MarshalBinary()
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
pk, err = pubKey.MarshalBinary()
|
||||
}
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
sks[i] = sk
|
||||
pks[i] = pk
|
||||
sigs[i] = make([]byte, p.SignatureSize)
|
||||
}
|
||||
a := GetAccelerator()
|
||||
if err := a.SignBatch(mode, msgs, sks, sigs); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if _, err := a.VerifyBatch(mode, msgs, sigs, pks); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkVerifyBatch_44_1(b *testing.B) { benchVerifyBatch(b, ModeMLDSA44, 1) }
|
||||
func BenchmarkVerifyBatch_44_16(b *testing.B) { benchVerifyBatch(b, ModeMLDSA44, 16) }
|
||||
func BenchmarkVerifyBatch_44_64(b *testing.B) { benchVerifyBatch(b, ModeMLDSA44, 64) }
|
||||
|
||||
func BenchmarkVerifyBatch_65_1(b *testing.B) { benchVerifyBatch(b, ModeMLDSA65, 1) }
|
||||
func BenchmarkVerifyBatch_65_16(b *testing.B) { benchVerifyBatch(b, ModeMLDSA65, 16) }
|
||||
func BenchmarkVerifyBatch_65_64(b *testing.B) { benchVerifyBatch(b, ModeMLDSA65, 64) }
|
||||
|
||||
func BenchmarkVerifyBatch_87_1(b *testing.B) { benchVerifyBatch(b, ModeMLDSA87, 1) }
|
||||
func BenchmarkVerifyBatch_87_16(b *testing.B) { benchVerifyBatch(b, ModeMLDSA87, 16) }
|
||||
func BenchmarkVerifyBatch_87_64(b *testing.B) { benchVerifyBatch(b, ModeMLDSA87, 64) }
|
||||
@@ -0,0 +1,248 @@
|
||||
// Copyright (C) 2019-2026, Lux Industries Inc. All rights reserved.
|
||||
// See the file LICENSE for licensing terms.
|
||||
|
||||
package gpu
|
||||
|
||||
// FIPS 204 ML-DSA pure-Go NTT over R_q = Z_q[X]/(X^256+1).
|
||||
//
|
||||
// Implements the standard negacyclic Cooley-Tukey forward and
|
||||
// Gentleman-Sande inverse described in FIPS 204 Algorithms 41 (NTT)
|
||||
// and 42 (NTT^{-1}), with the canonical zeta-power table from FIPS 204
|
||||
// Appendix B (rows match the reference C / Python implementation
|
||||
// byte-for-byte).
|
||||
//
|
||||
// Constant-time discipline: all secret-dependent operations here are
|
||||
// branch-free modular adds/subs/mults that operate uniformly across the
|
||||
// input. The NTT itself is a fixed-shape butterfly network independent
|
||||
// of input values, so it leaks nothing about its operand. mulQ, reduce,
|
||||
// and the butterflies are written without secret-dependent branches.
|
||||
//
|
||||
// This file is *always* compiled (no build tags). The CGO and non-CGO
|
||||
// dispatchers in gpu_{cgo,nocgo}.go both fall back to these helpers
|
||||
// when the batch is below the GPU threshold or the GPU is unavailable.
|
||||
|
||||
// Zetas is the precomputed table of ζ^BitReverse(k) mod Q for k in [0,256),
|
||||
// taken from FIPS 204 Appendix B (Table 1). This is the canonical NTT
|
||||
// twiddle-factor table; deviating from it produces results that fail
|
||||
// the FIPS 204 KAT vectors.
|
||||
//
|
||||
//nolint:gochecknoglobals // FIPS 204 Appendix B constants
|
||||
var Zetas = [N]uint32{
|
||||
0, 4808194, 3765607, 3761513, 5178923, 5496691, 5234739, 5178987,
|
||||
7778734, 3542485, 2682288, 2129892, 3764867, 7375178, 557458, 7159240,
|
||||
5010068, 4317364, 2663378, 6705802, 4855975, 7946292, 676590, 7044481,
|
||||
5152541, 1714295, 2453983, 1460718, 7737789, 4795319, 2815639, 2283733,
|
||||
3602218, 3182878, 2740543, 4793971, 5269599, 2101410, 3704823, 1159875,
|
||||
394148, 928749, 1095468, 4874037, 2071829, 4361428, 3241972, 2156050,
|
||||
3415069, 1759347, 7562881, 4805951, 3756790, 6444618, 6663429, 4430364,
|
||||
5483103, 3192354, 556856, 3870317, 2917338, 1853806, 3345963, 1858416,
|
||||
3073009, 1277625, 5744944, 3852015, 4183372, 5157610, 5258977, 8106357,
|
||||
2508980, 2028118, 1937570, 4564692, 2811291, 5396636, 7270901, 4158088,
|
||||
1528066, 482649, 1148858, 5418153, 7814814, 169688, 2462444, 5046034,
|
||||
4213992, 4892034, 1987814, 5183169, 1736313, 235407, 5130263, 3258457,
|
||||
5801164, 1787943, 5989328, 6125690, 3482206, 4197502, 7080401, 6018354,
|
||||
7062739, 2461387, 3035980, 621164, 3901472, 7153756, 2925816, 3374250,
|
||||
1356448, 5604662, 2683270, 5601629, 4912752, 2312838, 7727142, 7921254,
|
||||
348812, 8052569, 1011223, 6026202, 4561790, 6458164, 6143691, 1744507,
|
||||
1753, 6444997, 5720892, 6924527, 2660408, 6600190, 8321269, 2772600,
|
||||
1182243, 87208, 636927, 4415111, 4423672, 6084020, 5095502, 4663471,
|
||||
8352605, 822541, 1009365, 5926272, 6400920, 1596822, 4423473, 4620952,
|
||||
6695264, 4969849, 2678278, 4611469, 4829411, 635956, 8129971, 5925040,
|
||||
4234153, 6607829, 2192938, 6653329, 2387513, 4768667, 8111961, 5199961,
|
||||
3747250, 2296099, 1239911, 4541938, 3195676, 2642980, 1254190, 8368000,
|
||||
2998219, 141835, 8291116, 2513018, 7025525, 613238, 7070156, 6161950,
|
||||
7921677, 6458423, 4040196, 4908348, 2039144, 6500539, 7561656, 6201452,
|
||||
6757063, 2105286, 6006015, 6346610, 586241, 7200804, 527981, 5637006,
|
||||
6903432, 1994046, 2491325, 6987258, 507927, 7192532, 7655613, 6545891,
|
||||
5346675, 8041997, 2647994, 3009748, 5767564, 4148469, 749577, 4357667,
|
||||
3980599, 2569011, 6764887, 1723229, 1665318, 2028038, 1163598, 5011144,
|
||||
3994671, 8368538, 7009900, 3020393, 3363542, 214880, 545376, 7609976,
|
||||
3105558, 7277073, 508145, 7826699, 860144, 3430436, 140244, 6866265,
|
||||
6195333, 3123762, 2358373, 6187330, 5365997, 6663603, 2926054, 7987710,
|
||||
8077412, 3531229, 4405932, 4606686, 1900052, 7598542, 1054478, 7648983,
|
||||
}
|
||||
|
||||
// reduce32 computes a Barrett-style 32-bit reduction modulo Q.
|
||||
// Input range: full uint32. Output range: [0, Q).
|
||||
//
|
||||
// Constant-time: no branches, no table lookups dependent on a.
|
||||
func reduce32(a uint32) uint32 {
|
||||
// Reduce a in [0, 2^32) to [0, 2Q) via Barrett:
|
||||
// t = floor(a / Q)
|
||||
// then a - t*Q in [0, 2Q).
|
||||
// A final conditional subtract drops to [0, Q).
|
||||
// The mulhi factor mu = floor(2^45 / Q) = 4193792.
|
||||
// Implemented branch-free with the standard mask idiom.
|
||||
|
||||
// Use uint64 arithmetic to compute (a * mu) >> 45.
|
||||
const mu uint64 = 4193792 // floor(2^45 / Q)
|
||||
t := uint32((uint64(a) * mu) >> 45)
|
||||
r := a - t*Q
|
||||
// Branch-free conditional subtract: r -= Q & -(r >= Q).
|
||||
mask := uint32(0) - uint32(boolToU32(r >= Q))
|
||||
r -= Q & mask
|
||||
return r
|
||||
}
|
||||
|
||||
// boolToU32 converts a bool to 1 (true) or 0 (false) without a branch on
|
||||
// secret data. Go guarantees this lowers to a setcc on amd64.
|
||||
func boolToU32(b bool) uint32 {
|
||||
if b {
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// montReduce returns (a * R^{-1}) mod Q where R = 2^32 (Montgomery
|
||||
// reduction). Input a is a 64-bit product of two reduced-form
|
||||
// 32-bit values, so a < Q * 2^32 < 2^55.
|
||||
//
|
||||
// Constant-time: branch-free.
|
||||
func montReduce(a uint64) uint32 {
|
||||
// t = (a * QInv) mod 2^32
|
||||
t := uint32(a) * QInv
|
||||
// u = (a + t * Q) / 2^32
|
||||
u := uint32((a + uint64(t)*uint64(Q)) >> 32)
|
||||
// u in [0, 2Q); conditional subtract.
|
||||
mask := uint32(0) - uint32(boolToU32(u >= Q))
|
||||
u -= Q & mask
|
||||
return u
|
||||
}
|
||||
|
||||
// mulModQ returns a * b mod Q.
|
||||
//
|
||||
// Constant-time: branch-free.
|
||||
func mulModQ(a, b uint32) uint32 {
|
||||
return uint32((uint64(a) * uint64(b)) % uint64(Q))
|
||||
}
|
||||
|
||||
// addModQ returns (a + b) mod Q, given a, b < Q.
|
||||
//
|
||||
// Constant-time: branch-free mask conditional subtract.
|
||||
func addModQ(a, b uint32) uint32 {
|
||||
r := a + b
|
||||
mask := uint32(0) - uint32(boolToU32(r >= Q))
|
||||
r -= Q & mask
|
||||
return r
|
||||
}
|
||||
|
||||
// subModQ returns (a - b) mod Q, given a, b < Q.
|
||||
//
|
||||
// Constant-time: branch-free mask conditional add.
|
||||
func subModQ(a, b uint32) uint32 {
|
||||
r := a - b
|
||||
mask := uint32(0) - uint32(boolToU32(int32(r) < 0))
|
||||
r += Q & mask
|
||||
return r
|
||||
}
|
||||
|
||||
// NTTForwardCPU performs the FIPS 204 §4.3 forward NTT in-place on a
|
||||
// length-N coefficient slice over R_q. p is mutated to NTT-domain form.
|
||||
//
|
||||
// FIPS 204 Algorithm 41 (NTT) — Cooley-Tukey decimation-in-time.
|
||||
// Twiddle factors come from Zetas[] in bit-reversed order, exactly
|
||||
// matching the reference implementation. Output coefficients are in
|
||||
// [0, Q).
|
||||
//
|
||||
// Constant-time: the iteration structure depends only on N, not on p.
|
||||
func NTTForwardCPU(p *[N]uint32) {
|
||||
k := 1
|
||||
for length := 128; length >= 1; length >>= 1 {
|
||||
for start := 0; start < N; start += 2 * length {
|
||||
zeta := Zetas[k]
|
||||
k++
|
||||
for j := start; j < start+length; j++ {
|
||||
// Standard CT butterfly:
|
||||
// t = zeta * p[j+length]
|
||||
// p[j+length] = p[j] - t
|
||||
// p[j] = p[j] + t
|
||||
t := mulModQ(zeta, p[j+length])
|
||||
p[j+length] = subModQ(p[j], t)
|
||||
p[j] = addModQ(p[j], t)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// NTTInverseCPU performs the FIPS 204 §4.3 inverse NTT in-place on a
|
||||
// length-N evaluation-domain slice over R_q. p is mutated to coefficient
|
||||
// form.
|
||||
//
|
||||
// FIPS 204 Algorithm 42 (NTT^{-1}) — Gentleman-Sande decimation-in-
|
||||
// frequency, followed by the canonical scaling factor (1 - Q) (= N^{-1}
|
||||
// mod Q packed into the same loop). Output coefficients are in [0, Q).
|
||||
//
|
||||
// Constant-time: structure depends only on N.
|
||||
func NTTInverseCPU(p *[N]uint32) {
|
||||
k := N - 1
|
||||
for length := 1; length < N; length <<= 1 {
|
||||
for start := 0; start < N; start += 2 * length {
|
||||
// FIPS 204 reference uses -zeta in the inverse loop.
|
||||
zeta := Q - Zetas[k]
|
||||
k--
|
||||
for j := start; j < start+length; j++ {
|
||||
// Gentleman-Sande butterfly:
|
||||
// t = p[j]
|
||||
// p[j] = t + p[j+length]
|
||||
// p[j+length] = zeta * (t - p[j+length])
|
||||
t := p[j]
|
||||
p[j] = addModQ(t, p[j+length])
|
||||
p[j+length] = mulModQ(zeta, subModQ(t, p[j+length]))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Final scaling: multiply by f = N^{-1} mod Q = 8347681 (FIPS 204
|
||||
// reference constant; also equals (Q+1)/N reduced).
|
||||
const fInv uint32 = 8347681
|
||||
for i := 0; i < N; i++ {
|
||||
p[i] = mulModQ(fInv, p[i])
|
||||
}
|
||||
}
|
||||
|
||||
// PolyMulCPU computes c = a * b in R_q via the standard NTT-domain
|
||||
// pointwise multiplication: NTT(c) = NTT(a) ⊙ NTT(b), then INTT.
|
||||
//
|
||||
// All three slices must be length N. Inputs are not mutated.
|
||||
func PolyMulCPU(a, b *[N]uint32) [N]uint32 {
|
||||
var aN, bN [N]uint32
|
||||
aN = *a
|
||||
bN = *b
|
||||
NTTForwardCPU(&aN)
|
||||
NTTForwardCPU(&bN)
|
||||
var c [N]uint32
|
||||
for i := 0; i < N; i++ {
|
||||
c[i] = mulModQ(aN[i], bN[i])
|
||||
}
|
||||
NTTInverseCPU(&c)
|
||||
return c
|
||||
}
|
||||
|
||||
// PolyMulNTTDomainCPU computes c = a ⊙ b pointwise modulo Q. Both
|
||||
// inputs are assumed to already be in NTT (evaluation) domain. Useful
|
||||
// when many multiplications share the same NTT of one operand.
|
||||
func PolyMulNTTDomainCPU(a, b *[N]uint32) [N]uint32 {
|
||||
var c [N]uint32
|
||||
for i := 0; i < N; i++ {
|
||||
c[i] = mulModQ(a[i], b[i])
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
// PolyAddCPU returns a + b mod Q componentwise.
|
||||
func PolyAddCPU(a, b *[N]uint32) [N]uint32 {
|
||||
var c [N]uint32
|
||||
for i := 0; i < N; i++ {
|
||||
c[i] = addModQ(a[i], b[i])
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
// PolySubCPU returns a - b mod Q componentwise.
|
||||
func PolySubCPU(a, b *[N]uint32) [N]uint32 {
|
||||
var c [N]uint32
|
||||
for i := 0; i < N; i++ {
|
||||
c[i] = subModQ(a[i], b[i])
|
||||
}
|
||||
return c
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
// Copyright (C) 2019-2026, Lux Industries Inc. All rights reserved.
|
||||
// See the file LICENSE for licensing terms.
|
||||
|
||||
package gpu
|
||||
|
||||
import "sync"
|
||||
|
||||
// capWorkers clamps the worker count to a sane minimum/maximum. We never
|
||||
// want to fan out to fewer than 1 goroutine, and beyond ~64 the
|
||||
// scheduler contention dominates the work.
|
||||
//
|
||||
// Lives in a build-tag-free helper file so both the cgo and !cgo
|
||||
// engines can share it without duplication.
|
||||
func capWorkers(n int) int {
|
||||
if n < 1 {
|
||||
return 1
|
||||
}
|
||||
if n > 64 {
|
||||
return 64
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// parallelDo dispatches f over indices [0, total) across at most
|
||||
// workers goroutines. Returns when every index has completed.
|
||||
//
|
||||
// Lives in a build-tag-free helper file so both the cgo and !cgo
|
||||
// engines can share it.
|
||||
func parallelDo(total, workers int, f func(int)) {
|
||||
if total <= 0 {
|
||||
return
|
||||
}
|
||||
if workers < 1 {
|
||||
workers = 1
|
||||
}
|
||||
if workers > total {
|
||||
workers = total
|
||||
}
|
||||
if workers == 1 {
|
||||
for i := 0; i < total; i++ {
|
||||
f(i)
|
||||
}
|
||||
return
|
||||
}
|
||||
var wg sync.WaitGroup
|
||||
ch := make(chan int, total)
|
||||
for i := 0; i < total; i++ {
|
||||
ch <- i
|
||||
}
|
||||
close(ch)
|
||||
wg.Add(workers)
|
||||
for w := 0; w < workers; w++ {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for i := range ch {
|
||||
f(i)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
@@ -0,0 +1,225 @@
|
||||
// Copyright (C) 2019-2026, Lux Industries Inc. All rights reserved.
|
||||
// See the file LICENSE for licensing terms.
|
||||
|
||||
// Package gpu provides GPU/CPU-accelerated FIPS 204 (ML-DSA) primitives.
|
||||
//
|
||||
// All three parameter sets (ML-DSA-44, ML-DSA-65, ML-DSA-87) share the
|
||||
// same polynomial ring R_q = Z_q[X]/(X^N+1) with
|
||||
//
|
||||
// N = 256
|
||||
// Q = 8380417 = 2^23 - 2^13 + 1
|
||||
//
|
||||
// Only the matrix dimensions (k, ℓ), secret-key width η, challenge weight
|
||||
// τ, signature norms (β, γ_1, γ_2), and hint bound ω vary between modes.
|
||||
// One NTT kernel and one parameter table therefore serve all three
|
||||
// security levels — exactly the FIPS 204 §4 design.
|
||||
//
|
||||
// The package exposes:
|
||||
//
|
||||
// - A pure-Go CPU implementation of the negacyclic NTT and INTT over
|
||||
// R_q, with constant-time arithmetic on the verify side.
|
||||
// - A CGO-backed dispatch to github.com/luxfi/accel which uses the
|
||||
// Metal / CUDA / WebGPU backends from luxfi/accel for batched NTT
|
||||
// and batched ML-DSA sign/verify across modes {2, 3, 5}.
|
||||
//
|
||||
// The CPU and GPU paths produce byte-identical output (verify is
|
||||
// deterministic per FIPS 204 §5.3; sign is deterministic in non-hedged
|
||||
// mode). KAT-level equivalence is asserted in the test suite.
|
||||
package gpu
|
||||
|
||||
// FIPS 204 ring constants. These are fixed across all three parameter
|
||||
// sets and are the only inputs to the NTT kernel.
|
||||
const (
|
||||
// N is the polynomial degree (FIPS 204 §2.3).
|
||||
N = 256
|
||||
|
||||
// Q is the prime modulus 2^23 - 2^13 + 1 = 8380417 (FIPS 204 §2.3).
|
||||
//
|
||||
// Q is a 23-bit NTT-friendly prime: Q ≡ 1 (mod 2N), so a primitive
|
||||
// 2N-th root of unity exists in Z_Q. The chosen root is ζ = 1753
|
||||
// (FIPS 204 Appendix B).
|
||||
Q uint32 = 8380417
|
||||
|
||||
// Zeta is the primitive 512-th root of unity used by the FIPS 204
|
||||
// NTT (FIPS 204 Appendix B). ζ^N = -1 mod Q, ζ^(2N) = 1 mod Q.
|
||||
Zeta uint32 = 1753
|
||||
|
||||
// QInv is -Q^{-1} mod 2^32, used by Montgomery reduction
|
||||
// (FIPS 204 Appendix B). QInv * Q ≡ -1 (mod 2^32).
|
||||
QInv uint32 = 58728449
|
||||
|
||||
// MontR is 2^32 mod Q. Used as the constant "1" in Montgomery
|
||||
// representation.
|
||||
MontR uint32 = 4193792
|
||||
)
|
||||
|
||||
// Mode identifies a FIPS 204 parameter set by its NIST security level.
|
||||
// The integer encoding matches the standard's "mode" naming (2 = Level 2,
|
||||
// 3 = Level 3, 5 = Level 5) so callers can pass it directly to the
|
||||
// accel.LatticeOps.MLDSASignBatch / MLDSAVerifyBatch API.
|
||||
type Mode uint8
|
||||
|
||||
const (
|
||||
// ModeMLDSA44 is FIPS 204 ML-DSA-44 (NIST Level 2).
|
||||
ModeMLDSA44 Mode = 2
|
||||
// ModeMLDSA65 is FIPS 204 ML-DSA-65 (NIST Level 3, canonical Lux PQ).
|
||||
ModeMLDSA65 Mode = 3
|
||||
// ModeMLDSA87 is FIPS 204 ML-DSA-87 (NIST Level 5, high-value tier).
|
||||
ModeMLDSA87 Mode = 5
|
||||
)
|
||||
|
||||
// String returns the canonical FIPS 204 name of the mode.
|
||||
func (m Mode) String() string {
|
||||
switch m {
|
||||
case ModeMLDSA44:
|
||||
return "ML-DSA-44"
|
||||
case ModeMLDSA65:
|
||||
return "ML-DSA-65"
|
||||
case ModeMLDSA87:
|
||||
return "ML-DSA-87"
|
||||
default:
|
||||
return "ML-DSA-invalid"
|
||||
}
|
||||
}
|
||||
|
||||
// Params holds the FIPS 204 parameter set for one security level.
|
||||
// Per FIPS 204 §4, only the matrix shape (K × L), secret-key width Eta,
|
||||
// challenge weight Tau, signature norms Gamma1Bits / Gamma2 / Beta, and
|
||||
// hint bound Omega vary between modes. All schemes share the same ring
|
||||
// (N=256, Q=8380417), so a single NTT kernel serves all three.
|
||||
type Params struct {
|
||||
// Mode is the NIST level identifier (2, 3, or 5).
|
||||
Mode Mode
|
||||
|
||||
// Name is the canonical FIPS 204 algorithm name.
|
||||
Name string
|
||||
|
||||
// K is the number of rows in matrix A and length of vectors t, w, r
|
||||
// (FIPS 204 §4 Table 1).
|
||||
K int
|
||||
|
||||
// L is the number of columns in matrix A and length of vectors s, y, z
|
||||
// (FIPS 204 §4 Table 1).
|
||||
L int
|
||||
|
||||
// Eta is the bound on secret-key coefficients (s ∈ S_η^ℓ × S_η^k).
|
||||
Eta int
|
||||
|
||||
// Tau is the Hamming weight of the challenge polynomial c.
|
||||
Tau int
|
||||
|
||||
// Beta is the bound on |c·s|: Beta = Tau * Eta.
|
||||
Beta int
|
||||
|
||||
// Gamma1Bits is log_2(γ_1) — controls the y-coefficient sampling range
|
||||
// (FIPS 204 §4 Table 1).
|
||||
Gamma1Bits int
|
||||
|
||||
// Gamma1 is the bound on coefficients of y (FIPS 204 §4).
|
||||
// Gamma1 = 2^Gamma1Bits.
|
||||
Gamma1 int
|
||||
|
||||
// Gamma2 is the low-order rounding threshold (FIPS 204 §4 Table 1).
|
||||
Gamma2 int
|
||||
|
||||
// Omega is the maximum number of 1's in the hint h
|
||||
// (FIPS 204 §4 Table 1).
|
||||
Omega int
|
||||
|
||||
// PublicKeySize is the encoded public-key length in bytes.
|
||||
PublicKeySize int
|
||||
|
||||
// PrivateKeySize is the encoded private-key length in bytes.
|
||||
PrivateKeySize int
|
||||
|
||||
// SignatureSize is the encoded signature length in bytes.
|
||||
SignatureSize int
|
||||
|
||||
// CTildeSize is the length of the commitment hash c~ in bytes
|
||||
// (FIPS 204 §4 Table 1).
|
||||
CTildeSize int
|
||||
}
|
||||
|
||||
// allParams holds the canonical FIPS 204 parameter sets, indexed by mode.
|
||||
// Values match FIPS 204 §4 Table 1 (NIST FIPS 204 2024-08-13).
|
||||
//
|
||||
//nolint:gochecknoglobals // FIPS 204 standard constants
|
||||
var allParams = map[Mode]Params{
|
||||
ModeMLDSA44: {
|
||||
Mode: ModeMLDSA44,
|
||||
Name: "ML-DSA-44",
|
||||
K: 4,
|
||||
L: 4,
|
||||
Eta: 2,
|
||||
Tau: 39,
|
||||
Beta: 39 * 2, // tau * eta
|
||||
Gamma1Bits: 17,
|
||||
Gamma1: 1 << 17, // 131072
|
||||
Gamma2: 95232, // (Q-1)/88
|
||||
Omega: 80,
|
||||
PublicKeySize: 1312,
|
||||
PrivateKeySize: 2560,
|
||||
SignatureSize: 2420,
|
||||
CTildeSize: 32,
|
||||
},
|
||||
ModeMLDSA65: {
|
||||
Mode: ModeMLDSA65,
|
||||
Name: "ML-DSA-65",
|
||||
K: 6,
|
||||
L: 5,
|
||||
Eta: 4,
|
||||
Tau: 49,
|
||||
Beta: 49 * 4,
|
||||
Gamma1Bits: 19,
|
||||
Gamma1: 1 << 19, // 524288
|
||||
Gamma2: 261888, // (Q-1)/32
|
||||
Omega: 55,
|
||||
PublicKeySize: 1952,
|
||||
PrivateKeySize: 4032,
|
||||
SignatureSize: 3309,
|
||||
CTildeSize: 48,
|
||||
},
|
||||
ModeMLDSA87: {
|
||||
Mode: ModeMLDSA87,
|
||||
Name: "ML-DSA-87",
|
||||
K: 8,
|
||||
L: 7,
|
||||
Eta: 2,
|
||||
Tau: 60,
|
||||
Beta: 60 * 2,
|
||||
Gamma1Bits: 19,
|
||||
Gamma1: 1 << 19,
|
||||
Gamma2: 261888,
|
||||
Omega: 75,
|
||||
PublicKeySize: 2592,
|
||||
PrivateKeySize: 4896,
|
||||
SignatureSize: 4627,
|
||||
CTildeSize: 64,
|
||||
},
|
||||
}
|
||||
|
||||
// ParamsFor returns the FIPS 204 parameter set for the given mode.
|
||||
// Returns the zero Params and false for an unknown mode; callers
|
||||
// MUST check the boolean before consuming the value.
|
||||
func ParamsFor(m Mode) (Params, bool) {
|
||||
p, ok := allParams[m]
|
||||
return p, ok
|
||||
}
|
||||
|
||||
// MustParamsFor returns the FIPS 204 parameter set for the given mode
|
||||
// and panics on an unknown mode. Use in test code or static call sites
|
||||
// where the mode is a compile-time constant.
|
||||
func MustParamsFor(m Mode) Params {
|
||||
p, ok := allParams[m]
|
||||
if !ok {
|
||||
panic("mldsa/gpu: unknown ML-DSA mode " + m.String())
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// AllModes returns the canonical FIPS 204 modes in ascending NIST-level
|
||||
// order: {ML-DSA-44, ML-DSA-65, ML-DSA-87}. Useful for iterating cross-
|
||||
// mode test vectors.
|
||||
func AllModes() []Mode {
|
||||
return []Mode{ModeMLDSA44, ModeMLDSA65, ModeMLDSA87}
|
||||
}
|
||||
Reference in New Issue
Block a user