prio3: Adds polynomial multiplication using NTT

vdaf/prio3/arith/arith.go

@@ -89,9 +89,9 @@ type Vec[Vec ~[]E, E Elt] interface {
 	// DotProduct calculates z[i] = z[i] * x[i].
 	DotProduct(x Vec) E
 	// NTT calculates the number theoretic transform of the vector.
-	NTT(Vec)
+	NTT(Vec, uint)
 	// InvNTT calculates the inverse number theoretic transform on values.
-	InvNTT(Vec)
+	InvNTT(Vec, uint)
 	// SplitBits sets the vector of elements corresponding to the bits of n.
 	// The receiving vector sets v[i] = SetUint64(n(i)), where n(i) is the i-th
 	// bit of n, and len(v) >= log2(n).
@@ -133,8 +133,4 @@ type Poly[Poly ~[]E, E Elt] interface {
 	Evaluate(x *E) E
 	// Strip removes the higher-degree zero terms.
 	Strip() Poly
-	// Interpolate a polynomial passing through the points (x[i], y[i]), where
-	// x are the powers of an N-th root of unity, y are the values, and
-	// N = len(y) must be a power of two.
-	Interpolate(y []E)
 }

vdaf/prio3/arith/poly_test.go

@@ -14,6 +14,7 @@ func testPoly[P Poly[P, E], V Vec[V, E], E EltTest, F Fp[E]](t *testing.T) {
 	t.Run("interpolate", interpolate[P, V, E, F])
 	t.Run("strip", strip[P, E, F])
 	t.Run("mulSqr", mulSqrPoly[P, V, E, F])
+	t.Run("mulNTT", mulNTT[P, V, E, F])
 }
 
 func mulSqrPoly[P Poly[P, E], V Vec[V, E], E EltTest, F Fp[E]](t *testing.T) {
@@ -48,10 +49,10 @@ func mulSqrPoly[P Poly[P, E], V Vec[V, E], E EltTest, F Fp[E]](t *testing.T) {
 	}
 }
 
-func evalRootsUnity[P Poly[P, E], V Vec[V, E], E Elt, F Fp[E]](p P) V {
+func evalRootsUnity[P Poly[P, E], V Vec[V, E], E Elt, F Fp[E]](p P, n uint) V {
 	// evaluate p on the powers of the root of unity.
 	// p(w^0), p(w^1), p(w^2), ...
-	N, logN := math.NextPow2(uint(len(p)))
+	N, logN := math.NextPow2(n)
 	var wi, wn F = new(E), new(E)
 	wi.SetOne()
 	wn.SetRootOfUnityTwoN(logN)
@@ -71,16 +72,20 @@ func interpolate[P Poly[P, E], V Vec[V, E], E EltTest, F Fp[E]](t *testing.T) {
 		N := uint(1) << logN
 		p := NewPoly[P](N - 1)
 		mustRead(t, V(p))
-		values := evalRootsUnity[P, V, E, F](p)
+		values := evalRootsUnity[P, V, E, F](p, N)
 
 		y := NewVec[V](N)
-		y.NTT(V(p))
+		y.NTT(V(p), N)
 		if !slices.Equal(y, values) {
 			test.ReportError(t, y, values)
 		}
 
+		var invN F = new(E)
+		invN.InvTwoN(uint(logN))
+
 		p2 := NewPoly[P](N - 1)
-		p2.Interpolate(values)
+		V(p2).InvNTT(values, N)
+		V(p2).ScalarMul(invN)
 		if !slices.Equal(p, p2) {
 			test.ReportError(t, p, p2)
 		}
@@ -101,6 +106,38 @@ func strip[P Poly[P, E], E Elt, F Fp[E]](t *testing.T) {
 	}
 }
 
+type polyTest[P any] interface {
+	MulNlogN(x, y P)
+	MulNSquare(x, y P)
+}
+
+func mulNTT[P Poly[P, E], V Vec[V, E], E EltTest, F Fp[E]](t *testing.T) {
+	const DegX uint = 16
+	const DegY uint = 16
+
+	for degX := range DegX {
+		for degY := range DegY {
+			degZ := degX + degY
+			x := NewPoly[P](degX)
+			y := NewPoly[P](degY)
+			got := NewPoly[P](degZ)
+			want := NewPoly[P](degZ)
+
+			mustRead(t, V(x))
+			mustRead(t, V(y))
+
+			any(got).(polyTest[P]).MulNlogN(x, y)
+			any(want).(polyTest[P]).MulNSquare(x, y)
+
+			if !slices.EqualFunc(got, want,
+				func(x, y E) bool { return F(&x).IsEqual(&y) },
+			) {
+				test.ReportError(t, got, want, degX, degY)
+			}
+		}
+	}
+}
+
 func benchmarkPoly[P Poly[P, E], V Vec[V, E], E Elt, F Fp[E]](b *testing.B) {
 	x := F(new(E))
 	p := NewPoly[P](Degree)
@@ -113,7 +150,7 @@ func benchmarkPoly[P Poly[P, E], V Vec[V, E], E Elt, F Fp[E]](b *testing.B) {
 	N, _ := math.NextPow2(Degree)
 	pol := NewPoly[P](N - 1)
 	mustRead(b, V(pol))
-	values := evalRootsUnity[P, V, E, F](pol)
+	values := evalRootsUnity[P, V, E, F](pol, N)
 
 	b.Run("AddAssign", func(b *testing.B) {
 		for i := 0; i < b.N; i++ {
@@ -137,17 +174,12 @@ func benchmarkPoly[P Poly[P, E], V Vec[V, E], E Elt, F Fp[E]](b *testing.B) {
 	})
 	b.Run("NTT", func(b *testing.B) {
 		for i := 0; i < b.N; i++ {
-			V(p).NTT(values)
+			V(p).NTT(values, N)
 		}
 	})
 	b.Run("InvNTT", func(b *testing.B) {
 		for i := 0; i < b.N; i++ {
-			V(p).InvNTT(values)
-		}
-	})
-	b.Run("Interpolate", func(b *testing.B) {
-		for i := 0; i < b.N; i++ {
-			pol.Interpolate(values)
+			V(p).InvNTT(values, N)
 		}
 	})
 }

vdaf/prio3/arith/templates/poly.go.tmpl

@@ -13,6 +13,15 @@ type Poly []Fp
 func (p Poly) AddAssign(x Poly) { Vec(p).AddAssign(Vec(x)) }
 func (p Poly) SubAssign(x Poly) { Vec(p).SubAssign(Vec(x)) }
 func (p Poly) Mul(x, y Poly) {
+	const thresholdPolyMul = 128
+	if len(x) + len(y) - 1 < thresholdPolyMul {
+		p.MulNSquare(x, y)
+	} else{
+		p.MulNlogN(x, y)
+	}
+}
+
+func (p Poly) MulNSquare(x, y Poly) {
 	mustSumLen(p, x, y)
 	clear(p)
 	var xiyj Fp
@@ -24,6 +33,24 @@ func (p Poly) Mul(x, y Poly) {
 	}
 }
 
+func (p Poly) MulNlogN(x, y Poly) {
+	mustSumLen(p, x, y)
+	N, logN := math.NextPow2(uint(len(x) + len(y) - 1))
+	buf := make(Vec, 2*N)
+	lx, ly := buf[:N], buf[N:]
+	lx.NTT(Vec(x), N)
+	ly.NTT(Vec(y), N)
+	for i := range lx {
+		lx[i].MulAssign(&ly[i])
+	}
+
+	ly.InvNTT(lx, N)
+	var invN Fp
+	invN.InvTwoN(logN)
+	copy(p, ly)
+	Vec(p).ScalarMul(&invN)
+}
+
 func (p Poly) Sqr(x Poly) {
 	mustSumLen(p, x, x)
 	clear(p)
@@ -63,10 +90,3 @@ func (p Poly) Strip() Poly {
 	return p[:0]
 }
 
-func (p Poly) Interpolate(values []Fp) {
-	_, logN := math.NextPow2(uint(len(values)))
-	Vec(p).InvNTT(values)
-	var invN Fp
-	invN.InvTwoN(logN)
-	Vec(p).ScalarMul(&invN)
-}