speedup falcon signature verification

codec.c

@@ -72,7 +72,16 @@ Zf(modq_encode)(
 	return out_len;
 }
 
+#if FALCON_AVX2
+static inline uint32_t load4(const uint8_t *buf) {
+	uint32_t r;
+	memcpy(&r, buf, 4);
+	return r;
+}
+#endif
+
 /* see inner.h */
+TARGET_AVX2
 size_t
 Zf(modq_decode)(
 	uint16_t *x, unsigned logn,
@@ -88,10 +97,113 @@ Zf(modq_decode)(
 	if (in_len > max_in_len) {
 		return 0;
 	}
+
+	u = 0;
 	buf = in;
+#if FALCON_AVX2 // yyyAVX2+1
+	if (logn >= 3) { // We need at least 8 elements to decode at a time.
+		const __m256i mask     = _mm256_set1_epi32((1 << 14) - 1);
+		const __m256i Q 	   = _mm256_set1_epi32(12289);
+		const __m256i offsets  = _mm256_setr_epi32(18,4,14,0, 18,4,14,0);
+		/* Mask that byteswaps each 32-bit element of the __m256i */
+		const __m256i bytemask = _mm256_setr_epi8(
+			3,  2,  1,  0,
+			7,  6,  5,  4,
+			11, 10,  9,  8,
+			15, 14, 13, 12,
+			3,  2,  1,  0,
+			7,  6,  5,  4,
+			11, 10,  9,  8,
+			15, 14, 13, 12
+		);
+
+		uint16_t *out;
+		out = x;
+		/* 
+			The size of the polynomial in the encoded polynomial in bytes 
+			is given by (14 * (1 << logn) / 8). The encoded form represents
+			(1 << logn) polynomial coefficients, bit-packed as 14-bits each.
+
+			The least-common-multiple of 14 and 8 is 56, meaning the lowest
+			amount we could comfortably extract in parallel is 56 / 14 = 4 elements.
+
+			This means our parallel strategy could also be effective with 128-bit vectors,
+			but as we're targetting AVX2, we can double all of the numbers to process
+			not 7 bytes at a time, but 14.
+		*/
+		while (u < (n / 8)) {
+			/* 
+				We know that an element could never be in more than 3 bytes at a time.
+				The worst case is that the first bit is at index 7 of byte 0,
+				the takes up the entire byte 1, and the last few bits are in byte 2.
+
+				We can effectively work around this by performing fast 32-bit loads,
+				and simply reading an extra byte each time. The loads are indifferent
+				to which byte is the unused one, and we trivially make up for it with
+				the shifts later on.
+
+				We perform 4 movs to load words at buf, buf + 3, buf + 7, buf + 10.
+				Each pair of elements has the same offset, since the first element
+				will not use the 4th byte but the second element will.
+
+				The vector will contain 8 compressed elements (end-exclusive ranges):
+				1. 00..14 (bytes 0, 1)
+				2. 14..28 (bytes 1, 2, 3)
+				3. 28..42 (bytes 3, 4, 5)
+				4. 42..56 (bytes 5, 6)
+				...
+			*/
+			__m256i compressed = _mm256_setr_epi32(
+				load4(buf + 0),
+				load4(buf + 0),
+				load4(buf + 3),
+				load4(buf + 3),
+				load4(buf + 7),
+				load4(buf + 7),
+				load4(buf + 10),
+				load4(buf + 10)
+			);
+
+			/* 
+				We first perform the byteswap, which can be trivially done as a single vpshufb.
+				This shuffle will byteswap each of the 32-bit elements within the vector.
+			*/
+			__m256i swapped = _mm256_shuffle_epi8(compressed, bytemask);
+			/* 
+				We perform shifts that align each of the elements, whereever they are within
+				their byte-swapped 32-bit representation, to start at the first bit of the element.
+				This makes up for the overlapping bytes that occur when having bit-packed elements. 
+			*/
+			__m256i shifted = _mm256_srlv_epi32(swapped, offsets);
+			/* After aligning the elements, we simply mask them off to 14-bits */
+			__m256i masked = _mm256_and_si256(shifted, mask);
+
+			/* 
+				Here we perform a little trick, which is equivalent to reduce(.or, masked >= Q) 
+			   	If any of our elements is greater-than-or-equal to Q, the predicate is true
+			   	and we exist the decoding process.
+			*/
+			__m256i max = _mm256_max_epu32(masked, Q);
+			__m256i cmp = _mm256_cmpeq_epi32(masked, max);
+			if (!_mm256_testz_si256(cmp, cmp)) return 0;
+
+			/* 
+				Our final step is to pack the extracted elements from being in 32-bits 
+				(but masked off to 14), into actual 16-bit elements as the API uses. 
+			*/
+			__m256i packed = _mm256_packus_epi32(masked, masked);
+			__m256i perm = _mm256_permute4x64_epi64(packed, 0xD8);
+			_mm_storeu_si128((__m128i*)out, _mm256_castsi256_si128(perm));
+
+			u += 1;
+			buf += 14;
+			out += 8;
+		}
+		return in_len;
+	}
+#endif // yyyAVX2-
 	acc = 0;
 	acc_len = 0;
-	u = 0;
 	while (u < n) {
 		acc = (acc << 8) | (*buf ++);
 		acc_len += 8;

common.c

@@ -47,15 +47,18 @@ Zf(hash_to_point_vartime)(
 	 * nonce, the hashed output cannot be matched against potential
 	 * plaintexts).
 	 */
-	size_t n;
+	uint8_t buf[128];
+	uint8_t offset = 128;
 
+	size_t n;
 	n = (size_t)1 << logn;
 	while (n > 0) {
-		uint8_t buf[2];
-		uint32_t w;
-
-		inner_shake256_extract(sc, (void *)buf, sizeof buf);
-		w = ((unsigned)buf[0] << 8) | (unsigned)buf[1];
+		if (offset >= 128) {
+			inner_shake256_extract(sc, (void *)buf, sizeof buf);
+			offset = 0;
+		}
+		uint32_t w = ((unsigned)buf[offset] << 8) | (unsigned)buf[offset + 1];
+		offset += 2;
 		if (w < 61445) {
 			while (w >= 12289) {
 				w -= 12289;

vrfy.c

@@ -498,9 +498,39 @@ mq_div_12289(uint32_t x, uint32_t y)
 	return mq_montymul(y18, x);
 }
 
+
+
+#if FALCON_AVX2
+
+#define VEC __m256i
+#define VER(name) mq_8x_##name
+#define ADD(a, b)   _mm256_add_epi32((a), (b))
+#define AND(a, b)   _mm256_and_si256((a), (b))
+#define MUL(a, b)   _mm256_mullo_epi32((a), (b))
+#define NEG(a)      _mm256_sub_epi32(_mm256_setzero_si256(), (a))
+#define SHR(a, imm) _mm256_srli_epi32((a), (imm))
+#define SPLAT(a)    _mm256_set1_epi32((a))
+#define SUB(a, b)   _mm256_sub_epi32((a), (b))
+#include "vrfy_simd.h"
+
+#define VEC __m128i
+#define VER(name) mq_4x_##name
+#define ADD(a, b)   _mm_add_epi32((a), (b))
+#define AND(a, b)   _mm_and_si128((a), (b))
+#define MUL(a, b)   _mm_mullo_epi32((a), (b))
+#define NEG(a)      _mm_sub_epi32(_mm_setzero_si128(), (a))
+#define SHR(a, imm) _mm_srli_epi32((a), (imm))
+#define SPLAT(a)    _mm_set1_epi32((a))
+#define SUB(a, b)   _mm_sub_epi32((a), (b))
+#include "vrfy_simd.h"
+
+#endif
+
+
 /*
  * Compute NTT on a ring element.
  */
+TARGET_AVX2
 static void
 mq_NTT(uint16_t *a, unsigned logn)
 {
@@ -513,9 +543,71 @@ mq_NTT(uint16_t *a, unsigned logn)
 
 		ht = t >> 1;
 		for (i = 0, j1 = 0; i < m; i ++, j1 += t) {
-			size_t j, j2;
-			uint32_t s;
-
+			size_t j;
+
+#if FALCON_AVX2
+			switch (ht) {
+				case 8: {
+					/* Load 8 16-bit field elements from our two sources. */
+					__m128i t0  = _mm_loadu_si128((__m128i*)&a[j1]);
+					__m128i t1  = _mm_loadu_si128((__m128i*)&a[j1 + ht]);
+					/* Extend each 16-bit element to 32-bits to perform the multiplications. */
+					__m256i u_0 = _mm256_cvtepu16_epi32(t0);
+					__m256i u_1 = _mm256_cvtepu16_epi32(t1);
+
+					/* We're able to reuse the omega across the whole butterfly loop. */
+					__m256i s = _mm256_set1_epi32(GMb[m + i]);
+
+					/* V = a[j+t] * S */
+					__m256i v  = mq_8x_montymul(u_1, s);
+					/* a[j] = U + V mod q */
+					__m256i r0 = mq_8x_add(u_0, v);
+					/* a[j+t] = U - V mod q */
+					__m256i r1 = mq_8x_sub(u_0, v);
+
+					/* 
+						r0 and r1 contain 8 32-bit elements, where each element is reduced to 16-bits.
+						our goal is to shuffle around the elements so that we combine them and end
+						up with a single register containing 16, 16-bit elements, where the first 8
+						are from the r0 and the second are from r1.
+					*/
+					__m256i result = _mm256_permute4x64_epi64(
+						_mm256_packus_epi32(r0, r1),
+						0xD8
+					);
+					_mm256_storeu_si256((__m256i*)&a[j1], result);
+					continue;
+				}
+				case 4: {
+					/* Sets the first 4 elements of each vector, keeping the rest undefined. */
+					__m128i t0 = _mm_loadl_epi64((__m128i*)&a[j1]);
+					__m128i t1 = _mm_loadl_epi64((__m128i*)&a[j1 + ht]);
+
+					/* Extend each 16-bit element to 32-bits to perform the multiplications. */
+					__m128i u_0 = _mm_cvtepu16_epi32(t0);
+					__m128i u_1 = _mm_cvtepu16_epi32(t1);
+
+					/* We're able to reuse the omega across the whole butterfly loop. */
+					__m128i s = _mm_set1_epi32(GMb[m + i]);
+
+					/* V = a[j+t] * S */
+					__m128i v  = mq_4x_montymul(u_1, s);
+					/* a[j] = U + V mod q */
+					__m128i r0 = mq_4x_add(u_0, v);
+					/* a[j+t] = U - V mod q */
+					__m128i r1 = mq_4x_sub(u_0, v);
+
+					__m128i zero = _mm_setzero_si128();
+					__m128i lo = _mm_blend_epi16(r1, zero, 0xAA);
+					__m128i hi = _mm_blend_epi16(r0, zero, 0xAA);
+					__m128i result = _mm_packus_epi32(hi, lo);
+					_mm_storeu_si128((__m128i*)&a[j1], result);
+					continue;
+				}
+				default: {}
+			}
+#endif
+			uint32_t s, j2;
 			s = GMb[m + i];
 			j2 = j1 + ht;
 			for (j = j1; j < j2; j ++) {
@@ -534,6 +626,7 @@ mq_NTT(uint16_t *a, unsigned logn)
 /*
  * Compute the inverse NTT on a ring element, binary case.
  */
+TARGET_AVX2
 static void
 mq_iNTT(uint16_t *a, unsigned logn)
 {
@@ -549,9 +642,61 @@ mq_iNTT(uint16_t *a, unsigned logn)
 		hm = m >> 1;
 		dt = t << 1;
 		for (i = 0, j1 = 0; i < hm; i ++, j1 += dt) {
-			size_t j, j2;
-			uint32_t s;
+			size_t j;
+
+#if FALCON_AVX2
+			switch (t) {
+				case 8: {
+					/* Load 8 16-bit field elements from our two sources. */
+					__m128i t0  = _mm_loadu_si128((__m128i*)&a[j1]);
+					__m128i t1  = _mm_loadu_si128((__m128i*)&a[j1 + t]);
+					/* Extend each 16-bit element to 32-bits to perform the multiplications. */
+					__m256i u = _mm256_cvtepu16_epi32(t0);
+					__m256i v = _mm256_cvtepu16_epi32(t1);
+
+					/* We're able to reuse the omega across the whole butterfly loop. */
+					__m256i s = _mm256_set1_epi32(iGMb[hm + i]);
+
+					/* a[j] = U + V mod q */
+					__m256i r0  = mq_8x_add(u, v);
+					/* a[j+t] = (U - V) * S mod q */
+					__m256i r1 = mq_8x_montymul(mq_8x_sub(u, v), s);
+
+					__m256i result = _mm256_permute4x64_epi64(
+						_mm256_packus_epi32(r0, r1),
+						0xD8
+					);
+					_mm256_storeu_si256((__m256i*)&a[j1], result);
+					continue;
+				}
+				case 4: {
+					/* Sets the first 4 elements of each vector, keeping the rest undefined. */
+					__m128i t0 = _mm_loadl_epi64((__m128i*)&a[j1]);
+					__m128i t1 = _mm_loadl_epi64((__m128i*)&a[j1 + t]);
+					/* Extend each 16-bit element to 32-bits to perform the multiplications. */
+					__m128i u = _mm_cvtepu16_epi32(t0);
+					__m128i v = _mm_cvtepu16_epi32(t1);
+
+					/* We're able to reuse the omega across the whole butterfly loop. */
+					__m128i s = _mm_set1_epi32(iGMb[hm + i]);
+
+					/* a[j] = U + V mod q */
+					__m128i r0  = mq_4x_add(u, v);
+					/* a[j+t] = (U - V) * S mod q */
+					__m128i r1 = mq_4x_montymul(mq_4x_sub(u, v), s);
+
+					__m128i zero = _mm_setzero_si128();
+					__m128i lo = _mm_blend_epi16(r1, zero, 0xAA);
+					__m128i hi = _mm_blend_epi16(r0, zero, 0xAA);
+					__m128i result = _mm_packus_epi32(hi, lo);
+					_mm_storeu_si128((__m128i*)&a[j1], result);
+					continue;
+				}
+				default: {}
+			}
+#endif
 
+			uint32_t s, j2;
 			j2 = j1 + t;
 			s = iGMb[hm + i];
 			for (j = j1; j < j2; j ++) {

vrfy_simd.h

@@ -0,0 +1,50 @@
+TARGET_AVX2
+static inline VEC
+VER(montymul)(VEC x, VEC y) {
+    VEC z, w;
+    const VEC vQ   = SPLAT(12289);
+    const VEC vQ0I = SPLAT(12287); // -1/q mod 2^16
+    const VEC mask = SPLAT(0xFFFF);
+
+    z = MUL(x, y);
+    w = MUL(AND(MUL(z, vQ0I), mask), vQ);
+    z = SHR(ADD(z, w), 16);
+
+    z = SUB(z, vQ);
+    z = ADD(z, AND(vQ, NEG(SHR(z, 31))));
+    return z;
+
+}
+
+TARGET_AVX2
+static inline VEC
+VER(add)(VEC x, VEC y) {
+    VEC d;
+    const VEC vQ   = SPLAT(12289);
+
+    d = SUB(ADD(x, y), vQ);
+    d = ADD(d, AND(vQ, NEG(SHR(d, 31))));
+    return d;
+}
+
+TARGET_AVX2
+static inline VEC
+VER(sub)(VEC x, VEC y) {
+    VEC d;
+    const VEC vQ   = SPLAT(12289);
+
+    d = SUB(x, y);
+    d = ADD(d, AND(vQ, NEG(SHR(d, 31))));
+    return d;
+}
+
+#undef VEC
+#undef VER
+#undef NAME
+#undef SPLAT
+#undef ADD
+#undef SUB
+#undef MUL
+#undef SHR
+#undef AND
+#undef NEG