cpu

12 files changed, 338 insertions, 189 deletions

diff --git a/gentestvec/gentestvec_alu.c b/gentestvec/gentestvec_alu.c
index 39de43e..4a11a74 100644
--- a/gentestvec/gentestvec_alu.c
+++ b/gentestvec/gentestvec_alu.c
@@ -6,7 +6,8 @@
 typedef enum {
   ADD = 0b0000,
   SUB = 0b0001,
-  SLT = 0b0011,
+  SLT = 0b0010,
+  SLTU = 0b0011,
 
   AND = 0b0100,
   OR = 0b0101,
@@ -31,6 +32,9 @@ void test_op(OP op, uint32_t a, uint32_t b) {
   case SLT:
     result = (int32_t)a < (int32_t)b;
     break;
+  case SLTU:
+    result = a < b;
+    break;
 
   case AND:
     result = a & b;
@@ -79,6 +83,7 @@ int main(int argc, const char *argv[]) {
   test_op(SLT, 0x8fffffff, 0xffffffff);
   test_op(SLT, 0xffffffff, 0x00000001);
   test_op(SLT, 0x00000001, 0xffffffff);
+  test_op_random(SLTU, 1000);
 
   test_op_random(OR, 1000);
   test_op(OR, 0x00000000, 0x00000000);
diff --git a/sim/testbench_alu.v b/sim/testbench_alu.v
index 6403cd7..f011ed2 100644
--- a/sim/testbench_alu.v
+++ b/sim/testbench_alu.v
@@ -39,7 +39,7 @@ module testbench_alu();
   assign exp_zero = exp_flags[0];
 
   reg [31:0] alu_test_count, alu_error_count;
-  reg [103:0] alu_testvec [0:9999];
+  reg [103:0] alu_testvec [0:20000];
 
   initial begin
     #5;
@@ -62,7 +62,7 @@ module testbench_alu();
 
     alu_test_count = alu_test_count + 1;
 
-    if ((alu_test_count == 9027)) begin
+    if ((alu_test_count == 10027)) begin
       $display("FINISHED (ALU), with %d errors out of %d tests.", alu_error_count, alu_test_count);
       #16;
 
@@ -73,9 +73,9 @@ module testbench_alu();
   
   
 
-  alu #(.N(32)) alu (
-    .src0(a),
-    .src1(b),
+  alu alu (
+    .a(a),
+    .b(b),
     .op(op),
     .result(result),
     .zero(zero)
diff --git a/sim/testbench_register_file.v b/sim/testbench_register_file.v
index 79825dc..a22f3e1 100644
--- a/sim/testbench_register_file.v
+++ b/sim/testbench_register_file.v
@@ -14,12 +14,12 @@ register_file uut (
   .clk(clk),
   .rst(rst),
   .we(we),
-  .addr_read0(addr_rs0),
-  .addr_read1(addr_rs1),
-  .addr_write2(addr_rd2),
-  .data_read0(data_rs0),
-  .data_read1(data_rs1),
-  .data_write2(data_rd2)
+  .rs1(addr_rs0),
+  .rs2(addr_rs1),
+  .rd(addr_rd2),
+  .rs1_data(data_rs0),
+  .rs2_data(data_rs1),
+  .rd_data(data_rd2)
 );
 
 integer file, r, eof;
diff --git a/src/alu.v b/src/alu.v
index e90694b..6962129 100644
--- a/src/alu.v
+++ b/src/alu.v
@@ -1,43 +1,42 @@
-module alu #(
-  parameter N = 32
-)(
-  input [N-1:0] src0, src1,
-  input [3:0] op, // alu_op[3:2] 00: ARITHMETIC, 01: LOGIC, 10: SHIFT 
-  output reg [N-1:0] result,
+module alu (
+  input [31:0] a, b,
+  input [3:0] op,
+  output reg [31:0] result,
   output zero
 );
 
-wire [N-1:0] arithmetic_result, logic_result, shift_result;
+wire [31:0] arithmetic_result, logic_result, shift_result;
 
-arithmetic_unit #(.N(N)) au (
-  .src0(src0),
-  .src1(src1),
+arithmetic_unit au (
+  .a(a),
+  .b(b),
   .op(op[1:0]),
   .result(arithmetic_result)
 );
 
-logic_unit #(.N(N)) lu (
-  .src0(src0),
-  .src1(src1),
+logic_unit lu (
+  .a(a),
+  .b(b),
   .op(op[1:0]),
   .result(logic_result)
 );
 
-shift_unit #(.N(N)) su (
-  .src0(src0),
-  .shamt(src1),
+shift_unit su (
+  .a(a),
+  .b(b[4:0]),
   .op(op[1:0]),
   .result(shift_result)
 );
 
 always @ (*) begin
   case (op[3:2])
-    2'b00: result <= arithmetic_result;
-    2'b01: result <= logic_result;
-    2'b10: result <= shift_result;
+    2'b00: result <= arithmetic_result; // ARITHMETIC
+    2'b01: result <= logic_result;      // LOGIC
+    2'b10: result <= shift_result;      // SHIFT
+    default: result <= 31'b0;
   endcase
 end
 
-assign zero = ~|result;
+assign zero = result == 32'b0;
 
 endmodule
diff --git a/src/arithmetic_unit.v b/src/arithmetic_unit.v
index 59f5cae..7902e8c 100644
--- a/src/arithmetic_unit.v
+++ b/src/arithmetic_unit.v
@@ -1,23 +1,21 @@
-module arithmetic_unit #(
-  parameter N = 32
-)(
-  input [N-1:0] src0, src1,
-  input [1:0] op, // 00: ADD, 01: SUB, 11: SLT
-  output [N-1:0] result
+module arithmetic_unit (
+  input [31:0] a, b,
+  input [1:0] op,
+  output reg [31:0] result
 );
 
-wire [N-1:0] src1_inv, sum;
-wire cin, src0_lt_src1, overflow;
+wire signed [31:0] a_signed, b_signed;
 
-assign src1_inv = op[0] ? ~src1 : src1;
-assign cin = op[0];
+assign a_signed = a;
+assign b_signed = b;
 
-assign sum = src0 + src1_inv + cin;
-
-assign overflow = ~(src0[N-1] ^ src1[N-1] ^ op[0]) & (src0[N-1] ^ sum[N-1]);
-
-assign src0_lt_src1 = overflow ^ sum[N-1];
-
-assign result = op[1] ? {{(N-1){1'b0}}, src0_lt_src1} : sum;
+always @ (*) begin
+  case (op)
+    2'b00: result <= a + b;                               // ADD
+    2'b01: result <= a - b;                               // SUB
+    2'b10: result <= { {31{1'b0}}, a_signed < b_signed }; // SLT
+    2'b11: result <= { {31{1'b0}}, a < b };               // SLTU
+  endcase
+end
 
 endmodule
diff --git a/src/control_unit.v b/src/control_unit.v
new file mode 100644
index 0000000..35fc32c
--- /dev/null
+++ b/src/control_unit.v
@@ -0,0 +1,82 @@
+module control_unit (
+ input clk, rst,
+ input [6:0] opcode,
+ input [2:0] funct3,
+ input [6:0] funct7,
+ input zero,
+ output pc_we,
+ output mem_addr_src,
+ output mem_we,
+ output instr_we,
+ output [1:0] result_src,
+ output [3:0] alu_op,
+ output [1:0] alu_src0_src,
+ output [1:0] alu_src1_src,
+ output [1:0] imm_src,
+ output rf_we
+);
+
+parameter s00_fetch     = 4'h0,
+          s01_decode    = 4'h1,
+          s02_mem_addr  = 4'h2,
+          s03_mem_read  = 4'h3,
+          s04_mem_wb    = 4'h4,
+          s05_mem_write = 4'h5,
+          s06_execute_r = 4'h6,
+          s07_alu_wb    = 4'h7,
+          s08_execute_i = 4'h8,
+          s09_jal       = 4'h9,
+          s10_beq       = 4'ha;
+
+reg [3:0] state, next_state;
+
+always @ (posedge clk or posedge rst) begin
+  if (rst) state <= s00_fetch;
+  else state <= next_state;
+end
+
+always @ (*) begin
+  case(state)
+    s00_fetch:     next_state <= s01_decode;
+    s01_decode: case(opcode)
+      7'b0000011:  next_state <= s02_mem_addr;
+      7'b0100011:  next_state <= s02_mem_addr;
+      7'b0110011:  next_state <= s06_execute_r;
+      7'b0010011:  next_state <= s08_execute_i;
+      7'b1101111:  next_state <= s09_jal;
+      7'b1100011:  next_state <= s10_beq;
+      
+    endcase
+    s02_mem_addr:  case(opcode)
+      7'b0000011:  next_state <= s03_mem_read;
+      7'b0100011:  next_state <= s05_mem_write; 
+    endcase
+    s03_mem_read:  next_state <= s04_mem_wb;
+    s04_mem_wb:    next_state <= s00_fetch;
+    s05_mem_write: next_state <= s00_fetch;
+    s06_execute_r: next_state <= s07_alu_wb;
+    s07_alu_wb:    next_state <= s00_fetch;
+    s08_execute_i: next_state <= s07_alu_wb;
+    s09_jal:       next_state <= s07_alu_wb;
+    s10_beq:       next_state <= s00_fetch;
+  endcase
+end
+/*
+always @ (*) begin
+ case(state)
+    s00_fetch:
+    s01_decode:
+    s02_mem_addr:
+    s03_mem_read:
+    s04_mem_wb:
+    s05_mem_write:
+    s06_execute_r:
+    s07_alu_wb:
+    s08_execute_i:
+    s09_jal:
+    s10_beq:
+  endcase
+end
+*/
+
+endmodule
diff --git a/src/cpu.v b/src/cpu.v
index e420252..bd282d0 100644
--- a/src/cpu.v
+++ b/src/cpu.v
@@ -1,142 +1,159 @@
-module cpu #(
-  parameter N = 32,
-  parameter XLEN = 32
-)(
+module cpu (
   input clk,
   input rst
 );
 
 
-wire pc_we;
+// Control Unit
+
 wire mem_addr_src;
-wire mem_we;
+wire pc_we;
 wire instr_we;
 wire rf_we;
-wire [1:0] imm_src;
-wire [1:0] alu_src0_src;
-wire [1:0] alu_src1_src;
-wire [3:0] alu_op;
 wire alu_zero;
+wire [3:0] alu_op;
+wire [1:0] alu_src_a;
+wire [1:0] alu_src_b;
 wire [1:0] result_src;
 
-reg [N-1:0] result;
-
-
 // Fetch
-reg [N-1:0] pc, pc_old, instr;
 
+reg [31:0] pc;
 
 always @ (posedge clk or posedge rst) begin
- if (rst) pc <= 32'h0001_0000;
- else if (pc_we) pc <= result;
+  if (rst) pc <= 32'h0001_0000;
+  else if (pc_we) pc <= result;
 end
 
-always @ (posedge clk) begin
-  if (instr_we) begin
-    pc_old <= pc;
-    instr <= mem_data_read;
-  end
-end
+wire [31:0] mem_addr = mem_addr_src ? result : pc;
 
-// Memory
-wire [N-1:0] mem_addr;
-assign mem_addr = mem_addr_src ? pc : result;
+wire [31:0] mem_read_data;
 
-wire [N-1:0] mem_data_write;
-wire [N-1:0] mem_data_read;
-reg [N-1:0] data;
-
-memory_unit #(.N(N)) mu (
+memory_unit mu (
   .clk(clk),
   .rst(rst),
   .we(mem_we),
   .addr(mem_addr),
-  .data_read(mem_data_read),
-  .data_write(mem_data_write)
+  .read_data(mem_read_data),
+  .write_data(b_buf)
 );
 
-always @ (posedge clk) begin
-  data <= mem_data_read;
+reg [31:0] instruction;
+reg [31:0] pc_buf;
+
+always @ (posedge clk or posedge rst) begin
+  if (rst) begin
+    pc_buf <= 32'b0;
+    instruction <= 32'b0;
+  end else begin
+    if (instr_we) begin
+      pc_buf <= pc;
+      instruction <= mem_read_data; 
+    end
+  end
+end
+
+reg [31:0] data;
+
+always @(posedge clk or posedge rst) begin
+  if (rst) data <= 32'b0;
+  else data <= mem_read_data;
 end
 
-// Register file
+// Instruction Decode
+
+wire [6:0] opcode;
+wire [2:0] funct3;
+wire [6:0] funct7;
+wire [31:0] immediate;
+wire [4:0] rs1, rs2, rd;
+
+instruction_decode id (
+  .instruction(instruction),
+  .opcode(opcode),
+  .funct3(funct3),
+  .funct7(funct7),
+  .immediate(immediate),
+  .rs1(rs1),
+  .rs2(rs2),
+  .rd(rd)
+);
+
+// Register File
 
-reg [N-1:0] rf_data0_old, rf_data1_old;
-wire [N-1:0] rf_data0, rf_data1;
+wire [31:0] rs1_data, rs2_data;
 
-register_file #(.N(N), .XLEN(XLEN)) rf(
+register_file rf (
   .clk(clk),
   .rst(rst),
   .we(rf_we),
-  .addr_read0(instr[19:15]),
-  .addr_read1(instr[24:20]),
-  .addr_write2(instr[11:7]),
-  .data_read0(rf_data0),
-  .data_read1(rf_data1),
-  .data_write2(result)
+  .rs1(rs1),
+  .rs2(rs2),
+  .rd(rd),
+  .rs1_data(rs1_data),
+  .rs2_data(rs2_data),
+  .rd_data(result)
 );
 
-always @ (posedge clk) begin
-  rf_data0_old <= rf_data0;
-  rf_data1_old <= rf_data1;
-end
-
-assign mem_data_write = rf_data1_old;
+reg [31:0] a_buf, b_buf;
 
-reg [N-1:0] imm;
-
-always @ (*) begin
-  case (imm_src)
-    2'b00: imm <= {{20{instr[31]}}, instr[31:20]};
-    2'b01: imm <= {{20{instr[31]}}, instr[31:25], instr[11:7]};
-    2'b10: imm <= {{20{instr[31]}}, instr[7], instr[30:25], instr[11:8], 1'b0};
-  endcase
+always @ (posedge clk or posedge rst) begin
+  if (rst) begin
+    a_buf <= 32'b0;
+    b_buf <= 32'b0;
+  end else begin
+    a_buf <= rs1;
+    b_buf <= rs2;
+  end
 end
 
-reg [N-1:0] alu_src0, alu_src1;
+// Execute
+
+reg [31:0] a, b;
+wire [31:0] alu_result;
 
 always @ (*) begin
-  case(alu_src0_src)
-    2'b00: alu_src0 <= pc;
-    2'b01: alu_src0 <= pc_old;
-    2'b10: alu_src0 <= rf_data0_old;
+  case(alu_src_a)
+    2'b00: a <= pc;
+    2'b01: a <= pc_buf;
+    2'b10: a <= a_buf;
+    default: a <= 32'b0;
   endcase
 end
 
 always @ (*) begin
-  case(alu_src1_src)
-    2'b00: alu_src1 <= rf_data1_old;
-    2'b01: alu_src1 <= imm;
-    2'b10: alu_src1 <= {{(N-3){1'b0}}, 3'h4};
+  case(alu_src_b)
+    2'b00: b <= b_buf;
+    2'b01: b <= immediate;
+    2'b10: b <= 32'h4;
+    default: b <= 32'b0;
   endcase
 end
 
-// Execute
-
-wire [N-1:0] alu_result;
-
-alu #(.N(N)) alu (
-  .src0(alu_src0),
-  .src1(alu_src1),
+alu alu (
+  .a(a),
+  .b(b),
   .op(alu_op),
   .result(alu_result),
   .zero(alu_zero)
 );
 
-reg [N-1:0] alu_out;
+reg [31:0] result_buf;
 
-always @ (posedge clk) begin
-  alu_out <= alu_result;
+always @ (posedge clk or posedge rst) begin
+  if (rst) result_buf <= 32'b0;
+  else result_buf <= alu_result;
 end
 
+reg [31:0] result;
+
 always @ (*) begin
-  case (result_src)
-    2'b00: result <= alu_out;
+  case(result_src)
+    2'b00: result <= result_buf;
     2'b01: result <= data;
     2'b10: result <= alu_result;
+    default: result <= 32'b0;
   endcase
 end
 
-
-
 endmodule
diff --git a/src/instruction_decode.v b/src/instruction_decode.v
new file mode 100644
index 0000000..2a7eba7
--- /dev/null
+++ b/src/instruction_decode.v
@@ -0,0 +1,56 @@
+module instruction_decode(
+  input [31:0] instruction,
+  output [6:0] opcode,
+  output [2:0] funct3,
+  output [6:0] funct7,
+  output reg [31:0] immediate,
+  output [4:0] rs1, rs2, rd
+);
+
+parameter INSTRUCTION_FORMAT_UNKNOWN  = 3'b000,
+          INSTRUCTION_FORMAT_R        = 3'b001,
+          INSTRUCTION_FORMAT_I        = 3'b010,
+          INSTRUCTION_FORMAT_S        = 3'b011,
+          INSTRUCTION_FORMAT_B        = 3'b100,
+          INSTRUCTION_FORMAT_U        = 3'b101,
+          INSTRUCTION_FORMAT_J        = 3'b110;
+
+reg [2:0] instruction_format;
+
+always @ (*) begin
+  case (opcode)
+    7'b0110011: instruction_format <= INSTRUCTION_FORMAT_R; // ADD, SUB, SLL, SLT, SLTU, XOR, SRL, SRA, OR, AND
+    7'b0010011: instruction_format <= INSTRUCTION_FORMAT_I; // ADDI, SLTI, SLTIU, XORI, ORI, ANDI, SLLI, SRLI, SRAI
+    7'b0000011: instruction_format <= INSTRUCTION_FORMAT_I; // LB, LH, LW, LBU, LHU
+    7'b1100111: instruction_format <= INSTRUCTION_FORMAT_I; // JALR
+    7'b0100011: instruction_format <= INSTRUCTION_FORMAT_S; // SB, SH, SW
+    7'b1100011: instruction_format <= INSTRUCTION_FORMAT_B; // BEQ, BNE, BLT, BGE, BLTU, BGEU
+    7'b0110111: instruction_format <= INSTRUCTION_FORMAT_U; // LUI
+    7'b0010111: instruction_format <= INSTRUCTION_FORMAT_U; // AUIPC
+    7'b1101111: instruction_format <= INSTRUCTION_FORMAT_J; // JAL
+    7'b0001111: instruction_format <= INSTRUCTION_FORMAT_I; // FENCE, FENCE.I
+    7'b1110011: instruction_format <= INSTRUCTION_FORMAT_I; // ECALL, EBREAK, CSRRW, CSRRS, CSRRC, CSRRWI, CSRRSI, CSRRCI
+    default:    instruction_format <= INSTRUCTION_FORMAT_UNKNOWN;
+  endcase
+end
+
+always @ (*) begin
+  case (instruction_format)
+    INSTRUCTION_FORMAT_I: immediate <= { {21{instruction[31]}}, instruction[30:20] };
+    INSTRUCTION_FORMAT_S: immediate <= { {21{instruction[31]}}, instruction[30:25], instruction[11:8], instruction[7] };
+    INSTRUCTION_FORMAT_B: immediate <= { {20{instruction[31]}}, instruction[7], instruction[30:25], instruction[11:8], 1'b0 };
+    INSTRUCTION_FORMAT_U: immediate <= { instruction[31], instruction[30:12], 12'b0 };
+    INSTRUCTION_FORMAT_J: immediate <= { {12{instruction[31]}}, instruction[19:12], instruction[20], instruction[30:21], 1'b0 };
+    default: immediate <= 32'b0;
+  endcase
+end
+
+assign opcode = instruction[6:0];
+assign funct3 = instruction[14:12];
+assign funct7 = instruction[31:25];
+assign rs1    = instruction[19:15];
+assign rs2    = instruction[24:20];
+assign rd     = instruction[11:7];
+
+
+endmodule
diff --git a/src/logic_unit.v b/src/logic_unit.v
index 0bbd642..f9f62a2 100644
--- a/src/logic_unit.v
+++ b/src/logic_unit.v
@@ -1,16 +1,15 @@
-module logic_unit #(
-  parameter N = 32
-)(
-  input [N-1:0] src0, src1,
-  input [1:0] op, // 00: AND, 01: OR, 10: XOR
-  output reg [N-1:0] result
+module logic_unit (
+  input [31:0] a, b,
+  input [1:0] op,
+  output reg [31:0] result
 );
   
 always @ (*) begin
   case (op)
-  2'b00: result <= src0 & src1;
-  2'b01: result <= src0 | src1;
-  2'b10: result <= src0 ^ src1;
+    2'b00: result <= a & b; // AND
+    2'b01: result <= a | b; // OR
+    2'b10: result <= a ^ b; // XOR
+    default: result <= 32'b0;
   endcase
 end
 
diff --git a/src/memory_unit.v b/src/memory_unit.v
index b07c895..b2d7434 100644
--- a/src/memory_unit.v
+++ b/src/memory_unit.v
@@ -1,30 +1,28 @@
-module memory_unit #(
-  parameter N = 32
-)(
+module memory_unit (
   input clk,
   input rst,
   input we,
-  input [N-1:0] addr,
-  output reg [N-1:0] data_read,
-  input [N-1:0] data_write
+  input [31:0] addr,
+  output reg [31:0] read_data,
+  input [31:0] write_data
 );
 
-reg we_ram;
-wire [N-1:0] data_read_ram, data_read_rom;
+reg ram_we;
+wire [31:0] ram_read_data, rom_read_data;
 
-ram #(.N(N), .SIZE(1024)) ram(
+ram #(.N(32), .SIZE(1024)) ram(
   .clk(clk),
   .rst(rst),
   .we(we_ram),
-  .addr(addr[N-17:0]),
-  .data_read(data_read_ram),
-  .data_write(data_write)
+  .addr(addr),
+  .data_read(ram_read_data),
+  .data_write(write_data)
 );
 
-rom #(.N(N), .SIZE(1024)) rom(
+rom #(.N(32), .SIZE(1024)) rom(
   .clk(clk),
-  .addr(addr[N-17:0]),
-  .data_read(data_read_rom)
+  .addr(addr),
+  .data_read(rom_read_data)
 );
 
 // 0000 0000 Reserved
@@ -41,16 +39,16 @@ rom #(.N(N), .SIZE(1024)) rom(
 
 
 always @(*) begin
-  we_ram = 0;
-  if (addr[N-1:N-16] == 16'h0000) begin
-    data_read <= 0;
-  end else if (addr[N-1:N-16] >= 16'h0001 && addr[N-1:N-16] <= 16'h000F) begin
-    data_read <= data_read_rom;
-    we_ram = we;
-  end else if (addr[N-1:N-16] >= 16'h0010 && addr[N-1:N-16] <= 16'hFF0F) begin
-    data_read <= data_read_ram;
-  end else if (addr[N-1:N-16] >= 16'hFF10 && addr[N-1:N-16] <= 16'hFFFF) begin
-    data_read <= 0;
+  ram_we = 0;
+  if (addr[31:16] == 16'h0000) begin
+    read_data <= 0;
+  end else if (addr[31:16] >= 16'h0001 && addr[31:16] <= 16'h000F) begin
+    read_data <= rom_read_data;
+    ram_we = we;
+  end else if (addr[31:16] >= 16'h0010 && addr[31:16] <= 16'hFF0F) begin
+    read_data <= ram_read_data;
+  end else if (addr[31:16] >= 16'hFF10 && addr[31:16] <= 16'hFFFF) begin
+    read_data <= 0;
   end
 end
 
diff --git a/src/register_file.v b/src/register_file.v
index de697ab..5142a69 100644
--- a/src/register_file.v
+++ b/src/register_file.v
@@ -1,27 +1,23 @@
-module register_file #(
-  parameter N = 32,
-  parameter XLEN = 32
-)(
+module register_file (
   input clk, rst, we,
-  input [log2(XLEN)-1:0] addr_read0, addr_read1, addr_write2,
-  input [N-1:0] data_write2,
-  output reg [N-1:0] data_read0, data_read1
+  input [4:0] rs1, rs2, rd,
+  input [31:0] rd_data,
+  output reg [31:0] rs1_data, rs2_data
 );
 
-`include "include/log2.vh"
 
-reg [N-1:0] registers[XLEN-1:1];
+reg [31:0] registers[31:1];
 
 integer i;
 always @(posedge clk or rst) begin
   if (rst) begin
-    for (i = 1; i < XLEN; i = i + 1)
-      registers[i] <= 0;
+    for (i = 1; i < 32; i = i + 1)
+      registers[i] <= 32'b0;
   end else begin
-    data_read0 = (addr_read0 == 0) ? 0 : registers[addr_read0];
-    data_read1 = (addr_read1 == 0) ? 0 : registers[addr_read1];
-    if (we && (addr_write2 != 0)) begin
-      registers[addr_write2] <= data_write2;
+    rs1_data = (rs1 == 0) ? 32'b0 : registers[rs1];
+    rs2_data = (rs2 == 0) ? 32'b0 : registers[rs2];
+    if (we && (rd != 0)) begin
+      registers[rd] <= rd_data;
     end
   end
 end
diff --git a/src/shift_unit.v b/src/shift_unit.v
index d0aa9d3..df4b01d 100644
--- a/src/shift_unit.v
+++ b/src/shift_unit.v
@@ -1,17 +1,16 @@
-module shift_unit #(
-  parameter N = 32
-)(
-  input signed [N-1:0] src0,
-  input  [N-1:0] shamt,
-  input [1:0] op, // 00: SLL, 01: SRL, 11: SRA
-  output reg [N-1:0] result
+module shift_unit (
+  input signed [31:0] a,
+  input  [4:0] b,
+  input [1:0] op,
+  output reg [31:0] result
 );
 
 always @ (*) begin
   case (op)
-  2'b00: result <= src0 << shamt % N;
-  2'b01: result <= src0 >> shamt % N;
-  2'b11: result <= src0 >>> shamt % N;
+    2'b00: result <= a << b;   // SLL
+    2'b01: result <= a >> b;   // SRL
+    2'b11: result <= a >>> b;  // SRA
+    default: result <= 32'b0;
   endcase
 end