perf: adds (some) arm NEON assembly for koalabear and babybear (#588)

Author: gbotrel

field/asm/element_31b_arm64.s

@@ -0,0 +1,90 @@
+// Code generated by gnark-crypto/generator. DO NOT EDIT.
+#include "textflag.h"
+#include "funcdata.h"
+#include "go_asm.h"
+
+// addVec(res, a, b *Element, n uint64)
+// n is the number of blocks of 4 uint32 to process
+TEXT ·addVec(SB), NOFRAME|NOSPLIT, $0-32
+	LDP   res+0(FP), (R0, R1)
+	LDP   b+16(FP), (R2, R3)
+	VMOVS $const_q, V3
+	VDUP  V3.S[0], V3.S4      // broadcast q into V3
+
+loop1:
+	CBZ    R3, done2
+	VLD1.P 16(R1), [V0.S4]
+	VLD1.P 16(R2), [V1.S4]
+	VADD   V0.S4, V1.S4, V1.S4 // b = a + b
+	VSUB   V3.S4, V1.S4, V2.S4 // t = b - q
+	VUMIN  V2.S4, V1.S4, V1.S4 // b = min(t, b)
+	VST1.P [V1.S4], 16(R0)     // res = b
+	SUB    $1, R3, R3
+	JMP    loop1
+
+done2:
+	RET
+
+// subVec(res, a, b *Element, n uint64)
+// n is the number of blocks of 4 uint32 to process
+TEXT ·subVec(SB), NOFRAME|NOSPLIT, $0-32
+	LDP   res+0(FP), (R0, R1)
+	LDP   b+16(FP), (R2, R3)
+	VMOVS $const_q, V3
+	VDUP  V3.S[0], V3.S4      // broadcast q into V3
+
+loop3:
+	CBZ    R3, done4
+	VLD1.P 16(R1), [V0.S4]
+	VLD1.P 16(R2), [V1.S4]
+	VSUB   V1.S4, V0.S4, V1.S4 // b = a - b
+	VADD   V1.S4, V3.S4, V2.S4 // t = b + q
+	VUMIN  V2.S4, V1.S4, V1.S4 // b = min(t, b)
+	VST1.P [V1.S4], 16(R0)     // res = b
+	SUB    $1, R3, R3
+	JMP    loop3
+
+done4:
+	RET
+
+// sumVec(t *uint64, a *[]uint32, n uint64) res = sum(a[0...n])
+// n is the number of blocks of 16 uint32 to process
+TEXT ·sumVec(SB), NOFRAME|NOSPLIT, $0-24
+	// zeroing accumulators
+	VMOVQ $0, $0, V4
+	VMOVQ $0, $0, V5
+	VMOVQ $0, $0, V6
+	VMOVQ $0, $0, V7
+	LDP   t+0(FP), (R1, R0)
+	MOVD  n+16(FP), R2
+
+loop5:
+	CBZ R2, done6
+
+	// blockSize is 16 uint32; we load 4 vectors of 4 uint32 at a time
+	// (4*4)*4 = 64 bytes ~= 1 cache line
+	// since our values are 31 bits, we can add 2 by 2 these vectors
+	// we are left with 2 vectors of 4x32 bits values
+	// that we accumulate in 4*2*64bits accumulators
+	// the caller will reduce mod q the accumulators.
+
+	VLD2.P  32(R0), [V0.S4, V1.S4]
+	VADD    V0.S4, V1.S4, V0.S4    // a1 += a2
+	VLD2.P  32(R0), [V2.S4, V3.S4]
+	VADD    V2.S4, V3.S4, V2.S4    // a3 += a4
+	VUSHLL  $0, V0.S2, V1.D2       // convert low words to 64 bits
+	VADD    V1.D2, V5.D2, V5.D2    // acc2 += a2
+	VUSHLL2 $0, V0.S4, V0.D2       // convert high words to 64 bits
+	VADD    V0.D2, V4.D2, V4.D2    // acc1 += a1
+	VUSHLL  $0, V2.S2, V3.D2       // convert low words to 64 bits
+	VADD    V3.D2, V7.D2, V7.D2    // acc4 += a4
+	VUSHLL2 $0, V2.S4, V2.D2       // convert high words to 64 bits
+	VADD    V2.D2, V6.D2, V6.D2    // acc3 += a3
+	SUB     $1, R2, R2
+	JMP     loop5
+
+done6:
+	VADD   V4.D2, V6.D2, V4.D2   // acc1 += acc3
+	VADD   V5.D2, V7.D2, V5.D2   // acc2 += acc4
+	VST2.P [V4.D2, V5.D2], 0(R1) // store acc1 and acc2
+	RET

field/babybear/doc.go

@@ -0,0 +1,10 @@
+//go:build  !purego
+
+// Copyright 2020-2024 Consensys Software Inc.
+// Licensed under the Apache License, Version 2.0. See the LICENSE file for details.
+
+// Code generated by consensys/gnark-crypto DO NOT EDIT
+
+// We include the hash to force the Go compiler to recompile: 8620676634583589757
+#include "../asm/element_31b_arm64.s"
+

field/babybear/element_arm64.s

@@ -221,10 +221,8 @@ func GenerateCommonASM(w io.Writer, nbWords, nbBits int, hasVector bool) error {
 		if nbBits == 31 {
 			return GenerateF31ASM(f, hasVector)
 		} else {
-			fmt.Printf("nbWords: %d, nbBits: %d\n", nbWords, nbBits)
 			panic("not implemented")
 		}
-
 	}
 
 	f.GenerateReduceDefine()

field/babybear/vector_amd64.go

@@ -55,6 +55,14 @@ func GenerateCommonASM(w io.Writer, nbWords, nbBits int, hasVector bool) error {
 	f.WriteLn("#include \"go_asm.h\"")
 	f.WriteLn("")
 
+	if nbWords == 1 {
+		if nbBits == 31 {
+			return GenerateF31ASM(f, hasVector)
+		} else {
+			panic("not implemented")
+		}
+	}
+
 	if f.NbWords%2 != 0 {
 		panic("NbWords must be even")
 	}
@@ -216,3 +224,15 @@ func ElementASMFileName(nbWords, nbBits int) string {
 	}
 	return fmt.Sprintf(nameWN, nbWords)
 }
+
+func GenerateF31ASM(f *FFArm64, hasVector bool) error {
+	if !hasVector {
+		return nil // nothing for now.
+	}
+
+	f.generateAddVecF31()
+	f.generateSubVecF31()
+	f.generateSumVecF31()
+
+	return nil
+}