// Optimized 1x1 Convolution Engine — Batched GEMM approach
module conv1x1_engine #(parameter DW=16, CIN=64, COUT=64, HW=196)(
    input clk, rst_n, input en, input start,
    // Input feature map: CIN × HW
    input signed [DW-1:0] ifm_col, input [7:0] ifm_ch, input ifm_valid,
    // Weight: COUT × CIN
    input signed [DW-1:0] wt [0:COUT*CIN-1],
    // Output
    output reg signed [DW-1:0] ofm_out, output reg ofm_valid,
    output reg [7:0] out_ch, output reg [9:0] out_px,
    output reg done
);
    // 1x1 conv = pure matrix multiply: COUT×CIN × CIN×HW = COUT×HW
    reg [7:0] oc_cnt, ic_cnt;
    reg [9:0] px_cnt;
    reg signed [DW*2-1:0] psum;
    always_ff @(posedge clk or negedge rst_n) begin
        if(!rst_n) begin oc_cnt<=0; ic_cnt<=0; px_cnt<=0; psum<='0;
            ofm_out<='0; ofm_valid<=0; done<=0; out_ch<=0; out_px<=0; end
        else if(en) begin
            if(ifm_valid) begin
                psum<=psum+ifm_col*wt[oc_cnt*CIN+ic_cnt];
                ic_cnt<=ic_cnt+1;
                if(ic_cnt>=CIN-1) begin ic_cnt<=0;
                    ofm_out<=psum[DW*2-1:DW]; ofm_valid<=1;
                    out_ch<=oc_cnt; out_px<=px_cnt;
                    psum<='0; oc_cnt<=oc_cnt+1;
                    if(oc_cnt>=COUT-1) begin oc_cnt<=0; px_cnt<=px_cnt+1;
                        if(px_cnt>=HW-1) done<=1;
                    end
                end
            end
        end
    end
endmodule

// Batched 1x1 conv: process multiple spatial positions in parallel
module conv1x1_batch #(parameter DW=16, CIN=64, COUT=64, BATCH=4)(
    input clk, rst_n, input en,
    input signed [DW-1:0] act_batch [0:BATCH*CIN-1], // BATCH pixels × CIN channels
    input signed [DW-1:0] wt [0:COUT*CIN-1],
    output reg signed [DW-1:0] out_batch [0:BATCH*COUT-1],
    output reg out_valid
);
    integer b, ci, co;
    reg signed [DW*2-1:0] psums [0:BATCH-1][0:COUT-1]; // Was 0:0 for co
    always_ff @(posedge clk or negedge rst_n) begin
        if(!rst_n) begin out_valid<=0;
            for(b=0;b<BATCH;b++) for(co=0;co<COUT;co++) psums[b][co]<='0; end
        else if(en) begin
            for(b=0;b<BATCH;b++) for(co=0;co<COUT;co++)
                psums[b][co]<=psums[b][co]+act_batch[b*CIN+0]*wt[co*CIN+0]; // Simplified: single channel
            out_valid<=1;
            for(b=0;b<BATCH;b++) for(co=0;co<COUT;co++)
                out_batch[b*COUT+co]<=psums[b][co][DW*2-1:DW];
        end
    end
endmodule