From dd9b470b87a44c869b937ca41fed9b915e61bb67 Mon Sep 17 00:00:00 2001
From: gamzeisl <gamzeisl@hotmail.com>
Date: Thu, 10 Oct 2024 14:07:19 +0200
Subject: [PATCH] [fix] Correct pipeline stages in activation unit

---
 src/ita.sv            | 22 +++++++++++-----------
 src/ita_activation.sv | 10 +++-------
 2 files changed, 14 insertions(+), 18 deletions(-)

diff --git a/src/ita.sv b/src/ita.sv
index f164ac9..bd03c42 100644
--- a/src/ita.sv
+++ b/src/ita.sv
@@ -288,19 +288,19 @@ module ita
   );
 
   ita_activation i_activation (
-    .clk_i         (clk_i       ),
-    .rst_ni        (rst_ni      ),
-    .activation_i  (activation_q10),
-    .calc_en_i     (calc_en_q6 && last_inner_tile_q6  ),
-    .calc_en_q_i  (calc_en_q7 && last_inner_tile_q7  ),
-    .b_i           (ctrl_i.gelu_b  ),
-    .c_i           (ctrl_i.gelu_c  ),
+    .clk_i           (clk_i        ),
+    .rst_ni          (rst_ni       ),
+    .activation_i    (activation_q9),
+    .calc_en_i       (calc_en_q6 && last_inner_tile_q6),
+    .calc_en_q_i     (calc_en_q7 && last_inner_tile_q7),
+    .b_i             (ctrl_i.gelu_b),
+    .c_i             (ctrl_i.gelu_c),
     .requant_mode_i  (activation_requant_mode),
-    .requant_mult_i    (activation_requant_mult),
+    .requant_mult_i  (activation_requant_mult),
     .requant_shift_i (activation_requant_shift),
-    .requant_add_i         (activation_requant_add),
-    .data_i        (requant_oup),
-    .data_o        (post_activation)
+    .requant_add_i   (activation_requant_add),
+    .data_i          (requant_oup),
+    .data_o          (post_activation)
   );
 
   ita_fifo_controller i_fifo_controller (
diff --git a/src/ita_activation.sv b/src/ita_activation.sv
index 95813f4..0f54070 100644
--- a/src/ita_activation.sv
+++ b/src/ita_activation.sv
@@ -23,7 +23,7 @@ module ita_activation
   requant_oup_t data_q1, data_q2, data_q3, data_q4;
   activation_e activation_q1, activation_q2;
   oup_t gelu_out, requant_in;
-  requant_oup_t relu_out_d, relu_out_q1, relu_out_q2, requant_out;
+  requant_oup_t relu_out, requant_out;
   logic calc_en_q2, calc_en_q3;
 
   ita_requantizer i_requantizer (
@@ -43,7 +43,7 @@ module ita_activation
     for (genvar i = 0; i < N; i++) begin: relu_instances
       ita_relu i_relu (
         .data_i(data_q2[i]),
-        .data_o(relu_out_d[i])
+        .data_o(relu_out[i])
       );
     end
   endgenerate
@@ -70,7 +70,7 @@ module ita_activation
       end
       Relu: begin
         for (int i = 0; i < N; i++) begin
-          requant_in[i] = {{(WO-WI){relu_out_q2[i][WI-1]}}, relu_out_q2[i]};
+          requant_in[i] = {{(WO-WI){relu_out[i][WI-1]}}, relu_out[i]};
         end
       end
       default: begin
@@ -101,8 +101,6 @@ module ita_activation
       data_q4 <= '0;
       calc_en_q2 <= 0;
       calc_en_q3 <= 0;
-      relu_out_q1 <= '0;
-      relu_out_q2 <= '0;
     end else begin
       activation_q1 <= activation_i;
       activation_q2 <= activation_q1;
@@ -112,8 +110,6 @@ module ita_activation
       data_q4 <= data_q3;
       calc_en_q2 <= calc_en_q_i;
       calc_en_q3 <= calc_en_q2;
-      relu_out_q1 <= relu_out_d;
-      relu_out_q2 <= relu_out_q1;
     end
   end
 endmodule
\ No newline at end of file