fallensigh/mlkem-sync
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
171ffd91d3eca6841db4a900e0cab99117638b1b
mlkem-sync/sync_rtl/sample_ntt/TB/gen_vectors.py
FallenSigh d4c3fc86fc feat(tb): add Vivado XSIM Verilog testbenches for all 10 sync modules 
Add file-based vector testbenches ( + ) for:
- mod_add_sync, rng_sync, poly_arith_sync, comp_decomp_sync
- s_bram/sd_bram, sha3_chain_top
- ntt_core, poly_mul_sync
- sample_cbd_sync, sample_ntt_sync

Each module includes:
- tb_<module>_xsim.v: Vivado XSIM testbench
- gen_vectors.py: Python vector generator (stdlib only)
- vectors/<module>_input.hex: test input vectors
- xsim_run.tcl: compile + elaborate + simulate script
2026-06-25 20:48:38 +08:00

504 lines
16 KiB
Python
Raw Blame History

#!/usr/bin/env python3
"""gen_vectors.py - Test vector generator for sample_ntt_sync module.
Generates random rho+k+i+j test vectors, computes expected coefficients using
Keccak-p[1600, 24] matching the RTL's per-permutation squeeze pattern.
The RTL sample_ntt_sync absorbs rho || j || i, then repeatedly permutes
the keccak state, extracting d1,d2 from bits [23:0] of each permuted state
with rejection sampling (d < Q=3329).
Algorithm (matching RTL exactly):
1. Build absorb state: pad10*1 padding for SHAKE-128 (rate=1344) on msg(272b)
2. state = keccak_p(absorb_state)
3. Extract d1 = (state[11:8] << 8) | state[7:0]
d2 = (state[23:16] << 4) | state[15:12]
4. If d1 < Q: output d1 (12-bit unsigned)
5. If d2 < Q: output d2
6. state = keccak_p(state)
7. Repeat 3-6 until 256 coefficients collected
Bit ordering (FIPS 202 / RTL match):
The RTL feeds rho_i[0] as the first message bit into SHA3.
$readmemh stores hex MSB-first → rho_i[255:0] in MSB-first order.
Python: input bytes = reversed(bytes.fromhex(rho_hex)) + bytes([j, i])
Usage:
python3 gen_vectors.py # Generate vectors
python3 gen_vectors.py --verify # Verify results against expected
python3 gen_vectors.py --selftest # Run Keccak-p self-test
"""
import os
import random
import sys
Q = 3329
N_COEFFS = 256
# ======================================================================
# Keccak-p[1600, 24] implementation (stdlib only, matches RTL exactly)
# ======================================================================
# Rotation offsets for rho step: RHO[x][y]
RHO = [
[ 0, 36, 3, 41, 18],
[ 1, 44, 10, 45, 2],
[62, 6, 43, 15, 61],
[28, 55, 25, 21, 56],
[27, 20, 39, 8, 14],
]
# Round constants for iota step
RC = [
0x0000000000000001, 0x0000000000008082, 0x800000000000808A, 0x8000000080008000,
0x000000000000808B, 0x0000000080000001, 0x8000000080008081, 0x8000000000008009,
0x000000000000008A, 0x0000000000000088, 0x0000000080008009, 0x000000008000000A,
0x000000008000808B, 0x800000000000008B, 0x8000000000008089, 0x8000000000008003,
0x8000000000008002, 0x8000000000000080, 0x000000000000800A, 0x800000008000000A,
0x8000000080008081, 0x8000000000008080, 0x0000000080000001, 0x8000000080008008,
]
def ROTL64(val, n):
"""Rotate 64-bit value left by n bits."""
n = n & 63
return ((val << n) | (val >> (64 - n))) & 0xFFFFFFFFFFFFFFFF
def keccak_round(lanes, rnd_idx):
"""Single Keccak-f round: theta, rho, pi, chi, iota.
lanes: list of 25 64-bit integers, indexed as lane_idx = 5*y + x.
Returns new list of 25 lanes.
"""
# === Theta ===
C = [0] * 5 # C[x] = A[x,0] ^ A[x,1] ^ A[x,2] ^ A[x,3] ^ A[x,4]
for x in range(5):
for y in range(5):
C[x] ^= lanes[5 * y + x]
# D[x] = C[x-1] ^ ROTL(C[x+1], 1)
D = [0] * 5
for x in range(5):
D[x] = C[(x - 1) % 5] ^ ROTL64(C[(x + 1) % 5], 1)
new_lanes = [0] * 25
for y in range(5):
for x in range(5):
idx = 5 * y + x
new_lanes[idx] = lanes[idx] ^ D[x]
# === Rho + Pi ===
B = [0] * 25
for x in range(5):
for y in range(5):
src_idx = 5 * y + x
dst_row = (2 * x + 3 * y) % 5
dst_col = y
dst_idx = 5 * dst_row + dst_col
B[dst_idx] = ROTL64(new_lanes[src_idx], RHO[x][y])
# === Chi ===
result = [0] * 25
for y in range(5):
for x in range(5):
idx = 5 * y + x
idx1 = 5 * y + ((x + 1) % 5)
idx2 = 5 * y + ((x + 2) % 5)
result[idx] = B[idx] ^ ((~B[idx1] & 0xFFFFFFFFFFFFFFFF) & B[idx2])
# === Iota ===
result[0] ^= RC[rnd_idx]
return result
def keccak_p(state_int):
"""Keccak-p[1600, 24] permutation of a 1600-bit integer.
Returns 1600-bit integer.
"""
# Convert 1600-bit integer to 25 lanes of 64 bits
lanes = [0] * 25
for y in range(5):
for x in range(5):
idx = 5 * y + x
lanes[idx] = (state_int >> (64 * idx)) & 0xFFFFFFFFFFFFFFFF
for rnd in range(24):
lanes = keccak_round(lanes, rnd)
# Convert back to 1600-bit integer
result = 0
for y in range(5):
for x in range(5):
idx = 5 * y + x
result |= lanes[idx] << (64 * idx)
return result & ((1 << 1600) - 1)
def keccak_p_selftest():
"""Verify Keccak-p against known test vectors.
Validates by computing SHA3-256('') and SHA3-256('abc')
against Python's hashlib reference.
"""
import hashlib
def sha3_256_via_keccak(data):
"""Compute SHA3-256 using our keccak_p, compare with hashlib.
SHA3-256(M) = Keccak[c=512](M || 01, 256)
Rate = 1088 bits, capacity = 512 bits.
"""
# Convert message to bits, LSB-first per byte
msg_bits = ''.join(format(b, '08b')[::-1] for b in data)
# Append SHA-3 suffix '01' (bit 0, then bit 1)
msg_bits += '01'
# Apply pad10*1: '1' + required zeros + '1'
pad_len = (1088 - (len(msg_bits) % 1088)) % 1088
if pad_len == 0:
pad_len = 1088
# pad_len >= 2 always for pad10*1
msg_bits += '1' + '0' * (pad_len - 2) + '1'
# Build state by XOR-ing padded message into rate portion
state = 0
for i, bit in enumerate(msg_bits):
if bit == '1':
state ^= (1 << i)
# Absorb (single block for short messages)
state = keccak_p(state)
# Squeeze 256 bits from rate portion [255:0]
output_bytes = bytes(
(state >> (8 * i)) & 0xFF for i in range(32)
)
return output_bytes.hex()
# Test SHA3-256("") against hashlib
expected_empty = hashlib.sha3_256(b"").hexdigest()
got_empty = sha3_256_via_keccak(b"")
assert got_empty == expected_empty, \
f"SHA3-256('') mismatch:\n got: {got_empty}\n expected: {expected_empty}"
# Test SHA3-256("abc") against hashlib
expected_abc = hashlib.sha3_256(b"abc").hexdigest()
got_abc = sha3_256_via_keccak(b"abc")
assert got_abc == expected_abc, \
f"SHA3-256('abc') mismatch:\n got: {got_abc}\n expected: {expected_abc}"
# Also test that our keccak_p is not identity
assert keccak_p(0) != 0, "keccak_p(0) must not equal 0"
print("Keccak-p self-test PASSED (SHA3-256 verified against hashlib)")
return True
# ======================================================================
# Message construction (matching RTL pad10*1 for SHAKE-128)
# ======================================================================
def build_absorb_state(rho_hex, j, i):
"""Build the 1600-bit absorb state matching RTL sample_ntt_sync.
RTL absorb_state = {capacity(256b), pad10*1, suffix(1111), msg(272b)}
Message order (bit 0 = state[0]):
rho_i[0], ..., rho_i[255], j[0], ..., j[7], i[0], ..., i[7]
Args:
rho_hex: 64-char hex string (MSB-first).
j: 2-bit index (j_idx).
i: 2-bit index (i_idx).
Returns:
1600-bit integer representing the padded absorb state.
"""
# Build message bytes
rho_bytes = bytes.fromhex(rho_hex) # MSB-first → byte 0 = rho_i[255:248]
rho_rev = rho_bytes[::-1] # Reversed: byte 0 = rho_i[7:0]
msg_bytes = rho_rev + bytes([j & 0xFF, i & 0xFF])
# msg_bytes[0] = rho_i[7:0], msg_bytes[1] = rho_i[15:8], ...
# msg_bytes[32] = j, msg_bytes[33] = i
# Build 1600-bit state
# state[7:0] = msg_bytes[0]
# state[15:8] = msg_bytes[1]
# ...
# state[275:272] = 4'b1111 (SHAKE suffix)
# state[276] = 1'b1 (pad10*1 first 1)
# state[1342:277] = 0
# state[1343] = 1'b1 (pad10*1 final 1)
# state[1599:1344] = 0 (capacity)
state = 0
# Message bytes at bits [0:271]
for idx, b in enumerate(msg_bytes):
state |= (b & 0xFF) << (8 * idx)
# SHAKE suffix (1111) at bits [272:275]
state |= 0xF << 272
# pad10*1: 1 at bit 276
state |= 1 << 276
# pad10*1: 1 at bit 1343
state |= 1 << 1343
return state & ((1 << 1600) - 1)
def extract_d1_d2(state):
"""Extract d1, d2 from state[23:0] matching RTL.
Returns (d1, d2) as 12-bit unsigned integers [0, 4095].
"""
c0 = (state >> 0) & 0xFF
c1 = (state >> 8) & 0xFF
c2 = (state >> 16) & 0xFF
d1 = ((c1 & 0xF) << 8) | c0
d2 = (c2 << 4) | (c1 >> 4)
return d1, d2
def sample_ntt_compute(rho_hex, k, i, j):
"""Compute expected coefficients matching RTL sample_ntt_sync.
The RTL absorbs rho || j || i, then repeatedly permutes,
extracting d1,d2 from each permuted state with rejection sampling.
Args:
rho_hex: 64-char hex string.
k: 3-bit k value (unused in computation, passed through).
i: 2-bit i_idx.
j: 2-bit j_idx.
Returns:
list of 256 integers in [0, Q-1].
"""
state = build_absorb_state(rho_hex, j, i)
state = keccak_p(state)
coeffs = []
while len(coeffs) < N_COEFFS:
d1, d2 = extract_d1_d2(state)
if d1 < Q:
coeffs.append(d1)
if len(coeffs) == N_COEFFS:
break
if d2 < Q:
coeffs.append(d2)
if len(coeffs) == N_COEFFS:
break
# Re-permute for next squeeze
state = keccak_p(state)
return coeffs
# ======================================================================
# Helper functions
# ======================================================================
def random_hex(bits):
"""Generate a random hex string (MSB-first) of the given bit length."""
val = random.getrandbits(bits)
num_nibbles = (bits + 3) // 4
return f"{val:0{num_nibbles}X}"
def coeff_to_hex(val):
"""Convert 12-bit unsigned coefficient to 3-char hex string."""
return f"{val & 0xFFF:03X}"
# ======================================================================
# Vector generation and file I/O
# ======================================================================
def generate_one(k, i, j):
"""Generate a single test vector.
Returns dict with 'rho_hex', 'k', 'i', 'j', 'coeffs'.
"""
rho_hex = random_hex(256)
coeffs = sample_ntt_compute(rho_hex, k, i, j)
return {
"rho_hex": rho_hex,
"k": k,
"i": i,
"j": j,
"coeffs": coeffs,
}
def write_input_hex(vectors, filepath):
"""Write input vectors as packed hex for $readmemh.
Each line: {1'b0, j_idx[1:0], i_idx[1:0], k_i[2:0], rho_i[255:0]}
= 264 bits = 66 hex chars.
"""
os.makedirs(os.path.dirname(filepath), exist_ok=True)
with open(filepath, "w") as f:
for v in vectors:
rho_hex = v["rho_hex"]
k = v["k"] & 0x7
i = v["i"] & 0x3
j = v["j"] & 0x3
# Pack: {1'b0, j[1:0], i[1:0], k[2:0], rho[255:0]}
header = (j << 5) | (i << 3) | k
# header occupies bits [262:256] (when zero-padded properly)
# rho occupies bits [255:0]
# Total: 256 + 7 = 263 bits → 66 hex chars (264 bits, top bit = 0)
# header_hex: 2 chars
header_hex = f"{header:02X}"
packed = header_hex + rho_hex # 2 + 64 = 66 chars
f.write(packed + "\n")
def write_expected_hex(vectors, filepath):
"""Write expected coefficients: one 12-bit hex value per line."""
os.makedirs(os.path.dirname(filepath), exist_ok=True)
with open(filepath, "w") as f:
for idx, v in enumerate(vectors):
f.write(f"# VECTOR_{idx} k={v['k']} i={v['i']} j={v['j']} "
f"rho={v['rho_hex']}\n")
for c in v["coeffs"]:
f.write(coeff_to_hex(c) + "\n")
def verify_results(result_file, vectors):
"""Verify RTL output against expected values."""
with open(result_file, "r") as f:
lines = [line.strip() for line in f
if line.strip() and not line.strip().startswith("#")]
# Remove trailing comments
cleaned = []
for line in lines:
comment_idx = line.find(" #")
if comment_idx >= 0:
line = line[:comment_idx].strip()
cleaned.append(line)
expected_all = []
for v in vectors:
for c in v["coeffs"]:
expected_all.append(coeff_to_hex(c))
if len(cleaned) != len(expected_all):
print(f" COUNT MISMATCH: got {len(cleaned)}, expected {len(expected_all)}")
return False
mismatches = 0
for idx, (g, e) in enumerate(zip(cleaned, expected_all)):
if g.upper() != e.upper():
if mismatches < 10:
print(f" MISMATCH[{idx}]: got={g.upper()}, expected={e.upper()}")
mismatches += 1
if mismatches > 0:
print(f" Total mismatches: {mismatches}")
return False
return True
# ======================================================================
# Main
# ======================================================================
def main():
base_dir = os.path.dirname(os.path.abspath(__file__))
vectors_dir = os.path.join(base_dir, "vectors")
input_file = os.path.join(vectors_dir, "sample_ntt_input.hex")
expected_file = os.path.join(vectors_dir, "sample_ntt_expected.hex")
result_file = os.path.join(vectors_dir, "sample_ntt_result.hex")
if "--selftest" in sys.argv:
keccak_p_selftest()
print("Running quick smoke test on sample_ntt computation...")
rho = "000102030405060708090A0B0C0D0E0F101112131415161718191A1B1C1D1E1F"
coeffs = sample_ntt_compute(rho, k=4, i=0, j=0)
print(f" rho={rho[:16]}...: {len(coeffs)} coefficients")
print(f" first 4: {coeffs[:4]}")
assert len(coeffs) == 256, f"Expected 256 coefficients, got {len(coeffs)}"
assert all(0 <= c < Q for c in coeffs), "Coefficient out of range"
print("Smoke test PASSED")
return
verify_mode = "--verify" in sys.argv
if verify_mode:
if not os.path.exists(result_file):
print(f"ERROR: Result file not found: {result_file}")
sys.exit(1)
if not os.path.exists(input_file):
print(f"ERROR: Input file not found: {input_file}")
sys.exit(1)
print(f"Verifying results from {result_file}...")
# Recompute expected from input file
with open(input_file, "r") as f:
input_lines = [l.strip() for l in f if l.strip()]
vectors = []
for line in input_lines:
if len(line) != 66:
print(f"WARNING: Skipping line with unexpected length {len(line)}: {line[:20]}...")
continue
header_hex = line[:2]
rho_hex = line[2:]
header = int(header_hex, 16)
j = (header >> 5) & 0x3
i = (header >> 3) & 0x3
k = header & 0x7
coeffs = sample_ntt_compute(rho_hex, k, i, j)
vectors.append({"rho_hex": rho_hex, "k": k, "i": i, "j": j, "coeffs": coeffs})
ok = verify_results(result_file, vectors)
if ok:
print("ALL VECTORS PASSED")
else:
print("VERIFICATION FAILED")
sys.exit(1)
else:
# Generate mode
vector_count = 4
print(f"Generating {vector_count} test vectors...")
print(f"Running Keccak-p self-test first...")
keccak_p_selftest()
vectors = []
for idx in range(vector_count):
k = random.choice([2, 3, 4])
i = random.randint(0, 3)
j = random.randint(0, 3)
v = generate_one(k, i, j)
vectors.append(v)
print(f" Vector {idx}: k={k}, i={i}, j={j}, "
f"rho={v['rho_hex'][:8]}..., coeffs[0]={v['coeffs'][0]}")
write_input_hex(vectors, input_file)
print(f"Wrote {len(vectors)} vectors to {input_file}")
write_expected_hex(vectors, expected_file)
print(f"Wrote expected coefficients to {expected_file}")
# Sanity checks
for v in vectors:
for c in v["coeffs"]:
assert 0 <= c < Q, f"Coefficient {c} out of range [0, {Q-1}]"
assert len(v["coeffs"]) == N_COEFFS
print("All sanity checks passed.")
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink