Preparation for ARM64 Implementation of poly operations for dilithium package.

sign/internal/dilithium/arm64.go

@@ -0,0 +1,84 @@
+//go:build arm64 && !purego
+// +build arm64,!purego
+
+package dilithium
+
+// Execute an in-place forward NTT on as.
+//
+// Assumes the coefficients are in Montgomery representation and bounded
+// by 2*Q.  The resulting coefficients are again in Montgomery representation,
+// but are only bounded bt 18*Q.
+func (p *Poly) NTT() {
+	p.nttGeneric()
+}
+
+// Execute an in-place inverse NTT and multiply by Montgomery factor R
+//
+// Assumes the coefficients are in Montgomery representation and bounded
+// by 2*Q.  The resulting coefficients are again in Montgomery representation
+// and bounded by 2*Q.
+func (p *Poly) InvNTT() {
+	p.invNttGeneric()
+}
+
+// Sets p to the polynomial whose coefficients are the pointwise multiplication
+// of those of a and b.  The coefficients of p are bounded by 2q.
+//
+// Assumes a and b are in Montgomery form and that the pointwise product
+// of each coefficient is below 2³² q.
+func (p *Poly) MulHat(a, b *Poly) {
+	p.mulHatGeneric(a, b)
+}
+
+// Sets p to a + b.  Does not normalize polynomials.
+func (p *Poly) Add(a, b *Poly) {
+	polyAdd(p, a, b)
+}
+
+// Sets p to a - b.
+//
+// Warning: assumes coefficients of b are less than 2q.
+// Sets p to a + b.  Does not normalize polynomials.
+func (p *Poly) Sub(a, b *Poly) {
+	p.subGeneric(a, b)
+}
+
+// Writes p whose coefficients are in [0, 16) to buf, which must be of
+// length N/2.
+func (p *Poly) PackLe16(buf []byte) {
+	p.packLe16Generic(buf)
+}
+
+// Reduces each of the coefficients to <2q.
+func (p *Poly) ReduceLe2Q() {
+	p.reduceLe2QGeneric()
+}
+
+// Reduce each of the coefficients to <q.
+func (p *Poly) Normalize() {
+	p.normalizeGeneric()
+}
+
+// Normalize the coefficients in this polynomial assuming they are already
+// bounded by 2q.
+func (p *Poly) NormalizeAssumingLe2Q() {
+	p.normalizeAssumingLe2QGeneric()
+}
+
+// Checks whether the "supnorm" (see sec 2.1 of the spec) of p is equal
+// or greater than the given bound.
+//
+// Requires the coefficients of p to be normalized.
+func (p *Poly) Exceeds(bound uint32) bool {
+	return p.exceedsGeneric(bound)
+}
+
+// Sets p to 2ᵈ q without reducing.
+//
+// So it requires the coefficients of p  to be less than 2³²⁻ᴰ.
+func (p *Poly) MulBy2toD(q *Poly) {
+	p.mulBy2toDGeneric(q)
+}
+
+//go:noescape
+func polyAdd(p, a, b *Poly)

sign/internal/dilithium/arm64.s

@@ -0,0 +1,29 @@
+//go:build arm64 && !purego
+
+#include "textflag.h"
+
+// func polyAdd(p, a, b *Poly)
+TEXT ·polyAdd(SB), NOSPLIT|NOFRAME, $0-24
+    MOVD    p+0(FP), R0
+    MOVD    a+8(FP), R1
+    MOVD    b+16(FP), R2
+
+    // loop counter (each iteration processes 16 elements so 16 * 16 = 256 = N)
+    // manually unrolling could also be done, for now skipped
+    MOVW    $16, R3
+
+add:
+    VLD1.P  (64)(R1), [V0.S4, V1.S4, V2.S4, V3.S4]
+    VLD1.P  (64)(R2), [V4.S4, V5.S4, V6.S4, V7.S4]
+
+    VADD    V4.S4, V0.S4, V8.S4
+    VADD    V5.S4, V1.S4, V9.S4
+    VADD    V6.S4, V2.S4, V10.S4
+    VADD    V7.S4, V3.S4, V11.S4
+
+    VST1.P  [V8.S4, V9.S4, V10.S4, V11.S4], (64)(R0)
+
+    SUBS    $1, R3, R3
+    BGT     add
+
+    RET

sign/internal/dilithium/arm64_test.go

@@ -0,0 +1,21 @@
+//go:build arm64 && !purego
+// +build arm64,!purego
+
+package dilithium
+
+import (
+	"testing"
+	"unsafe"
+)
+
+func TestAssemblyImplementationAssumptions(t *testing.T) {
+	var p Poly
+
+	if len(p) != 256 {
+		t.Fatal("length of p must be 256. (assumption of the arm64 implementations)")
+	}
+
+	if unsafe.Sizeof(p) != 1024 {
+		t.Fatal("sizeof p be must be 1024 bytes. (assumption of the arm64 implementations)")
+	}
+}

sign/internal/dilithium/generic.go

@@ -1,5 +1,5 @@
-//go:build !amd64 || purego
-// +build !amd64 purego
+//go:build (!amd64 && !arm64) || purego
+// +build !amd64,!arm64 purego
 
 package dilithium