src/nmutil/popcount.py

   1 # SPDX-License-Identifier: LGPL-3-or-later
   2 # Copyright 2022 Jacob Lifshay programmerjake@gmail.com
   3
   4 # Funded by NLnet Assure Programme 2021-02-052, https://nlnet.nl/assure part
   5 # of Horizon 2020 EU Programme 957073.
   6
   7 from nmigen import Module
   8 from nmigen.hdl.ast import Value, Const, Signal
   9 from nmutil.plain_data import plain_data
  10 from nmutil.prefix_sum import tree_reduction
  11 from nmigen.cli import rtlil
  12
  13
  14 def pop_count(v, *, width=None, process_temporary=lambda v: v):
  15     """return the population count (number of 1 bits) of `v`.
  16     Arguments:
  17     v: nmigen.Value | int
  18         the value to calculate the pop-count of.
  19     width: int | None
  20         the bit-width of `v`.
  21         If `width` is None, then `v` must be a nmigen Value or
  22         match `v`'s width.
  23     process_temporary: function of (type(v)) -> type(v)
  24         called after every addition operation, can be used to introduce
  25         `Signal`s for the intermediate values in the pop-count computation
  26         like so:
  27
  28         ```
  29         def process_temporary(v):
  30             sig = Signal.like(v)
  31             m.d.comb += sig.eq(v)
  32             return sig
  33         ```
  34     """
  35     if isinstance(v, Value):
  36         if width is None:
  37             width = len(v)
  38         assert width == len(v)
  39         bits = [v[i] for i in range(width)]
  40         if len(bits) == 0:
  41             return Const(0)
  42     else:
  43         assert width is not None, "width must be given"
  44         # v and width are ints
  45         bits = [(v & (1 << i)) != 0 for i in range(width)]
  46         if len(bits) == 0:
  47             return 0
  48     return tree_reduction(bits, fn=lambda a, b: process_temporary(a + b))
  49
  50
  51 # run this as simply "python3 popcount.py" to create an ilang file that
  52 # can be viewed with yosys "read_ilang test_popcount.il; show top"
  53 if __name__ == "__main__":
  54     m = Module()
  55     v = Signal(8)
  56     x = Signal(8)
  57     pc = pop_count(v, width=8)
  58     m.d.comb += v.eq(pc)
  59     vl = rtlil.convert(m)
  60     with open("test_popcount.il", "w") as f:
  61         f.write(vl)