Port test/test_bounds.py to confirm bound for Montgomery

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

Author: jammychiou1

dev/x86_64/src/intt.S

@@ -323,6 +323,9 @@ vpblendd	$0xAA,%ymm9,%ymm7,%ymm7
  * int32_t (thus |a| <= R/2), we still have |montmul(a, b)| <= 3q/4. This can be
  * strengthened to |montmul_pos(a, b)| <= floor(3q/4) < ceil(3q/4) since LHS is
  * an integer and 3q/4 isn't.
+ *
+ * See test/test_bounds.py for more empirical evidence (and some minor technical
+ * details).
  */
 
 /* 4, 5, 6, 7: abs bound < ceil(3q/4) */

test/test_bounds.py

@@ -0,0 +1,124 @@
+# Copyright (c) The mlkem-native project authors
+# Copyright (c) The mldsa-native project authors
+# SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT
+
+#
+# The purpose of this script is to provide either brute-force proof
+# or empirical evidence to arithmetic bounds for the modular
+# arithmetic primitives used in this repository.
+#
+
+import random
+from functools import lru_cache
+from fractions import Fraction
+from math import ceil
+
+# Global constants
+R = 2**32
+Q = 8380417
+Qinv = pow(Q, -1, R)
+NQinv = pow(-Q, -1, R)
+
+
+#
+# Montgomery multiplication
+#
+
+
+def lift_signed_i32(x):
+    """Returns signed canonical representative modulo R=2^32."""
+    x = x % R
+    if x >= R // 2:
+        x -= R
+    return x
+
+
+@lru_cache(maxsize=None)
+def montmul_neg_twiddle(b):
+    return (b * NQinv) % R
+
+
+@lru_cache(maxsize=None)
+def montmul_pos_twiddle(b):
+    return (b * Qinv) % R
+
+
+def montmul_neg(a, b):
+    b_twiddle = montmul_neg_twiddle(b)
+    return (a * b + Q * lift_signed_i32(a * b_twiddle)) // R
+
+
+def montmul_pos(a, b):
+    b_twiddle = montmul_pos_twiddle(b)
+    return (a * b - Q * lift_signed_i32(a * b_twiddle)) // R
+
+
+#
+# Generic test functions
+#
+
+
+def test_random(f, test_name, num_tests=10000000, bound_a=R // 2, bound_b=Q // 2):
+    print(f"Randomly checking {test_name} ({num_tests} tests)...")
+    for i in range(num_tests):
+        if i % 100000 == 0:
+            print(f"... run {i} tests ({((i * 1000) // num_tests)/10}%)")
+        a = random.randrange(-bound_a, bound_a)
+        b = random.randrange(-bound_b, bound_b)
+        f(a, b)
+
+
+#
+# Test bound on "Montgomery multiplication with signed canonical constant", as
+# used in AVX2 [I]NTT
+#
+
+"""
+In @[Survey_Hwang23, Section 2.2], the author noted the bound*
+
+  |montmul(a, b)| <= (q/2) (1 + |a|/R).
+
+In particular, knowing that a fits inside int32_t (thus |a| <= R/2) already
+implies |montmul(a, b)| <= 3q/4 < ceil(3q/4).
+
+(*) Strictly speaking, they considered the negative/additive variant
+    montmul_neg(a, b), but the exact same bound and proof also work for the
+    positive/subtractive variant montmul_pos(a, b).
+"""
+
+
+def montmul_pos_const_bound(a):
+    return Fraction(Q, 2) * (1 + Fraction(abs(a), R))
+
+
+def montmul_pos_const_bound_test(a, b):
+    ab = montmul_pos(a, b)
+    bound = montmul_pos_const_bound(a)
+    if abs(ab) > bound:
+        print(f"montmul_pos_const_bound_test failure for (a,b)={(a,b)}")
+        print(f"montmul_pos(a,b): {ab}")
+        print(f"bound: {bound}")
+        assert False
+
+
+def montmul_pos_const_bound_test_random():
+    test_random(
+        montmul_pos_const_bound_test,
+        "bound on Montgomery multiplication with constant, as used in AVX2 [I]NTT",
+    )
+
+
+def montmul_pos_const_bound_tight():
+    """
+    This example shows that, unless we know more about a or b, the bound
+    |montmul(a, b)| < ceil(3q/4) is the tightest exclusive bound.
+    """
+    a_worst = -R // 2
+    b_worst = -(Q - 3) // 2
+    ab_worst = montmul_pos(a_worst, b_worst)
+    bound = ceil(Fraction(3 * Q, 4))
+    assert ab_worst == bound - 1
+
+
+montmul_pos_const_bound_test_random()
+montmul_pos_const_bound_tight()