Write Back stage unit and tests

3 files changed, 240 insertions, 0 deletions

diff --git a/rv64cpu/rtl/core/wb_stage.sv b/rv64cpu/rtl/core/wb_stage.sv
new file mode 100644
index 0000000..13dd417
--- /dev/null
+++ b/rv64cpu/rtl/core/wb_stage.sv
@@ -0,0 +1,116 @@
+//     Richard CPU
+//     Copyright (C) 2026  Richard Qin
+
+//     This program is free software: you can redistribute it and/or modify
+//     it under the terms of the GNU General Public License as published by
+//     the Free Software Foundation, either version 3 of the License, or
+//     (at your option) any later version.
+
+//     This program is distributed in the hope that it will be useful,
+//     but WITHOUT ANY WARRANTY; without even the implied warranty of
+//     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+//     GNU General Public License for more details.
+
+//     You should have received a copy of the GNU General Public License
+//     along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+// File path: rtl/core/wb_stage.sv
+// Description: Write Back Stage
+// Responsible for writing back to the general-purpose register file after multiplexing the final data (ALU result, memory read data, or PC+4)
+
+module wb_stage
+    import rv64_pkg::*;
+(
+    // Inputs from MEM/WB pipeline register
+    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                   mem_wb_reg_write_en,
+    input  logic [4:0]             mem_wb_rd_addr,
+    input  logic [1:0]             mem_wb_wb_sel,
+
+    input  logic                   mem_wb_is_csr,
+    input  logic [2:0]             mem_wb_csr_op,
+    input  logic [XLEN-1:0]        mem_wb_rs1_rdata,
+    input  logic [4:0]             mem_wb_rs1_addr,
+    input  logic [XLEN-1:0]        mem_wb_imm,
+    
+    input  logic                   mem_wb_illegal_instr,
+    input  logic                   mem_wb_is_ecall,
+    input  logic                   mem_wb_is_ebreak,
+    input  logic                   mem_wb_is_mret,
+    input  logic                   mem_wb_is_sret,
+
+    // Inputs from CSR Unit
+    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                   wb_reg_write_en,
+
+    // Outputs to CSR Unit
+    output logic                   csr_req,           // CSR request valid
+    output logic [2:0]             csr_op,            // CSR operation type
+    output logic [11:0]            csr_addr,          // CSR address
+    output logic [XLEN-1:0]        csr_rs1_data,      // Source data (from rs1_rdata or zimm)
+
+    // Outputs to Exception Handler (Trap Controller)
+    output logic                   trap_illegal_instr,
+    output logic                   trap_is_ecall,
+    output logic                   trap_is_ebreak,
+    output logic                   trap_is_mret,
+    output logic                   trap_is_sret,
+    output logic [XLEN-1:0]        trap_epc           // PC of the instruction causing the trap
+);
+
+    // --------------------------
+    // CSR Signal Routing
+    // --------------------------
+    assign csr_req  = mem_wb_is_csr;
+    assign csr_op   = mem_wb_csr_op;
+    assign csr_addr = mem_wb_imm[11:0];
+    
+    // For immediate CSR instructions (CSRRWI, CSRRSI, CSRRCI), 
+    // the source data is the zero-extended rs1_addr field.
+    // mem_wb_csr_op[2] is 1 for immediate variants.
+    assign csr_rs1_data = (mem_wb_csr_op[2]) ? {59'b0, mem_wb_rs1_addr} : mem_wb_rs1_rdata;
+
+    // --------------------------
+    // Data Multiplexer for Register File
+    // --------------------------
+    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;
+        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;
+            endcase
+        end
+    end
+
+    // --------------------------
+    // 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;
+
+    // --------------------------
+    // Trap / Exception Signals
+    // --------------------------
+    assign trap_illegal_instr = mem_wb_illegal_instr;
+    assign trap_is_ecall      = mem_wb_is_ecall;
+    assign trap_is_ebreak     = mem_wb_is_ebreak;
+    assign trap_is_mret       = mem_wb_is_mret;
+    assign trap_is_sret       = mem_wb_is_sret;
+    assign trap_epc           = mem_wb_pc;
+
+endmodule
diff --git a/rv64cpu/tb/CMakeLists.txt b/rv64cpu/tb/CMakeLists.txt
index 7ca5bf0..ecccd02 100644
--- a/rv64cpu/tb/CMakeLists.txt
+++ b/rv64cpu/tb/CMakeLists.txt
@@ -141,6 +141,15 @@ add_verilator_target(
     INCLUDE_DIRS ${RTL_COMMON_INCLUDES}
     TB_SOURCES  ${CMAKE_CURRENT_SOURCE_DIR}/unit/tb_mem_wb_reg.cc
 )
+
+add_verilator_target(
+    NAME        tb_wb_stage
+    TOP_MODULE  wb_stage
+    SOURCES     ${RTL_CORE_DIR}/include/rv64_pkg.sv
+                ${RTL_CORE_DIR}/wb_stage.sv
+    INCLUDE_DIRS ${RTL_COMMON_INCLUDES}
+    TB_SOURCES  ${CMAKE_CURRENT_SOURCE_DIR}/unit/tb_wb_stage.cc
+)
 #
 # add_verilator_target(
 #     NAME        tb_csr
diff --git a/rv64cpu/tb/unit/tb_wb_stage.cc b/rv64cpu/tb/unit/tb_wb_stage.cc
new file mode 100644
index 0000000..156d8e3
--- /dev/null
+++ b/rv64cpu/tb/unit/tb_wb_stage.cc
@@ -0,0 +1,115 @@
+//     Richard CPU
+//     Copyright (C) 2026  Richard Qin
+
+// File path: tb/unit/tb_wb_stage.cc
+// Description: WB Stage unit test
+
+#include <iostream>
+#include <cassert>
+#include <verilated.h>
+#include "Vwb_stage.h"
+
+int pass_count = 0;
+
+void tick(Vwb_stage* dut) {
+    dut->eval();
+}
+
+void check(bool cond, const char* name) {
+    if (cond) {
+        std::cout << "\033[32m[PASS] " << name << "\033[0m" << std::endl;
+        pass_count++;
+    } else {
+        std::cout << "\033[31m[FAIL] " << name << "\033[0m" << std::endl;
+        assert(false);
+    }
+}
+
+int main(int argc, char** argv) {
+    Verilated::commandArgs(argc, argv);
+    Vwb_stage* dut = new Vwb_stage;
+
+    std::cout << "----------------------------------" << std::endl;
+    std::cout << "        WB Stage Test" << std::endl;
+    std::cout << "----------------------------------" << std::endl;
+
+    // 1. Basic ALU writeback
+    dut->mem_wb_alu_result = 0xAA55AA55;
+    dut->mem_wb_mem_rdata = 0x11223344;
+    dut->mem_wb_pc = 0x80000000;
+    dut->mem_wb_reg_write_en = 1;
+    dut->mem_wb_rd_addr = 5;
+    dut->mem_wb_wb_sel = 0; // ALU
+    dut->mem_wb_is_csr = 0;
+    dut->mem_wb_csr_op = 0;
+    dut->mem_wb_rs1_rdata = 0x99;
+    dut->mem_wb_rs1_addr = 1;
+    dut->mem_wb_imm = 0x123;
+    dut->mem_wb_illegal_instr = 0;
+    dut->mem_wb_is_ecall = 0;
+    dut->mem_wb_is_ebreak = 0;
+    dut->mem_wb_is_mret = 0;
+    dut->mem_wb_is_sret = 0;
+    dut->csr_rdata = 0x0;
+
+    tick(dut);
+
+    check(dut->wb_reg_write_data == 0xAA55AA55, "WB_Sel_ALU");
+    check(dut->wb_reg_write_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");
+
+    // 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");
+
+    // 4. CSR Read/Write (e.g. CSRRS)
+    dut->mem_wb_is_csr = 1;
+    dut->mem_wb_wb_sel = 0; // Doesn't matter because is_csr forces CSR path
+    dut->mem_wb_csr_op = 2; // CSRRS
+    dut->mem_wb_imm = 0x300; // mstatus
+    dut->mem_wb_rs1_rdata = 0x00000008; // Set MIE
+    dut->mem_wb_rs1_addr = 10;
+    dut->csr_rdata = 0x00001800; // Old mstatus
+
+    tick(dut);
+
+    check(dut->wb_reg_write_data == 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");
+    check(dut->csr_rs1_data == 0x00000008, "CSR_RS1_Data_Rdata");
+
+    // 5. CSR Immediate (e.g. CSRRSI)
+    dut->mem_wb_csr_op = 6; // CSRRSI (funct3 = 3'b110)
+    tick(dut);
+    check(dut->csr_rs1_data == 10, "CSR_RS1_Data_Zimm");
+
+    // 6. Exception Signals Propagation
+    dut->mem_wb_illegal_instr = 1;
+    dut->mem_wb_is_ecall = 1;
+    dut->mem_wb_is_ebreak = 1;
+    dut->mem_wb_is_mret = 1;
+    dut->mem_wb_is_sret = 1;
+    tick(dut);
+
+    check(dut->trap_illegal_instr == 1, "Trap_Illegal_Instr");
+    check(dut->trap_is_ecall == 1, "Trap_ECALL");
+    check(dut->trap_is_ebreak == 1, "Trap_EBREAK");
+    check(dut->trap_is_mret == 1, "Trap_MRET");
+    check(dut->trap_is_sret == 1, "Trap_SRET");
+    check(dut->trap_epc == 0x80000000, "Trap_EPC");
+
+    std::cout << "----------------------------------" << std::endl;
+    std::cout << "\033[32mAll " << pass_count << " tests passed!\033[0m" << std::endl;
+    std::cout << "----------------------------------" << std::endl;
+
+    delete dut;
+    return 0;
+}