speed up ==, hash, <, >, <=, and >= for plain_data
[nmutil.git] / src / nmutil / picker.py
index 0babc074a9eeb9663793bcc4fd1bfdb259d631bc..7aab175d8d60c031020b7b208727156123b8c09f 100644 (file)
@@ -1,4 +1,5 @@
-""" Priority Picker: optimised back-to-back PriorityEncoder and Decoder
+# SPDX-License-Identifier: LGPL-3-or-later
+""" Priority Picker: optimized back-to-back PriorityEncoder and Decoder
     and MultiPriorityPicker: cascading mutually-exclusive pickers
 
     This work is funded through NLnet under Grant 2019-02-012
 """
 
 from nmigen import Module, Signal, Cat, Elaboratable, Array, Const, Mux
-from nmigen.cli import verilog, rtlil
+from nmigen.utils import bits_for
+from nmigen.cli import rtlil
 import math
+from nmutil.prefix_sum import prefix_sum
 
 
 class PriorityPicker(Elaboratable):
     """ implements a priority-picker.  input: N bits, output: N bits
 
-        * lsb_mode is for a LSB-priority picker
-        * reverse_i=True is for convenient reverseal of the input bits
+        * msb_mode is for a MSB-priority picker
+        * reverse_i=True is for convenient reversal of the input bits
         * reverse_o=True is for convenient reversal of the output bits
+        * `msb_mode=True` is redundant with `reverse_i=True, reverse_o=True`
+            but is allowed for backwards compatibility.
     """
-    def __init__(self, wid, lsb_mode=False, reverse_i=False, reverse_o=False):
+
+    def __init__(self, wid, msb_mode=False, reverse_i=False, reverse_o=False):
         self.wid = wid
         # inputs
-        self.lsb_mode = lsb_mode
+        self.msb_mode = msb_mode
         self.reverse_i = reverse_i
         self.reverse_o = reverse_o
         self.i = Signal(wid, reset_less=True)
         self.o = Signal(wid, reset_less=True)
-        self.en_o = Signal(reset_less=True) # true if any output is true
+
+        self.en_o = Signal(reset_less=True)
+        "true if any output is true"
 
     def elaborate(self, platform):
         m = Module()
 
         # works by saying, "if all previous bits were zero, we get a chance"
         res = []
-        ni = Signal(self.wid, reset_less = True)
+        ni = Signal(self.wid, reset_less=True)
         i = list(self.i)
         if self.reverse_i:
             i.reverse()
+        if self.msb_mode:
+            i.reverse()
         m.d.comb += ni.eq(~Cat(*i))
         prange = list(range(0, self.wid))
-        if self.lsb_mode:
+        if self.msb_mode:
             prange.reverse()
         for n in prange:
-            t = Signal(name="t%d" % n, reset_less = True)
+            t = Signal(name="t%d" % n, reset_less=True)
             res.append(t)
             if n == 0:
                 m.d.comb += t.eq(i[n])
@@ -71,7 +81,7 @@ class PriorityPicker(Elaboratable):
         # we like Cat(*xxx).  turn lists into concatenated bits
         m.d.comb += self.o.eq(Cat(*res))
         # useful "is any output enabled" signal
-        m.d.comb += self.en_o.eq(self.o.bool()) # true if 1 input is true
+        m.d.comb += self.en_o.eq(self.o.bool())  # true if 1 input is true
 
         return m
 
@@ -95,16 +105,16 @@ class MultiPriorityPicker(Elaboratable):
 
         Also outputted (optional): an index for each picked "thing".
     """
-    def __init__(self, wid, levels, indices=False, multiin=False):
+
+    def __init__(self, wid, levels, indices=False, multi_in=False):
         self.levels = levels
         self.wid = wid
         self.indices = indices
-        self.multiin = multiin
+        self.multi_in = multi_in
 
-
-        if multiin:
+        if multi_in:
             # multiple inputs, multiple outputs.
-            i_l = [] # array of picker outputs
+            i_l = []  # array of picker outputs
             for j in range(self.levels):
                 i = Signal(self.wid, name="i_%d" % j, reset_less=True)
                 i_l.append(i)
@@ -114,7 +124,7 @@ class MultiPriorityPicker(Elaboratable):
             self.i = Signal(self.wid, reset_less=True)
 
         # create array of (single-bit) outputs (unary)
-        o_l = [] # array of picker outputs
+        o_l = []  # array of picker outputs
         for j in range(self.levels):
             o = Signal(self.wid, name="o_%d" % j, reset_less=True)
             o_l.append(o)
@@ -128,7 +138,7 @@ class MultiPriorityPicker(Elaboratable):
 
         # add an array of indices
         lidx = math.ceil(math.log2(self.levels))
-        idx_o = [] # store the array of indices
+        idx_o = []  # store the array of indices
         for j in range(self.levels):
             i = Signal(lidx, name="idxo_%d" % j, reset_less=True)
             idx_o.append(i)
@@ -146,7 +156,7 @@ class MultiPriorityPicker(Elaboratable):
         p_mask = None
         pp_l = []
         for j in range(self.levels):
-            if self.multiin:
+            if self.multi_in:
                 i = self.i[j]
             else:
                 i = self.i
@@ -160,10 +170,10 @@ class MultiPriorityPicker(Elaboratable):
                 p_mask = Const(0, self.wid)
             else:
                 mask = Signal(self.wid, name="m_%d" % j, reset_less=True)
-                comb += mask.eq(prev_pp.o | p_mask) # accumulate output bits
+                comb += mask.eq(prev_pp.o | p_mask)  # accumulate output bits
                 comb += pp.i.eq(i & ~mask)          # mask out input
                 p_mask = mask
-            i = pp.i # for input to next round
+            i = pp.i  # for input to next round
             prev_pp = pp
 
         # accumulate the enables
@@ -178,34 +188,89 @@ class MultiPriorityPicker(Elaboratable):
 
         # for each picker enabled, pass that out and set a cascading index
         lidx = math.ceil(math.log2(self.levels))
-        prev_count = None
+        prev_count = 0
         for j in range(self.levels):
             en_o = pp_l[j].en_o
-            if prev_count is None:
-                comb += self.idx_o[j].eq(0)
-            else:
-                count1 = Signal(lidx, name="count_%d" % j, reset_less=True)
-                comb += count1.eq(prev_count + Const(1, lidx))
-                comb += self.idx_o[j].eq(Mux(en_o, count1, prev_count))
-            prev_count = self.idx_o[j]
+            count1 = Signal(lidx, name="count_%d" % j, reset_less=True)
+            comb += count1.eq(prev_count + Const(1, lidx))
+            comb += self.idx_o[j].eq(prev_count)
+            prev_count = Mux(en_o, count1, prev_count)
 
         return m
 
     def __iter__(self):
-        if self.multiin:
+        if self.multi_in:
             yield from self.i
         else:
             yield self.i
         yield from self.o
+        yield self.en_o
         if not self.indices:
             return
-        yield self.en_o
         yield from self.idx_o
 
     def ports(self):
         return list(self)
 
 
+class BetterMultiPriorityPicker(Elaboratable):
+    """A better replacement for MultiPriorityPicker that has O(log levels)
+        latency, rather than > O(levels) latency.
+    """
+
+    def __init__(self, width, levels, *, work_efficient=False):
+        assert isinstance(width, int) and width >= 1
+        assert isinstance(levels, int) and 1 <= levels <= width
+        assert isinstance(work_efficient, bool)
+        self.width = width
+        self.levels = levels
+        self.work_efficient = work_efficient
+        assert self.__index_sat > self.levels - 1
+        self.i = Signal(width)
+        self.o = [Signal(width, name=f"o_{i}") for i in range(levels)]
+        self.en_o = Signal(levels)
+
+    @property
+    def __index_width(self):
+        return bits_for(self.levels)
+
+    @property
+    def __index_sat(self):
+        return (1 << self.__index_width) - 1
+
+    def elaborate(self, platform):
+        m = Module()
+
+        def sat_add(a, b):
+            sum = Signal(self.__index_width + 1)
+            m.d.comb += sum.eq(a + b)
+            retval = Signal(self.__index_width)
+            m.d.comb += retval.eq(Mux(sum[-1], self.__index_sat, sum))
+            return retval
+        indexes = prefix_sum((self.i[i] for i in range(self.width - 1)),
+                             sat_add, work_efficient=self.work_efficient)
+        indexes.insert(0, 0)
+        for i in range(self.width):
+            sig = Signal(self.__index_width, name=f"index_{i}")
+            m.d.comb += sig.eq(indexes[i])
+            indexes[i] = sig
+        for level in range(self.levels):
+            m.d.comb += self.en_o[level].eq(self.o[level].bool())
+            for i in range(self.width):
+                index_matches = indexes[i] == level
+                m.d.comb += self.o[level][i].eq(index_matches & self.i[i])
+
+        return m
+
+    def __iter__(self):
+        yield self.i
+        yield from self.o
+        yield self.en_o
+
+    def ports(self):
+        return list(self)
+
+
 if __name__ == '__main__':
     dut = PriorityPicker(16)
     vl = rtlil.convert(dut, ports=dut.ports())