rtio-nmigen/rtio/sed/output_network.py

105 lines
4.3 KiB
Python

from nmigen import *
from nmigen.utils import *
from nmigen.asserts import *
from rtio.sed import layouts
__all__ = ["latency", "OutputNetwork"]
# Based on: https://github.com/Bekbolatov/SortingNetworks/blob/master/src/main/js/gr.js
def boms_get_partner(n, l, p):
if p == 1:
return n ^ (1 << (l - 1))
scale = 1 << (l - p)
box = 1 << p
sn = n//scale - n//scale//box*box
if sn == 0 or sn == (box - 1):
return n
if (sn % 2) == 0:
return n - scale
return n + scale
def boms_steps_pairs(lane_count):
d = log2_int(lane_count)
steps = []
for l in range(1, d+1):
for p in range(1, l+1):
pairs = []
for n in range(2**d):
partner = boms_get_partner(n, l, p)
if partner != n:
if partner > n:
pair = (n, partner)
else:
pair = (partner, n)
if pair not in pairs:
pairs.append(pair)
steps.append(pairs)
return steps
def latency(lane_count):
d = log2_int(lane_count)
return sum(l for l in range(1, d+1))
def cmp_wrap(a, b):
return Mux((a[-2] == a[-1]) & (b[-2] == b[-1]) & (a[-1] != b[-1]), a[-1], a < b)
class OutputNetwork(Elaboratable):
def __init__(self, lane_count, seqn_width, layout_payload):
self.lane_count = lane_count
self.seqn_width = seqn_width
self.layout_payload = layout_payload
self.steps = boms_steps_pairs(lane_count)
self.network = [[Record(layouts.output_network_node(seqn_width, layout_payload)) for _ in range(lane_count)] for _ in range(len(self.steps) + 1)]
for i in range(1, len(self.steps) + 1):
for rec in self.network[i]:
rec.seqn.reset_less = True
rec.replace_occured.reset_less = True
rec.nondata_replace_occured.reset_less = True
for field, _ in layout_payload:
getattr(rec.payload, field).reset_less = True
self.input = self.network[0]
self.output = self.network[-1]
def elaborate(self, platform):
m = Module()
for i in range(len(self.steps)):
for node1, node2 in self.steps[i]:
nondata_difference = Signal()
for field, _ in self.layout_payload:
if field != 'data':
f1 = getattr(self.network[i][node1].payload, field)
f2 = getattr(self.network[i][node2].payload, field)
with m.If(f1 != f2):
m.d.comb += nondata_difference.eq(1)
k1 = Cat(self.network[i][node1].payload.channel, ~self.network[i][node1].valid)
k2 = Cat(self.network[i][node2].payload.channel, ~self.network[i][node2].valid)
with m.If(k1 == k2):
with m.If(cmp_wrap(self.network[i][node1].seqn, self.network[i][node2].seqn)):
m.d.sync += self.network[i + 1][node1].eq(self.network[i][node2])
m.d.sync += self.network[i + 1][node2].eq(self.network[i][node1])
with m.Else():
m.d.sync += self.network[i + 1][node1].eq(self.network[i][node1])
m.d.sync += self.network[i + 1][node2].eq(self.network[i][node2])
m.d.sync += self.network[i + 1][node1].replace_occured.eq(1)
m.d.sync += self.network[i + 1][node1].nondata_replace_occured.eq(nondata_difference)
m.d.sync += self.network[i + 1][node2].valid.eq(0)
with m.Elif(k1 < k2):
m.d.sync += self.network[i + 1][node1].eq(self.network[i][node1])
m.d.sync += self.network[i + 1][node2].eq(self.network[i][node2])
with m.Else():
m.d.sync += self.network[i + 1][node1].eq(self.network[i][node2])
m.d.sync += self.network[i + 1][node2].eq(self.network[i][node1])
unchanged = list(range(self.lane_count))
for node1, node2 in self.steps[i]:
unchanged.remove(node1)
unchanged.remove(node2)
for node in unchanged:
m.d.sync += self.network[i + 1][node].eq(self.network[i][node])
return m