relu works

Author: WilliamZhang20

src/PE.v

@@ -4,9 +4,7 @@ module PE #(
     input wire clk,
     input wire rst,
     input wire clear,
-    input wire transpose_en,  // Enables fused transpose
-    input wire relu_en,       // Enables fused ReLU
-
+    input wire relu_en,
     input wire signed [WIDTH-1:0] a_in,
     input wire signed [WIDTH-1:0] b_in,
 
@@ -16,8 +14,6 @@ module PE #(
     output reg signed [WIDTH-1:0] c_out
 );
 
-    reg signed [WIDTH-1:0] product;
-
     always @(posedge clk or posedge rst) begin
         if (rst) begin
             a_out <= 0;
@@ -28,18 +24,10 @@ module PE #(
             b_out <= 0;
             c_out <= 0;
         end else begin
-            // Fused transpose
-            a_out <= transpose_en ? b_in : a_in;
-            b_out <= transpose_en ? a_in : b_in;
-
-            // Compute product
-            product = a_in * b_in;
+            a_out <= a_in;
+            b_out <= b_in;
 
-            // Apply ReLU if enabled
-            if (relu_en && product < 0)
-                c_out <= 0;
-            else
-                c_out <= c_out + product;
+            c_out <= c_out + (a_in * b_in);
         end
     end
 

src/systolic_array_2x2.v

@@ -22,7 +22,7 @@ module systolic_array_2x2 #(
     // Internal signals between PEs
     wire [WIDTH-1:0] a_wire [0:1][0:2];
     wire [WIDTH-1:0] b_wire [0:2][0:1];
-    wire [7:0] c_array [0:1][0:1];
+    wire signed [WIDTH-1:0] c_array [0:1][0:1];
 
     // Input loading at top-left
     assign a_wire[0][0] = a_data0;
@@ -34,11 +34,10 @@ module systolic_array_2x2 #(
     generate
         for (i = 0; i < 2; i = i + 1) begin : row
             for (j = 0; j < 2; j = j + 1) begin : col
-                PE #(.WIDTH(8)) pe_inst (
+                PE #(.WIDTH(WIDTH)) pe_inst (
                     .clk(clk),
                     .rst(rst),
                     .clear(clear),
-                    .transpose_en(transpose),
                     .relu_en(activation),
                     .a_in(a_wire[i][j]),
                     .b_in(b_wire[i][j]),
@@ -50,8 +49,10 @@ module systolic_array_2x2 #(
         end
     endgenerate
 
-    assign c00 = c_array[0][0];
-    assign c01 = c_array[0][1];
-    assign c10 = c_array[1][0];
-    assign c11 = c_array[1][1];
-endmodule
+    // Combinational logic for output with optional ReLU
+    assign c00 = activation ? (c_array[0][0] < 0 ? 0 : c_array[0][0]) : c_array[0][0];
+    assign c01 = activation ? (c_array[0][1] < 0 ? 0 : c_array[0][1]) : c_array[0][1];
+    assign c10 = activation ? (c_array[1][0] < 0 ? 0 : c_array[1][0]) : c_array[1][0];
+    assign c11 = activation ? (c_array[1][1] < 0 ? 0 : c_array[1][1]) : c_array[1][1];
+
+endmodule

test/test.py

@@ -29,16 +29,15 @@ async def load_matrix(dut, matrix, sel):
         dut.uio_in.value = 0
         await RisingEdge(dut.clk)
 
-async def read_signed_output(dut, transpose=False, relu=False):
+async def read_signed_output(dut, transpose=0, relu=0):
     # Apply instruction signal just before reading
-    dut.uio_in.value = 0b00000001 | (transpose << 1) | (relu << 2)
-    await ClockCycles(dut.clk, 1)
-    dut.uio_in.value = 0
-    await ClockCycles(dut.clk, 2)  # allow systolic array to compute
+    for i in range(3):
+        dut.uio_in.value = (transpose << 1) | (relu << 2)
+        await ClockCycles(dut.clk, 1)
 
     results = []
     for i in range(4):
-        dut.uio_in.value = 0  # Read mode
+        dut.uio_in.value = (transpose << 1) | (relu << 2)
         await ClockCycles(dut.clk, 1)
         val_unsigned = dut.uo_out.value.integer
         val_signed = val_unsigned if val_unsigned < 128 else val_unsigned - 256
@@ -61,14 +60,14 @@ async def test_relu_transpose(dut):
     dut.rst_n.value = 1
     await ClockCycles(dut.clk, 5)
 
-    A = [1, -2, 3, -4]   # row-major
-    B = [5, 6, -7, 8]    # row-major
+    A = [5, -6, 7, 8]  # row-major
+    B = [8, 9, 6, 8]  # row-major: [B00, B01, B10, B11]
 
     await load_matrix(dut, A, sel=0)
     await load_matrix(dut, B, sel=1)
 
-    expected = get_expected_matmul(A, B, transpose=True, relu=True)
-    results = await read_signed_output(dut)
+    expected = get_expected_matmul(A, B, transpose=False, relu=True)
+    results = await read_signed_output(dut, relu=1)
 
     for i in range(4):
         assert results[i] == expected[i], f"C[{i//2}][{i%2}] = {results[i]} != expected {expected[i]}"