fallensigh/mlkem-sync
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(dec): Decaps D0 - op_i widen + dk/c load + parse

Browse Source 
Scaffolding for ML-KEM Decaps (FIPS 203 Alg 18):
- op_i widened to 2-bit: 00=KeyGen, 01=Encaps, 10=Decaps (op_r too).
- New ST_DEC_LOAD state (D0: settles to DONE so load/parse is dbg-checkable).
- dk (=sk) streamed via dk_in_*; load logic routes each byte by region:
  [0,384K)->dk_pke (dkp_bram), [384K,768K+32)->ek_pke (ek_bram),
  [768K+32,+32)->H(ek) (hek_r), [768K+64,+32)->z (z_r). Routing uses the
  LIVE k_i input, not start-captured k_r (dk is streamed before start_i).
- c (=ct) streamed via c_in_* into a SEPARATE c_in_bram, so the computed c'
  (ct_bram) can later be compared against original c and J(z||c) can read c.
- New dbg taps: dbg_mprime_o/dbg_kbar_o/dbg_decz_o/dbg_dech_o.

TB: tb_mlkem_dec_katK_xsim verifies dk parse (H(ek), z, ek_pke/dk_pke BRAM
round-trip). gen_decaps_vectors.py emits dec_k{K}_c{N}_{dk,ct,ss,ctn,ssn}.hex
from the NIST KAT. run_tb.sh gains a 'dec' module (mirrors 'enc').

Regression fix: old KeyGen/Encaps TBs didn't connect the new input ports,
floating them to X and corrupting the ek/dkp write muxes -> tied off
dk_in_*/c_in_*/new dbg taps in both.

Verified: dec D0 K=2/3/4 PASS; KeyGen K=2 + Encaps K=2 unregressed.
This commit is contained in:
FallenSigh
2026-06-29 15:22:34 +08:00
parent 4091fd0676
commit 030931d4e5
52 changed files with 43220 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

127
sync_rtl/top/mlkem_top.v
View File

@@ -28,7 +28,7 @@ module mlkem_top #(
input clk,
input rst_n,
input [2:0] k_i, // RUNTIME ML-KEM param: 2=512, 3=768, 4=1024
input op_i, // 0 = KeyGen, 1 = Encaps (captured at start_i)
input [1:0] op_i, // 0=KeyGen, 1=Encaps, 2=Decaps (captured at start_i)
input [255:0] d_i, // KeyGen seed d (byte 0 in d_i[7:0])
input [255:0] z_i, // implicit-rejection seed z
input [255:0] msg_i, // Encaps message m (byte 0 in msg_i[7:0])
@@ -41,6 +41,21 @@ module mlkem_top #(
input ek_in_we, // write one ek byte (only honored in ST_IDLE for Encaps preload)
input [10:0] ek_in_addr, // ek byte address 0..ek_bytes-1
input [7:0] ek_in_byte, // ek byte value
// Decaps dk input: stream dk bytes (=sk, 768K+96) into the design before
// start_i with op_i=2. The load logic routes each byte by region:
// [0,384K) -> dk_pke (dkp_bram) = s_hat encoding
// [384K, 768K+32) -> ek_pke (ek_bram) = t_hat encoding || rho
// [768K+32, 768K+64) -> H(ek) (hek_r)
// [768K+64, 768K+96) -> z (z_r)
input dk_in_we, // write one dk byte (ST_IDLE Decaps preload)
input [11:0] dk_in_addr, // dk byte address 0..(768K+96)-1
input [7:0] dk_in_byte, // dk byte value
// Decaps ciphertext input: stream c bytes (=ct) into c_in_bram before
// start_i. Kept separate from ct_bram so the computed c' can be compared
// against the original c (and J(z||c) can read c) without aliasing.
input c_in_we, // write one c byte (ST_IDLE Decaps preload)
input [10:0] c_in_addr, // c byte address 0..ct_bytes-1
input [7:0] c_in_byte, // c byte value
// Encaps shared secret output (= K), valid at done_o
output [255:0] ss_o,
// Encaps ciphertext readback tap: ct byte by index 0..ct_bytes-1
@@ -65,7 +80,12 @@ module mlkem_top #(
output [255:0] dbg_sigma_o,
// Encaps debug taps: r (G high half) and H(ek)
output [255:0] dbg_r_o,
output [255:0] dbg_hek_o
output [255:0] dbg_hek_o,
// Decaps debug taps: m' (Decrypt output), K-bar (J output), parsed z/h.
output [255:0] dbg_mprime_o,
output [255:0] dbg_kbar_o,
output [255:0] dbg_decz_o, // parsed z (dk[768K+64:+32])
output [255:0] dbg_dech_o // parsed H(ek) (dk[768K+32:+32])
);
localparam Q = `Q; // 3329
@@ -82,15 +102,28 @@ module mlkem_top #(
wire [5:0] slot_t_rt = kk_rt + {1'b0, k_r} + {1'b0, k_r}; // t_hat base = kk+2k
wire [11:0] ek_bytes_rt = 12'd384 * {9'b0, k_r} + 12'd32; // 800/1184/1568
wire [11:0] dk_bytes_rt = 12'd384 * {9'b0, k_r}; // 768/1152/1536
// full dk = dk_pke(dk_bytes) || ek(ek_bytes) || H(ek)(32) || z(32). Region
// boundaries (used by Decaps dk load routing and dk readback).
wire [11:0] dk_ek_end = dk_bytes_rt + ek_bytes_rt; // ek_pke region end
wire [11:0] dk_hek_end = dk_ek_end + 12'd32; // H(ek) region end
// LOAD-TIME boundaries: dk is streamed in BEFORE start_i, so k_r is not yet
// captured. Route the dk load by the LIVE k_i input instead of k_r.
wire [11:0] dkp_bytes_ld = 12'd384 * {9'b0, k_i}; // 768/1152/1536
wire [11:0] ek_bytes_ld = 12'd384 * {9'b0, k_i} + 12'd32; // 800/1184/1568
wire [11:0] dk_ek_end_ld = dkp_bytes_ld + ek_bytes_ld; // ek_pke region end
wire [11:0] dk_hek_end_ld= dk_ek_end_ld + 12'd32; // H(ek) region end
// H(ek) block count = ceil((ek_bytes+1)/136): 6/9/12 for k=2/3/4 (table)
wire [3:0] h_nblk_rt = (k_r == 3'd2) ? 4'd6 : (k_r == 3'd3) ? 4'd9 : 4'd12;
wire [11:0] h_last_rt = {6'b0, h_nblk_rt} * 12'd136 - 12'd1; // final padded byte index
// ---- Encaps runtime params ----
reg op_r; // 0=KeyGen 1=Encaps (captured at start)
reg [255:0] m_r; // Encaps message m (captured at start)
reg [1:0] op_r; // 0=KeyGen 1=Encaps 2=Decaps (captured at start)
reg [255:0] m_r; // Encaps message m (captured at start); Decaps reuses for m'
reg [255:0] ss_r; // Encaps shared secret K (= G output low half)
reg [255:0] r_r; // Encaps PRF seed r (= G output high half)
// ---- Decaps runtime params (parsed from dk during load) ----
reg [255:0] z_r; // implicit-rejection seed z (dk tail), captured at load
reg [255:0] kbar_r; // K-bar = J(z||c) (D5)
// FIPS 203: eta2 = 2 for all parameter sets.
wire [1:0] eta2_rt = 2'd2;
// Compression params: (du,dv) = (10,4) for k=2/3, (11,5) for k=4.
@@ -100,6 +133,7 @@ module mlkem_top #(
wire [11:0] c1_bytes_rt = 12'd32 * {7'b0, du_rt} * {9'b0, k_r}; // 640/960/1408
wire [11:0] c2_bytes_rt = 12'd32 * {7'b0, dv_rt}; // 128/128/160
wire [11:0] ct_bytes_rt = c1_bytes_rt + c2_bytes_rt; // 768/1088/1568
wire [11:0] ct_bytes_rt = c1_bytes_rt + c2_bytes_rt; // 768/1088/1568
assign ss_o = ss_r;
// ---- E1: rho-load + byteDecode12 bookkeeping ----
@@ -145,6 +179,23 @@ module mlkem_top #(
assign ct_rd_addr = dbg_ct_idx_i;
assign dbg_ct_o = ct_rd_data;
// ---- c_in_bram: Decaps input ciphertext c (<=1568 B). Preloaded via
// c_in_* in ST_IDLE. Kept separate from ct_bram (which holds the computed
// c') so D7 can compare c' vs c and D5 J(z||c) can read c. Read port is
// muxed (D5 J-feed / D7 compare); for D0 only the load write path exists.
wire [10:0] cin_rd_addr;
wire [7:0] cin_rd_data;
reg [10:0] cin_rd_addr_r; // D5/D7 read address (tied 0 until then)
sd_bram #(.W(8), .D(2048), .A(11)) u_c_in_bram (
.clk(clk),
.rd_addr(cin_rd_addr), .rd_data(cin_rd_data),
.wr_en(c_in_we), .wr_addr(c_in_addr), .wr_data(c_in_byte)
);
assign cin_rd_addr = cin_rd_addr_r;
/* verilator lint_off UNUSEDSIGNAL */
wire [7:0] cin_rd_data_unused = cin_rd_data; // consumed in D5/D7
/* verilator lint_on UNUSEDSIGNAL */
// ================================================================
// Polynomial storage, sized for KMAX (worst case). Runtime k uses a
// sub-range. Slot layout (each slot = 256 coeffs):
@@ -330,26 +381,39 @@ module mlkem_top #(
reg [7:0] ek_wd, dkp_wd;
// ek BRAM write port: KeyGen ST_E drives ek_we/ek_wa/ek_wd; Encaps preloads
// ek from the external ek_in_* port (TB streams ek=pk before start_i). The
// two never overlap (preload happens in ST_IDLE before an Encaps run).
wire ek_we_mux = ek_in_we ? 1'b1 : ek_we;
wire [10:0] ek_wa_mux = ek_in_we ? ek_in_addr : ek_wa;
wire [7:0] ek_wd_mux = ek_in_we ? ek_in_byte : ek_wd;
// ek from the external ek_in_* port (TB streams ek=pk before start_i). Decaps
// routes the dk's ek_pke region (dk[384K:768K+32]) here at load time. The
// three never overlap (preloads happen in ST_IDLE before a run).
wire dk_ld_ekpke = dk_in_we && (dk_in_addr >= dkp_bytes_ld)
&& (dk_in_addr < dk_ek_end_ld); // ek_pke region
wire dk_ld_dkpke = dk_in_we && (dk_in_addr < dkp_bytes_ld); // dk_pke region
wire [10:0] dk_ek_off = (dk_in_addr - dkp_bytes_ld); // offset within ek
wire ek_we_mux = ek_in_we ? 1'b1 :
dk_ld_ekpke ? 1'b1 : ek_we;
wire [10:0] ek_wa_mux = ek_in_we ? ek_in_addr :
dk_ld_ekpke ? dk_ek_off : ek_wa;
wire [7:0] ek_wd_mux = ek_in_we ? ek_in_byte :
dk_ld_ekpke ? dk_in_byte : ek_wd;
sd_bram #(.W(8), .D(2048), .A(11)) u_ek_bram (
.clk(clk),
.rd_addr(ek_rd_addr), .rd_data(ek_rd_data),
.wr_en(ek_we_mux), .wr_addr(ek_wa_mux), .wr_data(ek_wd_mux)
);
// dk_pke BRAM write port: KeyGen ST_E (s_hat byteEncode) drives dkp_we/wa/wd;
// Decaps routes the dk's dk_pke region (dk[0:384K]) here at load time. The
// two never overlap (preload in ST_IDLE before a Decaps run).
wire dkp_we_mux = dk_ld_dkpke ? 1'b1 : dkp_we;
wire [10:0] dkp_wa_mux = dk_ld_dkpke ? dk_in_addr[10:0] : dkp_wa;
wire [7:0] dkp_wd_mux = dk_ld_dkpke ? dk_in_byte : dkp_wd;
sd_bram #(.W(8), .D(2048), .A(11)) u_dkp_bram (
.clk(clk),
.rd_addr(dkp_rd_addr), .rd_data(dkp_rd_data),
.wr_en(dkp_we), .wr_addr(dkp_wa), .wr_data(dkp_wd)
.wr_en(dkp_we_mux), .wr_addr(dkp_wa_mux), .wr_data(dkp_wd_mux)
);
// full dk = dk_pke(dk_bytes) || ek(ek_bytes) || H(ek)(32) || z(32)
wire [11:0] dk_ek_end = dk_bytes_rt + ek_bytes_rt; // ek region end
wire [11:0] dk_hek_end = dk_ek_end + 12'd32; // H(ek) region end
// full dk = dk_pke(dk_bytes) || ek(ek_bytes) || H(ek)(32) || z(32).
// (dk_ek_end / dk_hek_end declared near the top with the other size wires.)
// Debug-region selects for dk readback (combinational region decode).
wire dbgdk_in_dkp = (dbg_dk_idx_i < dk_bytes_rt);
@@ -390,6 +454,8 @@ module mlkem_top #(
localparam ST_ENC_V = 5'd17; // v = INTT(sum t_hat o y_hat) + e2 + mu
localparam ST_ENC_C2 = 5'd18; // Compress_dv + byteEncode_dv -> ct c2
localparam ST_ENC_E2MV = 5'd19; // relocate e2 bank_t[0] -> bank_a[E2_ASLOT]
// ---- Decaps states ----
localparam ST_DEC_LOAD = 5'd20; // dk/c already streamed in; parse/settle (D0)
localparam ST_DONE = 5'd31;
reg [4:0] st, st_next;
@@ -418,6 +484,11 @@ module mlkem_top #(
assign dbg_sigma_o = sigma_r;
assign dbg_r_o = r_r;
assign dbg_hek_o = hek_r;
// Decaps taps: m' reuses m_r (Decrypt writes it), z/h parsed from dk at load.
assign dbg_mprime_o = m_r;
assign dbg_kbar_o = kbar_r;
assign dbg_decz_o = z_r;
assign dbg_dech_o = hek_r; // Decaps parses dk's H(ek) into hek_r at load
// ---- sha3_top in G mode: data_i = {K_byte, d} (d byte0 in [7:0]) ----
reg sha3_valid;
@@ -897,7 +968,11 @@ module mlkem_top #(
always @(*) begin
st_next = st;
case (st)
ST_IDLE: if (start_i) st_next = op_i ? ST_ENC_H : ST_G;
ST_IDLE: if (start_i) st_next = (op_i == 2'd2) ? ST_DEC_LOAD :
(op_i == 2'd1) ? ST_ENC_H : ST_G;
// D0: settle after dk/c parse, then (D1+) proceed to Decrypt. For now
// ST_DEC_LOAD just lands in DONE so the load/parse can be dbg-checked.
ST_DEC_LOAD: st_next = ST_DONE;
ST_G: if (sha3_vo) st_next = ST_A;
ST_A: if (a_pair >= kk_rt) st_next = ST_C;
ST_C: if (c_poly >= {1'b0, k_r, 1'b0}) st_next = ST_N;
@@ -927,7 +1002,7 @@ module mlkem_top #(
if (!rst_n) begin
st <= ST_IDLE;
k_r <= 3'd0;
op_r <= 1'b0;
op_r <= 2'd0;
m_r <= 256'd0;
ss_r <= 256'd0;
r_r <= 256'd0;
@@ -1024,6 +1099,9 @@ module mlkem_top #(
h_mblast <= 1'b0;
h_ack <= 1'b0;
hek_r <= 256'd0;
z_r <= 256'd0;
kbar_r <= 256'd0;
cin_rd_addr_r <= 11'd0;
h_blk <= 3'd0;
h_byte <= 8'd0;
h_phase <= 2'd0;
@@ -1043,11 +1121,12 @@ module mlkem_top #(
em_we <= 1'b0; // e2-relocate bank_a write default low (E6)
// Kick off when entering from IDLE: KeyGen starts G; Encaps captures
// op/m and arms the H(ek) machinery (ST_ENC_H reuses the ST_H FSM).
// op/m and arms the H(ek) machinery (ST_ENC_H reuses the ST_H FSM);
// Decaps captures op (dk/c already streamed; z/h captured at load).
if (st == ST_IDLE && start_i) begin
k_r <= k_i; // capture runtime ML-KEM param
op_r <= op_i;
if (op_i) begin
if (op_i == 2'd1) begin
m_r <= msg_i; // capture Encaps message
// arm H(ek) (same fields the ST_E->ST_H arming sets)
h_blk <= 3'd0;
@@ -1057,11 +1136,25 @@ module mlkem_top #(
h_mblast <= 1'b0;
h_ack <= 1'b1; // ready to consume final digest
h_wb_vld <= 1'b0;
end else if (op_i == 2'd2) begin
// Decaps: dk_pke/ek_pke already in BRAM; z/H(ek) captured into
// z_r/hek_r during load (below). Nothing else to arm in D0.
end else begin
sha3_valid <= 1'b1;
sha3_ack <= 1'b1;
end
end
// Decaps dk load: capture the H(ek) and z byte regions into registers
// as they stream in (the dk_pke/ek_pke regions go to BRAM via the
// write muxes). Byte i within a region lands in [8i +: 8]. Uses the
// LIVE k_i boundaries (k_r not yet captured during preload).
if (dk_in_we) begin
if (dk_in_addr >= dk_ek_end_ld && dk_in_addr < dk_hek_end_ld)
hek_r[(dk_in_addr - dk_ek_end_ld)*8 +: 8] <= dk_in_byte; // H(ek)
else if (dk_in_addr >= dk_hek_end_ld)
z_r[(dk_in_addr - dk_hek_end_ld)*8 +: 8] <= dk_in_byte; // z
end
// Drop valid once accepted
if (sha3_valid && sha3_ready) sha3_valid <= 1'b0;