`timescale 1ns/1ps
// Synthesizable output-stationary PE. One job between clear pulses.
// Inputs are sampled before the rising edge. Forwarding is registered.
// reset > clear > advance. No multiplier pipeline is inserted here.
module signed_mac_pe #(
  parameter integer W = 8,
  parameter integer K = 2,
  parameter integer AW = 2*W + $clog2(K),
  parameter integer KW = (K > 1) ? $clog2(K) : 1,
  parameter integer CW = (K > 1) ? $clog2(K+1) : 1
) (
  input  logic clk, reset, clear, advance,
  input  logic signed [W-1:0] a_in, b_in,
  input  logic a_valid_in, b_valid_in,
  input  logic [KW-1:0] a_k_in, b_k_in,
  output logic signed [W-1:0] a_out, b_out,
  output logic a_valid_out, b_valid_out,
  output logic [KW-1:0] a_k_out, b_k_out,
  output logic signed [AW-1:0] acc,
  output logic [CW-1:0] count,
  output wire done,
  output logic protocol_error
);
  localparam integer PW = 2*W;
  wire signed [PW-1:0] product;
  wire signed [AW-1:0] addend;
  assign product = $signed(a_in) * $signed(b_in);
  assign addend = {{(AW-PW){product[PW-1]}}, product};
  assign done = (count == K);

  always_ff @(posedge clk) begin
    if (reset || clear) begin
      a_out <= '0; b_out <= '0;
      a_valid_out <= 1'b0; b_valid_out <= 1'b0;
      a_k_out <= '0; b_k_out <= '0;
      acc <= '0; count <= '0; protocol_error <= 1'b0;
    end else if (advance) begin
      a_out <= a_in; b_out <= b_in;
      a_valid_out <= a_valid_in; b_valid_out <= b_valid_in;
      a_k_out <= a_k_in; b_k_out <= b_k_in;
      if (a_valid_in && b_valid_in) begin
        if ((a_k_in != b_k_in) || (count >= K) || (a_k_in != count)) begin
          protocol_error <= 1'b1;
        end else begin
          acc <= acc + addend;
          count <= count + 1'b1;
        end
      end
    end
  end
endmodule