Remove tracing elaboratable and replace with Fragment.origins

Author: rroohhh

naps/cores/debug/fsm_status_reg.py

@@ -10,7 +10,7 @@
 def fsm_status_reg(platform, m, fsm):
     if isinstance(platform, SocPlatform):
         fsm_state = StatusSignal(name=f"{fsm.state.name}_reg")  # TODO: use meaningful shape value here (needs deferring)
-        def signal_fixup_hook(platform, top_fragment: Fragment, sames):
+        def signal_fixup_hook(platform, top_fragment: Fragment):
             fsm_state.width = fsm.state.width
             fsm_state.decoder = fsm.state.decoder
         platform.prepare_hooks.insert(0, signal_fixup_hook)

naps/cores/peripherals/csr_bank.py

@@ -9,8 +9,9 @@
 
 
 class CsrBank(Elaboratable):
-    def __init__(self):
+    def __init__(self, names):
         self.memorymap = MemoryMap()
+        self.name = ".".join(names)
 
     def reg(self, name: str, signal: _Csr):
         assert isinstance(signal, _Csr)
@@ -67,5 +68,5 @@ def handle_write(m, addr, data, write_done):
                 write_done(Response.ERR)
 
         m = Module()
-        m.submodules += Peripheral(handle_read, handle_write, self.memorymap)
+        m.submodules += Peripheral(handle_read, handle_write, self.memorymap, self.name)
         return m

naps/cores/peripherals/csr_bank_zynq_test.py

@@ -5,7 +5,7 @@
 class TestAxiSlave(unittest.TestCase):
     def check_csr_bank(self, num_csr=10, testdata=0x12345678, use_axi_interconnect=False):
         platform = ZynqSocPlatform(SimPlatform(), use_axi_interconnect)
-        csr_bank = CsrBank()
+        csr_bank = CsrBank("test")
         for i in range(num_csr):
             csr_bank.reg("csr#{}".format(i), ControlSignal(32))
 
@@ -26,7 +26,7 @@ def test_csr_bank_interconnect(self):
 
     def test_simple_test_csr_bank(self):
         platform = ZynqSocPlatform(SimPlatform())
-        csr_bank = CsrBank()
+        csr_bank = CsrBank("test")
         csr_bank.reg("csr", ControlSignal(32))
 
         def testbench():

naps/cores/stream/stream_memory.py

@@ -1,4 +1,5 @@
 from amaranth import *
+from amaranth.lib.memory import Memory
 from amaranth.utils import bits_for
 
 from naps import BasicStream, stream_transformer

naps/soc/devicetree_overlay.py

@@ -17,7 +17,7 @@ def devicetree_overlay(platform, overlay_name, overlay_content, placeholder_subs
     if not hasattr(platform, "devicetree_overlays"):
         platform.devicetree_overlays = dict()
 
-        def overlay_hook(platform, top_fragment: Fragment, sames):
+        def overlay_hook(platform, top_fragment: Fragment):
             assert hasattr(top_fragment, "memorymap")
             memorymap = top_fragment.memorymap
 

naps/soc/hooks.py

@@ -4,104 +4,27 @@
 from amaranth.hdl._ast import Assign, Property, Switch, Print, Operator, Slice, Part, Concat, SwitchValue, ClockSignal, ResetSignal, Initial, ValueCastable
 from amaranth.hdl._ir import RequirePosedge
 
+from naps.soc.tracing_elaborate import get_elaboratable, get_module
+
 from .csr_types import _Csr, ControlSignal, StatusSignal, EventReg
 from .memorymap import MemoryMap
 from .pydriver.driver_items import DriverItem, DriverData
-from .tracing_elaborate import ElaboratableSames
 from ..util.py_serialize import is_py_serializable
 
 
-def csr_and_driver_item_hook(platform, top_fragment: Fragment, sames: ElaboratableSames):
+def csr_and_driver_item_hook(platform, top_fragment: Fragment):
     from naps.cores.peripherals.csr_bank import CsrBank
     already_done = []
 
-    def inner(fragment):
-        elaboratable = sames.get_elaboratable(fragment)
-        if elaboratable:
-            class_members = [(s, getattr(elaboratable, s)) for s in dir(elaboratable)]
-            csr_signals = [(name, member) for name, member in class_members if isinstance(member, _Csr)]
-            
-            def get_statement_csrs(stmt):
-                csrs = set()
-                if stmt is None:
-                    pass
-                # statements
-                elif isinstance(stmt, Assign):
-                    csrs |= get_statement_csrs(stmt.lhs)
-                    csrs |= get_statement_csrs(stmt.rhs)
-                elif isinstance(stmt, Property):
-                    csrs |= get_statement_csrs(stmt.message)
-                    csrs |= get_statement_csrs(stmt.test)
-                elif isinstance(stmt, Switch):
-                    csrs |= get_statement_csrs(stmt.test)
-                    for _patterns, statements, _src_loc  in stmt.cases:
-                        for statement in statements:
-                            csrs |= get_statement_csrs(statement)
-                elif isinstance(stmt, Print):
-                    for chunk in stmt.message._chunks:
-                        if isinstance(chunk, tuple):
-                            value, _format_spec = chunk
-                            csrs |= get_statement_csrs(value)
-                # Values
-                elif isinstance(stmt, Operator):
-                    for operand in stmt.operands:
-                        csrs |= get_statement_csrs(operand)
-                elif isinstance(stmt, Slice):
-                    csrs |= get_statement_csrs(stmt.value)
-                elif isinstance(stmt, Part):
-                    csrs |= get_statement_csrs(stmt.value)
-                    csrs |= get_statement_csrs(stmt.offset)
-                elif isinstance(stmt, Concat):
-                    for part in stmt.parts:
-                        csrs |= get_statement_csrs(part)
-                elif isinstance(stmt, SwitchValue):
-                    csrs |= get_statement_csrs(stmt.test)
-                    for pattern, value in stmt.cases:
-                        csrs |= get_statement_csrs(value)
-                elif isinstance(stmt, (ClockSignal, ResetSignal, Initial)):
-                    pass
-                elif isinstance(stmt, _Csr):
-                    csrs.add(stmt)
-                elif isinstance(stmt, (Signal, Const, ValueCastable)):
-                    pass
-                else:
-                    raise AssertionError("unknown object {} of type {} in statement", stmt, type(stmt))
-                return csrs
-
-            fragment_signals = set()
-            for _domain, statements in fragment.statements.items():
-                for stmt in statements:
-                    fragment_signals |= get_statement_csrs(stmt)
-
-            csr_signals += [
-                (signal.name, signal) for signal in fragment_signals
-                if isinstance(signal, _Csr)
-                   and signal.name != "$signal"
-                   and not any(signal is cmp_signal for name, cmp_signal in csr_signals)
-            ]
-
-            new_csr_signals = [(name, signal) for name, signal in csr_signals if not any(signal is done for done in already_done)]
-            old_csr_signals = [(name, signal) for name, signal in csr_signals if any(signal is done for done in already_done)]
-            for name, signal in new_csr_signals:
-                already_done.append(signal)
-
-            mmap = fragment.memorymap = MemoryMap()
-
-            if new_csr_signals:
-                m = Module()
-                csr_bank = m.submodules.csr_bank = CsrBank()
-                for name, signal in new_csr_signals:
-                    if isinstance(signal, (ControlSignal, StatusSignal, EventReg)):
-                        csr_bank.reg(name, signal)
-                        signal._MustUse__used = True
-
-                mmap.allocate_subrange(csr_bank.memorymap, name=None)  # name=None means that the Memorymap will be inlined
-                platform.to_inject_subfragments.append((m, "ignore"))
+    def inner(fragment, names):
+        elaboratables = get_elaboratable(fragment) or ()
 
-            for name, signal in old_csr_signals:
-                mmap.add_alias(name, signal)
+        class_members = []
+        driver_items = []
 
-            driver_items = [
+        for elaboratable in elaboratables:
+            class_members += list(elaboratable.__dict__.items())
+            driver_items += [
                 (name, getattr(elaboratable, name))
                 for name in dir(elaboratable)
                 if isinstance(getattr(elaboratable, name), DriverItem)
@@ -111,19 +34,103 @@ def get_statement_csrs(stmt):
                 for name in dir(elaboratable)
                 if is_py_serializable(getattr(elaboratable, name)) and not name.startswith("_")
             ]
-            for name, driver_item in driver_items:
-                fragment.memorymap.add_driver_item(name, driver_item)
+
+        csr_signals = [(name, member) for name, member in class_members if isinstance(member, _Csr)]
+
+        def get_statement_csrs(stmt):
+            csrs = set()
+            if stmt is None:
+                pass
+            # statements
+            elif isinstance(stmt, Assign):
+                csrs |= get_statement_csrs(stmt.lhs)
+                csrs |= get_statement_csrs(stmt.rhs)
+            elif isinstance(stmt, Property):
+                csrs |= get_statement_csrs(stmt.message)
+                csrs |= get_statement_csrs(stmt.test)
+            elif isinstance(stmt, Switch):
+                csrs |= get_statement_csrs(stmt.test)
+                for _patterns, statements, _src_loc  in stmt.cases:
+                    for statement in statements:
+                        csrs |= get_statement_csrs(statement)
+            elif isinstance(stmt, Print):
+                for chunk in stmt.message._chunks:
+                    if isinstance(chunk, tuple):
+                        value, _format_spec = chunk
+                        csrs |= get_statement_csrs(value)
+            # Values
+            elif isinstance(stmt, Operator):
+                for operand in stmt.operands:
+                    csrs |= get_statement_csrs(operand)
+            elif isinstance(stmt, Slice):
+                csrs |= get_statement_csrs(stmt.value)
+            elif isinstance(stmt, Part):
+                csrs |= get_statement_csrs(stmt.value)
+                csrs |= get_statement_csrs(stmt.offset)
+            elif isinstance(stmt, Concat):
+                for part in stmt.parts:
+                    csrs |= get_statement_csrs(part)
+            elif isinstance(stmt, SwitchValue):
+                csrs |= get_statement_csrs(stmt.test)
+                for pattern, value in stmt.cases:
+                    csrs |= get_statement_csrs(value)
+            elif isinstance(stmt, (ClockSignal, ResetSignal, Initial)):
+                pass
+            elif isinstance(stmt, _Csr):
+                csrs.add(stmt)
+            elif isinstance(stmt, (Signal, Const, ValueCastable)):
+                pass
+            else:
+                raise AssertionError("unknown object {} of type {} in statement", stmt, type(stmt))
+            return csrs
+
+        fragment_signals = set()
+        for _domain, statements in fragment.statements.items():
+            for stmt in statements:
+                fragment_signals |= get_statement_csrs(stmt)
+
+        csr_signals += [
+            (signal.name, signal) for signal in fragment_signals
+            if isinstance(signal, _Csr)
+                and signal.name != "$signal"
+                and not any(signal is cmp_signal for name, cmp_signal in csr_signals)
+        ]
+
+        new_csr_signals = [(name, signal) for name, signal in csr_signals if not any(signal is done for done in already_done)]
+        old_csr_signals = [(name, signal) for name, signal in csr_signals if any(signal is done for done in already_done)]
+        for name, signal in new_csr_signals:
+            already_done.append(signal)
+
+        mmap = fragment.memorymap = MemoryMap()
+
+        if new_csr_signals:
+            m = Module()
+            csr_bank = m.submodules.csr_bank = CsrBank(names)
+            print(f"-> adding csr bank for {'.'.join(names)}")
+            for name, signal in new_csr_signals:
+                if isinstance(signal, (ControlSignal, StatusSignal, EventReg)):
+                    csr_bank.reg(name, signal)
+                    signal._MustUse__used = True
+
+            mmap.allocate_subrange(csr_bank.memorymap, name=None)  # name=None means that the Memorymap will be inlined
+            platform.to_inject_subfragments.append((m, "ignore"))
+
+        for name, signal in old_csr_signals:
+            mmap.add_alias(name, signal)
+
+        for name, driver_item in driver_items:
+            fragment.memorymap.add_driver_item(name, driver_item)
 
         for subfragment, name, _src_loc in fragment.subfragments:
             if isinstance(subfragment, RequirePosedge):
                 continue
-            inner(subfragment)
-    inner(top_fragment)
+            inner(subfragment, [*names, str(name)])
+    inner(top_fragment, ["top"])
 
 
-def address_assignment_hook(platform, top_fragment: Fragment, sames: ElaboratableSames):
-    def inner(fragment):
-        module = sames.get_module(fragment)
+def address_assignment_hook(platform, top_fragment: Fragment):
+    def inner(fragment, names):
+        module = get_module(fragment)
         if hasattr(module, "peripheral"):  # we have the fragment of a marker module for a peripheral
             fragment.memorymap = module.peripheral.memorymap
             return
@@ -132,17 +139,18 @@ def inner(fragment):
         for sub_fragment, sub_name, _src_loc in fragment.subfragments:
             if isinstance(sub_fragment, RequirePosedge) or sub_name == "ignore":
                 continue
-            inner(sub_fragment)
+            inner(sub_fragment, [*names, str(sub_name)])
 
         # add everything to the own memorymap
         if not hasattr(fragment, "memorymap"):
             fragment.memorymap = MemoryMap()
+        print(f"-> assigning address for {'.'.join(names)}")
         for sub_fragment, sub_name, _src_loc in fragment.subfragments:
             if isinstance(sub_fragment, RequirePosedge) or sub_name == "ignore":
                 continue
             assert hasattr(sub_fragment, "memorymap")  # this holds because we did depth first recursion
             fragment.memorymap.allocate_subrange(sub_fragment.memorymap, sub_name)
-    inner(top_fragment)
+    inner(top_fragment, ["top"])
 
     # prepare and finalize the memorymap
     top_memorymap: MemoryMap = top_fragment.memorymap
@@ -173,20 +181,25 @@ def get_child_bits(memorymap: MemoryMap):
     platform.memorymap = top_memorymap
 
 
-def peripherals_collect_hook(platform, top_fragment: Fragment, sames: ElaboratableSames):
+def peripherals_collect_hook(platform, top_fragment: Fragment):
     platform.peripherals = []
 
-    def collect_peripherals(platform, fragment: Fragment, sames):
-        module = sames.get_module(fragment)
+    def collect_peripherals(platform, fragment: Fragment, names):
+        module = get_module(fragment)
         if module:
             if hasattr(module, "peripheral"):
+                try:
+                    for elab in get_elaboratable(fragment) or ():
+                        print(f"-> collected peripheral {elab.name}")
+                except:
+                    print(f"-> collected peripheral at {'.'.join(names)}")
                 platform.peripherals.append(module.peripheral)
         for f, name, _src_loc in fragment.subfragments:
             if isinstance(f, RequirePosedge):
                 continue
-            collect_peripherals(platform, f, sames)
+            collect_peripherals(platform, f, [*names, str(name)])
 
-    collect_peripherals(platform, top_fragment, sames)
+    collect_peripherals(platform, top_fragment, ["top"])
 
     ranges = [(peripheral.range(), peripheral) for peripheral in platform.peripherals
               if not peripheral.memorymap.is_empty and not peripheral.memorymap.was_inlined]

naps/soc/peripheral.py

@@ -22,7 +22,8 @@ def __init__(
             self,
             handle_read: HandleRead,
             handle_write: HandleWrite,
-            memorymap: MemoryMap
+            memorymap: MemoryMap,
+            name: str | None = None
     ):
         """
         A `Peripheral` is a thing that is memorymaped in the SOC.
@@ -38,6 +39,7 @@ def __init__(
         self.handle_read = handle_read
         self.handle_write = handle_write
         self.memorymap = memorymap
+        self.name = name
 
     def range(self):
         return self.memorymap.absolute_range_of_direct_children.range()

naps/soc/platform/jtag/jtag_soc_platform.py

@@ -27,7 +27,7 @@ def __init__(self, platform):
         self.jtag_active = Signal()
         self.jtag_debug_signals = Signal(32)
 
-        def peripherals_connect_hook(platform, top_fragment: Fragment, sames):
+        def peripherals_connect_hook(platform, top_fragment: Fragment):
             from naps import JTAGPeripheralConnector
             if platform.peripherals:
                 aggregator = PeripheralsAggregator()

naps/soc/platform/sim/sim_soc_platform.py

@@ -20,7 +20,7 @@ def __init__(self, platform):
         assert isinstance(platform, SimPlatform)
         super().__init__(platform)
 
-        def peripherals_connect_hook(platform, top_fragment: Fragment, sames):
+        def peripherals_connect_hook(platform, top_fragment: Fragment):
             from naps.cores.axi import AxiEndpoint, AxiLitePeripheralConnector
 
             if platform.peripherals:

naps/soc/platform/zynq/zynq_soc_platform.py

@@ -24,7 +24,7 @@ def __init__(self, platform, use_axi_interconnect=False):
         self.ps7 = PS7(here_is_the_only_place_that_instanciates_ps7=True)
         self.final_to_inject_subfragments.append((self.ps7, "ps7"))
 
-        def peripherals_connect_hook(platform, top_fragment: Fragment, sames):
+        def peripherals_connect_hook(platform, top_fragment: Fragment):
             from naps.cores.axi import AxiEndpoint, AxiLitePeripheralConnector, AxiFullToLiteBridge, AxiInterconnect
 
             if platform.peripherals: