soc/interconnect/stream: add BufferizeEndpoints

diff --git a/litex/soc/interconnect/stream.py b/litex/soc/interconnect/stream.py
index f6b1bb0213b4f398376afeb0b8d0b8f569a0edaf..0647af94b451974deefc15d44e50592406b9a5ef 100644 (file)
--- a/litex/soc/interconnect/stream.py
+++ b/litex/soc/interconnect/stream.py
@@ -459,4 +459,34 @@ class Pipeline(Module):
             self.source = m.source
 
 
+# Add buffers on Endpoints (can be used to improve timings)
+class BufferizeEndpoints(ModuleTransformer):
+    def __init__(self, *names):
+        self.names = names
+
+    def transform_instance(self, submodule):
+        endpoints = get_endpoints(submodule)
+        sinks = {}
+        sources = {}
+        for name, endpoint in endpoints.items():
+            if not self.names or name in self.names:
+                if isinstance(endpoint, Sink):
+                    sinks.update({name: endpoint})
+                elif isinstance(endpoint, Source):
+                    sources.update({name: endpoint})
+
+        # add buffer on sinks
+        for name, sink in sinks.items():
+            buf = Buffer(sink.description)
+            submodule.submodules += buf
+            setattr(submodule, name, buf.d)
+            submodule.comb += Record.connect(buf.q, sink)
+
+        # add buffer on sources
+        for name, source in sources.items():
+            buf = Buffer(source.description)
+            submodule.submodules += buf
+            submodule.comb += Record.connect(source, buf.d)
+            setattr(submodule, name, buf.q)
+
 # XXX