// mlkem_axi_slave.v — AXI4 (burst) slave wrapper for the ML-KEM core. // Makes the ML-KEM accelerator (KeyGen / Encaps / Decaps) CPU-controllable. // // The bus runs at sys_clk (50 MHz). The ML-KEM core cannot close timing at // 50 MHz, so it is clocked by mlkem_clk = sys_clk / 2 (25 MHz), produced by a // toggle flip-flop in this wrapper. mlkem_clk is a phase-aligned synchronous // divide of sys_clk; every signal crossing between the two domains still goes // through an explicit synchronizer or toggle handshake. // NOTE (FPGA): for real silicon route mlkem_clk through a BUFG/BUFGCE rather // than a logic-generated net; kept as a plain toggle here for portability. // // Register map (byte offsets from base, e.g. virtual 0xbf400000): // 0x000 CTRL W bit0 = start (launches op selected in CONFIG) // 0x004 STATUS R bit0 busy, bit1 done, bit2 preload_busy // W1C write bit1=1 to clear the latched done flag // 0x008 CONFIG RW [2:0] k (2/3/4), [5:4] op (0=KeyGen 1=Encaps 2=Decaps) // 0x010..0x02C D RW KeyGen seed d (8 words; d byte0 in word0[7:0]) // 0x030..0x04C Z RW implicit-reject seed z (8 words) // 0x050..0x06C MSG RW Encaps message m (8 words) // 0x070..0x08C SS R shared secret out (8 words) // 0x100 PRELOAD_SEL RW 0=ek_in 1=dk_in 2=c_in // 0x104 PRELOAD_ADDR RW byte address (auto-increments after each data write) // 0x108 PRELOAD_DATA W write low byte -> streams to selected port (wait-stated) // 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) module mlkem_axi_slave ( input clk, // sys_clk (AXI bus clock, 50 MHz) input resetn, // sys_resetn (active low, already synchronized) // AXI Write Address Channel input [4:0] s_axi_awid, input [31:0] s_axi_awaddr, input [7:0] s_axi_awlen, input [2:0] s_axi_awsize, input [1:0] s_axi_awburst, input s_axi_awlock, input [3:0] s_axi_awcache, input [2:0] s_axi_awprot, input s_axi_awvalid, output reg s_axi_awready, // AXI Write Data Channel input [4:0] s_axi_wid, input [31:0] s_axi_wdata, input [3:0] s_axi_wstrb, input s_axi_wlast, input s_axi_wvalid, output reg s_axi_wready, // AXI Write Response Channel output [4:0] s_axi_bid, output [1:0] s_axi_bresp, output reg s_axi_bvalid, input s_axi_bready, // AXI Read Address Channel input [4:0] s_axi_arid, input [31:0] s_axi_araddr, input [7:0] s_axi_arlen, input [2:0] s_axi_arsize, input [1:0] s_axi_arburst, input s_axi_arlock, input [3:0] s_axi_arcache, input [2:0] s_axi_arprot, input s_axi_arvalid, output reg s_axi_arready, // AXI Read Data Channel output reg [4:0] s_axi_rid, output reg [31:0] s_axi_rdata, output [1:0] s_axi_rresp, output reg s_axi_rlast, output reg s_axi_rvalid, input s_axi_rready, // Interrupt output (active high level, asserted while done latched) output intr ); // ========================================== // Address map (byte offsets within the 1 MB slave window) // ========================================== localparam ADDR_CTRL = 12'h000; localparam ADDR_STATUS = 12'h004; localparam ADDR_CONFIG = 12'h008; localparam ADDR_D = 12'h010; // 0x010..0x02C (8 words) localparam ADDR_Z = 12'h030; // 0x030..0x04C localparam ADDR_MSG = 12'h050; // 0x050..0x06C localparam ADDR_SS = 12'h070; // 0x070..0x08C (read-only) localparam ADDR_PRE_SEL = 12'h100; localparam ADDR_PRE_ADR = 12'h104; localparam ADDR_PRE_DAT = 12'h108; localparam ADDR_RDB_SEL = 12'h10C; localparam ADDR_RDB_ADR = 12'h110; localparam ADDR_RDB_DAT = 12'h114; assign s_axi_bresp = 2'b00; assign s_axi_rresp = 2'b00; function [31:0] apply_wstrb; input [31:0] old_value; input [31:0] new_value; input [3:0] strobe; begin apply_wstrb = old_value; if (strobe[0]) apply_wstrb[7:0] = new_value[7:0]; if (strobe[1]) apply_wstrb[15:8] = new_value[15:8]; if (strobe[2]) apply_wstrb[23:16] = new_value[23:16]; if (strobe[3]) apply_wstrb[31:24] = new_value[31:24]; end endfunction // ========================================== // mlkem_clk = sys_clk / 2 (25 MHz) via toggle FF. // ========================================== reg mlkem_clk; always @(posedge clk or negedge resetn) begin if (!resetn) mlkem_clk <= 1'b0; else mlkem_clk <= ~mlkem_clk; end // Reset for the mlkem domain: reuse sys_resetn (async-assert, and the core's // own rst_n is sampled on mlkem_clk edges). wire mlkem_rstn = resetn; // ========================================== // sys_clk-domain config / data registers (quasi-static; fed to the core // combinationally — written long before start, stable across the op). // ========================================== reg [2:0] cfg_k; // 2/3/4 reg [1:0] cfg_op; // 0=KeyGen 1=Encaps 2=Decaps reg [31:0] d_word [0:7]; reg [31:0] z_word [0:7]; reg [31:0] msg_word [0:7]; wire [255:0] d_bus = {d_word[7], d_word[6], d_word[5], d_word[4], d_word[3], d_word[2], d_word[1], d_word[0]}; wire [255:0] z_bus = {z_word[7], z_word[6], z_word[5], z_word[4], z_word[3], z_word[2], z_word[1], z_word[0]}; wire [255:0] msg_bus = {msg_word[7],msg_word[6],msg_word[5],msg_word[4], msg_word[3],msg_word[2],msg_word[1],msg_word[0]}; // Preload / readback index registers (sys domain). reg [1:0] pre_sel; // 0=ek 1=dk 2=c reg [11:0] pre_addr; // byte address into the selected port reg [1:0] rdb_sel; // 0=ek 1=dk_pke 2=full-dk 3=ciphertext reg [11:0] rdb_addr; // output byte index // ========================================== // CDC: start pulse (sys toggle -> mlkem edge-detect -> 1-cyc start_i) // ========================================== reg start_tg; // sys-domain toggle, flips on CTRL.start write reg [2:0] start_sync; // mlkem-domain 2FF sync (+1 for edge) wire mlkem_start = start_sync[2] ^ start_sync[1]; // 1 mlkem_clk pulse // ========================================== // CDC: busy / done back to sys // ========================================== wire core_busy; // busy_o (mlkem domain) wire core_done; // done_o pulse (mlkem domain) reg done_latch; // sticky done in mlkem domain reg [1:0] busy_sync; // -> sys reg [1:0] done_sync; // -> sys // done-clear: sys toggle -> mlkem edge -> clears done_latch reg dclr_tg; // sys-domain toggle (STATUS W1C of done) reg [2:0] dclr_sync; // mlkem-domain 2FF sync (+1 for edge) wire mlkem_dclr = dclr_sync[2] ^ dclr_sync[1]; // ========================================== // CDC: preload byte write handshake. // sys: latch {sel,addr,byte}, flip pre_req_tg, hold wready low. // mlkem: 2FF sync pre_req_tg, on edge pulse the selected *_in_we 1 cycle, // flip pre_ack_tg. // sys: 2FF sync pre_ack_tg, on edge release the bus (addr++). // ========================================== reg pre_req_tg; // sys toggle: a byte is pending reg [2:0] pre_req_sync; // mlkem 2FF (+1 for edge) wire mlkem_pre_we = pre_req_sync[2] ^ pre_req_sync[1]; reg pre_ack_tg; // mlkem toggle: byte consumed reg [2:0] pre_ack_sync; // sys 2FF (+1 for edge) wire pre_ack_edge = pre_ack_sync[2] ^ pre_ack_sync[1]; reg pre_busy; // sys: a preload handshake is in flight // latched byte/port for the in-flight preload (sys, stable while pre_busy) reg [1:0] pre_sel_l; reg [11:0] pre_addr_l; reg [7:0] pre_byte_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 // 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; // 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). wire ek_in_we = mlkem_pre_we && (pre_sel_l == 2'd0); wire dk_in_we = mlkem_pre_we && (pre_sel_l == 2'd1); wire c_in_we = mlkem_pre_we && (pre_sel_l == 2'd2); mlkem_top #(.KMAX(4)) u_mlkem ( .clk (mlkem_clk), .rst_n (mlkem_rstn), .k_i (cfg_k), .op_i (cfg_op), .d_i (d_bus), .z_i (z_bus), .msg_i (msg_bus), .start_i (mlkem_start), .busy_o (core_busy), .done_o (core_done), // ek preload .ek_in_we (ek_in_we), .ek_in_addr (pre_addr_l[10:0]), .ek_in_byte (pre_byte_l), // dk preload .dk_in_we (dk_in_we), .dk_in_addr (pre_addr_l), .dk_in_byte (pre_byte_l), // ciphertext preload .c_in_we (c_in_we), .c_in_addr (pre_addr_l[10:0]), .c_in_byte (pre_byte_l), // shared secret out .ss_o (ss_o), // ciphertext readback tap .dbg_ct_idx_i (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), .dbg_idx_i (8'd0), .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_o (dbg_byte_o), // full-dk byte readback .dbg_dk_idx_i (dbg_dk_idx), .dbg_dk_o (dbg_dk_o), // remaining debug taps unused .dbg_rho_o (/* unused */), .dbg_sigma_o (/* unused */), .dbg_r_o (/* unused */), .dbg_hek_o (/* unused */), .dbg_mprime_o (/* unused */), .dbg_kbar_o (/* unused */), .dbg_decz_o (/* unused */), .dbg_dech_o (/* unused */) ); // ========================================== // mlkem_clk domain: synchronizers, done latch, preload-ack toggle. // ========================================== always @(posedge mlkem_clk or negedge mlkem_rstn) begin if (!mlkem_rstn) begin start_sync <= 3'd0; dclr_sync <= 3'd0; pre_req_sync <= 3'd0; pre_ack_tg <= 1'b0; done_latch <= 1'b0; end else begin // 2FF (+edge) synchronizers of sys-domain toggles start_sync <= {start_sync[1:0], start_tg}; dclr_sync <= {dclr_sync[1:0], dclr_tg}; pre_req_sync <= {pre_req_sync[1:0], pre_req_tg}; // acknowledge a consumed preload byte (write-enable pulsed this cycle) if (mlkem_pre_we) pre_ack_tg <= ~pre_ack_tg; // sticky done: set on the core's done_o pulse, cleared by synced W1C. // (mlkem_start also clears it so a new run starts with done low.) if (core_done) done_latch <= 1'b1; else if (mlkem_dclr) done_latch <= 1'b0; else if (mlkem_start) done_latch <= 1'b0; end end // ========================================== // sys_clk domain: sync busy/done/preload-ack back from mlkem. // ========================================== always @(posedge clk or negedge resetn) begin if (!resetn) begin busy_sync <= 2'd0; done_sync <= 2'd0; pre_ack_sync <= 3'd0; end else begin busy_sync <= {busy_sync[0], core_busy}; done_sync <= {done_sync[0], done_latch}; pre_ack_sync <= {pre_ack_sync[1:0], pre_ack_tg}; end end wire status_busy = busy_sync[1]; wire status_done = done_sync[1]; assign intr = status_done; // ========================================== // 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; reg [1: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]; assign s_axi_bid = w_id; integer wi; always @(posedge clk or negedge resetn) begin if (!resetn) begin w_state <= W_IDLE; s_axi_awready <= 1'b1; s_axi_wready <= 1'b0; s_axi_bvalid <= 1'b0; cfg_k <= 3'd4; // default ML-KEM-1024 cfg_op <= 2'd0; pre_sel <= 2'd0; pre_addr <= 12'd0; rdb_sel <= 2'd0; rdb_addr <= 12'd0; start_tg <= 1'b0; dclr_tg <= 1'b0; pre_req_tg <= 1'b0; pre_busy <= 1'b0; pre_sel_l <= 2'd0; pre_addr_l <= 12'd0; pre_byte_l <= 8'd0; w_last_l <= 1'b0; for (wi = 0; wi < 8; wi = wi + 1) begin d_word[wi] <= 32'd0; z_word[wi] <= 32'd0; msg_word[wi] <= 32'd0; end end else begin // 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 end case (w_state) W_IDLE: begin s_axi_awready <= 1'b1; if (s_axi_awvalid && s_axi_awready) begin s_axi_awready <= 1'b0; s_axi_wready <= 1'b1; w_addr <= s_axi_awaddr; w_id <= s_axi_awid; w_burst <= s_axi_awburst; w_state <= W_DATA; end end W_DATA: begin if (s_axi_wvalid && s_axi_wready) begin 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_req_tg <= ~pre_req_tg; pre_busy <= 1'b1; w_last_l <= s_axi_wlast; s_axi_wready <= 1'b0; w_state <= W_PRELOAD; end else begin // normal single-cycle register write case (1'b1) (w_offset == ADDR_CTRL): if (s_axi_wdata[0]) start_tg <= ~start_tg; (w_offset == ADDR_STATUS): if (s_axi_wdata[1]) dclr_tg <= ~dclr_tg; // clear done (w_offset == ADDR_CONFIG): begin cfg_k <= s_axi_wdata[2:0]; cfg_op <= s_axi_wdata[5:4]; end (w_offset == ADDR_PRE_SEL): pre_sel <= s_axi_wdata[1:0]; (w_offset == ADDR_PRE_ADR): pre_addr <= s_axi_wdata[11:0]; (w_offset == ADDR_RDB_SEL): rdb_sel <= s_axi_wdata[1:0]; (w_offset == ADDR_RDB_ADR): rdb_addr <= s_axi_wdata[11:0]; (w_offset >= ADDR_D && w_offset < ADDR_D + 12'h20): d_word[(w_offset-ADDR_D)>>2] <= apply_wstrb(d_word[(w_offset-ADDR_D)>>2], s_axi_wdata, s_axi_wstrb); (w_offset >= ADDR_Z && w_offset < ADDR_Z + 12'h20): z_word[(w_offset-ADDR_Z)>>2] <= apply_wstrb(z_word[(w_offset-ADDR_Z)>>2], s_axi_wdata, s_axi_wstrb); (w_offset >= ADDR_MSG && w_offset < ADDR_MSG + 12'h20): msg_word[(w_offset-ADDR_MSG)>>2] <= apply_wstrb(msg_word[(w_offset-ADDR_MSG)>>2], s_axi_wdata, s_axi_wstrb); default: ; // SS / RDB_DAT / unmapped: ignore endcase if (w_burst == 2'b01) w_addr <= w_addr + 32'd4; if (s_axi_wlast) begin s_axi_wready <= 1'b0; s_axi_bvalid <= 1'b1; w_state <= W_RESP; end end end end W_PRELOAD: begin // wait for the synced ack (pre_ack_edge clears pre_busy above) if (!pre_busy) 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 W_RESP: begin if (s_axi_bvalid && s_axi_bready) begin s_axi_bvalid <= 1'b0; s_axi_awready <= 1'b1; w_state <= W_IDLE; end end default: w_state <= W_IDLE; endcase end end // ========================================== // AXI Read channel (INCR burst). // ========================================== localparam R_IDLE = 1'b0; localparam R_DATA = 1'b1; reg 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]; always @(posedge clk or negedge resetn) begin if (!resetn) begin r_state <= R_IDLE; s_axi_arready <= 1'b1; s_axi_rvalid <= 1'b0; s_axi_rlast <= 1'b0; s_axi_rid <= 5'd0; r_addr <= 32'd0; r_cnt <= 8'd0; r_len <= 8'd0; r_burst <= 2'd0; end else begin case (r_state) R_IDLE: begin s_axi_arready <= 1'b1; if (s_axi_arvalid && s_axi_arready) begin s_axi_arready <= 1'b0; r_addr <= s_axi_araddr; 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; end end R_DATA: begin if (s_axi_rvalid && s_axi_rready) begin if (r_cnt == r_len) begin s_axi_rvalid <= 1'b0; s_axi_rlast <= 1'b0; s_axi_arready <= 1'b1; 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; end end end default: r_state <= R_IDLE; endcase end end // Read-data mux. Use the live address while in R_IDLE accepting AR so the // 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; always @(*) begin s_axi_rdata = 32'd0; 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}; (cur_off == ADDR_CONFIG): s_axi_rdata = {26'd0, cfg_op, 1'b0, cfg_k}; (cur_off == ADDR_PRE_SEL): s_axi_rdata = {30'd0, pre_sel}; (cur_off == ADDR_PRE_ADR): s_axi_rdata = {20'd0, pre_addr}; (cur_off == ADDR_RDB_SEL): s_axi_rdata = {30'd0, rdb_sel}; (cur_off == ADDR_RDB_ADR): s_axi_rdata = {20'd0, rdb_addr}; (cur_off == ADDR_RDB_DAT): s_axi_rdata = {24'd0, rdb_byte}; (cur_off >= ADDR_D && cur_off < ADDR_D + 12'h20): s_axi_rdata = d_word[(cur_off-ADDR_D)>>2]; (cur_off >= ADDR_Z && cur_off < ADDR_Z + 12'h20): s_axi_rdata = z_word[(cur_off-ADDR_Z)>>2]; (cur_off >= ADDR_MSG && cur_off < ADDR_MSG + 12'h20): s_axi_rdata = msg_word[(cur_off-ADDR_MSG)>>2]; (cur_off >= ADDR_SS && cur_off < ADDR_SS + 12'h20): s_axi_rdata = ss_o[((cur_off-ADDR_SS)>>2)*32 +: 32]; default: s_axi_rdata = 32'd0; endcase end end endmodule