Add DMA-friendly MLKEM windows

This commit is contained in:
2026-07-05 20:07:14 +08:00
parent d13f7933d8
commit 2b7c43edcb
7 changed files with 281 additions and 69 deletions

View File

@@ -1,8 +1,8 @@
`define FFT_CSR_ADDR 16'hf000
`define FFT_IN_RE_BASE 16'h1000
`define FFT_IN_IM_BASE 16'h2000
`define FFT_OUT_RE_BASE 16'h3000
`define FFT_OUT_IM_BASE 16'h4000
`define FFT_CSR_ADDR 16'hc000
`define FFT_IN_RE_BASE 16'h8000
`define FFT_IN_IM_BASE 16'h9000
`define FFT_OUT_RE_BASE 16'ha000
`define FFT_OUT_IM_BASE 16'hb000
module fft_controller #(
parameter FFT_LEN = 1024, // FFT处理长度
@@ -201,10 +201,10 @@ module fft_controller #(
assign s_rresp = 2'b00;
// 地址译码
wire hit_in_re = (aw_addr_reg[15:12] == 4'h1);
wire hit_in_im = (aw_addr_reg[15:12] == 4'h2);
wire hit_out_re = (ar_addr_reg[15:12] == 4'h3);
wire hit_out_im = (ar_addr_reg[15:12] == 4'h4);
wire hit_in_re = (aw_addr_reg[15:12] == 4'h8);
wire hit_in_im = (aw_addr_reg[15:12] == 4'h9);
wire hit_out_re = (ar_addr_reg[15:12] == 4'ha);
wire hit_out_im = (ar_addr_reg[15:12] == 4'hb);
// 寄存器控制
reg [3:0] fft_ctrl;
@@ -380,4 +380,4 @@ module fft_controller #(
s_rdata = rdata_d;
end
endmodule
endmodule

View File

@@ -2,9 +2,9 @@
//
// Both software stacks use the same base address, 0xbf400000. Their internal
// maps do not overlap:
// ML-KEM: 0x0000..0x0114
// FFT: 0x1000, 0x2000, 0x3000, 0x4000, 0xf000 windows
// Route 0x0xxx to ML-KEM and all other offsets to FFT.
// ML-KEM: 0x0000..0x3fff
// FFT: 0x8000..0xcfff
// Route the low ML-KEM block to ML-KEM and the high FFT block to FFT.
module fft_mlkem_axi_slave (
input clk,
input resetn,
@@ -54,8 +54,8 @@ module fft_mlkem_axi_slave (
output mlkem_intr
);
wire aw_to_mlkem = (s_axi_awaddr[15:12] == 4'h0);
wire ar_to_mlkem = (s_axi_araddr[15:12] == 4'h0);
wire aw_to_mlkem = (s_axi_awaddr[15:12] <= 4'h3);
wire ar_to_mlkem = (s_axi_araddr[15:12] <= 4'h3);
reg write_sel_mlkem;
reg write_active;

View File

@@ -24,6 +24,9 @@
// 0x10C RDBK_SEL RW 0=ek 1=dk_pke 2=full-dk 3=ciphertext
// 0x110 RDBK_ADDR RW output byte index
// 0x114 RDBK_DATA R selected readback-tap byte (low 8 bits)
// 0x1000..0x1FFF EK W: preload ek_in, R: read generated ek
// 0x2000..0x2FFF DK W: preload dk_in, R: read generated full dk
// 0x3000..0x3FFF CT W: preload c_in, R: read generated ciphertext
module mlkem_axi_slave (
input clk, // sys_clk (AXI bus clock, 50 MHz)
input resetn, // sys_resetn (active low, already synchronized)
@@ -94,6 +97,9 @@ module mlkem_axi_slave (
localparam ADDR_RDB_SEL = 12'h10C;
localparam ADDR_RDB_ADR = 12'h110;
localparam ADDR_RDB_DAT = 12'h114;
localparam WIN_EK = 4'h1;
localparam WIN_DK = 4'h2;
localparam WIN_CT = 4'h3;
assign s_axi_bresp = 2'b00;
assign s_axi_rresp = 2'b00;
@@ -111,6 +117,41 @@ module mlkem_axi_slave (
end
endfunction
function [7:0] select_byte;
input [31:0] word;
input [1:0] lane;
begin
case (lane)
2'd0: select_byte = word[7:0];
2'd1: select_byte = word[15:8];
2'd2: select_byte = word[23:16];
default: select_byte = word[31:24];
endcase
end
endfunction
function [1:0] window_preload_sel;
input [3:0] window;
begin
case (window)
WIN_DK: window_preload_sel = 2'd1;
WIN_CT: window_preload_sel = 2'd2;
default: window_preload_sel = 2'd0;
endcase
end
endfunction
function [1:0] window_readback_sel;
input [3:0] window;
begin
case (window)
WIN_DK: window_readback_sel = 2'd2; // full dk
WIN_CT: window_readback_sel = 2'd3; // ciphertext
default: window_readback_sel = 2'd0; // ek
endcase
end
endfunction
// ==========================================
// mlkem_clk = sys_clk / 2 (25 MHz) via toggle FF.
// ==========================================
@@ -184,19 +225,25 @@ module mlkem_axi_slave (
reg [1:0] pre_sel_l;
reg [11:0] pre_addr_l;
reg [7:0] pre_byte_l;
reg pre_auto_inc_l;
// ==========================================
// mlkem_top instance
// ==========================================
wire [7:0] dbg_byte_o, dbg_dk_o, dbg_ct_o;
wire [255:0] ss_o;
wire dbg_byte_sel = (rdb_sel == 2'd1); // 0=ek, 1=dk_pke
wire bulk_read_active;
wire [1:0] bulk_rdb_sel;
wire [11:0] bulk_rdb_addr;
wire [1:0] eff_rdb_sel = bulk_read_active ? bulk_rdb_sel : rdb_sel;
wire [11:0] eff_rdb_addr = bulk_read_active ? bulk_rdb_addr : rdb_addr;
wire dbg_byte_sel = (eff_rdb_sel == 2'd1); // 0=ek, 1=dk_pke
// The core's ek/dkp read-address muxes give dbg_dk_idx_i priority over
// dbg_byte_idx_i whenever it lands in the ek (768..1567) or dk_pke (<768)
// region (mlkem_top: ek_rd_addr / dkp_rd_addr). So only drive dbg_dk_idx_i
// with the real index during a full-dk readback (rdb_sel==2); otherwise hold
// it out of range (0xFFF) so an ek/dk_pke readback isn't corrupted.
wire [11:0] dbg_dk_idx = (rdb_sel == 2'd2) ? rdb_addr : 12'hFFF;
wire [11:0] dbg_dk_idx = (eff_rdb_sel == 2'd2) ? eff_rdb_addr : 12'hFFF;
// one-cycle write-enable pulses into the core (combinational off the synced
// edge; pre_sel_l/pre_addr_l/pre_byte_l are stable in sys, sampled by core).
@@ -230,7 +277,7 @@ module mlkem_axi_slave (
// shared secret out
.ss_o (ss_o),
// ciphertext readback tap
.dbg_ct_idx_i (rdb_addr[10:0]),
.dbg_ct_idx_i (eff_rdb_addr[10:0]),
.dbg_ct_o (dbg_ct_o),
// coefficient debug tap (unused by CPU path; tie inputs off)
.dbg_slot_i (6'd0),
@@ -238,7 +285,7 @@ module mlkem_axi_slave (
.dbg_coeff_o (/* unused */),
// ek / dk_pke byte readback
.dbg_byte_sel_i (dbg_byte_sel),
.dbg_byte_idx_i (rdb_addr[10:0]),
.dbg_byte_idx_i (eff_rdb_addr[10:0]),
.dbg_byte_o (dbg_byte_o),
// full-dk byte readback
.dbg_dk_idx_i (dbg_dk_idx),
@@ -304,17 +351,27 @@ module mlkem_axi_slave (
// AXI Write channel (INCR burst). Preload-data beats are wait-stated until
// the cross-domain byte handshake completes (W_PRELOAD).
// ==========================================
localparam W_IDLE = 2'd0;
localparam W_DATA = 2'd1;
localparam W_PRELOAD = 2'd2;
localparam W_RESP = 2'd3;
localparam W_IDLE = 3'd0;
localparam W_DATA = 3'd1;
localparam W_PRELOAD = 3'd2;
localparam W_STREAM = 3'd3;
localparam W_RESP = 3'd4;
reg [1:0] w_state;
reg [2:0] w_state;
reg [31:0] w_addr;
reg [4:0] w_id;
reg [1:0] w_burst;
reg w_last_l; // saved wlast across a preload wait
wire [11:0] w_offset = w_addr[11:0];
wire w_stream_window = (w_addr[15:12] == WIN_EK) ||
(w_addr[15:12] == WIN_DK) ||
(w_addr[15:12] == WIN_CT);
reg [31:0] stream_wdata;
reg [3:0] stream_wstrb;
reg [11:0] stream_byte_addr;
reg [2:0] stream_lane;
reg [1:0] stream_sel;
assign s_axi_bid = w_id;
@@ -338,7 +395,13 @@ module mlkem_axi_slave (
pre_sel_l <= 2'd0;
pre_addr_l <= 12'd0;
pre_byte_l <= 8'd0;
pre_auto_inc_l <= 1'b0;
w_last_l <= 1'b0;
stream_wdata <= 32'd0;
stream_wstrb <= 4'd0;
stream_byte_addr <= 12'd0;
stream_lane <= 3'd0;
stream_sel <= 2'd0;
for (wi = 0; wi < 8; wi = wi + 1) begin
d_word[wi] = 32'd0;
z_word[wi] = 32'd0;
@@ -348,7 +411,8 @@ module mlkem_axi_slave (
// clear the preload-busy flag when the mlkem domain acks the byte
if (pre_ack_edge) begin
pre_busy <= 1'b0;
pre_addr <= pre_addr + 12'd1; // auto-increment for next byte
if (pre_auto_inc_l)
pre_addr <= pre_addr + 12'd1; // legacy PRE_DAT auto-increment
end
case (w_state)
@@ -366,11 +430,21 @@ module mlkem_axi_slave (
W_DATA: begin
if (s_axi_wvalid && s_axi_wready) begin
if (w_offset == ADDR_PRE_DAT) begin
if (w_stream_window) begin
stream_wdata <= s_axi_wdata;
stream_wstrb <= s_axi_wstrb;
stream_byte_addr <= w_offset;
stream_lane <= 3'd0;
stream_sel <= window_preload_sel(w_addr[15:12]);
w_last_l <= s_axi_wlast;
s_axi_wready <= 1'b0;
w_state <= W_STREAM;
end else if (w_offset == ADDR_PRE_DAT) begin
// launch cross-domain byte write; stall the beat.
pre_sel_l <= pre_sel;
pre_addr_l <= pre_addr;
pre_byte_l <= s_axi_wdata[7:0];
pre_auto_inc_l <= 1'b1;
pre_req_tg <= ~pre_req_tg;
pre_busy <= 1'b1;
w_last_l <= s_axi_wlast;
@@ -424,6 +498,31 @@ module mlkem_axi_slave (
end
end
W_STREAM: begin
if (!pre_busy) begin
if (stream_lane < 3'd4 && !stream_wstrb[stream_lane[1:0]]) begin
stream_lane <= stream_lane + 3'd1;
end else if (stream_lane < 3'd4) begin
pre_sel_l <= stream_sel;
pre_addr_l <= stream_byte_addr + {10'd0, stream_lane[1:0]};
pre_byte_l <= select_byte(stream_wdata, stream_lane[1:0]);
pre_auto_inc_l <= 1'b0;
pre_req_tg <= ~pre_req_tg;
pre_busy <= 1'b1;
stream_lane <= stream_lane + 3'd1;
end else begin
if (w_burst == 2'b01) w_addr <= w_addr + 32'd4;
if (w_last_l) begin
s_axi_bvalid <= 1'b1;
w_state <= W_RESP;
end else begin
s_axi_wready <= 1'b1;
w_state <= W_DATA;
end
end
end
end
W_RESP: begin
if (s_axi_bvalid && s_axi_bready) begin
s_axi_bvalid <= 1'b0;
@@ -439,15 +538,31 @@ module mlkem_axi_slave (
// ==========================================
// AXI Read channel (INCR burst).
// ==========================================
localparam R_IDLE = 1'b0;
localparam R_DATA = 1'b1;
localparam R_IDLE = 3'd0;
localparam R_DATA = 3'd1;
localparam R_BULK_WAIT = 3'd2;
localparam RDB_WAIT_CYCLES = 3'd4;
reg r_state;
reg [2:0] r_state;
reg [31:0] r_addr;
reg [7:0] r_cnt;
reg [7:0] r_len;
reg [1:0] r_burst;
wire [11:0] r_offset = r_addr[11:0];
wire ar_stream_window = (s_axi_araddr[15:12] == WIN_EK) ||
(s_axi_araddr[15:12] == WIN_DK) ||
(s_axi_araddr[15:12] == WIN_CT);
reg r_bulk;
reg [1:0] bulk_sel_r;
reg [11:0] bulk_addr_r;
reg [1:0] bulk_lane;
reg [2:0] bulk_wait;
reg [31:0] bulk_word;
assign bulk_read_active = r_bulk && (r_state == R_BULK_WAIT);
assign bulk_rdb_sel = bulk_sel_r;
assign bulk_rdb_addr = bulk_addr_r;
always @(posedge clk or negedge resetn) begin
if (!resetn) begin
@@ -460,6 +575,12 @@ module mlkem_axi_slave (
r_cnt <= 8'd0;
r_len <= 8'd0;
r_burst <= 2'd0;
r_bulk <= 1'b0;
bulk_sel_r <= 2'd0;
bulk_addr_r <= 12'd0;
bulk_lane <= 2'd0;
bulk_wait <= 3'd0;
bulk_word <= 32'd0;
end else begin
case (r_state)
R_IDLE: begin
@@ -470,10 +591,44 @@ module mlkem_axi_slave (
r_len <= s_axi_arlen;
r_burst <= s_axi_arburst;
r_cnt <= 8'd0;
s_axi_rvalid <= 1'b1;
s_axi_rlast <= (s_axi_arlen == 8'd0);
s_axi_rid <= s_axi_arid;
r_state <= R_DATA;
r_bulk <= ar_stream_window;
if (ar_stream_window) begin
bulk_sel_r <= window_readback_sel(s_axi_araddr[15:12]);
bulk_addr_r <= s_axi_araddr[11:0];
bulk_lane <= 2'd0;
bulk_wait <= RDB_WAIT_CYCLES;
bulk_word <= 32'd0;
s_axi_rvalid <= 1'b0;
s_axi_rlast <= 1'b0;
r_state <= R_BULK_WAIT;
end else begin
s_axi_rvalid <= 1'b1;
s_axi_rlast <= (s_axi_arlen == 8'd0);
r_state <= R_DATA;
end
end
end
R_BULK_WAIT: begin
if (bulk_wait != 3'd0) begin
bulk_wait <= bulk_wait - 3'd1;
end else begin
case (bulk_lane)
2'd0: bulk_word[7:0] <= rdb_byte;
2'd1: bulk_word[15:8] <= rdb_byte;
2'd2: bulk_word[23:16] <= rdb_byte;
default: bulk_word[31:24] <= rdb_byte;
endcase
if (bulk_lane == 2'd3) begin
s_axi_rvalid <= 1'b1;
s_axi_rlast <= (r_cnt == r_len);
r_state <= R_DATA;
end else begin
bulk_lane <= bulk_lane + 2'd1;
bulk_addr_r <= bulk_addr_r + 12'd1;
bulk_wait <= RDB_WAIT_CYCLES;
end
end
end
R_DATA: begin
@@ -482,11 +637,22 @@ module mlkem_axi_slave (
s_axi_rvalid <= 1'b0;
s_axi_rlast <= 1'b0;
s_axi_arready <= 1'b1;
r_bulk <= 1'b0;
r_state <= R_IDLE;
end else begin
r_cnt <= r_cnt + 8'd1;
s_axi_rlast <= (r_cnt + 8'd1 == r_len);
if (r_burst == 2'b01) r_addr <= r_addr + 32'd4;
if (r_bulk) begin
bulk_addr_r <= (r_burst == 2'b01) ? (r_offset + 12'd4) : r_offset;
bulk_lane <= 2'd0;
bulk_wait <= RDB_WAIT_CYCLES;
bulk_word <= 32'd0;
s_axi_rvalid <= 1'b0;
s_axi_rlast <= 1'b0;
r_state <= R_BULK_WAIT;
end else begin
s_axi_rlast <= (r_cnt + 8'd1 == r_len);
end
end
end
end
@@ -499,12 +665,14 @@ module mlkem_axi_slave (
// first beat is valid the cycle rvalid asserts (matches matmul style).
wire [11:0] cur_off = (r_state == R_DATA) ? r_offset : s_axi_araddr[11:0];
// Selected readback byte (taps are combinational over the core's BRAMs).
wire [7:0] rdb_byte = (rdb_sel == 2'd3) ? dbg_ct_o :
(rdb_sel == 2'd2) ? dbg_dk_o : dbg_byte_o;
wire [7:0] rdb_byte = (eff_rdb_sel == 2'd3) ? dbg_ct_o :
(eff_rdb_sel == 2'd2) ? dbg_dk_o : dbg_byte_o;
always @(*) begin
s_axi_rdata = 32'd0;
if (r_state == R_DATA || (r_state == R_IDLE && s_axi_arvalid)) begin
if (r_bulk && r_state == R_DATA) begin
s_axi_rdata = bulk_word;
end else if (r_state == R_DATA || (r_state == R_IDLE && s_axi_arvalid)) begin
case (1'b1)
(cur_off == ADDR_CTRL): s_axi_rdata = 32'd0;
(cur_off == ADDR_STATUS): s_axi_rdata = {29'd0, pre_busy, status_done, status_busy};