// 4x4 Matrix Multiply Array (Output-Stationary)
module matmul_array #(parameter DW=16, AN=4, AW=32)(
    input clk, rst_n, start, output done,
    input signed [DW-1:0] a_in[0:AN-1], input a_valid,
    input signed [DW-1:0] b_in[0:AN-1], input b_valid,
    output signed [AW-1:0] c_out[0:AN-1][0:AN-1], output c_valid
);
    wire signed [AW-1:0]   acc [0:AN-1][0:AN-1];
    wire signed [DW-1:0]  aw [0:AN][0:AN-1];
    wire signed [DW-1:0]  bw [0:AN-1][0:AN];
    genvar i,j; generate
        for(i=0;i<AN;i++) begin:fa assign aw[0][i]=a_valid?a_in[i]:'0; end
        for(j=0;j<AN;j++) begin:fb assign bw[j][0]=b_valid?b_in[j]:'0; end
        for(i=0;i<AN;i++) for(j=0;j<AN;j++) begin:pe
            mm_pe #(.DW(DW),.AW(AW)) u(.clk(clk),.rst_n(rst_n),
                .a(aw[i][j]),.b(bw[i][j]),.ao(aw[i+1][j]),.bo(bw[i][j+1]),.s(acc[i][j]));
        end
        for(i=0;i<AN;i++) for(j=0;j<AN;j++) assign c_out[i][j]=acc[i][j];
    endgenerate
    reg [7:0] cnt; always_ff @(posedge clk or negedge rst_n)
        if(!rst_n) cnt<=0; else if(start) cnt<=0; else if(cnt<255) cnt<=cnt+1;
    assign done=(cnt==AN+2); assign c_valid=done;
endmodule

module mm_pe #(parameter DW=16, AW=32)(
    input clk, rst_n, input signed [DW-1:0] a,b, output reg signed [DW-1:0] ao,bo, output reg signed [AW-1:0] s
);
    always_ff @(posedge clk or negedge rst_n)
        if(!rst_n) begin ao<='0; bo<='0; s<='0; end
        else begin ao<=a; bo<=b; s<=s+a*b; end
endmodule