Add bounds reasoning comments to AVX2 ntt/intt

Signed-off-by: jammychiou1 <[email protected]>

Author: jammychiou1

dev/x86_64/src/intt.S

@@ -43,6 +43,12 @@ vpsrlq		$32,%ymm\r0,%ymm\r0
 vpblendd	$0xAA,%ymm\r1,%ymm\r0,%ymm\r3
 .endm
 
+/*
+ * Compute l + h, montmul(h - l, zh) then store the results back to l, h
+ * respectively.
+ *
+ * The general abs bound of Montgomery multiplication is 3q/4.
+ */
 .macro butterfly l,h,zl0=1,zl1=1,zh0=2,zh1=2
 vpsubd		%ymm\l,%ymm\h,%ymm12
 vpaddd		%ymm\h,%ymm\l,%ymm\l
@@ -74,6 +80,8 @@ vmovdqa		256*\off+160(%rdi),%ymm9
 vmovdqa		256*\off+192(%rdi),%ymm10
 vmovdqa	 	256*\off+224(%rdi),%ymm11
 
+/* All: abs bound < q */
+
 /* level 0 */
 vpermq		$0x1B,(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS_QINV+296-8*\off-8)*4(%rsi),%ymm3
 vpermq		$0x1B,(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS+296-8*\off-8)*4(%rsi),%ymm15
@@ -99,6 +107,19 @@ vmovshdup	%ymm3,%ymm1
 vmovshdup	%ymm15,%ymm2
 butterfly	10,11,1,3,2,15
 
+/* 4, 6, 8, 10: abs bound < 2q; 5, 7, 9, 11: abs bound < 3q/4 */
+/*
+ * Note that since 2^31 / q > 256, the sum of all 256 coefficients does not
+ * overflow. This allows us to greatly simplify the range analysis by relaxing
+ * and unifying the bounds of all coefficients on the same layer. As a concrete
+ * example, here we relax the bounds on 5, 7, 9, 11 and conclude that
+ *
+ * All: abs bound < 2q
+ *
+ * In all but last of the following layers, we do the same relaxation without
+ * explicit mention.
+ */
+
 /* level 1 */
 vpermq		$0x1B,(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS_QINV+168-8*\off-8)*4(%rsi),%ymm3
 vpermq		$0x1B,(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS+168-8*\off-8)*4(%rsi),%ymm15
@@ -114,6 +135,8 @@ vmovshdup	%ymm15,%ymm2
 butterfly	8,10,1,3,2,15
 butterfly	9,11,1,3,2,15
 
+/* All: abs bound < 4q */
+
 /* level 2 */
 vpermq		$0x1B,(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS_QINV+104-8*\off-8)*4(%rsi),%ymm3
 vpermq		$0x1B,(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS+104-8*\off-8)*4(%rsi),%ymm15
@@ -124,6 +147,8 @@ butterfly	5,9,1,3,2,15
 butterfly	6,10,1,3,2,15
 butterfly	7,11,1,3,2,15
 
+/* All: abs bound < 8q */
+
 /* level 3 */
 shuffle2	4,5,3,5
 shuffle2	6,7,4,7
@@ -137,6 +162,8 @@ butterfly	4,7
 butterfly	6,9
 butterfly	8,11
 
+/* All: abs bound < 16q */
+
 /* level 4 */
 shuffle4	3,4,10,4
 shuffle4	6,8,3,8
@@ -150,6 +177,8 @@ butterfly	3,8
 butterfly	6,7
 butterfly	5,11
 
+/* All: abs bound < 32q */
+
 /* level 5 */
 shuffle8	10,3,9,3
 shuffle8	6,5,10,5
@@ -163,6 +192,8 @@ butterfly	10,5
 butterfly	6,8
 butterfly	4,11
 
+/* All: abs bound < 64q */
+
 vmovdqa		%ymm9,256*\off+  0(%rdi)
 vmovdqa		%ymm10,256*\off+ 32(%rdi)
 vmovdqa		%ymm6,256*\off+ 64(%rdi)
@@ -194,6 +225,8 @@ vpbroadcastd	(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS+2)*4(%rsi),%ymm2
 butterfly	8,10
 butterfly	9,11
 
+/* All: abs bound < 128q */
+
 /* level 7 */
 vpbroadcastd	(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS_QINV+0)*4(%rsi),%ymm1
 vpbroadcastd	(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS+0)*4(%rsi),%ymm2
@@ -203,11 +236,27 @@ butterfly	5,9
 butterfly	6,10
 butterfly	7,11
 
+/* 4, 5, 6, 7: abs bound < 256q; 8, 9, 10, 11: abs bound < 3q/4 */
+
 vmovdqa         %ymm8,512+32*\off(%rdi)
 vmovdqa         %ymm9,640+32*\off(%rdi)
 vmovdqa         %ymm10,768+32*\off(%rdi)
 vmovdqa         %ymm11,896+32*\off(%rdi)
 
+/*
+ * In order to (a) remove the factor of 256 arising from the 256-point INTT
+ * butterflies and (b) transform the output into Montgomery domain, we need to
+ * multiply all coefficients by 2^32/256.
+ *
+ * For ymm{8,9,10,11}, the scaling has been merged into the last butterfly, so
+ * only ymm{4,5,6,7} need to be scaled explicitly.
+ *
+ * The scaling is achieved by computing montmul(-, MLD_AVX2_DIV), so the output
+ * will have an abs bound of 3q/4.
+ *
+ * 4, 5, 6, 7: abs bound < 256q
+ */
+
 vmovdqa		(MLD_AVX2_BACKEND_DATA_OFFSET_8XDIV_QINV)*4(%rsi),%ymm1
 vmovdqa		(MLD_AVX2_BACKEND_DATA_OFFSET_8XDIV)*4(%rsi),%ymm2
 vpmuldq		%ymm1,%ymm4,%ymm12
@@ -256,6 +305,8 @@ vmovshdup	%ymm7,%ymm7
 vpblendd	$0xAA,%ymm8,%ymm6,%ymm6
 vpblendd	$0xAA,%ymm9,%ymm7,%ymm7
 
+/* 4, 5, 6, 7: abs bound < 3q/4 */
+
 vmovdqa         %ymm4,  0+32*\off(%rdi)
 vmovdqa         %ymm5,128+32*\off(%rdi)
 vmovdqa         %ymm6,256+32*\off(%rdi)

dev/x86_64/src/ntt.S

@@ -44,6 +44,16 @@ vpsrlq		$32,%ymm\r0,%ymm\r0
 vpblendd	$0xAA,%ymm\r1,%ymm\r0,%ymm\r3
 .endm
 
+/*
+ * Compute l + montmul(h, zh), l - montmul(h, zh) then store the results back to
+ * l, h respectively.
+ *
+ * Although the general abs bound of Montgomery multiplication is 3q/4, we use
+ * the more convenient bound q here.
+ *
+ * In conclusion, the magnitudes of all coefficients grow by at most q after
+ * each layer.
+ */
 .macro butterfly l,h,zl0=1,zl1=1,zh0=2,zh1=2
 vpmuldq		%ymm\zl0,%ymm\h,%ymm13
 vmovshdup	%ymm\h,%ymm12
@@ -56,16 +66,30 @@ vpmuldq		%ymm0,%ymm13,%ymm13
 vpmuldq		%ymm0,%ymm14,%ymm14
 
 vmovshdup	%ymm\h,%ymm\h
-vpblendd	$0xAA,%ymm12,%ymm\h,%ymm\h
+vpblendd	$0xAA,%ymm12,%ymm\h,%ymm\h /* mulhi(h, zh) */
 
-vpsubd		%ymm\h,%ymm\l,%ymm12
-vpaddd		%ymm\h,%ymm\l,%ymm\l
+/*
+ * Originally, mulhi(h, zh) should be subtracted by mulhi(q, mullo(h, zl)) in
+ * order to complete computing
+ *
+ *   montmul(h, zh) = mulhi(h, zh) - mulhi(q, mullo(h, zl)).
+ *
+ * Here, since mulhi(q, mullo(h, zl)) has not been computed yet, this task is
+ * delayed until after add/sub.
+ */
+vpsubd		%ymm\h,%ymm\l,%ymm12 /* l - mulhi(h, zh)
+                                  * = l - montmul(h, zh)
+                                  *   - mulhi(q, mullo(h, zl)) */
+vpaddd		%ymm\h,%ymm\l,%ymm\l /* l + mulhi(h, zh)
+                                  * = l + montmul(h, zh)
+                                  *   + mulhi(q, mullo(h, zl)) */
 
 vmovshdup	%ymm13,%ymm13
-vpblendd	$0xAA,%ymm14,%ymm13,%ymm13
+vpblendd	$0xAA,%ymm14,%ymm13,%ymm13 /* mulhi(q, mullo(h, zl)) */
 
-vpaddd		%ymm13,%ymm12,%ymm\h
-vpsubd		%ymm13,%ymm\l,%ymm\l
+/* Finish the delayed task mentioned above */
+vpaddd		%ymm13,%ymm12,%ymm\h /* l - montmul(h, zh) */
+vpsubd		%ymm13,%ymm\l,%ymm\l /* l + montmul(h, zh) */
 .endm
 
 .macro levels0t1 off
@@ -82,11 +106,15 @@ vmovdqa		640+32*\off(%rdi),%ymm9
 vmovdqa		768+32*\off(%rdi),%ymm10
 vmovdqa	 	896+32*\off(%rdi),%ymm11
 
+/* All: abs bound < q */
+
 butterfly	4,8
 butterfly	5,9
 butterfly	6,10
 butterfly	7,11
 
+/* All: abs bound < 2q */
+
 /* level 1 */
 vpbroadcastd	(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS_QINV+2)*4(%rsi),%ymm1
 vpbroadcastd	(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS+2)*4(%rsi),%ymm2
@@ -98,6 +126,8 @@ vpbroadcastd	(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS+3)*4(%rsi),%ymm2
 butterfly	8,10
 butterfly	9,11
 
+/* All: abs bound < 3q */
+
 vmovdqa		%ymm4,  0+32*\off(%rdi)
 vmovdqa		%ymm5,128+32*\off(%rdi)
 vmovdqa		%ymm6,256+32*\off(%rdi)
@@ -132,6 +162,8 @@ shuffle8	5,9,4,9
 shuffle8	6,10,5,10
 shuffle8	7,11,6,11
 
+/* All: abs bound < 4q */
+
 /* level 3 */
 vmovdqa		(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS_QINV+8+8*\off)*4(%rsi),%ymm1
 vmovdqa		(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS+8+8*\off)*4(%rsi),%ymm2
@@ -146,6 +178,8 @@ shuffle4	8,10,3,10
 shuffle4	4,6,8,6
 shuffle4	9,11,4,11
 
+/* All: abs bound < 5q */
+
 /* level 4 */
 vmovdqa		(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS_QINV+40+8*\off)*4(%rsi),%ymm1
 vmovdqa		(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS+40+8*\off)*4(%rsi),%ymm2
@@ -160,6 +194,8 @@ shuffle2	5,6,7,6
 shuffle2	3,4,5,4
 shuffle2	10,11,3,11
 
+/* All: abs bound < 6q */
+
 /* level 5 */
 vmovdqa		(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS_QINV+72+8*\off)*4(%rsi),%ymm1
 vmovdqa		(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS+72+8*\off)*4(%rsi),%ymm2
@@ -171,6 +207,8 @@ butterfly	8,4,1,10,2,15
 butterfly	7,3,1,10,2,15
 butterfly	6,11,1,10,2,15
 
+/* All: abs bound < 7q */
+
 /* level 6 */
 vmovdqa		(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS_QINV+104+8*\off)*4(%rsi),%ymm1
 vmovdqa		(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS+104+8*\off)*4(%rsi),%ymm2
@@ -186,6 +224,8 @@ vmovshdup	%ymm2,%ymm15
 butterfly	5,3,1,10,2,15
 butterfly	4,11,1,10,2,15
 
+/* All: abs bound < 8q */
+
 /* level 7 */
 vmovdqa		(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS_QINV+168+8*\off)*4(%rsi),%ymm1
 vmovdqa		(MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS+168+8*\off)*4(%rsi),%ymm2
@@ -211,6 +251,8 @@ vpsrlq		$32,%ymm1,%ymm10
 vmovshdup	%ymm2,%ymm15
 butterfly	3,11,1,10,2,15
 
+/* All: abs bound < 9q */
+
 vmovdqa		%ymm9,256*\off+  0(%rdi)
 vmovdqa		%ymm8,256*\off+ 32(%rdi)
 vmovdqa		%ymm7,256*\off+ 64(%rdi)