Simplify main_stage.py

diff --git a/src/soc/branch/main_stage.py b/src/soc/branch/main_stage.py
index 02731da68f7d673743fc177ff6a535fd87bd29c5..1e2b3e7192613a0419327d326193c078626b6c8b 100644 (file)
--- a/src/soc/branch/main_stage.py
+++ b/src/soc/branch/main_stage.py
@@ -44,9 +44,16 @@ class BranchMainStage(PipeModBase):
         aa = Signal(i_fields['AA'][0:-1].shape())
         comb += aa.eq(i_fields['AA'][0:-1])
 
+        branch_imm_addr = Signal(64, reset_less=True)
         branch_addr = Signal(64, reset_less=True)
         branch_taken = Signal(reset_less=True)
         comb += branch_taken.eq(0)
+
+        with m.If(aa):
+            comb += branch_addr.eq(branch_imm_addr)
+        with m.Else():
+            comb += branch_addr.eq(branch_imm_addr + self.i.nia)
+            
         
 
         ##########################
@@ -56,12 +63,8 @@ class BranchMainStage(PipeModBase):
             with m.Case(InternalOp.OP_B):
                 li = Signal(i_fields['LI'][0:-1].shape())
                 comb += li.eq(i_fields['LI'][0:-1])
-                with m.If(aa):
-                    comb += branch_addr.eq(Cat(Const(0, 2), li))
-                    comb += branch_taken.eq(1)
-                with m.Else():
-                    comb += branch_addr.eq(Cat(Const(0, 2), li) + self.i.nia)
-                    comb += branch_taken.eq(1)
+                comb += branch_imm_addr.eq(Cat(Const(0, 2), li))
+                comb += branch_taken.eq(1)
 
         comb += self.o.nia_out.data.eq(branch_addr)
         comb += self.o.nia_out.ok.eq(branch_taken)

diff --git a/src/soc/branch/pipe_data.py b/src/soc/branch/pipe_data.py
index 2b3f7c4f4492f65d3ec749bcdadc370592c01b9f..4904ba7e4671e4e20a2f14247444801383e74e83 100644 (file)
--- a/src/soc/branch/pipe_data.py
+++ b/src/soc/branch/pipe_data.py
@@ -43,14 +43,14 @@ class BranchInputData(IntegerData):
 class BranchOutputData(IntegerData):
     def __init__(self, pspec):
         super().__init__(pspec)
-        self.lr = Signal(64, reset_less=True)
-        self.spr = Signal(64, reset_less=True)
+        self.lr = Data(64, name="lr")
+        self.spr = Data(64, name="spr")
         self.nia_out = Data(64, name="nia_out")
 
     def __iter__(self):
         yield from super().__iter__()
-        yield self.lr
-        yield self.spr
+        yield from self.lr
+        yield from self.spr
         yield from self.nia_out
 
     def eq(self, i):