Flush the decode/execute stages if a branch is taken

rv32.sv

@@ -29,6 +29,8 @@ module rv32 (
     logic mem_flush;
 
     rv32_hazard hazard (
+        .branch_taken_in(mem_branch_taken),
+
         .fetch_stall_out(fetch_stall),
         .fetch_flush_out(fetch_flush),
 
@@ -37,7 +39,7 @@ module rv32 (
 
         .execute_stall_out(execute_stall),
         .execute_flush_out(execute_flush),
-        
+
         .mem_stall_out(mem_stall),
         .mem_flush_out(mem_flush)
     );

rv32_hazard.sv

@@ -2,6 +2,8 @@
 `define RV32_HAZARD
 
 module rv32_hazard (
+    input branch_taken_in,
+
     output fetch_stall_out,
     output fetch_flush_out,
 
@@ -18,10 +20,10 @@ module rv32_hazard (
     assign fetch_flush_out = 0;
 
     assign decode_stall_out = execute_stall_out;
-    assign decode_flush_out = fetch_stall_out;
+    assign decode_flush_out = fetch_stall_out || branch_taken_in;
 
     assign execute_stall_out = mem_stall_out;
-    assign execute_flush_out = decode_stall_out;
+    assign execute_flush_out = decode_stall_out || branch_taken_in;
 
     assign mem_stall_out = 0;
     assign mem_flush_out = execute_stall_out;