//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Matthew Johns <mrj1g17@soton.ac.uk>
// Byron Theobald <bt4g16@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

// `include "defs.sv"
import accelerator_pkg::*;

module accelerator_top (
    output logic  [31:0] apu_result,
    output logic  [4:0]  apu_flags_o,
    output logic         apu_gnt,
    output logic         apu_rvalid,
    input  wire         clk,
    input  wire         n_reset,
    input  wire         apu_req,
    input  wire  [2:0][31:0] apu_operands_i,
    input  wire  [5:0]  apu_op,
    input  wire  [14:0] apu_flags_i,
    output wire         data_req_o,
    input  wire         data_gnt_i,
    input  wire         data_rvalid_i,
    output wire         data_we_o,
    output wire  [3:0]  data_be_o,
    output wire  [31:0] data_addr_o,
    output wire  [31:0] data_wdata_o,
    input  wire  [31:0] data_rdata_i,
    output wire         core_halt_o
);

////////////////////////////////////////////////////////////////////////////////
// OUTPUT VARIABLE DECLARATIONS
////////////////////////////////////////////////////////////////////////////////

// CSR OUTPUTS
wire [4:0] vl;
wire [1:0] vsew;
wire [1:0] vlmul;

// DECODER OUTPUTS
wire [31:0] scalar_operand1;
wire [31:0] scalar_operand2;
wire [10:0] immediate_operand;
wire [4:0] vs1_addr;
wire [4:0] vs2_addr;
logic [4:0] vd_addr;
wire csr_write;
wire preserve_vl;
wire set_vl_max;
wire [1:0] elements_to_write;
wire [1:0] cycle_count;
wire vec_reg_write;
vreg_wb_src_t vd_data_src;
vreg_addr_src_t vs3_addr_src;
pe_arith_op_t pe_op;
pe_saturate_mode_t saturate_mode;
pe_output_mode_t output_mode;
pe_operand_t operand_select;
wire [1:0] pe_mul_us;
wire [1:0] widening;
apu_result_src_t apu_result_select;
wire unsigned_immediate;
wire wide_vs1;

// VLSU OUTPUTS
wire [127:0] vlsu_wdata;
logic vec_reg_write_lsu;
logic vlsu_done;
logic [4:0] vl_next_comb;

logic vlsu_en;
logic vlsu_load;
logic vlsu_store;
logic vlsu_strided;
logic vlsu_ready;

// VECTOR REGISTERS OUTPUTS
wire [127:0] vs1_data;
wire [127:0] vs2_data;
wire [127:0] vs3_data;

// ARITHMETIC STAGE OUTPUTS
wire [127:0] arith_output;
wire [127:0] replicated_scalar;

////////////////////////////////////////////////////////////////////////////////
// MODULE INSTANTIATION
////////////////////////////////////////////////////////////////////////////////

wire [31:0] apu_operands [2:0];
assign apu_operands[0] = apu_operands_i[0];
assign apu_operands[1] = apu_operands_i[1];
assign apu_operands[2] = apu_operands_i[2];

////////////////////////////////////////
// CSRs
vector_csrs vcsrs0 (
    .vl(vl),
    .vsew(vsew),
    .vlmul(vlmul),
    .vl_next_comb(vl_next_comb),
    .clk(clk),
    .n_reset(n_reset),
    .avl_in(scalar_operand1),
    .vtype_in(immediate_operand[4:0]),
    .write(csr_write),
    .saturate_flag(1'b0),
    .preserve_vl(preserve_vl),
    .set_vl_max(set_vl_max)
);

////////////////////////////////////////
// DECODER
vector_decoder vdec0 (
    .apu_rvalid(apu_rvalid),
    .apu_gnt(apu_gnt),
    .scalar_operand1(scalar_operand1),
    .scalar_operand2(scalar_operand2),
    .immediate_operand(immediate_operand),
    .vs1_addr(vs1_addr),
    .vs2_addr(vs2_addr),
    .vd_addr(vd_addr),
    .csr_write(csr_write),
    .preserve_vl(preserve_vl),
    .set_vl_max(set_vl_max),
    .elements_to_write(elements_to_write),
    .cycle_count(cycle_count),
    .vec_reg_write(vec_reg_write),
    .vd_data_src(vd_data_src),
    .vs3_addr_src(vs3_addr_src),
    .pe_op(pe_op),
    .saturate_mode(saturate_mode),
    .output_mode(output_mode),
    .operand_select(operand_select),
    .pe_mul_us(pe_mul_us),
    .widening(widening),
    .apu_result_select(apu_result_select),
    .unsigned_immediate(unsigned_immediate),
    .wide_vs1(wide_vs1),
    .clk(clk),
    .n_reset(n_reset),
    .apu_req(apu_req),
    .apu_operands(apu_operands),
    .apu_op(apu_op),
    .apu_flags_i(apu_flags_i),
    .vl(vl),
    .vsew(vsew),
    .vlsu_en_o(vlsu_en),
    .vlsu_load_o(vlsu_load),
    .vlsu_store_o(vlsu_store),
    .vlsu_strided_o(vlsu_strided),
    .vlsu_ready_i(vlsu_ready),
    .vlsu_done_i(vlsu_done),
    .core_halt_o(core_halt_o)
);

////////////////////////////////////////
// VECTOR REGISTERS
logic [127:0] vd_data;
logic [4:0] vs3_addr, vs3_addr_vlsu;
always_comb begin
    case (vd_data_src)
        VREG_WB_SRC_MEMORY:

    case (vs3_addr_src)
        VS3_ADDR_SRC_DECODE:
            vs3_addr = vd_addr;
        VS3_ADDR_SRC_VLSU:
            vs3_addr = vs3_addr_vlsu;
        default:
            vs3_addr = '0;
    endcase
end

vector_registers vreg0 (
    .vs1_data(vs1_data),
    .vs2_data(vs2_data),
    .vs3_data(vs3_data),
    .vd_data(vd_data),
    .vs1_addr(vs1_addr),
    .vs2_addr(vs2_addr),
    .vd_addr(vs3_addr),
    .vsew(vsew),
    .vlmul(vlmul),
    .elements_to_write(elements_to_write),
    .clk(clk),
    .n_reset(n_reset),
    .write(vec_reg_write | vec_reg_write_lsu ),
    .widening_op(widening[0]),
    .wide_vs1(wide_vs1),
    .load_operation(vlsu_load)
);

////////////////////////////////////////
// PEs CONTAINED IN ARITHMETIC STAGE WRAPPER
arith_stage arith_stage0 (
    .arith_output(arith_output),
    .replicated_scalar(replicated_scalar),
    .clk(clk),
    .n_reset(n_reset),
    .vs1_data(vs1_data),
    .vs2_data(vs2_data),
    .vs3_data(vs3_data),
    .scalar_operand(scalar_operand1),
    .imm_operand(immediate_operand[4:0]),
    .elements_to_write(elements_to_write),
    .cycle_count(cycle_count),
    .op(pe_op),
    .saturate_mode(saturate_mode),
    .output_mode(output_mode),
    .operand_select(operand_select),
    .widening(widening),
    .mul_us(pe_mul_us),
    .unsigned_immediate(unsigned_immediate),
    .wide_vs1(wide_vs1),
    .vl(vl),
    .vsew(vsew)
);

////////////////////////////////////////
// VLSU
vector_lsu vlsu0 (
    .clk(clk),
    .n_reset(n_reset),

    .vl_i(vl),
    .vsew_i(vsew),
    .vlmul_i(vlmul),

    .vlsu_en_i(vlsu_en),
    .vlsu_load_i(vlsu_load),
    .vlsu_store_i(vlsu_store),
    .vlsu_strided_i(vlsu_strided),
    .vlsu_ready_o(vlsu_ready),
    .vlsu_done_o(vlsu_done),

    .data_req_o(data_req_o),
    .data_gnt_i(data_gnt_i),
    .data_rvalid_i(data_rvalid_i),
    .data_addr_o(data_addr_o),
    .data_we_o(data_we_o),
    .data_be_o(data_be_o),
    .data_rdata_i(data_rdata_i),
    .data_wdata_o(data_wdata_o),

    .cycle_count_i(cycle_count),

    .op0_data_i(scalar_operand1),
    .op1_data_i(scalar_operand2),

    .vs_wdata_o(vlsu_wdata),
    .vs_rdata_i(vs3_data),
    .vr_addr_i(vd_addr),
    .vs3_addr_o(vs3_addr_vlsu),
    .vr_we_o(vec_reg_write_lsu)
);

////////////////////////////////////////////////////////////////////////////////
// RESULT SELECTION - what value to return to CPU
////////////////////////////////////////////////////////////////////////////////
logic [31:0] reg_apu_result;
assign apu_flags_o = '0;
assign apu_result = reg_apu_result;

// Updated VL is arriving two cycles too late
always_comb begin
	 reg_apu_result = '0;
    case (apu_result_select)
        APU_RESULT_SRC_VL:
            reg_apu_result = {'0, vl_next_comb};
        APU_RESULT_SRC_VS2_0:
            case (vsew)
                2'd0: // 8b
                    reg_apu_result = { {24{vs2_data[7]}}, vs2_data[7:0] };
                2'd1: // 16b
                    reg_apu_result = { {16{vs2_data[15]}}, vs2_data[15:0] };
                2'd2:// 32b
                    reg_apu_result = vs2_data[31:0];
            endcase
    endcase
end

endmodule

            vd_data = vlsu_wdata;
        VREG_WB_SRC_ARITH:
            vd_data = arith_output;
        VREG_WB_SRC_SCALAR:
            vd_data = replicated_scalar;


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Byron Theobald <bt4g16@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

/* verilator lint_off ALWCOMBORDER */
module address_unit(
    input logic clk_i, n_rst_i,
    input logic [31:0] base_addr_i,
    input logic [31:0] stride_i,
    input logic [4:0] vl_i,
    input logic [1:0] vsew_i,
    input logic au_start_i,
    input logic au_next_i,
    output logic [3:0] au_be_o,
    output logic [6:0] au_bc_o,
    output logic [31:0] au_addr_o,
    output logic [6:0] vd_offset_o,
    output logic au_valid_o,
    output logic au_ready_o,
    output logic au_final_o);

    logic [1:0] ib_select; // Low 2 bits of initial address
    logic [3:0] be_gen;

    logic [31:0] next_el_pre, next_el_addr;
    logic [31:0] cycle_addr;
    logic [6:0] cycle_bytes;

    typedef enum {RESET, FIRST, CYCLE, WAIT, FINAL} be_state;
    be_state current_state, next_state;

    logic signed [6:0] byte_track, byte_track_next;
    logic cycle_load;

    assign vd_offset_o = (vl_i << vsew_i) - byte_track;

    always_comb begin
        if(au_valid_o) begin // Make sure that all our addresses are word aligned
            au_addr_o = {cycle_addr[31:2], 2'd0};
            au_be_o = be_gen;
        end else begin
            au_addr_o = 32'd0;
            au_be_o = 4'b0000;
        end
    end

    always_comb begin
        if(au_start_i)
            byte_track_next = {2'd0, vl_i} << vsew_i; // Bytes dependent on element size
        else if(current_state != CYCLE)
            byte_track_next = byte_track;
        else
            byte_track_next = (byte_track >= cycle_bytes) ? (byte_track - cycle_bytes) : 7'd0;
    end

    always_ff @(posedge clk_i, negedge n_rst_i) begin
        if(~n_rst_i)
            byte_track <= 7'd0;
        else
            byte_track <= byte_track_next;
    end

    always_ff @(posedge clk_i, negedge n_rst_i) begin
        if(~n_rst_i)
            cycle_addr <= 32'd0;
        else if(au_start_i)
            cycle_addr <= base_addr_i;
        else if(current_state != CYCLE)
            cycle_addr <= cycle_addr;
        else
            cycle_addr <= next_el_addr;
    end

    always_ff @(posedge clk_i, negedge n_rst_i) begin
        if(~n_rst_i)
            current_state <= RESET;
        else
            current_state <= next_state;
    end

    always_comb begin
        be_gen = 4'b0000;
        case(vsew_i)
            2'b00 : begin // 8 Bit
                ib_select = cycle_addr[1:0];

                if(stride_i > 1) begin
                    be_gen[ib_select] = 1'b1;

                    // Where is our next byte?
                    next_el_pre = cycle_addr + stride_i;
                    if(next_el_pre[31:2] == cycle_addr[31:2] && byte_track > 1) begin
                        be_gen[next_el_pre[1:0]] = 1'b1;
                        next_el_addr = next_el_pre + stride_i; // Stride by second element
                    end else begin
                        next_el_addr = next_el_pre;
                    end

                    // Calculate the number of bytes for cycle
                    cycle_bytes = {5'd0, be_gen[3]} + {5'd0, be_gen[2]} + {5'd0, be_gen[1]} + {5'd0, be_gen[0]};
                end else if(stride_i == 1) begin
                    be_gen[0] = (ib_select == 0) ? 1 : 0;
                    be_gen[1] = (ib_select == 1 || byte_track > 1) ? 1'b1 : 1'b0;
                    be_gen[2] = (ib_select == 2 || byte_track > 2) ? 1'b1 : 1'b0;
                    be_gen[3] = (ib_select == 3 || byte_track > 3) ? 1'b1 : 1'b0;
                    next_el_addr = {cycle_addr[31:2], 2'b0} + 32'd4;

                    // Calculate the number of bytes for cycle
                    cycle_bytes = {5'd0, be_gen[3]} + {5'd0, be_gen[2]} + {5'd0, be_gen[1]} + {5'd0, be_gen[0]};
                end else if(stride_i == 0) begin
                    be_gen[ib_select] = 1'b1;
                    cycle_bytes = {2'b0, vl_i}; // Read all bytes in 1 cycle
                end
            end
            2'b01 : begin // 16 Bit
                ib_select = {cycle_addr[1], 1'b0}; // Force alignment byte 0 or 2

                if(stride_i > 2) begin // Always 1 element
                    // Always set 2 bytes
                    be_gen[ib_select] = 1'b1;
                    be_gen[ib_select+1] = 1'b1;
                    next_el_addr = {cycle_addr[31:1], 1'b0} + {stride_i[31:1], 1'b0};

                    // Calculate the number of bytes for cycle
<TARGET>
                end else if (stride_i == 2) begin // Up to 2 Elements
                    be_gen[1:0] = (ib_select == 0) ? 2'b11 : 2'b00;
                    be_gen[3:2] = (ib_select == 2 || byte_track > 2) ? 2'b11 : 2'b00;
                    next_el_addr = {cycle_addr[31:2], 2'b0} + 32'd4;

                    // Calculate the number of bytes for cycle

                end else if (stride_i == 0) begin
                    be_gen[ib_select] = 1'b1;
                    be_gen[ib_select+1] = 1'b1;
                    cycle_bytes = {1'b0, vl_i, 1'b0}; // Read all bytes in 1 cycle
                end
            end
            2'b10 : begin // 32 Bit
                ib_select = 2'd0; // Force alignment to byte 0

                if(stride_i >= 4) begin // Always 1 element
                    be_gen = 4'b1111;
                    next_el_addr = {cycle_addr[31:2], 2'b0} + {stride_i[31:2], 2'b0}; // stride_i is always a multiple of 4

                    // Calculate the number of bytes for cycle
                    cycle_bytes = {5'd0, be_gen[3]} + {5'd0, be_gen[2]} + {5'd0, be_gen[1]} + {5'd0, be_gen[0]};
                end else if(stride_i == 0) begin
                    be_gen = 4'b1111;
                    cycle_bytes = {1'b0, vl_i, 1'b0}; // Read all bytes in 1 cycle
                end
            end
            default : $error("Invalid VSEW");
        endcase
    end

    always_comb begin
        cycle_load = 1'b0;
        au_valid_o = 1'b0;
        au_ready_o = 1'b0;
        au_final_o = 1'b0;
        case(current_state)
            RESET: begin
                au_ready_o = 1'b1;
                if(au_start_i)
                    next_state = FIRST;
                else
                    next_state = RESET;
            end
            FIRST: begin
                au_valid_o = 1'b1;
                if(stride_i != 0) begin
                    next_state = WAIT;
                end else begin
                    next_state = RESET;
                end
            end
            CYCLE: begin
                au_valid_o = 1'b1;
                if(byte_track_next == 0) begin
                    next_state = WAIT;
                end else begin
                    cycle_load = 1'b1;
                    next_state = WAIT;
                end
            end
            WAIT: begin
                if(au_next_i)
                    next_state = CYCLE;
                else if(byte_track_next == 0)
                    next_state = FINAL;
                else
                    next_state = WAIT;
            end
            FINAL: begin
                au_final_o = 1'b1;
                next_state = RESET;
            end
        endcase
    end

    initial begin
      $dumpfile("test.vcd");
      $dumpvars;
      $display(" Model running...\n");
   end
endmodule
/* verilator lint_on ALWCOMBORDER */


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Matthew Johns <mrj1g17@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

// Module instance to contain the PEs and supporting logic such as input/output
// selection logic. This is just to tidy up the top-level a bit.

import accelerator_pkg::*;

module arith_stage (
    output logic [127:0] arith_output,
    output logic [127:0] replicated_scalar, // Want this output for vmv.v.i
    input wire clk,
    input wire n_reset,
    input wire [127:0] vs1_data,
    input wire [127:0] vs2_data,
    input wire [127:0] vs3_data,
    input wire [31:0] scalar_operand,
    input wire [4:0] imm_operand,
    input wire [1:0] elements_to_write,
    input wire [1:0] cycle_count,
    input pe_arith_op_t op,
    input pe_saturate_mode_t saturate_mode,
    input pe_output_mode_t output_mode,
    input pe_operand_t operand_select,
    input wire [1:0] widening,
    input wire [1:0] mul_us,
    input wire unsigned_immediate,
    input wire wide_vs1,
    input wire [4:0] vl,
    input wire [1:0] vsew
);

logic [31:0] reduction_intermediate_reg;

// logic [127:0] replicated_scalar;

wire [31:0] pe0_out;
wire [31:0] pe1_out;
wire [31:0] pe2_out;
wire [31:0] pe3_out;

logic [31:0] pe0_b_data;
logic [31:0] pe1_b_data;
logic [31:0] pe2_b_data;
logic [31:0] pe3_b_data;

logic [31:0] scalar_to_replicate;

pe_32b pe0 (
    .out(pe0_out),
    .a(vs2_data[31:0]),
    .b(pe0_b_data),
    .c(vs3_data[31:0]),
    .op(op),
    .vsew(vsew),
    .widening(widening),
    .mul_us(mul_us),
    .saturate_mode(saturate_mode),
    .output_mode(output_mode),
    .wide_b(wide_vs1)
);

pe_32b pe1 (
    .out(pe1_out),
    .a(vs2_data[63:32]),
    .b(pe1_b_data),
    .c(vs3_data[63:32]),
    .op(op),
    .vsew(vsew),
    .widening(widening),
    .mul_us(mul_us),
    .saturate_mode(saturate_mode),
    .output_mode(output_mode),
    .wide_b(wide_vs1)
);

pe_32b pe2 (
    .out(pe2_out),
    .a(vs2_data[95:64]),
    .b(pe2_b_data),
    .c(vs3_data[95:64]),
    .op(op),
    .vsew(vsew),
    .widening(widening),
    .mul_us(mul_us),
    .saturate_mode(saturate_mode),
    .output_mode(output_mode),
    .wide_b(wide_vs1)
);

pe_32b pe3 (
    .out(pe3_out),
    .a(vs2_data[127:96]),
    .b(pe3_b_data),
    .c(vs3_data[127:96]),
    .op(op),
    .vsew(vsew),
    .widening(widening),
    .mul_us(mul_us),
    .saturate_mode(saturate_mode),
    .output_mode(output_mode),
    .wide_b(wide_vs1)
);

scalar_replicate scalar_rep0 (
    .replicated_out(replicated_scalar),
    .scalar_in(scalar_to_replicate),
    .vsew(vsew),
    .us(1'b0)
);

// Update the intermediate register used for reduction operations every cycle
always_ff @(posedge clk, negedge n_reset)
    if (~n_reset)
        reduction_intermediate_reg <= '0;
    else
        reduction_intermediate_reg <= pe3_out;

////////////////////////////////////////////////////////////////////////////////
// PE INPUT OPERAND SELECTION
////////////////////////////////////////////////////////////////////////////////
always_comb
begin
    scalar_to_replicate = scalar_operand;

    case (operand_select)
        PE_OPERAND_VS1:
        begin
            pe0_b_data = vs1_data[31:0];
            pe1_b_data = vs1_data[63:32];
            pe2_b_data = vs1_data[95:64];
            pe3_b_data = vs1_data[127:96];
        end
        PE_OPERAND_SCALAR:
        begin
            pe0_b_data = replicated_scalar[31:0];
            pe1_b_data = replicated_scalar[63:32];
            pe2_b_data = replicated_scalar[95:64];
            pe3_b_data = replicated_scalar[127:96];
        end

                pe3_b_data = {'0, imm_operand[4:0]};
            end
            else
            begin
                pe0_b_data = {{27{imm_operand[4]}}, imm_operand[4:0]};
                pe1_b_data = {{27{imm_operand[4]}}, imm_operand[4:0]};
                pe2_b_data = {{27{imm_operand[4]}}, imm_operand[4:0]};
                pe3_b_data = {{27{imm_operand[4]}}, imm_operand[4:0]};
            end
            // This line handles a special case - the scalar replicate logic is
            // used for vmv.v.i instructions but the PE is not used. In this
            // case the immediate operand needs to be replicated and returned.
            scalar_to_replicate = {{27{imm_operand[4]}}, imm_operand[4:0]};
        end
        PE_OPERAND_RIPPLE:
        begin
            // For first cycle of reduction operation want to look at vs1[0].
            // For later cycles need the intermediate value from last cycle.
            if (cycle_count == 2'd0)
                pe0_b_data = vs1_data[31:0];
            else
                pe0_b_data = reduction_intermediate_reg;
            pe1_b_data = pe0_out;
            pe2_b_data = pe1_out;
            pe3_b_data = pe2_out;
        end
    endcase
end

////////////////////////////////////////////////////////////////////////////////
// PE OUTPUT SELECTION
////////////////////////////////////////////////////////////////////////////////
always_comb
begin
    // For reduction operations the output comes from a single PE (the last in
    // the chain), but which PE is last depends on VL.
    if (operand_select == PE_OPERAND_RIPPLE)
        case (vl[1:0])
            2'd0:
                // Perhaps counterintuitively, if vl[1:0] is zero that means all
                // four elements are being written so want the last PE.
                arith_output = {'0, pe3_out};
            2'd1:
                arith_output = {'0, pe0_out};
            2'd2:
                arith_output = {'0, pe1_out};
            2'd3:
                arith_output = {'0, pe2_out};
            default:
                arith_output = {'0, pe3_out};
        endcase
    else
        arith_output = {pe3_out, pe2_out, pe1_out, pe0_out};
end

endmodule

        PE_OPERAND_IMMEDIATE:
        begin
            if (unsigned_immediate)
            begin


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Byron Theobald <bt4g16@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

// W_IN must be smaller than W_OUT
module bit_ext #(parameter W_IN = 8, parameter W_OUT = 12)

    assign a_out = {{(W_OUT-W_IN){a_in[W_IN-1]}}, a_in};

endmodule

    (input logic signed [W_IN-1:0] a_in,
     output logic signed [W_OUT-1:0] a_out);


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Byron Theobald <bt4g16@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

module mapping_unit
   (input logic [127:0] arith_format_o,
    output  logic [31:0] memory_format_i,
    input  logic [1:0] sew_i,
    input  logic [1:0] reg_select);

                vd_wr_data0 = '{
                    vd_data[103:96],
                    vd_data[71:64],
                    vd_data[39:32],
                    vd_data[7:0]
                };
            2'd1: // 16b
            begin
                vd_wr_data1 = {
                    vd_data[111:96],
                    vd_data[79:64]
                };
                vd_wr_data0 = {
                    vd_data[47:32],

                };
            end
            2'd2: // 32b
            begin
                vd_wr_data3 = vd_data[127:96];
                vd_wr_data2 = vd_data[95:64];
                vd_wr_data1 = vd_data[63:32];
                vd_wr_data0 = vd_data[31:0];*/
endmodule


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Matthew Johns <mrj1g17@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

// NOTE (Matthew Johns) - there is similarity between parts of this code and the
// proc_unit.sv module made by me in my third-year project. This is because the
// functionality is similar and therefore I'm using what I learnt previously.

// `include "defs.sv"
import accelerator_pkg::*;

module pe_32b (
    output logic [31:0] out,
    // output logic flag_saturated // TODO: add this flag for CSRs
    input wire [31:0] a,
    input wire [31:0] b,
    input wire [31:0] c,
    input pe_arith_op_t op,
    input wire [1:0] vsew,
    input wire [1:0] widening, // 2'd1 for widening, 2'd2 for quad widening
    input wire [1:0] mul_us, // Specifies each multiplier input as signed or unsigned
    input pe_saturate_mode_t saturate_mode,
    input pe_output_mode_t output_mode,
    input wire wide_b
);

// Usually input "a" is vs2, input "b" is vs1 and "c" is vs3/vd. (This is for
// most standard arithmetic operations)

logic signed [32:0] mult_a;
logic signed [32:0] mult_b;
logic signed [65:0] mult_wide;
logic signed [32:0] selected_mult_out;

logic [32:0] add_out;
logic [32:0] add_a;
logic [32:0] add_b;
logic [32:0] addend;

// Intermediate before saturation/ReLU
// Has to be 33 bits to give at least one bit for saturation
logic [32:0] arith_result;

logic macc;
logic subtract;

// Instantiate sign extension module for inputs a, b and c
wire [31:0] sign_ext_a;
wire [31:0] sign_ext_b;
wire [31:0] sign_ext_c;
vw_sign_ext se0 (
    .sign_ext_a(sign_ext_a),
    .sign_ext_b(sign_ext_b),
    .sign_ext_c(sign_ext_c),
    .a(a),
    .b(b),
    .c(c),
    .widening(widening),
    .vsew(vsew),
    .wide_b(wide_b)
);

////////////////////////////////////////////////////////////////////////////////
// ARITHMETIC STAGE
////////////////////////////////////////////////////////////////////////////////
always_comb
begin
    subtract = 1'b0;
    macc = 1'b0;
    arith_result = '0;

    case (op)
        // 4'h0: // Add
        PE_ARITH_ADD:
            arith_result = add_out;
        // 4'h1: // Sub
        PE_ARITH_SUB:
        begin
            arith_result = add_out;
            subtract = 1'b1;
        end
        // 4'h2: // Left-shift
        PE_ARITH_LSHIFT:
            arith_result = {1'b0, (a << b)};
        // 4'h3: // Multiply
        PE_ARITH_MUL:
            arith_result = selected_mult_out;
        // 4'h4: // Multiply-add
        PE_ARITH_MULADD:
        begin
            macc = 1'b1;
            arith_result = add_out;
        end
        // 4'h5: // XOR
        PE_ARITH_XOR:
            arith_result = {1'b0, (a ^ b)};
        // 4'h6: // Right-shift
        PE_ARITH_RSHIFT_LOG:
            arith_result = {1'b0, (a >> b)};
        // 4'h7: // Right-shift (arithmetic)
        PE_ARITH_RSHIFT_AR:
            arith_result = {1'b0, (a >>> b)};
        // 4'h8: // OR
        PE_ARITH_OR:
            arith_result = {1'b0, (a | b)};
        // 4'h9: // AND
        PE_ARITH_AND:
            arith_result = {1'b0, (a & b)};
    endcase
end

always_comb
begin
    // Multiplier needs to be able to be toggled between signed/unsigned for the
    // individual operands (depending on instructions). Can do this by adding an
    // extra bit to the inputs and sign-extending them for signed operations.
    // Then take the correct number of bits from the bottom.
    if (mul_us[1])
        mult_a = {1'b0, a};
    else
        mult_a = {sign_ext_a[31], sign_ext_a};
    if (mul_us[0])
        mult_b = {1'b0, b};
    else
        mult_b = {sign_ext_b[31], sign_ext_b};

    mult_wide = mult_a * mult_b;

    // Select adder inputs
    if (macc)
    begin
        add_a = {1'b0, mult_wide[31:0]};
        add_b = {1'b0, c};
    end
    else
    begin
        // Need to sign extend for saturated ops as another bit is gained to be
        // used for saturation.
        // This means instructions are fixed as signed. Would have to split up
        // if wanted to toggle signed/unsigned.
        // For some operations (such as regular addition) sign extension isn't
        // needed. But it doesn't do any harm and simplifies the logic.
        add_a = {sign_ext_a[31], sign_ext_a};
        add_b = {sign_ext_b[31], sign_ext_b};
    end

    if (subtract)
        addend = ~add_b + 1'b1;
    else
        addend = add_b;
    add_out = add_a + addend;

end

// Select multiplier output. V spec requires fractional saturating multiplies to
// take the upper bits for saturation instead of lower bits. Should be able to
// toggle this as we might need to saturate it and keep lower bits later on
always_comb
begin
    if (saturate_mode == PE_SAT_UPPER)
        case (vsew)
            2'd0: // 8b
                selected_mult_out = {{24{1'b0}}, mult_wide[8:0]};
            2'd1: // 16b
                selected_mult_out = {{16{1'b0}}, mult_wide[16:0]};
            2'd2: // 32b
                selected_mult_out = mult_wide[32:0];
            default:
                selected_mult_out = mult_wide[32:0];
        endcase
    else
        selected_mult_out = mult_wide[32:0];
end

////////////////////////////////////////////////////////////////////////////////
// SATURATE STAGE
////////////////////////////////////////////////////////////////////////////////
logic [31:0] sat_result;
// Instantiate the saturation blocks. One for each element width. Output is
// selected from one of them at a time.
wire [7:0] sat8_result;
sat_unit #(
    .W_IN(33),
    .W_OUT(8)
) sat8
(
    .a_in(arith_result),
    .a_out(sat8_result)
);
wire [15:0] sat16_result;
sat_unit #(
    .W_IN(33),
    .W_OUT(16)
) sat16
(
    .a_in(arith_result),
    .a_out(sat16_result)
);
wire [31:0] sat32_result;
sat_unit #(
    .W_IN(33),
    .W_OUT(32)
) sat32
(
    .a_in(arith_result),
    .a_out(sat32_result)
);

always_comb
begin
    sat_result = arith_result[31:0];
    // For widening, need to saturate to next larger element size
    if (widening[0])
        case (vsew)

    else
        case (vsew)
            2'd0: // 8b
                sat_result = {'0, sat8_result};
            2'd1: // 16b
                sat_result = {'0, sat16_result};
            2'd2: // 32b
                sat_result = sat32_result;
        endcase
end

////////////////////////////////////////////////////////////////////////////////
// OUTPUT MODE SELECT
////////////////////////////////////////////////////////////////////////////////
always_comb
begin
    out = arith_result[31:0];
    case (output_mode)
        PE_OP_MODE_RESULT:
            if (saturate_mode == PE_SAT_NONE)
                out = arith_result[31:0];
            else
                out = sat_result;
        PE_OP_MODE_PASS_MAX:
            // Will do arithmetic op of a-b. If negative, b is larger so pass b
            if (arith_result[32])
                out = b;
            else
                out = a;
        PE_OP_MODE_PASS_MIN:
            if (arith_result[32])
                out = a;
            else
                out = b;
            // Important point for this: am looking at bit 31 (not 32) because
            // the inputs aren't sign-extended for anything non-saturating.
    endcase
end

endmodule

            2'd0: // 8b -> 16b
                sat_result = {'0, sat16_result};
            2'd1: // 16b -> 32b
                sat_result = sat32_result;
        endcase


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Byron Theobald <bt4g16@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

module relu_bound
    #(parameter W=8, parameter N=(6<<4))
    (input logic signed [W-1:0] a,
     output logic [W-2:0] ar);

    logic signed [W-1:0] zero = 'd0;
    always_comb begin
        if(a < zero)

            ar = 0;
        else if (a > N)
            ar = N[6:0];
        else
            ar = a[6:0];


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Byron Theobald <bt4g16@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

// W_IN must be larger than W_OUT
module sat_unit #(parameter W_IN = 13, parameter W_OUT = 8)
    (input logic signed [W_IN-1:0] a_in,
     output logic signed [W_OUT-1:0] a_out);

    logic signed [W_IN-1:0] max_in = {{(W_IN-W_OUT+1){1'b0}},{(W_OUT-1){1'b1}}};
    logic signed [W_IN-1:0] min_in = {{(W_IN-W_OUT+1){1'b1}},{(W_OUT-1){1'b0}}};
    logic signed [W_OUT-1:0] max_out = {1'b0,{(W_OUT-1){1'b1}}};

    logic signed [W_OUT-1:0] min_out = {1'b1,{(W_OUT-1){1'b0}}};

    assign a_out = a_in < min_in ? min_out : (a_in > max_in) ? max_out : a_in[W_OUT-1:0];

    //initial $display("%d, %d, %d, %d", max_in, min_in, max_out, min_out);


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Matthew Johns <mrj1g17@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

// import accelerator_pkg::*;

module scalar_replicate (
    output logic [127:0] replicated_out,
    input wire [31:0] scalar_in,
    input wire [1:0] vsew,
    input wire us
);

// When a scalar operand comes into the accelerator there is only one copy of
// it. However, each PE needs a copy for it's calculation, so it must be
// replicated into the right position for each PE. This module does that.

always_comb
    case (vsew)
        2'd0: // 8b
            if (us)
                replicated_out = {
                    24'd0,
                    scalar_in[7:0],
                    24'd0,
                    scalar_in[7:0],
                    24'd0,
                    scalar_in[7:0],
                    24'd0,
                    scalar_in[7:0]
                };
            else
                replicated_out = {
                    {24{scalar_in[7]}},
                    scalar_in[7:0],
                    {24{scalar_in[7]}},
                    scalar_in[7:0],
                    {24{scalar_in[7]}},
                    scalar_in[7:0],
                    {24{scalar_in[7]}},
                    scalar_in[7:0]
                };
        2'd1: // 16b
            if (us)
                replicated_out = {
                    16'd0,
                    scalar_in[15:0],
                    16'd0,
                    scalar_in[15:0],
                    16'd0,
                    scalar_in[15:0],
                    16'd0,
                    scalar_in[15:0]
                };

                    scalar_in[15:0]
                };
        2'd2: // 32b
            replicated_out = {
                {4{scalar_in}}
            };
        default:
            replicated_out = {'0, scalar_in};
    endcase

endmodule

            else
                replicated_out = {
                    {16{scalar_in[15]}},
                    scalar_in[15:0],
                    {16{scalar_in[15]}},


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Byron Theobald <bt4g16@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

`timescale 1ns/10ps

`include "bit_ext.sv"

module tb_bit_ext;

logic signed [7:0] in;
logic signed [11:0] out;
bit_ext #(.W_IN(8), .W_OUT(12)) bit_ext(in, out);

initial begin
    `ifndef VERILATOR
        $dumpfile("bit_ext.vcd");
        $dumpvars;
    `endif

    for(in = -128; in < 127; in++) begin


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Byron Theobald <bt4g16@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

`timescale 1ns/10ps

`include "relu_bound.sv"

module tb_relu_bound;

logic signed [7:0] in;
logic [6:0] out;
relu_bound #(.W(8)) rb(in, out);

logic signed [7:0] i;

initial begin
    `ifndef VERILATOR
        $dumpfile("relu_bound.vcd");
        $dumpvars;
    `endif

    for(in = -128; in < 127; in++) begin


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Byron Theobald <bt4g16@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

`timescale 1ns/10ps

`include "sat_unit.sv"

module tb_sat_unit;

logic signed [12:0] in;
logic signed  [7:0] out;

logic signed [7:0] i;

initial begin
    `ifndef VERILATOR
        $dumpfile("sat_unit.vcd");
        $dumpvars;
    `endif

    for(in = -256; in < 256; in++) begin
        #1ns $display("%d, %d", in, out);
    end

    $finish;
end

endmodule

sat_unit #(.W_IN(13), .W_OUT(8)) satu(in, out);


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Byron Theobald <bt4g16@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

// 32-Bit input, byte_position selects the bytes to load.
// Selected bytes will be packed and loaded into the
// register starting from byte_loaded

module temporary_reg (
    input  logic clk_i, n_rst_i,
    input  logic byte_enable_valid,
    input  logic read_data_valid,
    input  logic clear_register,
    input  logic [31:0] memory_read_i,
    input  logic [3:0] byte_enable_i,
    input  logic [6:0] byte_select_i,
    output logic [127:0] wide_vd_o);

    logic [3:0] byte_enable_reg;

    always_ff @(posedge clk_i, negedge n_rst_i) begin
        if(~n_rst_i)
            byte_enable_reg <= 1'b0;
        else if(byte_enable_valid)
            byte_enable_reg <= byte_enable_i;
    end

    // Tempoary register, split into bytes.
    logic [7:0] temp_reg [15:0];

    // Split memory read into bytes
    logic [7:0] memory_read_bytes [3:0];

    always_comb begin // Split bytes out of word
        memory_read_bytes[0] = memory_read_i[7:0];
        memory_read_bytes[1] = memory_read_i[15:8];
        memory_read_bytes[2] = memory_read_i[23:16];
        memory_read_bytes[3] = memory_read_i[31:24];
    end

    // Packed read bytes, shift higher elements down
    logic [7:0] memory_read_packed [3:0];
    logic [3:0] packed_set;

    always_comb begin
		  memory_read_packed = '{default:0};
        packed_set = 4'b0000;
        if(byte_enable_reg[0]) begin
            packed_set[0] = 1'b1;
            memory_read_packed[0] = memory_read_bytes[0];
        end
        if(byte_enable_reg[1]) begin
            casez(packed_set)
                4'bzzz0 : begin
                    packed_set[0] = 1'b1;
                    memory_read_packed[0] = memory_read_bytes[1];
                end
                4'bzzz1 : begin
                    packed_set[1] = 1'b1;
                    memory_read_packed[1] = memory_read_bytes[1];
                end
					 default: begin
						  memory_read_packed[1] = '0;
					 end
            endcase
        end
        if(byte_enable_reg[2]) begin
            casez(packed_set)
                4'bzz00 : begin
                    packed_set[0] = 1'b1;
                    memory_read_packed[0] = memory_read_bytes[2];
                end
                4'bzz01 : begin

					 end
            endcase
        end
        if(byte_enable_reg[3]) begin
            casez(packed_set)
                4'bz000 : begin
                    packed_set[0] = 1'b1;
                    memory_read_packed[0] = memory_read_bytes[3];
                end
                4'bz001 : begin
                    packed_set[1] = 1'b1;
                    memory_read_packed[1] = memory_read_bytes[3];
                end
                4'bz011 : begin
                    packed_set[2] = 1'b1;
                    memory_read_packed[2] = memory_read_bytes[3];
                end
                4'bz111 : begin
                    packed_set[3] = 1'b1;
                    memory_read_packed[3] = memory_read_bytes[3];
                end
					 default: begin
						  memory_read_packed[3] = '0;
					 end
            endcase
        end
    end

    // Write elements into register
    always_ff @(posedge clk_i, negedge n_rst_i) begin
        if(~n_rst_i) begin
            temp_reg <= '{default: 8'd0};
        end else if(clear_register) begin
            temp_reg <= '{default: 8'd0};
        end else if(read_data_valid) begin
            if(packed_set[0]) temp_reg[byte_select_i+0] <= memory_read_packed[0];
            if(packed_set[1]) temp_reg[byte_select_i+1] <= memory_read_packed[1];
            if(packed_set[2]) temp_reg[byte_select_i+2] <= memory_read_packed[2];
            if(packed_set[3]) temp_reg[byte_select_i+3] <= memory_read_packed[3];
        end
    end

    assign wide_vd_o = {temp_reg[15], temp_reg[14], temp_reg[13], temp_reg[12],
                            temp_reg[11], temp_reg[10], temp_reg[ 9], temp_reg[ 8],
                            temp_reg[ 7], temp_reg[ 6], temp_reg[ 5], temp_reg[ 4],
                            temp_reg[ 3], temp_reg[ 2], temp_reg[ 1], temp_reg[ 0]};

endmodule

                    packed_set[1] = 1'b1;
                    memory_read_packed[1] = memory_read_bytes[2];
                end
                4'bzz11 : begin
                    packed_set[2] = 1'b1;


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Matthew Johns <mrj1g17@soton.ac.uk>
// Byron Theobald <bt4g16@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

// NOTE (Matthew Johns) - there is similarity between parts of this code and the
// csr.sv module made by me in my third-year project. This is because the
// functionality is similar and therefore I'm using what I learnt previously.

module vector_csrs (
    output logic [4:0] vl,
    output logic [1:0] vsew,
    output logic [1:0] vlmul,
    output logic [4:0] vl_next_comb,
    input wire clk,
    input wire n_reset,
    input wire [31:0] avl_in,
    input wire [4:0] vtype_in,
    input wire write,
    input wire saturate_flag,
    input wire preserve_vl,
    input wire set_vl_max
);

// CSRs included in this implementation:
// 0x009    vxsat   fixed-point saturate flag
// 0x00A    vxrm    fixed-point rounding mode (fixed to )
// 0xC20    vl      vector length
// 0xC21    vtype   vector data type register
// 0xC22    vlenb   vector register length in bytes
// For this accelerator they're separate from the standard CSR file - so we
// might as well put them in one block of their own, in this order

logic [31:0] csrs [4:0];

logic [4:0] vl_next;
logic [4:0] max_vl;
logic [2:0] per_reg;

always_ff @(posedge clk, negedge n_reset)
    if (~n_reset)
    begin
        for (int i=0; i<4; i++)
            csrs[i] <= '0;
            // vlenb is read-only, so can assign it at reset
        csrs[4] <= 32'd4;
    end
    else
    begin
        if (write)
        begin
            // vtype will be changed for every vsetvli instruction
            csrs[3] <= {'0, vtype_in};

            // Don't always want to write VL, eg. if rd == 0 and rs1 == 0
            // preserve_vl controls when it's left unchanged
            if (~preserve_vl)
                csrs[2] <= {'0 , vl_next};
        end
    end

always_comb
begin
    // Spec defines vsew as 3 bits of vtype, but our max element is 32b so the
    // top bit will always be zero and we can just look at the lower two
    vsew = csrs[3][3:2];
    vlmul = csrs[3][1:0];

    // If the AVL being suggested in the instruction is larger than max_vl, need
    // to set VL to max_vl. Also do this if set_vl_max asserted
    if ( set_vl_max | (avl_in > max_vl) )

assign vl = csrs[2][4:0];

// How many elements fit into a single register for each value of VSEW?
// Can work this out by dividing vlenb by vsew
assign per_reg = csrs[4][2:0] >> vtype_in[4:2];

// Max VL value equals the max number of elements per register * LMUL. LMUL is
// in powers of 2 so can use a shift
assign max_vl = per_reg << vtype_in[1:0];

assign vl_next_comb = vl_next;

endmodule

        vl_next = max_vl;
    else
        vl_next = avl_in[4:0];

end


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Matthew Johns <mrj1g17@soton.ac.uk>
// Byron Theobald <bt4g16@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

// `include "defs.sv"
import accelerator_pkg::*;

module vector_decoder (
    output logic apu_rvalid,
    output logic apu_gnt,
    output logic [31:0] scalar_operand1,
    output logic [31:0] scalar_operand2,
    output logic [10:0] immediate_operand,
    output logic [4:0] vs1_addr,
    output logic [4:0] vs2_addr,
    output logic [4:0] vd_addr,
    output logic csr_write,
    output logic preserve_vl,
    output logic set_vl_max,
    output logic [1:0] elements_to_write,
    output logic [1:0] cycle_count,
    output logic vec_reg_write,
    output vreg_wb_src_t vd_data_src,
    output vreg_addr_src_t vs3_addr_src,
    output pe_arith_op_t pe_op,
    output pe_saturate_mode_t saturate_mode,
    output pe_output_mode_t output_mode,
    output pe_operand_t operand_select,
    output logic [1:0] pe_mul_us,
    output logic [1:0] widening,
    output apu_result_src_t apu_result_select,
    output logic unsigned_immediate,
    output logic wide_vs1,
    input wire clk,
    input wire n_reset,
    input wire apu_req,
    input wire [31:0] apu_operands [2:0],
    input wire [5:0] apu_op,
    input wire [14:0] apu_flags_i,
    input wire [4:0] vl,
    input wire [1:0] vsew,
    output logic vlsu_en_o,
    output logic vlsu_load_o,
    output logic vlsu_store_o,
    output logic vlsu_strided_o,
    input logic vlsu_ready_i,
    input logic vlsu_done_i,
    output logic core_halt_o
);

enum {WAIT, EXEC, VALID} state, next_state;

logic [1:0] max_cycle_count;
logic multi_cycle_instr;
logic fix_vd_addr;

// Registers to store values from APU interface during instruction execution
logic [31:0] reg_apu_operands [2:0];
logic [5:0] reg_apu_op;
logic [14:0] reg_apu_flags_i;

// Assign variables for individual parts of instructions for readability
logic [2:0] funct3;
logic [6:0] major_opcode;
logic [5:0] funct6;
logic [4:0] source1;
logic [4:0] source2;
logic [4:0] destination;
logic [2:0] mop; // Vector Addressing Mode
assign funct3 = reg_apu_operands[0][14:12];
assign major_opcode = reg_apu_operands[0][6:0];
assign funct6 = reg_apu_operands[0][31:26];
assign source1 = reg_apu_operands[0][19:15];
assign source2 = reg_apu_operands[0][24:20];
assign destination = reg_apu_operands[0][11:7];
assign mop = funct6[2:0];

assign scalar_operand1 = reg_apu_operands[1];
assign scalar_operand2 = reg_apu_operands[2];

always_ff @(posedge clk, negedge n_reset)
    if(~n_reset)
    begin
        state <= WAIT;
        reg_apu_operands <= '{3{'0}};
        reg_apu_op <= '0;
        reg_apu_flags_i <= '0;
    end
    else
    begin
        state <= next_state;

        // In wait state, can load data from APU interface ready for the next
        // instruction. Only do this when it's valid, otherwise will screw any
        // invalid instruction checking code
        if ((state == WAIT) & apu_req)
        begin
            reg_apu_operands[0] <= apu_operands[0];
            reg_apu_operands[1] <= apu_operands[1];
            reg_apu_operands[2] <= apu_operands[2];
            reg_apu_op <= apu_op;
            reg_apu_flags_i <= apu_flags_i;
        end
    end

logic core_halt_ctrl;

assign core_halt_o = core_halt_ctrl;

/*always_ff @(posedge clk, negedge n_reset) begin
    if(~n_reset)
        core_halt_o <= 1'b0;
    else
        core_halt_o <= core_halt_ctrl;
end*/

always_comb
begin
    apu_rvalid = 1'b0;
    apu_gnt = 1'b0;
    next_state = state;
    core_halt_ctrl = 1'b0;

    case (state)
        WAIT:
        begin
            apu_gnt = 1'b1;
            if (apu_req)
                next_state = EXEC;

                next_state = WAIT;
        end
        EXEC:
        begin
            core_halt_ctrl = 1'b1;

            if (vlsu_load_o | vlsu_store_o) begin
                if(vlsu_done_i) begin
                    apu_rvalid = 1'b1;
                    next_state = WAIT;
                end
            end else if (cycle_count == max_cycle_count) begin
                apu_rvalid = 1'b1;
                next_state = WAIT;
            end
        end
    endcase
end

// VECTOR REGISTER ADDRESS GENERATION
always_ff @(posedge clk, negedge n_reset)
    if (~n_reset)
    begin
        cycle_count <= '0;
    end
    else
    begin
        if (state == WAIT || (vlsu_load_o | vlsu_store_o))
            cycle_count <= '0;
        else
            cycle_count <= cycle_count + 1'b1;

    end

logic [3:0] vl_zero_indexed;

always_comb
begin
    // Subtract 1 because if VL=4/8/16 it will want another cycle otherwise
    // Number of loads dependant on SEW (For contiguous 8-bit values)
    // TODO: Determine strided count
    // vl_zero_indexed = (vl - 1'b1) >> (2'd2 - vsew); // Used for memory
    vl_zero_indexed = vl - 1'b1;
    // Elements can be handled 4 at a time so divide VL by 4, or force 0
    max_cycle_count = multi_cycle_instr ? vl_zero_indexed[3:2] : 2'd0;

    case (vsew)
        2'd0: // 8b
        begin
            vs1_addr = source1 + cycle_count;
            vs2_addr = source2 + cycle_count;
            if (fix_vd_addr)
                vd_addr = destination;
            else
            begin
                if(widening[0])
                    vd_addr = destination + {cycle_count, 1'b0};
                else
                    vd_addr = destination + cycle_count;
            end
        end
        2'd1: // 16b
        begin
            vs1_addr = source1 + {cycle_count, 1'b0};
            vs2_addr = source2 + {cycle_count, 1'b0};
            if (fix_vd_addr)
                vd_addr = destination;
            else
                vd_addr = destination + {cycle_count, 1'b0};
        end
        default:
        begin
            vs1_addr = source1 + cycle_count;
            vs2_addr = source2 + cycle_count;
            if (fix_vd_addr)
                vd_addr = destination;
            else
                vd_addr = destination + cycle_count;
        end
    endcase

    if (funct3 == V_OPCFG)
        immediate_operand = reg_apu_operands[0][30:20];
    else
            immediate_operand = {'0, reg_apu_operands[0][19:15]};
end

always_comb
begin
    elements_to_write = 2'd0;

    if (multi_cycle_instr)
    begin
        if (cycle_count == max_cycle_count)
            if (operand_select == PE_OPERAND_RIPPLE)
            // Reductions only want to write in last cycle to only one element
                elements_to_write = 2'd1;
            else
            // On last cycle, work out how many elements remain
                elements_to_write = vl[1:0];
        else
            elements_to_write = 2'd0;
    end
end

////////////////////////////////////////////////////////////////////////////////
// ACCELERATOR CONTROL SIGNALS
always_comb
begin
    // Assign defaults for when not executing
    csr_write = 1'b0;
    preserve_vl = 1'b0;
    set_vl_max = 1'b0;
    vec_reg_write = 1'b0;
    vd_data_src = VREG_WB_SRC_ARITH;
    vs3_addr_src = VS3_ADDR_SRC_DECODE;
    pe_op = PE_ARITH_ADD;
    operand_select = PE_OPERAND_VS1;
    saturate_mode = PE_SAT_NONE;
    output_mode = PE_OP_MODE_RESULT;
    pe_mul_us = 2'b00;
    widening = 2'b00;
    apu_result_select = APU_RESULT_SRC_VL;
    multi_cycle_instr = 1'b0;
    unsigned_immediate = 1'b0;
    wide_vs1 = 1'b0;

    vlsu_en_o = 1'b0;
    vlsu_load_o = 1'b0;
    vlsu_store_o = 1'b0;
    vlsu_strided_o = 1'b0;

    // Used to control decoder module itself
    fix_vd_addr = 1'b0;

    // Control signals during instruction execution
    if (state == EXEC)
    begin
        if (major_opcode == V_MAJOR_LOAD_FP)
        begin
            if(funct3 == 3'b111) begin
                vd_data_src = VREG_WB_SRC_MEMORY;
                vlsu_en_o = 1'b1;
                vlsu_load_o = 1'b1;
                if(mop == 3'b010) vlsu_strided_o = 1'b1;
            end else $error("Unimplemented LOAD_FP instruction");
        end
        else if (major_opcode == V_MAJOR_STORE_FP)
        begin
            if(funct3 == 3'b111) begin
                vs3_addr_src = VS3_ADDR_SRC_VLSU;
                fix_vd_addr = 1'b1;
                vlsu_en_o = 1'b1;
                vlsu_store_o = 1'b1;
            end else $error("Unimplemented STORE_FP instruction");
        end
        else if (major_opcode == V_MAJOR_OP_V)
        begin
            // Consider vsetvli instructions separately (different format)
            if (funct3 == V_OPCFG)
            begin
                csr_write = 1'b1;
                apu_result_select = APU_RESULT_SRC_VL;
                if (source1 == '0)
                begin
                    if (destination == '0)
                        preserve_vl = 1'b1;
                    else
                        set_vl_max = 1'b1;
                end
            end
            else
            begin
                // Look for all other OP-V instructions
                case (funct6)

                    // vadd, vredsum
                    6'b000000:
                    begin
                        pe_op = PE_ARITH_ADD;
                        vec_reg_write = 1'b1;
                        multi_cycle_instr = 1'b1;
                        // vadd.vv
                        if (funct3 == V_OPIVV)
                            operand_select = PE_OPERAND_VS1;
                        // vadd.vx
                        else if (funct3 == V_OPIVX)
                            operand_select = PE_OPERAND_SCALAR;
                        else if (funct3 == V_OPMVV) // vredsum
                        begin
                            operand_select = PE_OPERAND_RIPPLE;
                            fix_vd_addr = 1'b1;
                        end
                    end

                    // vsub
                    6'b000010:
                    begin
                        pe_op = PE_ARITH_SUB;
                        vec_reg_write = 1'b1;
                        multi_cycle_instr = 1'b1;
                    end

                    // vmin
                    6'b000101:
                    begin
                        pe_op = PE_ARITH_SUB;
                        output_mode = PE_OP_MODE_PASS_MIN;
                        vec_reg_write = 1'b1;
                        multi_cycle_instr = 1'b1;
                        // Supports vmin.vv and vmin.vx
                        if (funct3 == V_OPIVV)
                            operand_select = PE_OPERAND_VS1;
                        else if (funct3 == V_OPIVX)
                            operand_select = PE_OPERAND_SCALAR;
                    end

                    // vmax, vredmax
                    6'b000111:
                    begin
                        pe_op = PE_ARITH_SUB;
                        vec_reg_write = 1'b1;
                        output_mode = PE_OP_MODE_PASS_MAX;
                        multi_cycle_instr = 1'b1;
                        // vredmax
                        if (funct3 == V_OPMVV)
                        begin
                            fix_vd_addr = 1'b1;
                            operand_select = PE_OPERAND_RIPPLE;
                        end
                        // Supports vmax.vv and vmax.vx
                        else if (funct3 == V_OPIVV)
                            operand_select = PE_OPERAND_VS1;
                        else if (funct3 == V_OPIVX)
                            operand_select = PE_OPERAND_SCALAR;
                    end

                    // vand
                    6'b001001:
                    begin
                        pe_op = PE_ARITH_AND;
                        vec_reg_write = 1'b1;
                        multi_cycle_instr = 1'b1;
                    end

                    // vor
                    6'b001010:
                    begin
                        pe_op = PE_ARITH_OR;
                        vec_reg_write = 1'b1;
                        multi_cycle_instr = 1'b1;
                    end

                    // vxor
                    6'b001011:
                    begin
                        pe_op = PE_ARITH_XOR;
                        vec_reg_write = 1'b1;
                        multi_cycle_instr = 1'b1;
                    end

                    // VWXUNARY0 (vmv.x.s)
                    6'b010000:
                    begin
                        apu_result_select = APU_RESULT_SRC_VS2_0;
                    end

                    // vmv.v
                    6'b010111:
                    begin
                        vec_reg_write = 1'b1;
                        multi_cycle_instr = 1'b1;
                        vd_data_src = VREG_WB_SRC_SCALAR;
                        // vmv.v.i
                        if (funct3 == V_OPIVI)
                            operand_select = PE_OPERAND_IMMEDIATE;
                        // vmv.v.x
                        else if (funct3 == V_OPIVX)
                            operand_select = PE_OPERAND_SCALAR;
                    end

                    // vsadd
                    6'b100001:
                    begin
                        pe_op = PE_ARITH_ADD;
                        saturate_mode = PE_SAT;
                        vec_reg_write = 1'b1;
                        multi_cycle_instr = 1'b1;
                    end

                    // vsll/vmul
                    6'b100101:
                    begin
                        vec_reg_write = 1'b1;
                        multi_cycle_instr = 1'b1;
                        if (funct3 == V_OPIVV)
                        begin
                            pe_op = PE_ARITH_LSHIFT;
                            operand_select = PE_OPERAND_VS1;
                        end
                        else if (funct3 == V_OPIVX)
                        begin
                            pe_op = PE_ARITH_LSHIFT;
                            operand_select = PE_OPERAND_SCALAR;
                        end
                        else if (funct3 == V_OPIVI)
                        begin
                            pe_op = PE_ARITH_LSHIFT;
                            operand_select = PE_OPERAND_IMMEDIATE;
                        end
                        else if (funct3 == V_OPMVV)
                        begin
                            pe_op = PE_ARITH_MUL;
                            operand_select = PE_OPERAND_VS1;
                        end

                    end

                    // vsmul
                    6'b100111:
                    begin
                        pe_op = PE_ARITH_MUL;
                        saturate_mode = PE_SAT_UPPER;
                        vec_reg_write = 1'b1;
                        multi_cycle_instr = 1'b1;
                        if (funct3 == V_OPIVV)
                            operand_select = PE_OPERAND_VS1;
                        else if (funct3 == V_OPIVX)
                            operand_select = PE_OPERAND_SCALAR;
                    end

                    // vsrl
                    6'b101000:
                    begin
                        pe_op = PE_ARITH_RSHIFT_LOG;
                        vec_reg_write = 1'b1;
                        multi_cycle_instr = 1'b1;
                        if (funct3 == V_OPIVV)
                            operand_select = PE_OPERAND_VS1;
                        else if (funct3 == V_OPIVX)
                            operand_select = PE_OPERAND_SCALAR;
                        else if (funct3 == V_OPIVI)
                        begin
                            unsigned_immediate = 1'b1;
                            operand_select = PE_OPERAND_IMMEDIATE;
                        end
                    end

                    // vsra
                    6'b101001:
                    begin
                        pe_op = PE_ARITH_RSHIFT_AR;
                        vec_reg_write = 1'b1;
                        multi_cycle_instr = 1'b1;
                        if (funct3 == V_OPIVV)
                            operand_select = PE_OPERAND_VS1;
                        else if (funct3 == V_OPIVX)
                            operand_select = PE_OPERAND_SCALAR;
                        else if (funct3 == V_OPIVI)
                        begin
                            unsigned_immediate = 1'b1;
                            operand_select = PE_OPERAND_IMMEDIATE;
                        end
                    end

                    // vwredsum
                    6'b110001:
                    begin
                        pe_op = PE_ARITH_ADD;
                        operand_select = PE_OPERAND_RIPPLE;
                        vec_reg_write = 1'b1;
                        fix_vd_addr = 1'b1;
                        multi_cycle_instr = 1'b1;
                        widening = 2'b01;
                        wide_vs1 = 1'b1;
                    end

                    // vwmul
                    6'b111011:
                    begin
                        pe_op = PE_ARITH_MUL;
                        vec_reg_write = 1'b1;
                        multi_cycle_instr = 1'b1;
                        widening = 2'b01;
                        if (funct3 == V_OPMVV)
                            operand_select = PE_OPERAND_VS1;
                        else if (funct3 == V_OPMVX)
                            operand_select = PE_OPERAND_SCALAR;
                    end

                    default:
                        $error("Unsupported vector instruction");

                endcase
            end
        end
        else
            $error("Unrecognised major opcode");
    end
end

endmodule


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Matthew Johns <mrj1g17@soton.ac.uk>
// Byron Theobald <bt4g16@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

import accelerator_pkg::*;

module vector_lsu (
    input  wire         clk,
    input  wire         n_reset,

    // Vector CSR
    input  wire [4:0]   vl_i,
    input  wire [1:0]   vsew_i,
    input  wire [1:0]   vlmul_i,

    // VLSU Decoded Control
    input  wire         vlsu_en_i,
    input  wire         vlsu_load_i,
    input  wire         vlsu_store_i,
    input  wire         vlsu_strided_i,
    output logic        vlsu_ready_o,
    output logic        vlsu_done_o,

    // OBI Memory Master
    output logic        data_req_o,
    input  logic        data_gnt_i,
    input  logic        data_rvalid_i,
    output logic [31:0] data_addr_o,
    output logic        data_we_o,
    output logic [3:0]  data_be_o,
    input  logic [31:0] data_rdata_i,
    output logic [31:0] data_wdata_o,

    input  logic [1:0]  cycle_count_i,

    // Target Data
    input  wire [31:0]  op0_data_i, // Source (Load) / Destination (Store)
    input  wire [31:0]  op1_data_i, // Stride

    // Wide vector register port
    output logic [127:0] vs_wdata_o,
    input  logic [127:0] vs_rdata_i,
    input  logic [4:0] vr_addr_i,
    output logic [4:0] vs3_addr_o, // Redirected vector register address
    output logic vr_we_o
);

    logic [31:0] vs_rdata_sel;
    logic [5:0] vsew_size;

    logic au_start;
    logic [3:0] au_be;
    logic [6:0] au_bc;
    logic [31:0] au_addr;
    logic au_valid, au_ready;
    logic [6:0] vd_offset;

    temporary_reg tr (
        .clk_i              (clk),
        .n_rst_i            (n_reset),
        .byte_enable_valid  (data_req_o),
        .read_data_valid    (data_rvalid_i),
        .clear_register     (au_start),
        .memory_read_i      (data_rdata_i),
        .byte_enable_i      (au_be),
        .byte_select_i      (vd_offset + {vr_addr_i[1:0], 2'b00}),
        .wide_vd_o          (vs_wdata_o)
    );

    always_comb begin
		  data_wdata_o = 'd0;
        case(vsew_i)
            2'd0 : begin
                data_wdata_o = {vs_rdata_i[103:96], vs_rdata_i[71:64], vs_rdata_i[39:32], vs_rdata_i[7:0]};
            end
            2'd1 : begin

                    2'd2 : data_wdata_o = vs_rdata_i[95:64];
                    2'd3 : data_wdata_o = vs_rdata_i[127:96];
                endcase
            end
        endcase
    end

    logic [1:0] ib_select; // Low 2 bits of initial address
    logic [3:0] be_gen;

    logic [31:0] next_el_pre, next_el_addr;
    logic [31:0] cycle_addr, stride;
    logic [6:0] cycle_bytes;

    typedef enum {RESET, LOAD_FIRST, LOAD_CYCLE, LOAD_WAIT, LOAD_FINAL, STORE_CYCLE, STORE_WAIT, STORE_FINAL} be_state;
    be_state current_state, next_state;

    logic signed [6:0] byte_track, byte_track_next;
    logic cycle_load, cycle_addr_inc, store_cycles_inc;

    logic [2:0] store_cycle_bytes;
    logic [3:0] store_cycle_be;
    logic [2:0] store_cycles, store_cycles_cnt;

    assign stride = vlsu_strided_i ? op1_data_i : (32'd1 << vsew_i);
    assign data_addr_o = vlsu_store_i ? ({cycle_addr[31:2], 2'd0} + (store_cycles_cnt << 2)) : {cycle_addr[31:2], 2'd0};
    assign au_be = be_gen;
    assign vd_offset = (vl_i << vsew_i) - byte_track;

    always_comb begin
        if(byte_track >= 4)
            store_cycle_be = 4'b1111;
        else if(byte_track >= 3)
            store_cycle_be = 4'b0111;
        else if(byte_track >= 2)
            store_cycle_be = 4'b0011;
        else
            store_cycle_be = 4'b0001;

        data_be_o = vlsu_store_i ? store_cycle_be : 4'b1111;
    end

    always_comb begin
        if(au_start)
            byte_track_next = {2'd0, vl_i} << vsew_i; // Bytes dependent on element size
        else if(cycle_addr_inc)
            byte_track_next = (byte_track >= cycle_bytes) ? (byte_track - cycle_bytes) : 7'd0;
        else if(store_cycles_inc)
            byte_track_next = byte_track - store_cycle_bytes;
        else
            byte_track_next = byte_track;
    end

    always_ff @(posedge clk, negedge n_reset) begin
        if(~n_reset)
            byte_track <= 7'd0;
        else
            byte_track <= byte_track_next;
    end

    always_ff @(posedge clk, negedge n_reset) begin
        if(~n_reset)
            cycle_addr <= 32'd0;
        else if(au_start)
            cycle_addr <= op0_data_i;
        else if(cycle_addr_inc)
            cycle_addr <= next_el_addr;
        else
            cycle_addr <= cycle_addr;
    end

    assign store_cycles = (vl_i >> 2-vsew_i)+1;
    assign vs3_addr_o = vr_addr_i + store_cycles_cnt;
    always_ff @(posedge clk, negedge n_reset) begin
        if(~n_reset)
            store_cycles_cnt <= 2'd0;
        else if(au_start)
            store_cycles_cnt <= 2'd0;
        else if (store_cycles_inc)
            store_cycles_cnt <= store_cycles_cnt + 2'd1;
    end

    always_ff @(posedge clk, negedge n_reset) begin
        if(~n_reset)
            current_state <= RESET;
        else
            current_state <= next_state;
    end

    always_comb begin
        be_gen = 4'b0000;
        next_el_pre = '0;
		  cycle_bytes = '0;
		  ib_select = '0;
		  next_el_addr = '0;
        case(vsew_i)
            2'b00 : begin // 8 Bit
                ib_select = cycle_addr[1:0];

                if(stride > 32'd1) begin
                    be_gen[ib_select] = 1'b1;

                    // Where is our next byte?
                    next_el_pre = cycle_addr + stride;
                    if(next_el_pre[31:2] == cycle_addr[31:2] && byte_track > 1) begin
                        be_gen[next_el_pre[1:0]] = 1'b1;
                        next_el_addr = next_el_pre + stride; // Stride by second element
                    end else begin
                        next_el_addr = next_el_pre;
                    end

                    // Calculate the number of bytes for LOAD_CYCLE
                    cycle_bytes = {5'd0, be_gen[3]} + {5'd0, be_gen[2]} + {5'd0, be_gen[1]} + {5'd0, be_gen[0]};
                end else if(stride == 1) begin
                    be_gen[0] = (ib_select == 32'd0) ? 1 : 0;
                    be_gen[1] = (ib_select == 1 || byte_track > 1) ? 1'b1 : 1'b0;
                    be_gen[2] = (ib_select == 2 || byte_track > 2) ? 1'b1 : 1'b0;
                    be_gen[3] = (ib_select == 3 || byte_track > 3) ? 1'b1 : 1'b0;
                    next_el_addr = {cycle_addr[31:2], 2'b0} + 32'd4;

                    // Calculate the number of bytes for LOAD_CYCLE
                    cycle_bytes = {5'd0, be_gen[3]} + {5'd0, be_gen[2]} + {5'd0, be_gen[1]} + {5'd0, be_gen[0]};
                end else if(stride == 32'd0) begin
                    be_gen[ib_select] = 1'b1;
                    cycle_bytes = {2'b0, vl_i}; // Read all bytes in 1 LOAD_CYCLE
                end
            end
            2'b01 : begin // 16 Bit
                ib_select = {cycle_addr[1], 1'b0}; // Force alignment byte 0 or 2

                if(stride > 32'd2) begin // Always 1 element
                    // Always set 2 bytes
                    be_gen[ib_select] = 1'b1;
                    be_gen[ib_select+1] = 1'b1;
                    next_el_addr = {cycle_addr[31:1], 1'b0} + {stride[31:1], 1'b0};

                    // Calculate the number of bytes for LOAD_CYCLE
                    cycle_bytes = {5'd0, be_gen[3]} + {5'd0, be_gen[2]} + {5'd0, be_gen[1]} + {5'd0, be_gen[0]};
                end else if (stride == 32'd2) begin // Up to 2 Elements
                    be_gen[1:0] = (ib_select == 0) ? 2'b11 : 2'b00;
                    be_gen[3:2] = (ib_select == 2 || byte_track > 2) ? 2'b11 : 2'b00;
                    next_el_addr = {cycle_addr[31:2], 2'b0} + 32'd4;

                    // Calculate the number of bytes for LOAD_CYCLE
                    cycle_bytes = {5'd0, be_gen[3]} + {5'd0, be_gen[2]} + {5'd0, be_gen[1]} + {5'd0, be_gen[0]};
                end else if (stride == 32'd0) begin
                    be_gen[ib_select] = 1'b1;
                    be_gen[ib_select+1] = 1'b1;
                    cycle_bytes = {1'b0, vl_i, 1'b0}; // Read all bytes in 1 LOAD_CYCLE
                end
            end
            2'b10 : begin // 32 Bit
                ib_select = 2'd0; // Force alignment to byte 0

                if(stride >= 32'd4) begin // Always 1 element
                    be_gen = 4'b1111;
                    next_el_addr = {cycle_addr[31:2], 2'b0} + {stride[31:2], 2'b0}; // stride is always a multiple of 4

                    // Calculate the number of bytes for LOAD_CYCLE
                    cycle_bytes = {5'd0, be_gen[3]} + {5'd0, be_gen[2]} + {5'd0, be_gen[1]} + {5'd0, be_gen[0]};
                end else if(stride == 32'd0) begin
                    be_gen = 4'b1111;
                    cycle_bytes = {1'b0, vl_i, 1'b0}; // Read all bytes in 1 LOAD_CYCLE
                end
            end
            default : $error("Invalid VSEW");
        endcase
    end

    always_comb begin
        cycle_load = 1'b0;
        data_req_o = 1'b0;
        data_we_o = 1'b0;
        au_start = 1'b0;
        au_ready = 1'b0;
        vlsu_done_o = 1'b0;
        vlsu_ready_o = 1'b0;
        cycle_addr_inc = 1'b0;
        store_cycles_inc = 1'b0;
        vr_we_o = 1'b0;
        case(current_state)
            RESET: begin
                vlsu_ready_o = 1'b1;
                if(vlsu_load_i) begin
                    au_start = 1'b1;
                    next_state = LOAD_FIRST;
                end else if (vlsu_store_i) begin
                    au_start = 1'b1;
                    next_state = STORE_CYCLE;
                end else begin
                    next_state = RESET;
                end
            end
            LOAD_FIRST: begin
                next_state = LOAD_CYCLE;
            end
            LOAD_CYCLE: begin
                if(byte_track_next == 0) begin
                    next_state = LOAD_WAIT;
                end else begin
                    data_req_o = 1'b1;
                    cycle_load = 1'b1;
                    next_state = LOAD_WAIT;
                end
            end
            LOAD_WAIT: begin
                if(data_rvalid_i) begin
                    cycle_addr_inc = 1'b1;
                    next_state = LOAD_CYCLE;
                end else if(byte_track_next == 0)
                    next_state = LOAD_FINAL;
                else
                    next_state = LOAD_WAIT;
            end
            LOAD_FINAL: begin
                next_state = RESET;
                vlsu_done_o = 1'b1;
                vr_we_o = 1'b1;
            end
            STORE_CYCLE: begin
                data_req_o = 1'b1;
                data_we_o = 1'b1;
                next_state = STORE_WAIT;
            end
            STORE_WAIT: begin
                if(data_rvalid_i) begin
                    if(store_cycles_cnt == store_cycles) begin
                        next_state = STORE_FINAL;
                    end else begin
                        store_cycles_inc = 1'b1;
                        next_state = STORE_CYCLE;
                    end
                end else begin
                    next_state = STORE_WAIT;
                end
            end
            STORE_FINAL: begin
                vlsu_done_o = 1'b1;
                next_state = RESET;
            end
        endcase
    end

endmodule

                case(vs3_addr_o[0])
                    1'd0 : data_wdata_o = {vs_rdata_i[47:32], vs_rdata_i[15:0]};
                    1'd1 : data_wdata_o = {vs_rdata_i[111:96], vs_rdata_i[79:64]};
                endcase
            end


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Matthew Johns <mrj1g17@soton.ac.uk>
// Byron Theobald <bt4g16@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

module vector_registers (
    output logic [127:0] vs1_data,
    output logic [127:0] vs2_data,
    output logic [127:0] vs3_data,
    input wire [127:0] vd_data,
    input wire [4:0] vs1_addr,
    input wire [4:0] vs2_addr,
    input wire [4:0] vd_addr,   // Generally this doubles up as vs3 address
    input wire [1:0] vsew,
    input wire [1:0] vlmul,
    input wire [1:0] elements_to_write,
    input wire clk,
    input wire n_reset,
    input wire write,
    input wire widening_op,
    input wire wide_vs1,
    input logic load_operation
);

localparam VLEN = 32;

logic [VLEN-1:0] vregs [31:0];

// Addresses for each of the read ports for each operand.
// Each operand may require up to four read ports.
// Reason: need 4 elements per operand per cycle to maintain throughput in
// the 4 PEs. For 8b elements, 4 elements are stored in a single register;
// however with 32b elements they will be spread across 4 registers (when
// LMUL > 1). Otherwise SIMD throughput would be tiny.
// These addresses will be consecutive for each operand. Just want a tidy
// efficient way of producing them without an adder for each one.
logic [4:0] vs1_addr0;
logic [4:0] vs1_addr1;
logic [4:0] vs1_addr2;
logic [4:0] vs1_addr3;

logic [4:0] vs2_addr0;
logic [4:0] vs2_addr1;
logic [4:0] vs2_addr2;
logic [4:0] vs2_addr3;

logic [4:0] vd_addr0;
logic [4:0] vd_addr1;
logic [4:0] vd_addr2;
logic [4:0] vd_addr3;

// Structured data to write back to registers
logic [VLEN-1:0] vd_wr_data0;
logic [VLEN-1:0] vd_wr_data1;
logic [VLEN-1:0] vd_wr_data2;
logic [VLEN-1:0] vd_wr_data3;

// Unstructured data read from registers
logic [VLEN-1:0] vs1_rd_data0;
logic [VLEN-1:0] vs1_rd_data1;
logic [VLEN-1:0] vs1_rd_data2;
logic [VLEN-1:0] vs1_rd_data3;

logic [VLEN-1:0] vs2_rd_data0;
logic [VLEN-1:0] vs2_rd_data1;
logic [VLEN-1:0] vs2_rd_data2;
logic [VLEN-1:0] vs2_rd_data3;

logic [VLEN-1:0] vs3_rd_data0;
logic [VLEN-1:0] vs3_rd_data1;
logic [VLEN-1:0] vs3_rd_data2;
logic [VLEN-1:0] vs3_rd_data3;

// Write-enable signals for each write port. Could get away with making wr_en1
// only 2 bits, as will only write 16b elements. Similarly could make wr_en2 and
// wr_en3 single bits. But I don't know what that would synthesise to if I did
logic [3:0] wr_en0;
logic [3:0] wr_en1;
logic [3:0] wr_en2;
logic [3:0] wr_en3;

// Effective vsew can be modified for widening ops
logic [1:0] eff_vsew;

// REGISTER WRITE
always_ff @(posedge clk, negedge n_reset)
    if (~n_reset)
        vregs <= '{VLEN{'0}};
    else
    begin
        // Don't want to write to v0 (reserved for vector mask)
        if (write & (vd_addr != '0))
        begin
            if (wr_en0[0])
                vregs[vd_addr0][7:0] <= vd_wr_data0[7:0];
            if (wr_en0[1])
                vregs[vd_addr0][15:8] <= vd_wr_data0[15:8];
            if (wr_en0[2])
                vregs[vd_addr0][23:16] <= vd_wr_data0[23:16];
            if (wr_en0[3])
                vregs[vd_addr0][31:24] <= vd_wr_data0[31:24];

            if (wr_en1[0])
                vregs[vd_addr1][7:0] <= vd_wr_data1[7:0];
            if (wr_en1[1])
                vregs[vd_addr1][15:8] <= vd_wr_data1[15:8];
            if (wr_en1[2])
                vregs[vd_addr1][23:16] <= vd_wr_data1[23:16];
            if (wr_en1[3])
                vregs[vd_addr1][31:24] <= vd_wr_data1[31:24];

            if (wr_en2[0])
                vregs[vd_addr2][7:0] <= vd_wr_data2[7:0];
            if (wr_en2[1])
                vregs[vd_addr2][15:8] <= vd_wr_data2[15:8];
            if (wr_en2[2])
                vregs[vd_addr2][23:16] <= vd_wr_data2[23:16];
            if (wr_en2[3])
                vregs[vd_addr2][31:24] <= vd_wr_data2[31:24];

            if (wr_en3[0])
                vregs[vd_addr3][7:0] <= vd_wr_data3[7:0];
            if (wr_en3[1])
                vregs[vd_addr3][15:8] <= vd_wr_data3[15:8];
            if (wr_en3[2])
                vregs[vd_addr3][23:16] <= vd_wr_data3[23:16];
            if (wr_en3[3])
                vregs[vd_addr3][31:24] <= vd_wr_data3[31:24];
        end
    end

// REGISTER READ
assign vs1_rd_data0 = vregs[vs1_addr0];
assign vs1_rd_data1 = vregs[vs1_addr1];
assign vs1_rd_data2 = vregs[vs1_addr2];
assign vs1_rd_data3 = vregs[vs1_addr3];
assign vs2_rd_data0 = vregs[vs2_addr0];
assign vs2_rd_data1 = vregs[vs2_addr1];
assign vs2_rd_data2 = vregs[vs2_addr2];
assign vs2_rd_data3 = vregs[vs2_addr3];
assign vs3_rd_data0 = vregs[vd_addr0];
assign vs3_rd_data1 = vregs[vd_addr1];
assign vs3_rd_data2 = vregs[vd_addr2];
assign vs3_rd_data3 = vregs[vd_addr3];

// ADDRESS CALCULATION
always_comb
begin
    // Logic behind this: Need 4 consecutive addresses, but don't want to just
    // have adders for each one to increment the address.
    // This is only useful for LMUL > 1, as for LMUL = 1, only one register will
    // be read anyway. If LMUL > 1, base addresses will always be even. So can
    // add one to it by making last bit 1. Add 2 by making second-last bit 1.
    // Add 3 by doing both.
    // If LMUL = 2 then adding 2 that way won't work, but also it won't be used
    // because only the first 2 registers will be used.
    vs1_addr0 = vs1_addr;
    vs1_addr1 = {vs1_addr[4:1], 1'b1};
    vs1_addr2 = {vs1_addr[4:2], 1'b1, vs1_addr[0]};
    vs1_addr3 = {vs1_addr[4:2], 2'b11};

    vs2_addr0 = vs2_addr;
    vs2_addr1 = {vs2_addr[4:1], 1'b1};
    vs2_addr2 = {vs2_addr[4:2], 1'b1, vs2_addr[0]};
    vs2_addr3 = {vs2_addr[4:2], 2'b11};

    vd_addr0 = vd_addr;
    vd_addr1 = {vd_addr[4:1], 1'b1};
    vd_addr2 = {vd_addr[4:2], 1'b1, vd_addr[0]};
    vd_addr3 = {vd_addr[4:2], 2'b11};
end

// WRITE-ENABLE GENERATION
always_comb
begin
    // Note: can ignore LMUL in below cases as LMUL will restrict the max of VL,
    // which will prevent from writing to higher registers than LMUL wants
    if (widening_op)
        case (vsew)
            2'd0: // 8b -> 16b
                eff_vsew = 2'd1;
            2'd1: // 16b -> 32b
                eff_vsew = 2'd2;
            default:
            begin
                // Shouldn't come to this
                eff_vsew = vsew;
                $error("Widening ops with VSEW=32b are not supported");
            end
        endcase
    else
        eff_vsew = vsew;

    wr_en0 = '0;
    wr_en1 = '0;
    wr_en2 = '0;
    wr_en3 = '0;

    if(load_operation) begin
        case (vlmul)
            2'd0: begin
                case(vd_addr[1:0])
                   2'b00 : wr_en0 = 4'b1111;
                   2'b01 : wr_en1 = 4'b1111;
                   2'b10 : wr_en2 = 4'b1111;
                   2'b11 : wr_en3 = 4'b1111;
                endcase
            end
            2'd1: begin
                if(vd_addr[1] == 1'b0) begin
                   wr_en0 = 4'b1111;
                   wr_en1 = 4'b1111;
                end else begin
                   wr_en2 = 4'b1111;
                   wr_en3 = 4'b1111;
                end
            end
            2'd2: begin
                wr_en0 = 4'b1111;
                wr_en1 = 4'b1111;
                wr_en2 = 4'b1111;
                wr_en3 = 4'b1111;
            end
        endcase
    end else begin
        case (eff_vsew)
            2'd0: // 8b
            begin
                // Only interested in first write port
                case (elements_to_write)
                    2'd0: // Write all elements
                        wr_en0 = 4'b1111;
                    2'd1:
                        wr_en0 = 4'b0001;
                    2'd2:
                        wr_en0 = 4'b0011;
                    2'd3:
                        wr_en0 = 4'b0111;
                endcase
            end
            2'd1: // 16b
            begin
                // Only interested in first 2 write ports
                case (elements_to_write)
                    2'd0: // Write all elements
                    begin
                        wr_en0 = 4'b1111;
                        wr_en1 = 4'b1111;
                    end
                    2'd1:
                    begin
                        wr_en0 = 4'b0011;
                    end
                    2'd2:
                    begin
                        wr_en0 = 4'b1111;
                    end
                    4'd3:
                    begin
                        wr_en0 = 4'b1111;
                        wr_en1 = 4'b0011;
                    end
                endcase
            end
            2'd2: // 32b
            begin
                // Need to consider all write ports
                // wr_en0 always enabled, otherwise would be writing no elements
                wr_en0 = 4'b1111;
                case (elements_to_write)
                    2'd0: // Write all elements
                    begin
                        wr_en1 = 4'b1111;
                        wr_en2 = 4'b1111;
                        wr_en3 = 4'b1111;

                    begin
                        wr_en1 = 4'b1111;
                        wr_en2 = 4'b1111;
                    end
                endcase
            end
        endcase
    end
end

// OUTPUT DATA MAP
// TODO: consider need for sign extension for VSEW<32b
// Have to pad the extra space when using smaller elements to give correct
// alignment into the PEs
always_comb
begin
    case (vsew)
        2'd0: // 8b
        begin
            // For vwredsum (theoretically other mixed-width instructions) the B
            // operand (which comes from vs1[0]) is 2*VSEW bits. So treat it as
            // for wider ones.
            if (wide_vs1)
                vs1_data = {
                    {16{1'b0}},
                    vs1_rd_data1[31:16],
                    {16{1'b0}},
                    vs1_rd_data1[15:0],
                    {16{1'b0}},
                    vs1_rd_data0[31:16],
                    {16{1'b0}},
                    vs1_rd_data0[15:0]
                };
            else
                vs1_data = {
                    {24{1'b0}},
                    vs1_rd_data0[31:24],
                    {24{1'b0}},
                    vs1_rd_data0[23:16],
                    {24{1'b0}},
                    vs1_rd_data0[15:8],
                    {24{1'b0}},
                    vs1_rd_data0[7:0]
                };

            vs2_data = {
                {24{1'b0}},
                vs2_rd_data0[31:24],
                {24{1'b0}},
                vs2_rd_data0[23:16],
                {24{1'b0}},
                vs2_rd_data0[15:8],
                {24{1'b0}},
                vs2_rd_data0[7:0]
            };
            // For widening ops such as MACC, the third operand needs to be the
            // widened width rather than VSEW. Copy the mapping of vs3 from the
            // next largest element. Does not apply to VSEW=32b as widening ops
            // are not supported.
            if (widening_op)
                vs3_data = {
                    {16{1'b0}},
                    vs3_rd_data1[31:16],
                    {16{1'b0}},
                    vs3_rd_data1[15:0],
                    {16{1'b0}},
                    vs3_rd_data0[31:16],
                    {16{1'b0}},
                    vs3_rd_data0[15:0]
                };
            else
                vs3_data = {
                    {24{1'b0}},
                    vs3_rd_data0[31:24],
                    {24{1'b0}},
                    vs3_rd_data0[23:16],
                    {24{1'b0}},
                    vs3_rd_data0[15:8],
                    {24{1'b0}},
                    vs3_rd_data0[7:0]
                };
        end
        2'd1: // 16b
        begin
            if (wide_vs1)
                vs1_data = {
                    vs1_rd_data3,
                    vs1_rd_data2,
                    vs1_rd_data1,
                    vs1_rd_data0
                };
            else
                vs1_data = {
                    {16{1'b0}},
                    vs1_rd_data1[31:16],
                    {16{1'b0}},
                    vs1_rd_data1[15:0],
                    {16{1'b0}},
                    vs1_rd_data0[31:16],
                    {16{1'b0}},
                    vs1_rd_data0[15:0]
                };

            vs2_data = {
                {16{1'b0}},
                vs2_rd_data1[31:16],
                {16{1'b0}},
                vs2_rd_data1[15:0],
                {16{1'b0}},
                vs2_rd_data0[31:16],
                {16{1'b0}},
                vs2_rd_data0[15:0]
            };
            if (widening_op)
                vs3_data = {
                    vs3_rd_data3,
                    vs3_rd_data2,
                    vs3_rd_data1,
                    vs3_rd_data0
                };
            else
                vs3_data = {
                    {16{1'b0}},
                    vs3_rd_data1[31:16],
                    {16{1'b0}},
                    vs3_rd_data1[15:0],
                    {16{1'b0}},
                    vs3_rd_data0[31:16],
                    {16{1'b0}},
                    vs3_rd_data0[15:0]
                };
        end
        2'd2: // 32b
        begin
            vs1_data = {
                vs1_rd_data3,
                vs1_rd_data2,
                vs1_rd_data1,
                vs1_rd_data0
            };
            vs2_data = {
                vs2_rd_data3,
                vs2_rd_data2,
                vs2_rd_data1,
                vs2_rd_data0
            };
            vs3_data = {
                vs3_rd_data3,
                vs3_rd_data2,
                vs3_rd_data1,
                vs3_rd_data0
            };
        end
        default:
        begin
            vs1_data = '0;
            vs2_data = '0;
            vs3_data = '0;
        end
    endcase
end

// INPUT DATA MAP
// Take wide combined result data from PEs and remove padding
always_comb
begin
    vd_wr_data3 = '0;
    vd_wr_data2 = '0;
    vd_wr_data1 = '0;
    vd_wr_data0 = '0;

    if(load_operation) begin
        case (vlmul)
            2'd0: begin
                case(vd_addr[1:0])
                   2'b00 : vd_wr_data0 = vd_data[31:0];
                   2'b01 : vd_wr_data1 = vd_data[63:32];
                   2'b10 : vd_wr_data2 = vd_data[95:64];
                   2'b11 : vd_wr_data3 = vd_data[127:96];
                endcase
            end
            2'd1: begin
                if(vd_addr[1] == 1'b0) begin
                    vd_wr_data1 = vd_data[63:32];
                    vd_wr_data0 = vd_data[31:0];
                end else begin
                    vd_wr_data3 = vd_data[127:96];
                    vd_wr_data2 = vd_data[95:64];
                end
            end
            2'd2: begin
                vd_wr_data3 = vd_data[127:96];
                vd_wr_data2 = vd_data[95:64];
                vd_wr_data1 = vd_data[63:32];
                vd_wr_data0 = vd_data[31:0];
            end
        endcase
    end else begin
        case (eff_vsew)
            2'd0: // 8b
                vd_wr_data0 = {
                    vd_data[103:96],
                    vd_data[71:64],
                    vd_data[39:32],
                    vd_data[7:0]
                };
            2'd1: // 16b
            begin
                vd_wr_data1 = {
                    vd_data[111:96],
                    vd_data[79:64]
                };
                vd_wr_data0 = {
                    vd_data[47:32],
                    vd_data[15:0]
                };
            end
            2'd2: // 32b
            begin
                vd_wr_data3 = vd_data[127:96];
                vd_wr_data2 = vd_data[95:64];
                vd_wr_data1 = vd_data[63:32];
                vd_wr_data0 = vd_data[31:0];
            end
        endcase
    end
end

endmodule

                    end
                    // 2'd1: // Not needed, just wr_en0 = '1
                    2'd2:
                    begin


//
// SPDX-License-Identifier: CERN-OHL-S-2.0+
//
// Copyright (C) 2020-21 Embecosm Limited <www.embecosm.com>
// Contributed by:
// Matthew Johns <mrj1g17@soton.ac.uk>
// Byron Theobald <bt4g16@soton.ac.uk>
//
// This source is distributed WITHOUT ANY EXPRESS OR IMPLIED WARRANTY,
// INCLUDING OF MERCHANTABILITY, SATISFACTORY QUALITY AND FITNESS FOR
// A PARTICULAR PURPOSE. Please see the CERN-OHL-S v2 for applicable
// conditions.
// Source location: https://github.com/AI-Vector-Accelerator
//

// Variable width sign extension module. Used to sign-extend 3 PE inputs for
// signed/widening multiplication

module vw_sign_ext (
    output logic [31:0] sign_ext_a,
    output logic [31:0] sign_ext_b,
    output logic [31:0] sign_ext_c,
    input wire [31:0] a,
    input wire [31:0] b,
    input wire [31:0] c,
    input wire [1:0] widening, // 2'd1 for 2*widening, 2'd2 for quad widening
    input wire [1:0] vsew,
    input wire wide_b
);

always_comb begin
	 sign_ext_a = '0;
	 sign_ext_b = '0;
	 sign_ext_c = '0;
    case(vsew)

            // operand for each PE is 2*VSEW bits because it is an intermediate
            // result. So treat it as if vsew was twice as large
            if (wide_b)
                sign_ext_b = {{16{b[15]}}, b[15:0]};
            else
                sign_ext_b = {{24{b[7]}}, b[7:0]};

            if (widening[0])
                sign_ext_c = {{16{c[15]}}, c[15:0]};
            else if (widening[1])
                sign_ext_c = c;
            else
                sign_ext_c = {{24{c[7]}}, c[7:0]};
        end
        2'd1: // 16b
        begin
            sign_ext_a = {{16{a[15]}}, a[15:0]};

            if (wide_b)
                sign_ext_b = b;
            else
                sign_ext_b = {{16{b[15]}}, b[15:0]};

            if (widening[0])
                sign_ext_c = c;
            else if (widening[1])
                $error("Trying to quad-widen 16b elements!");
            else
                sign_ext_c = {{16{c[15]}}, c[15:0]};
        end
        default:
        begin
            sign_ext_a = a;
            sign_ext_b = b;
            sign_ext_c = c;
        end
    endcase
end

endmodule