mlkem-sync/sync_rtl/poly_mul/poly_mul_sync.v

// poly_mul_sync.v - Synchronous NTT-domain polynomial multiplier
//
// Computes pointwise (Karatsuba-like base-case) multiplication of two
// 256-coefficient NTT-domain polynomials.
//
// Operation flow:
//   IDLE → LOAD (256× A+B pairs) → COMP_CALC (read+compute, 1 cycle)
//        → COMP_C0 (output c0) → COMP_C1 (output c1) → DONE → IDLE
//
// The LOAD phase accepts both A and B coefficients simultaneously
// (one pair per cycle) on coeff_a_in/coeff_b_in.
//
// The COMPUTE phase outputs the 256 result coefficients one per cycle
// via valid/ready handshake on coeff_out/valid_o/ready_i.
//
// Memory: 256×12-bit register arrays for A and B coefficients.
//
// Interface:
//   clk, rst_n     - Clock, active-low reset
//   coeff_a_in[11:0]- Polynomial A coefficient input
//   coeff_b_in[11:0]- Polynomial B coefficient input
//   valid_i         - Input valid
//   ready_o         - Ready to accept input (high in IDLE/LOAD)
//   coeff_out[11:0] - Result coefficient output
//   valid_o         - Output valid (high in COMP_C0/COMP_C1)
//   ready_i         - Output consumer ready

module poly_mul_sync (
    input         clk, rst_n,
    input  [11:0] coeff_a_in,
    input  [11:0] coeff_b_in,
    input         valid_i,
    output        ready_o,
    output [11:0] coeff_out,
    output        valid_o,
    input         ready_i
);

    // State definitions
    localparam S_IDLE     = 3'd0;
    localparam S_LOAD     = 3'd1;
    localparam S_COMP_CALC = 3'd2;
    localparam S_COMP_C0  = 3'd3;
    localparam S_COMP_C1  = 3'd4;
    localparam S_DONE     = 3'd5;

    reg [2:0] state, next_state;

    // Coefficient storage (register arrays)
    reg [11:0] mem_A [0:255];
    reg [11:0] mem_B [0:255];

    // Counters
    reg [7:0] load_cnt;       // 0..256 for loading 256 pairs
    reg [6:0] comp_k;         // 0..127, current base-case index

    // Registered basecase_mul results
    reg [11:0] c0_reg, c1_reg;

    // Combinational read signals for COMP_CALC
    wire [7:0] addr_even = {comp_k, 1'b0};   // comp_k * 2  (7+1 = 8 bits)
    wire [7:0] addr_odd  = {comp_k, 1'b1};   // comp_k * 2 + 1
    wire [11:0] mem_a0 = mem_A[addr_even];
    wire [11:0] mem_a1 = mem_A[addr_odd];
    wire [11:0] mem_b0 = mem_B[addr_even];
    wire [11:0] mem_b1 = mem_B[addr_odd];

    // Zeta ROM
    wire [11:0] zeta;
    poly_mul_zeta_rom u_zeta (
        .addr (comp_k),
        .zeta (zeta)
    );

    // Basecase multiply
    wire [11:0] bc_c0, bc_c1;
    basecase_mul u_bc (
        .a0  (mem_a0),
        .a1  (mem_a1),
        .b0  (mem_b0),
        .b1  (mem_b1),
        .zeta(zeta),
        .c0  (bc_c0),
        .c1  (bc_c1)
    );

    // Output interface
    assign ready_o = (state == S_IDLE) || (state == S_LOAD);
    assign valid_o = (state == S_COMP_C0) || (state == S_COMP_C1);
    assign coeff_out = (state == S_COMP_C0) ? c0_reg : c1_reg;

    // State transition logic (combinational)
    always @* begin
        next_state = state;
        case (state)
            S_IDLE:     if (valid_i && ready_o) next_state = S_LOAD;
            S_LOAD:     if (load_cnt >= 255 && valid_i && ready_o)
                            next_state = S_COMP_CALC;
            S_COMP_CALC: next_state = S_COMP_C0;
            S_COMP_C0:  if (valid_o && ready_i) next_state = S_COMP_C1;
            S_COMP_C1:  if (valid_o && ready_i) begin
                            if (comp_k >= 127) next_state = S_DONE;
                            else               next_state = S_COMP_CALC;
                        end
            S_DONE:     next_state = S_IDLE;
            default:    next_state = S_IDLE;
        endcase
    end

    // Sequential logic
    integer i;
    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            state    <= S_IDLE;
            load_cnt <= 8'd0;
            comp_k   <= 7'd0;
            c0_reg   <= 12'd0;
            c1_reg   <= 12'd0;
            for (i = 0; i < 256; i = i + 1) begin
                mem_A[i] <= 12'd0;
                mem_B[i] <= 12'd0;
            end
        end else begin
            state <= next_state;

            // ---- LOAD phase ----
            // First coefficient captured on IDLE → LOAD transition
            if (state == S_IDLE && valid_i && ready_o) begin
                mem_A[0] <= coeff_a_in;
                mem_B[0] <= coeff_b_in;
                load_cnt <= 8'd1;
            end

            // Subsequent coefficients in LOAD state
            if (state == S_LOAD && valid_i && ready_o) begin
                mem_A[load_cnt] <= coeff_a_in;
                mem_B[load_cnt] <= coeff_b_in;
                load_cnt <= load_cnt + 8'd1;
            end

            // ---- COMPUTE phase ----
            // COMP_CALC: capture basecase_mul results
            if (state == S_COMP_CALC) begin
                c0_reg <= bc_c0;
                c1_reg <= bc_c1;
            end

            // COMP_C0 → COMP_C1: c0 was consumed, increment comp_k
            if (state == S_COMP_C0 && valid_o && ready_i) begin
                // comp_k stays same, c1 still to output
            end

            // COMP_C1 → COMP_CALC: c1 was consumed, advance to next pair
            if (state == S_COMP_C1 && valid_o && ready_i) begin
                comp_k <= comp_k + 7'd1;
            end

            // ---- DONE → IDLE: reset counters ----
            if (state == S_DONE) begin
                load_cnt <= 8'd0;
                comp_k   <= 7'd0;
            end
        end
    end

endmodule