Merge pull request #3 from Richard-Qin-001/fix/wb-stage-review-9938572677202897816

Refactor WB Stage for Timing, and Precise Exceptions

2 files changed, 25 insertions, 17 deletions

diff --git a/rv64cpu/rtl/core/wb_stage.sv b/rv64cpu/rtl/core/wb_stage.sv
index 13dd417..26c9c7b 100644
--- a/rv64cpu/rtl/core/wb_stage.sv
+++ b/rv64cpu/rtl/core/wb_stage.sv
@@ -25,6 +25,7 @@ module wb_stage
     input  logic [XLEN-1:0]        mem_wb_alu_result,
     input  logic [XLEN-1:0]        mem_wb_mem_rdata,
     input  logic [XLEN-1:0]        mem_wb_pc,
+    input  logic [XLEN-1:0]        mem_wb_pc_plus_4,
     
     input  logic                   mem_wb_reg_write_en,
     input  logic [4:0]             mem_wb_rd_addr,
@@ -46,8 +47,8 @@ module wb_stage
     input  logic [XLEN-1:0]        csr_rdata,         // Data read from CSR unit
 
     // Outputs to Register File
-    output logic [XLEN-1:0]        wb_reg_write_data,
-    output logic [4:0]             wb_reg_write_addr,
+    output logic [XLEN-1:0]        wb_rd_wdata,
+    output logic [4:0]             wb_rd_addr,
     output logic                   wb_reg_write_en,
 
     // Outputs to CSR Unit
@@ -66,9 +67,15 @@ module wb_stage
 );
 
     // --------------------------
+    // Exception/Trap Detection
+    // --------------------------
+    logic has_trap;
+    assign has_trap = mem_wb_illegal_instr | mem_wb_is_ecall | mem_wb_is_ebreak;
+
+    // --------------------------
     // CSR Signal Routing
     // --------------------------
-    assign csr_req  = mem_wb_is_csr;
+    assign csr_req  = mem_wb_is_csr & ~has_trap;
     assign csr_op   = mem_wb_csr_op;
     assign csr_addr = mem_wb_imm[11:0];
     
@@ -83,13 +90,13 @@ module wb_stage
     always_comb begin
         if (mem_wb_is_csr) begin
             // CSR reads write the previous CSR value to the destination register
-            wb_reg_write_data = csr_rdata;
+            wb_rd_wdata = csr_rdata;
         end else begin
             case (mem_wb_wb_sel)
-                2'b00: wb_reg_write_data = mem_wb_alu_result;
-                2'b01: wb_reg_write_data = mem_wb_mem_rdata;
-                2'b10: wb_reg_write_data = mem_wb_pc + 64'd4;
-                default: wb_reg_write_data = '0;
+                2'b00: wb_rd_wdata = mem_wb_alu_result;
+                2'b01: wb_rd_wdata = mem_wb_mem_rdata;
+                2'b10: wb_rd_wdata = mem_wb_pc_plus_4;
+                default: wb_rd_wdata = '0;
             endcase
         end
     end
@@ -98,10 +105,8 @@ module wb_stage
     // Register File Write Control
     // --------------------------
     // We only write back if the instruction didn't trap.
-    // For now, assume pure write-enable propagation.
-    // A future trap controller might kill wb_reg_write_en if taking a trap.
-    assign wb_reg_write_addr  = mem_wb_rd_addr;
-    assign wb_reg_write_en    = mem_wb_reg_write_en;
+    assign wb_rd_addr      = mem_wb_rd_addr;
+    assign wb_reg_write_en = mem_wb_reg_write_en & ~has_trap;
 
     // --------------------------
     // Trap / Exception Signals
diff --git a/rv64cpu/tb/unit/tb_wb_stage.cc b/rv64cpu/tb/unit/tb_wb_stage.cc
index 156d8e3..715a47d 100644
--- a/rv64cpu/tb/unit/tb_wb_stage.cc
+++ b/rv64cpu/tb/unit/tb_wb_stage.cc
@@ -37,6 +37,7 @@ int main(int argc, char** argv) {
     dut->mem_wb_alu_result = 0xAA55AA55;
     dut->mem_wb_mem_rdata = 0x11223344;
     dut->mem_wb_pc = 0x80000000;
+    dut->mem_wb_pc_plus_4 = 0x80000004;
     dut->mem_wb_reg_write_en = 1;
     dut->mem_wb_rd_addr = 5;
     dut->mem_wb_wb_sel = 0; // ALU
@@ -54,20 +55,20 @@ int main(int argc, char** argv) {
 
     tick(dut);
 
-    check(dut->wb_reg_write_data == 0xAA55AA55, "WB_Sel_ALU");
-    check(dut->wb_reg_write_addr == 5, "WB_RdAddr");
+    check(dut->wb_rd_wdata == 0xAA55AA55, "WB_Sel_ALU");
+    check(dut->wb_rd_addr == 5, "WB_RdAddr");
     check(dut->wb_reg_write_en == 1, "WB_RegWriteEn");
     check(dut->csr_req == 0, "CSR_Req_0");
 
     // 2. Memory read writeback
     dut->mem_wb_wb_sel = 1; // MEM
     tick(dut);
-    check(dut->wb_reg_write_data == 0x11223344, "WB_Sel_MEM");
+    check(dut->wb_rd_wdata == 0x11223344, "WB_Sel_MEM");
 
     // 3. PC+4 writeback (JAL/JALR)
     dut->mem_wb_wb_sel = 2; // PC+4
     tick(dut);
-    check(dut->wb_reg_write_data == 0x80000004, "WB_Sel_PC_Plus_4");
+    check(dut->wb_rd_wdata == 0x80000004, "WB_Sel_PC_Plus_4");
 
     // 4. CSR Read/Write (e.g. CSRRS)
     dut->mem_wb_is_csr = 1;
@@ -80,7 +81,7 @@ int main(int argc, char** argv) {
 
     tick(dut);
 
-    check(dut->wb_reg_write_data == 0x00001800, "CSR_Read_To_Reg");
+    check(dut->wb_rd_wdata == 0x00001800, "CSR_Read_To_Reg");
     check(dut->csr_req == 1, "CSR_Req_1");
     check(dut->csr_op == 2, "CSR_Op_Propagated");
     check(dut->csr_addr == 0x300, "CSR_Addr_Extracted");
@@ -105,6 +106,8 @@ int main(int argc, char** argv) {
     check(dut->trap_is_mret == 1, "Trap_MRET");
     check(dut->trap_is_sret == 1, "Trap_SRET");
     check(dut->trap_epc == 0x80000000, "Trap_EPC");
+    check(dut->wb_reg_write_en == 0, "Trap_Gates_RegWriteEn");
+    check(dut->csr_req == 0, "Trap_Gates_CSR_Req");
 
     std::cout << "----------------------------------" << std::endl;
     std::cout << "\033[32mAll " << pass_count << " tests passed!\033[0m" << std::endl;