You are an expert Verilog hardware designer working on a complete FPGA or ASIC project. You are given the full source code of a Verilog project. One of the modules in the project is incomplete — its implementation is missing and marked with:

// >>> Module Implementation Begin

// <<< Module Implementation End

Your task is to complete the missing implementation so that the entire project functions correctly. You must infer what the missing module or logic is supposed to do by analyzing the rest of the codebase: module names, signal connections, instantiations, comments, and usage patterns.

Follow these rules:

1. Only generate Verilog code to go between the `// >>> Module Implementation Begin` and `// <<< Module Implementation End` comments.
2. Use Verilog-2001 syntax. Ensure the code is synthesizable.
3. Assume the goal is a working, integrated hardware design — your code must make the whole system function as intended.
4. Analyze the naming, I/O ports, and how the module is instantiated or used in other parts of the project.
5. Implement correct behavior based on design context. Do not guess randomly — reason it out.
6. Avoid vendor-specific features unless explicitly required.
7. Do not output comments, explanation, or anything outside the required code region unless explicitly asked.

Your output should contain **only the Verilog code** that fits inside the marked region.

/*
 * Copyright (c) 2024 Your Name
 * SPDX-License-Identifier: Apache-2.0
 */

`default_nettype none

module tt_um_revenantx86_tinytpu 
(
    input  wire [7:0] ui_in,    // Dedicated inputs
    output wire [7:0] uo_out,   // Dedicated outputs
    input  wire [7:0] uio_in,   // IOs: Input path
    output wire [7:0] uio_out,  // IOs: Output path
    output wire [7:0] uio_oe,   // IOs: Enable path (active high: 0=input, 1=output)
    input  wire       ena,      // always 1 when the design is powered, so you can ignore it
    input  wire       clk,      // clock
    input  wire       rst_n     // reset_n - low to reset
);
// >>> Module Implementation Begin

// <<< Module Implementation End
endmodule




module dff_mem #
(
    parameter D_W = 8, 
    parameter WORD = 8
)
(
    input wire clk,
    input wire rst,
    // Read Logic
    input wire rd_en,
    input wire [$clog2(WORD)-1:0] addr,
    // Write logic
    input wire wr_en,
    // IO
    input wire [D_W-1:0] data_in,
    output reg [D_W-1:0] data_out
);
//
// Memory Array
reg [D_W-1:0] mem [WORD-1:0];
//
// Write Logic
always @(posedge clk) begin
    if(wr_en) begin
        mem[addr] <= data_in;
    end
end

// Read Logic
always @(*) begin
    if(rd_en) begin
        data_out <= mem[addr];  // Update output only if rd_en is high
    end
    else if(!rd_en) begin
        data_out <= 0;  // Update output only if rd_en is high
    end 
    else if(rst) begin
        data_out <= 0;  // Optionally reset the output to 0 on reset
    end
end


endmodule

module input_control # 
(
    parameter D_W = 8,
    parameter N = 2,
    parameter WORD = 8
)
(
    input wire clk,
    input wire rst,
    input wire data_in_x,
    input wire data_in_y,
    input wire load_en,
    input wire init,
    output reg [(N*D_W)-1:0] out_x_flat,
    output reg [(N*D_W)-1:0] out_y_flat,
    output reg out_init
);
//
/*
    State Machine
*/
//enum {IDLE,LOAD,TRANSFER} STATE;
localparam IDLE     = 2'b00;
localparam LOAD     = 2'b01;
localparam TRANSFER = 2'b10;
reg [1:0] STATE;
//
wire [D_W-1:0] out_x [N-1:0];
wire [D_W-1:0] out_y [N-1:0];
/*
    Generate Memory Block
*/
reg  [$clog2(WORD)-1:0] addr_x_ram     [N-1:0];
reg                     rd_en_x_ram    [N-1:0];
reg                     wr_en_x_ram    [N-1:0];
reg  [D_W-1:0]          data_in_x_reg;
//
reg  [$clog2(WORD)-1:0] addr_y_ram     [N-1:0];
reg                     rd_en_y_ram    [N-1:0];
reg                     wr_en_y_ram    [N-1:0];
reg  [D_W-1:0]          data_in_y_reg;
//
genvar n;
generate 
    for(n=0;n<N;n=n+1) begin
        //Initalize memory X & Y
        dff_mem #( .D_W(D_W), .WORD(WORD)) 
                mem_x (
                        .clk(clk),
                        .rst(rst),
                        .addr(addr_x_ram[n]),
                        .rd_en(rd_en_x_ram[n]),
                        .wr_en(wr_en_x_ram[n]),
                        .data_in(data_in_x_reg),
                        .data_out(out_x[n])
                    );
        dff_mem #( .D_W(D_W), .WORD(WORD)) 
                mem_y (
                        .clk(clk),
                        .rst(rst),
                        .addr(addr_y_ram[n]),
                        .rd_en(rd_en_y_ram[n]),
                        .wr_en(wr_en_y_ram[n]),
                        .data_in(data_in_y_reg),
                        .data_out(out_y[n])
                    );
    end
endgenerate

always @(posedge clk) begin
    if(rst) begin
        STATE <= IDLE;
        out_init <= 0; 
    end
    else begin
        case(STATE) 
        IDLE: begin
            out_init <= 0; 
            if(load_en) 
                STATE <= LOAD;
            else if(init)begin
                out_init <= 1; 
                STATE <= TRANSFER;
            end
        end
        LOAD: begin
            if(!load_en)
                STATE <= IDLE;
        end
        TRANSFER: begin
            out_init <= 0; 
            if(rd_en_x_ram[N-1]) begin
                STATE <= IDLE;
            end
        end
        endcase
    end
end


reg [$clog2(D_W)-1:0]       bit_counter;
reg [$clog2(N)-1:0]         addr_counter;
reg [$clog2(N)-1:0]         ram_counter;

integer i;
always @(posedge clk) begin
    if(rst) begin
        //
        bit_counter   <= 0;
        addr_counter  <= 0;
        ram_counter   <= 0;
        data_in_x_reg <= 0;
        data_in_y_reg <= 0;
        //
        wr_en_x_ram[0] <= 0;
        rd_en_x_ram[0] <= 0;
        wr_en_y_ram[0] <= 0;
        rd_en_y_ram[0] <= 0;
        for(i=0;i<N;i=i+1) begin
            // Reset X
            addr_x_ram[i]  <= 0;
            // Reset Y
            addr_y_ram[i]  <= 0;                
        end
        //
    end
    else begin
        case(STATE)
            IDLE : begin
                addr_counter <= 0;
                // X Row
                wr_en_x_ram[N-1] <= 0;
                // Y Column
                wr_en_y_ram[N-1] <= 0;
            end
            LOAD: begin
                //
                bit_counter <= bit_counter + 1;
                // Load data for X row
                wr_en_x_ram[ram_counter] <= 0;
                data_in_x_reg <= {data_in_x,data_in_x_reg[D_W-1:1]};
                //
                // Load data for Y Column
                wr_en_y_ram[ram_counter] <= 0;
                data_in_y_reg <= {data_in_y,data_in_y_reg[D_W-1:1]};
                //
                if(bit_counter == D_W-1) begin
                    // X Row
                    wr_en_x_ram[ram_counter] <= 1;
                    // Y Column
                    wr_en_y_ram[ram_counter] <= 1;
                end
                //
                if(wr_en_x_ram[ram_counter] == 1) begin
                    // X Row
                    addr_x_ram[ram_counter]  <= addr_x_ram[ram_counter] + 1;
                    // X Row
                    addr_y_ram[ram_counter]  <= addr_y_ram[ram_counter] + 1;
                end
                if(addr_x_ram[ram_counter] == N-1 && wr_en_x_ram[ram_counter]) begin
                    // X Row & Y Row common ram word block counter
                    ram_counter <= ram_counter + 1;
                end
            end
            TRANSFER : begin
                // X & Y Common 
                wr_en_x_ram[ram_counter] <= 0;
                //
                if(rd_en_x_ram[N-1]) begin
                    rd_en_x_ram[0] <= 0;
                    rd_en_y_ram[0] <= 0;
                end
                else begin // Initalize transfer
                    rd_en_x_ram[0] <= 1;
                    rd_en_y_ram[0] <= 1;
                end
                addr_counter <= addr_counter + 1;
                // Assign updated addres X & Y
                addr_x_ram[0] <= addr_counter;
                addr_y_ram[0] <= addr_counter;
            end
        endcase
    end
end

// Assign reading
genvar x;
for (x=1;x<N;x=x+1)
begin
    always@(posedge clk)
    begin
        if (rst) begin
            addr_x_ram[x]  <= 0; 
            rd_en_x_ram[x] <= 0;
            // Assign for Y column 
            addr_y_ram[x] <= 0;
            rd_en_y_ram[x] <= 0;
        end
        else if(STATE == TRANSFER || STATE == IDLE) begin
            // Assign for X row
            addr_x_ram[x]  <= addr_x_ram[x-1]; 
            rd_en_x_ram[x] <= rd_en_x_ram[x-1];
            // Assign for Y column 
            addr_y_ram[x] <= addr_y_ram[x-1];
            rd_en_y_ram[x] <= rd_en_y_ram[x-1];
        end
        //else if (STATE == IDLE) begin
        //    rd_en_x_ram[N-1] <= 0;
        //    rd_en_y_ram[N-1] <= 0;
        //end
    end
end

// Flatten output signal
integer r;
// Packing the 2D array 'z' into the flat wire 'data_core_z_flat'
always @(*) begin
    for (r = 0; r < N; r = r + 1) begin
        // Compute the starting index of each signal in the flat wire
        out_x_flat[(r+1)*D_W-1 -: D_W] = out_x[r];
        out_y_flat[(r+1)*D_W-1 -: D_W] = out_y[r];
    end
end
endmodule

module mac #
(
    parameter D_W = 16
)
(   
    input  wire               clk,         // Clock
    input  wire               rst,         // reset
    input  wire               init,        // Initalize first op
    input  wire [D_W-1:0]     in_x,        // Input from X axis
    input  wire [D_W-1:0]     in_y,        // Input from Y axis
    output reg  [D_W-1:0]     out_x,       // x axis input passthrough
    output reg  [D_W-1:0]     out_y,       // y axis input passthrough
    output reg  [(2*D_W-1):0] out_z,       // Output for the sum
    output reg                out_init     // Init passthrough
);
//
always @(posedge clk) begin
    //
    out_x <= in_x;
    out_y <= in_y;
    out_init <= init;
    //
    if(rst) begin                       // Reset
        out_z <= 0;                     
        out_y <= 0;
        out_x <= 0;
        out_init <= 0;
    end
    //
    else if(init)
        out_z <= in_x * in_y;           // Multiply & Override 
    //
    else
        out_z <= out_z + (in_x* in_y);  // Multiply & Accumulate
    //
end
//
endmodule

module output_control #
(
    parameter D_W = 8,
    parameter N = 2
)
(
    // Global Signal
    input  wire                   clk,
    input  wire                   rst,
    // Systolic Output
    input  wire [(N*N*2*D_W)-1:0] core_out_z,
    input  wire                   init, 
    // IO Output
    output reg                    data_out_z,
    output reg                    tx_ready
);
/*
    STATE Machine
*/
//
localparam IDLE = 1'b0;
localparam TX   = 1'b1;
reg        STATE;
//enum {IDLE,TX} STATE;
//
// Counters
reg [$clog2(2*D_W)-1:0] bit_counter;
reg [$clog2(N)-1:0]     column_counter;
reg [$clog2(N)-1:0]     row_counter;
//
wire [2*D_W-1:0]        data_tx_reg;
//
//
assign data_tx_reg = core_out_z[(row_counter*N + column_counter)*2*D_W +: 2*D_W];
//
reg init_delay [N-1:0];
//
always @(posedge clk) begin
    if(rst) begin
        STATE <= IDLE;
    end
    else begin
        init_delay[0] <= init;
        case(STATE) 
            IDLE: begin
                if(init_delay[N-1])
                    STATE <= TX;
            end
            TX: begin
                if(row_counter == N-1 && column_counter == N-1 && bit_counter == 2*D_W-1) 
                    STATE <= IDLE;
            end
        endcase
    end
end

integer x,r,c;
always @(posedge clk) begin
    //
    if(rst) begin
        //
        bit_counter <= 0;
        column_counter <= 0;
        row_counter <= 0;
        tx_ready <= 0;
        //
    end
    //
    else begin
        case(STATE) 
            IDLE: begin
                //
                bit_counter     <= 0;
                column_counter  <= 0;
                row_counter     <= 0;
                tx_ready        <= 0;
                data_out_z      <= 0;
                //        
            end
            TX: begin
                //
                tx_ready    <= 1;
                bit_counter <= bit_counter + 1;
                data_out_z  <=  data_tx_reg[bit_counter];
                //
                if(bit_counter == (2*D_W-1)) begin
                    column_counter <= column_counter + 1;
                    bit_counter    <= 0;
                end
                if(column_counter == N-1 && bit_counter == (2*D_W-1)) begin
                    column_counter <= 0;
                    row_counter   <= row_counter + 1;
                end
                //
            end
        endcase
    end
end

// Assign delay counter
genvar i;
for (i=1;i<N;i=i+1) begin
    always @(posedge clk) begin
        init_delay[i] <= init_delay[i-1];
    end
end

endmodule

module systolic #
(
    parameter D_W = 8,
    parameter N = 2
)
(
    input   wire             clk,                   // Clock Input
    input   wire             rst,                   // Reset cores & Regs
    input   wire             init,                  // Init process
    //input   wire [D_W-1:0]   x        [N-1:0],       // Row input
    //input   wire [D_W-1:0]   y        [N-1:0],       // Column input
    //output  reg  [2*D_W-1:0] z        [N-1:0][N-1:0] // MAC unit indv outputs
    input wire [(N*D_W)-1:0] x_flat,  // Flat wire to hold the packed array
    input wire [(N*D_W)-1:0] y_flat,  // Flat wire to hold the packed array
    output reg [(N*N*2*D_W)-1:0] z_flat // Flat wire to hold the packed array
);
//
// Initlize wires
//
wire [D_W-1:0]   in_x     [N-1:0][N-1:0];
wire [D_W-1:0]   in_y     [N-1:0][N-1:0];
wire [D_W-1:0]   out_x    [N-1:0][N-1:0];
wire [D_W-1:0]   out_y    [N-1:0][N-1:0];
wire [2*D_W-1:0] out_z    [N-1:0][N-1:0];
wire             init_in  [N-1:0][N-1:0];
wire             init_out [N-1:0][N-1:0];
// wire assign to output
//assign z = out_z;
//
// flatten-unflatten
reg  [D_W-1:0]   x       [N-1:0];        // Row input
reg  [D_W-1:0]   y       [N-1:0];        // Column input
reg  [2*D_W-1:0] z       [N-1:0][N-1:0]; // MAC unit indv outputs
//

integer r, c;
// Assuming data_core_z_flat is already declared
always @(*) begin
    for (r = 0; r < N; r = r + 1) begin
        x[r] = x_flat[(r+1)*D_W-1 -: D_W];
        y[r] = y_flat[(r+1)*D_W-1 -: D_W];
        for (c = 0; c < N; c = c + 1) begin
            // Compute the starting index of each signal in the flat wire
            z_flat[((r * N + c) * (2 * D_W)) +: (2 * D_W)] = out_z[r][c];
        end
    end
end
//
//initalize counter
//
// Initalize MAC module COL >> ROW
assign init_in[0][0] = init;
//
// Loop Generate MAC matrix
genvar i,j,k;
generate 
    //
    for(i=0; i<N; i=i+1) begin
        //
        // Assign initialization wire to first row
        if(i==0) begin
            for(k=0; k<N-1;k++) begin
                assign init_in[0][k+1] = init_out[0][k];
            end 
        end
        //
        for(j=0; j<N; j = j + 1) begin
            //
            // Assign wires to mac modules
            for(k=0;k<N-1;k=k+1) begin
                assign in_x[i][k+1] = out_x[i][k];
                assign in_y[k+1][j] = out_y[k][j];
            end
            //
            // Asssign initalization signal to columns
            if(i==0) begin
                for (k=0;k<N-1;k++) begin
                    assign init_in[k+1][j]=init_out[k][j];
                end
            end
            //
            // Input Columns & Row Assignment
            if(j==0) assign in_x[i][0] = x[i];
            if(i==0) assign in_y[0][j] = y[j];
            //
            mac #(.D_W(D_W)) mac_inst 
                (
                    .clk      (clk),
                    .rst      (rst),
                    .init     (init_in[i][j]),
                    .out_init (init_out[i][j]),
                    .in_x     (in_x[i][j]),
                    .in_y     (in_y[i][j]),
                    .out_x    (out_x[i][j]),
                    .out_y    (out_y[i][j]),
                    .out_z    (out_z[i][j])
                );
            //
        end
    end
    //
endgenerate
//
endmodule

module tinytpu_top #
(
    parameter D_W = 8,
    parameter N = 2,
    parameter WORD = 2
)
(
    input  wire clk,
    input  wire rst,
    input  wire data_in_x,
    input  wire data_in_y,
    input  wire load_en,
    input  wire init,
    output wire data_out_z,
    output reg  tx_ready
);
//
// REG & WIRES
wire [(N*D_W)-1:0]     out_x_flat;
wire [(N*D_W)-1:0]     out_y_flat;
wire [(N*N*2*D_W)-1:0] out_z_flat;
wire out_init;
reg tx_ready_reg;

always @(posedge clk) begin
    tx_ready <= tx_ready_reg;
end 
//wire  [D_W-1:0]   data_core_x  [N-1:0];
//wire  [D_W-1:0]   data_core_y  [N-1:0];
//wire  [2*D_W-1:0] data_core_z   [N-1:0][N-1:0];
//
/*
    Input Control Module
*/
input_control #(.D_W(D_W), .N(N), .WORD(WORD)) 
        input_control_inst (
                        .clk(clk),
                        .rst(rst),
                        .data_in_x(data_in_x),
                        .data_in_y(data_in_y),
                        .load_en(load_en),
                        .init(init),
                        .out_x_flat(out_x_flat),
                        .out_y_flat(out_y_flat),
                        .out_init(out_init)
                     );
/*
    Output Control Module
*/
output_control #(.D_W(D_W), .N(N))
        output_control_inst (
                                .clk(clk),
                                .rst(rst),
                                .core_out_z(out_z_flat),
                                .init(init),
                                .data_out_z(data_out_z),
                                .tx_ready(tx_ready_reg)
                            );
/*
    Systolic Core
*/
systolic #(.D_W(D_W), .N(N)) 
        systolic_inst (
                        .clk(clk),
                        .rst(rst),
                        .init(out_init),
                        .x_flat(out_x_flat),
                        .y_flat(out_y_flat),
                        .z_flat(out_z_flat)
                      );
//
//
//
endmodule

Reference solution:
rtl/generated.v:   //
  wire rst;
  wire data_in_x;
  wire data_in_y;
  wire load_en;
  wire init;
  wire data_out_z;
  wire tx_ready;
  //
  assign rst = ~rst_n;
  //
  /*
    Assign IO
  */
  // Input Assign
  assign data_in_x = ui_in[0];
  assign data_in_y = ui_in[1];
  assign load_en   = ui_in[2];
  assign init      = ui_in[3];
  // Output Assign
  assign uo_out[0] = data_out_z;
  assign uo_out[1] = tx_ready;
  assign uo_out[2] = 0;
  assign uo_out[3] = 0;
  assign uo_out[4] = 0;
  assign uo_out[5] = 0;
  assign uo_out[6] = 0;
  assign uo_out[7] = 0;
  //
  assign uio_oe = 8'b00;
  assign uio_out = 8'b0;
  //
  tinytpu_top #(.D_W(8), .N(2), .WORD(2)) 
      tinytpu_top_inst (
                          .clk(clk),
                          .rst(rst),
                          .data_in_x(data_in_x),
                          .data_in_y(data_in_y),
                          .load_en(load_en),
                          .init(init),
                          .data_out_z(data_out_z),
                          .tx_ready(tx_ready)
                       );