#!/usr/bin/env python3 """gen_vectors.py - Generate poly_mul test vectors for Vivado XSIM testbench. Generates input vectors for poly_mul_sync module in hex format. Each line encodes {256x12-bit poly A, 256x12-bit poly B} = 6144 bits packed as a single hex number (1536 hex chars). The poly_mul_sync module uses basecase_mul (degree-1 Karatsuba multiplication) with zeta from poly_mul_zeta_rom. For each pair of degree-1 polynomials (A[2k], A[2k+1]) and (B[2k], B[2k+1]): c0 = (a0*b0 + a1*b1*zeta) mod Q c1 = (a0*b1 + a1*b0) mod Q Usage: python3 gen_vectors.py Output: vectors/poly_mul_input.hex """ import sys import os import random # ============================================================ # Constants # ============================================================ Q = 3329 # ML-KEM modulus N = 256 # Number of coefficients per polynomial # poly_mul_zeta_rom values (zeta = zeta_bitRev[k]^2 * 17 mod Q) # These match sync_rtl/poly_mul/poly_mul_zeta_rom.v exactly POLY_MUL_ZETA = [ 17, 3312, 2761, 568, 583, 2746, 2649, 680, 1637, 1692, 723, 2606, 2288, 1041, 1100, 2229, 1409, 1920, 2662, 667, 3281, 48, 233, 3096, 756, 2573, 2156, 1173, 3015, 314, 3050, 279, 1703, 1626, 1651, 1678, 2789, 540, 1789, 1540, 1847, 1482, 952, 2377, 1461, 1868, 2687, 642, 939, 2390, 2308, 1021, 2437, 892, 2388, 941, 733, 2596, 2337, 992, 268, 3061, 641, 2688, 1584, 1745, 2298, 1031, 2037, 1292, 3220, 109, 375, 2954, 2549, 780, 2090, 1239, 1645, 1684, 1063, 2266, 319, 3010, 2773, 556, 757, 2572, 2099, 1230, 561, 2768, 2466, 863, 2594, 735, 2804, 525, 1092, 2237, 403, 2926, 1026, 2303, 1143, 2186, 2150, 1179, 2775, 554, 886, 2443, 1722, 1607, 1212, 2117, 1874, 1455, 1029, 2300, 2110, 1219, 2935, 394, 885, 2444, 2154, 1175, ] def barrett_mul(a: int, b: int) -> int: """Barrett modular multiplication: (a * b) mod Q.""" return (a * b) % Q def compute_expected(a_coeffs: list, b_coeffs: list) -> list: """Compute expected output of poly_mul_sync. The poly_mul_sync module processes coefficient pairs using basecase_mul: For each k in 0..127: a0=A[2k], a1=A[2k+1], b0=B[2k], b1=B[2k+1], zeta=zeta_rom[k] c0 = (a0*b0 + a1*b1*zeta) mod Q c1 = (a0*b1 + a1*b0) mod Q """ result = [0] * N for k in range(N // 2): a0 = a_coeffs[2 * k] a1 = a_coeffs[2 * k + 1] b0 = b_coeffs[2 * k] b1 = b_coeffs[2 * k + 1] zeta = POLY_MUL_ZETA[k] # c0 = (a0*b0 + a1*b1*zeta) mod Q t1 = barrett_mul(a0, b0) t2 = barrett_mul(a1, b1) t2_zeta = barrett_mul(t2, zeta) c0 = (t1 + t2_zeta) % Q # c1 = (a0*b1 + a1*b0) mod Q t3 = barrett_mul(a0, b1) t4 = barrett_mul(a1, b0) c1 = (t3 + t4) % Q result[2 * k] = c0 result[2 * k + 1] = c1 return result def coeffs_pair_to_hex(a_coeffs: list, b_coeffs: list) -> str: """Convert 256+256 12-bit coefficients to a 1536-char hex string. A coeffs in the MSB half, B coeffs in the LSB half. A[0] at the very top, B[255] at the very bottom. """ result = 0 for c in a_coeffs: result = (result << 12) | (c & 0xFFF) for c in b_coeffs: result = (result << 12) | (c & 0xFFF) return f"{result:01536X}" def write_vector(f, a_coeffs: list, b_coeffs: list, label: str): """Write a single test vector to the hex file.""" hex_str = coeffs_pair_to_hex(a_coeffs, b_coeffs) f.write(f"// {label}\n") f.write(hex_str + "\n") def generate_vectors(): """Generate test vectors for poly_mul_sync.""" os.makedirs("vectors", exist_ok=True) hex_path = os.path.join("vectors", "poly_mul_input.hex") tests = [] # --- Test 0: All zeros --- zeros = [0] * N expected = [0] * N tests.append((zeros, zeros, expected, "A=0, B=0")) # --- Test 1: A = 1 at all positions, B = all zeros --- a_ones = [1] * N tests.append((a_ones, zeros, [0] * N, "A=1, B=0")) # --- Test 2: A = [1,0,1,0,...], B = [0,1,0,1,...] --- a_10 = [1 if i % 2 == 0 else 0 for i in range(N)] b_01 = [0 if i % 2 == 0 else 1 for i in range(N)] expected = compute_expected(a_10, b_01) tests.append((a_10, b_01, expected, "A=1010..., B=0101...")) # --- Test 3: impulse at index 0 for both A and B --- a_imp0 = [0] * N a_imp0[0] = 1 b_imp0 = [0] * N b_imp0[0] = 1 expected = compute_expected(a_imp0, b_imp0) tests.append((a_imp0, b_imp0, expected, "A=imp[0], B=imp[0]")) # --- Test 4: impulse at index 1 for both A and B --- a_imp1 = [0] * N a_imp1[1] = 1 b_imp1 = [0] * N b_imp1[1] = 1 expected = compute_expected(a_imp1, b_imp1) tests.append((a_imp1, b_imp1, expected, "A=imp[1], B=imp[1]")) # --- Test 5: impulse pair (0,1) --- a_imp01 = [0] * N a_imp01[0] = 1 a_imp01[1] = 1 b_imp01 = [0] * N b_imp01[0] = 1 b_imp01[1] = 1 expected = compute_expected(a_imp01, b_imp01) tests.append((a_imp01, b_imp01, expected, "A=imp[0,1], B=imp[0,1]")) # --- Test 6: all pairs identity --- a_id = [1, 0] * (N // 2) b_id = [0, 0] * (N // 2) b_id[0] = 1 # B = [1,0,0,0,...] expected = compute_expected(a_id, b_id) tests.append((a_id, b_id, expected, "A=all pairs (1,0), B=imp[0]")) # --- Tests 7-10: random vectors --- random.seed(0xBEEF) for i in range(4): a_rand = [random.randrange(0, Q) for _ in range(N)] b_rand = [random.randrange(0, Q) for _ in range(N)] expected = compute_expected(a_rand, b_rand) tests.append((a_rand, b_rand, expected, f"random {i}")) # Write input hex file with open(hex_path, "w") as f: f.write("// poly_mul_sync test vectors\n") f.write("// Format: {768 hex chars A coeffs}{768 hex chars B coeffs}\n") f.write("// A coeffs: 256 x 12-bit values, A[0] at MSB\n") f.write("// B coeffs: 256 x 12-bit values, B[0] after A[255]\n") f.write("\n") for a_coeffs, b_coeffs, expected, label in tests: write_vector(f, a_coeffs, b_coeffs, label) print(f"Generated {len(tests)} test vectors → {hex_path}") # Print expected results for reference print("\nExpected output summary (first 4 coeffs of each test):") for a_coeffs, b_coeffs, expected, label in tests: first4 = expected[:4] print(f" {label:35s} → [{first4[0]:4d}, {first4[1]:4d}, {first4[2]:4d}, {first4[3]:4d}, ...]") if __name__ == "__main__": generate_vectors()