Fix core freezing during flushing (#568)

coreblocks/params/keys.py

@@ -54,3 +54,8 @@ class AsyncInterruptInsertSignalKey(SimpleKey[Signal]):
 @dataclass(frozen=True)
 class MretKey(SimpleKey[Method]):
     pass
+
+
+@dataclass(frozen=True)
+class CoreStateKey(SimpleKey[Method]):
+    pass

coreblocks/params/layouts.py

@@ -333,6 +333,11 @@ def __init__(self, gen_params: GenParams):
 
         self.precommit: LayoutList = [fields.rob_id, fields.side_fx]
 
+        self.flushing = ("flushing", 1)
+        """ Core is currently flushed """
+
+        self.core_state: LayoutList = [self.flushing]
+
 
 class RSLayouts:
     """Layouts used in the reservation station."""

coreblocks/scheduler/scheduler.py

@@ -6,6 +6,8 @@
 from transactron.lib import FIFO, Forwarder
 from coreblocks.params import SchedulerLayouts, GenParams, OpType
 from transactron.utils import assign, AssignType
+from transactron.utils.dependencies import DependencyManager
+from coreblocks.params.keys import CoreStateKey
 from coreblocks.utils.protocols import FuncBlock
 
 
@@ -305,11 +307,16 @@ def elaborate(self, platform):
             source1 = self.rf_read1(m, {"reg_id": instr.regs_p.rp_s1})
             source2 = self.rf_read2(m, {"reg_id": instr.regs_p.rp_s2})
 
+            # when core is flushed, rp_dst are discarded.
+            # source operands may never become ready, skip waiting for them in any in RSes/FBs.
+            core_state = self.gen_params.get(DependencyManager).get_dependency(CoreStateKey())
+            flushing = core_state(m).flushing
+
             data = {
                 # when operand value is valid the convention is to set operand source to 0
                 "rs_data": {
-                    "rp_s1": Mux(source1.valid, 0, instr.regs_p.rp_s1),
-                    "rp_s2": Mux(source2.valid, 0, instr.regs_p.rp_s2),
+                    "rp_s1": Mux(source1.valid | flushing, 0, instr.regs_p.rp_s1),
+                    "rp_s2": Mux(source2.valid | flushing, 0, instr.regs_p.rp_s2),
                     "rp_s1_reg": instr.regs_p.rp_s1,
                     "rp_s2_reg": instr.regs_p.rp_s2,
                     "rp_dst": instr.regs_p.rp_dst,

coreblocks/stages/backend.py

@@ -56,9 +56,8 @@ def elaborate(self, platform):
             result = self.m_get_result(m)
             self.m_rob_mark_done(m, rob_id=result.rob_id, exception=result.exception)
 
-            with m.If(result.exception == 0):
-                self.m_rf_write_val(m, reg_id=result.rp_dst, reg_val=result.result)
-                with m.If(result.rp_dst != 0):
-                    self.m_rs_update(m, reg_id=result.rp_dst, reg_val=result.result)
+            self.m_rf_write_val(m, reg_id=result.rp_dst, reg_val=result.result)
+            with m.If(result.rp_dst != 0):
+                self.m_rs_update(m, reg_id=result.rp_dst, reg_val=result.result)
 
         return m

coreblocks/stages/retirement.py

@@ -1,12 +1,13 @@
 from amaranth import *
+from coreblocks.params.layouts import RetirementLayouts
 
-from transactron.core import Method, Transaction, TModule
+from transactron.core import Method, Transaction, TModule, def_method
 from transactron.lib.simultaneous import condition
 from transactron.utils.dependencies import DependencyManager
 
 from coreblocks.params.genparams import GenParams
 from coreblocks.params.isa import ExceptionCause
-from coreblocks.params.keys import GenericCSRRegistersKey
+from coreblocks.params.keys import CoreStateKey, GenericCSRRegistersKey
 from coreblocks.structs_common.csr_generic import CSRAddress, DoubleCounterCSR
 
 
@@ -48,10 +49,14 @@ def __init__(
 
         self.instret_csr = DoubleCounterCSR(gen_params, CSRAddress.INSTRET, CSRAddress.INSTRETH)
 
+        self.dependency_manager = gen_params.get(DependencyManager)
+        self.core_state = Method(o=self.gen_params.get(RetirementLayouts).core_state, nonexclusive=True)
+        self.dependency_manager.add_dependency(CoreStateKey(), self.core_state)
+
     def elaborate(self, platform):
         m = TModule()
 
-        m_csr = self.gen_params.get(DependencyManager).get_dependency(GenericCSRRegistersKey()).m_mode
+        m_csr = self.dependency_manager.get_dependency(GenericCSRRegistersKey()).m_mode
         m.submodules.instret_csr = self.instret_csr
 
         side_fx = Signal(reset=1)
@@ -98,6 +103,8 @@ def flush_instr(rob_entry):
                 ~rob_entry.exception | (rob_entry.exception & ecr_entry.valid & (ecr_entry.rob_id == rob_entry.rob_id))
             )
 
+        core_flushing = Signal()
+
         with m.FSM("NORMAL") as fsm:
             with m.State("NORMAL"):
                 with Transaction().body(m, request=retire_valid) as retire_transaction:
@@ -154,6 +161,8 @@ def flush_instr(rob_entry):
                             # Not using default condition, because we want to block if branch is not ready
                             flush_instr(rob_entry)
 
+                            m.d.comb += core_flushing.eq(1)
+
                     validate_transaction.schedule_before(retire_transaction)
 
             with m.State("TRAP_FLUSH"):
@@ -169,6 +178,8 @@ def flush_instr(rob_entry):
                     with m.If(core_empty):
                         m.next = "TRAP_RESUME"
 
+                m.d.comb += core_flushing.eq(1)
+
             with m.State("TRAP_RESUME"):
                 with Transaction().body(m):
                     # Resume core operation
@@ -185,4 +196,8 @@ def flush_instr(rob_entry):
         # Disable executing any side effects from instructions in core when it is flushed
         m.d.comb += side_fx.eq(~fsm.ongoing("TRAP_FLUSH"))
 
+        @def_method(m, self.core_state)
+        def _():
+            return {"flushing": core_flushing}
+
         return m

test/scheduler/test_scheduler.py

@@ -5,10 +5,13 @@
 from amaranth import *
 from amaranth.sim import Settle
 from parameterized import parameterized_class
+from coreblocks.params.keys import CoreStateKey
+from coreblocks.params.layouts import RetirementLayouts
 from coreblocks.stages.rs_func_block import RSBlockComponent
 
 from transactron.core import Method
 from transactron.lib import FIFO, AdapterTrans, Adapter
+from transactron.utils.dependencies import DependencyManager
 from coreblocks.scheduler.scheduler import Scheduler
 from coreblocks.structs_common.rf import RegisterFile
 from coreblocks.structs_common.rat import FRAT
@@ -82,6 +85,10 @@ def elaborate(self, platform):
         m.submodules.rob_retire = self.rob_retire = TestbenchIO(AdapterTrans(self.rob.retire))
         m.submodules.instr_input = self.instr_inp = TestbenchIO(AdapterTrans(instr_fifo.write))
         m.submodules.free_rf_inp = self.free_rf_inp = TestbenchIO(AdapterTrans(free_rf_fifo.write))
+        m.submodules.core_state = self.core_state = TestbenchIO(
+            Adapter(o=self.gen_params.get(RetirementLayouts).core_state)
+        )
+        self.gen_params.get(DependencyManager).add_dependency(CoreStateKey(), self.core_state.adapter.iface)
 
         # main scheduler
         m.submodules.scheduler = self.scheduler = Scheduler(
@@ -360,6 +367,11 @@ def process():
 
             return process
 
+        @def_method_mock(lambda: self.m.core_state)
+        def core_state_mock():
+            # TODO: flushing test
+            return {"flushing": 0}
+
         with self.run_simulation(self.m, max_cycles=1500) as sim:
             for i in range(self.rs_count):
                 sim.add_sync_process(
@@ -371,3 +383,4 @@ def process():
             )
             sim.add_sync_process(self.make_queue_process(io=self.m.free_rf_inp, input_queues=[self.free_regs_queue]))
             sim.add_sync_process(instr_input_process)
+            sim.add_sync_process(core_state_mock)