Apply ConsecutiveSlice optimizations to CompiledSimulation for ~46% reduction in compile time.

pyrtl/compilesim.py

@@ -15,7 +15,13 @@
 from pyrtl.helperfuncs import infer_val_and_bitwidth
 from pyrtl.memory import MemBlock, RomBlock
 from pyrtl.pyrtlexceptions import PyrtlError, PyrtlInternalError
-from pyrtl.simulation import SimulationTrace, _trace_sort_key
+from pyrtl.simulation import (
+    ConsecutiveSlice,
+    SimulationTrace,
+    _trace_sort_key,
+    make_consecutive_slices,
+    shift,
+)
 from pyrtl.wire import Const, Input, Output, Register, WireVector
 
 __all__ = ["CompiledSimulation"]
@@ -673,17 +679,85 @@ def _build_concat(self, write, _op, _param, args, dest):
                 )
             )
 
-    def _build_select(self, write, _op, param, args, dest):
-        for n in range(self._limbs(dest)):
-            bits = [
-                f"((1&({self.varname[args[0]]}[{b // 64}]>>{b % 64}))<<{en})"
-                for en, b in enumerate(param[64 * n : min(dest.bitwidth, 64 * (n + 1))])
-            ]
-            write(
-                "{dest}[{n}] = {bits};".format(
-                    dest=self.varname[dest], n=n, bits="|".join(bits)
+    def _next_limb_start(self, bit_position: int) -> int:
+        """Given a bit position, return the bit position of the first bit in the next
+        64-bit limb.
+
+        This rounds bit_position up to the next even multiple of 64. For example, all
+        integers in the range [0, 63] return 64, and all integers in the range [64, 127]
+        return 128.
+        """
+        return bit_position + (64 - bit_position % 64)
+
+    def _split_consecutive_slices(
+        self, dest: WireVector, op_param: list[int], n: int
+    ) -> list[ConsecutiveSlice]:
+        """Make ConsecutiveSlices for the ``n``th dest limb.
+
+        To simplify processing, split ConsecutiveSlices that fetch bits from multiple
+        arg limbs into ConsecutiveSlices that don't fetch bits from multiple arg limbs.
+        """
+        # Get op_params for the `n`th dest limb.
+        current_param = op_param[64 * n : min(dest.bitwidth, 64 * (n + 1))]
+
+        split_slices = []
+        for consecutive_slice in make_consecutive_slices(current_param):
+            next_limb_arg_start = self._next_limb_start(consecutive_slice.arg_start)
+            arg_end = consecutive_slice.arg_start + consecutive_slice.length
+
+            if next_limb_arg_start >= arg_end:
+                # consecutive_slice fetches bits from one arg limb.
+                split_slices.append(consecutive_slice)
+            else:
+                # consecutive_slice fetches bits from two arg limbs. Split it into two
+                # ConsecutiveSlices that each fetch bits from one arg limb.
+                first_limb_length = next_limb_arg_start - consecutive_slice.arg_start
+                split_slices.append(
+                    ConsecutiveSlice(
+                        length=first_limb_length,
+                        arg_start=consecutive_slice.arg_start,
+                        dest_start=consecutive_slice.dest_start,
+                    )
                 )
-            )
+                split_slices.append(
+                    ConsecutiveSlice(
+                        length=consecutive_slice.length - first_limb_length,
+                        arg_start=next_limb_arg_start,
+                        dest_start=consecutive_slice.dest_start + first_limb_length,
+                    )
+                )
+
+        return split_slices
+
+    def _build_slice(self, write, _op, param, args, dest):
+        # Iterate over the slice's dest limbs.
+        for n in range(self._limbs(dest)):
+            split_slices = self._split_consecutive_slices(dest, param, n)
+
+            expr_parts = []
+            for split_slice in split_slices:
+                # We only need to fetch bits from one arg limb because
+                # _split_consecutive_slices split up any slices that fetch bits from
+                # multiple arg limbs.
+                expr = f"{self.varname[args[0]]}[{split_slice.arg_start // 64}]"
+                expr = shift(expr, ">>", split_slice.arg_start % 64)
+
+                slice_arg_end = split_slice.arg_start + split_slice.length
+                next_limb_arg_start = self._next_limb_start(split_slice.arg_start)
+                assert slice_arg_end <= next_limb_arg_start
+                actual_arg_end = min(args[0].bitwidth, next_limb_arg_start)
+                # If this split_slice does not fetch all of the arg limb's remaining
+                # bits, we must mask the arg limb.
+                if slice_arg_end < actual_arg_end:
+                    expr = f"({expr} & 0x{(1 << split_slice.length) - 1:X})"
+
+                assert split_slice.dest_start < 64
+                assert split_slice.length <= 64 - split_slice.dest_start
+                expr_parts.append(shift(expr, "<<", split_slice.dest_start))
+
+            expr = "|".join(expr_parts)
+
+            write(f"{self.varname[dest]}[{n}] = {expr};")
 
     def _declare_mem_helpers(self, write):
         helpers = """
@@ -851,7 +925,7 @@ def _create_code(self, write):
             "-": self._build_sub,
             "*": self._build_mul,
             "c": self._build_concat,
-            "s": self._build_select,
+            "s": self._build_slice,
         }
         for net in self.block:  # topological order
             if net.op in "r@":

pyrtl/simulation.py

@@ -598,63 +598,65 @@ def _mem_update(self, net):
 #
 
 
+def shift(value: str, direction: str, amount: int) -> str:
+    """Return an expression that shifts ``value`` by ``amount``, in ``direction``."""
+    if amount == 0:
+        return value
+    return f"({value} {direction} {amount})"
+
+
 @dataclass
-class ContiguousSlice:
-    """Represents a contiguous range of bits mapped from a slice :class:`LogicNet`'s
-    input to output.
+class ConsecutiveSlice:
+    """Represents a consecutive range of bits mapped from a slice :class:`LogicNet`'s
+    ``arg`` to ``dest``.
     """
 
     length: int
-    """Length of this contiguous bit slice, in bits."""
+    """Length of this consecutive bit slice, in bits."""
+
     arg_start: int
     """Bit index of the start (least-significant) bit in the slice's input.
 
     The bit slice's input bits are in the range [arg_start, arg_start + length).
     """
+
     dest_start: int
     """Bit index of the start (least-significant) bit in the slice's output.
 
     The bit slice's output bits are in the range [dest_start, dest_start + length).
     """
 
 
-def make_contiguous_slices(net: LogicNet) -> list[ContiguousSlice]:
-    """Given a slice :class:`LogicNet`, return a run-length compressed version of its
-    :attr:`LogicNet.op_param`.
+def make_consecutive_slices(op_param: list[int]) -> list[ConsecutiveSlice]:
+    """Return a run-length compressed version of a slice :attr:`LogicNet.op_param`.
 
-    Slice :class:`LogicNets<LogicNet>` are very frequently used to extract a contiguous
-    range of bits from a :class:`WireVector`, for example ``wire[2:7]``. But the
-    :attr:`LogicNet.op_param` are difficult to process efficiently because they specify
-    the input-to-output mapping one bit at a time.
+    Slice :class:`LogicNets<LogicNet>` (:attr:`~LogicNet.op` ``s``) are very frequently
+    used to extract a consecutive range of bits from a :class:`WireVector`, for example
+    ``wire_vector[2:7]``. But the :attr:`LogicNet.op_param` are difficult to process
+    efficiently because they specify the ``arg``-to-``dest`` mapping one bit at a time.
 
     To process these common :class:`LogicNets<LogicNet>` more efficiently, this function
-    compresses contiguous ranges of bits in the input-to-output mapping into
-    :class:`ContiguousSlices<ContiguousSlice>`.
+    compresses consecutive ranges of bits in the ``arg``-to-``dest`` mapping into
+    :class:`ConsecutiveSlices<ConsecutiveSlice>`.
     """
-    if net.op != "s":
-        msg = f"Invalid LogicNet op ({net.op})"
-        raise PyrtlInternalError(msg)
-    if len(net.op_param) == 0:
-        return []
-
-    contiguous_slices = []
+    consecutive_slices = []
 
     length = 1
-    arg_start = net.op_param[0]
+    arg_start = op_param[0]
     dest_start = 0
-    for dest_index, arg_index in enumerate(net.op_param[1:], start=1):
+    for dest_index, arg_index in enumerate(op_param[1:], start=1):
         if arg_index == arg_start + length:
             # `arg_index` continues the current slice.
             length += 1
         else:
-            # `arg_index` ends the current slice because it is discontiguous.
-            contiguous_slices.append(ContiguousSlice(length, arg_start, dest_start))
+            # `arg_index` ends the current slice because it is not consecutive.
+            consecutive_slices.append(ConsecutiveSlice(length, arg_start, dest_start))
             length = 1
             arg_start = arg_index
             dest_start = dest_index
 
-    contiguous_slices.append(ContiguousSlice(length, arg_start, dest_start))
-    return contiguous_slices
+    consecutive_slices.append(ConsecutiveSlice(length, arg_start, dest_start))
+    return consecutive_slices
 
 
 class FastSimulation:
@@ -997,11 +999,6 @@ def _compiled(self) -> str:
             "x": lambda sel, f, t: f"({f}) if ({sel}==0) else ({t})",
         }
 
-        def shift(value, direction, shift_amt):
-            if shift_amt == 0:
-                return value
-            return f"({value} {direction} {shift_amt})"
-
         for net in self.block:
             if net.op in simple_func:
                 argvals = (self._arg_varname(arg) for arg in net.args)
@@ -1016,16 +1013,17 @@ def shift(value, direction, shift_amt):
                 expr = " | ".join(expr_parts)
             elif net.op == "s":
                 arg = self._arg_varname(net.args[0])
-                contiguous_slices = make_contiguous_slices(net)
+                consecutive_slices = make_consecutive_slices(net.op_param)
 
                 expr_parts = []
-                for contiguous_slice in contiguous_slices:
-                    expr = shift(arg, ">>", contiguous_slice.arg_start)
-                    arg_end = contiguous_slice.arg_start + contiguous_slice.length
+                for consecutive_slice in consecutive_slices:
+                    expr = shift(arg, ">>", consecutive_slice.arg_start)
+                    arg_end = consecutive_slice.arg_start + consecutive_slice.length
                     if net.args[0].bitwidth > arg_end:
-                        expr = f"({expr} & {(1 << contiguous_slice.length) - 1})"
+                        # Mask if we're not taking all the arg's remaining bits.
+                        expr = f"({expr} & 0x{(1 << consecutive_slice.length) - 1:X})"
 
-                    expr_parts.append(shift(expr, "<<", contiguous_slice.dest_start))
+                    expr_parts.append(shift(expr, "<<", consecutive_slice.dest_start))
 
                 expr = "|".join(expr_parts)
             elif net.op == "m":

tests/test_compilesim.py

@@ -123,6 +123,42 @@ def test_bitslice2_and_concat_simulation(self):
         self.r.next <<= pyrtl.concat(left, right)
         self.check_trace("o 01377777\n")
 
+    def test_multiple_limb_bitslice_simulation(self):
+        pyrtl.reset_working_block()
+
+        # `big_input`'s storage spans two 64-bit limbs.
+        big_input = pyrtl.Input(name="big_input", bitwidth=128)
+
+        out = pyrtl.Output(name="out", bitwidth=32)
+        # This slice must fetch and combine bits from both of `big_input`'s limbs.
+        out <<= big_input[48:80]
+
+        sim = self.sim()
+        sim.step({"big_input": 0xFFFF_EEEE_DDDD_1234_5678_CCCC_BBBB_AAAA})
+        self.assertEqual(sim.inspect("out"), 0x1234_5678)
+
+    def test_sparse_bitslice_simulation(self):
+        pyrtl.reset_working_block()
+
+        # `bigger_input`'s storage spans three 64-bit limbs.
+        bigger_input = pyrtl.Input(name="bigger_input", bitwidth=160)
+
+        even_bits = pyrtl.Output(name="even_bits", bitwidth=80)
+        even_bits <<= bigger_input[::2]
+
+        odd_bits = pyrtl.Output(name="odd_bits", bitwidth=80)
+        odd_bits <<= bigger_input[1::2]
+
+        sim = self.sim()
+        # 0xA == 0b1010
+        sim.step({"bigger_input": 0xAAAA_AAAA_AAAA_AAAA_AAAA_AAAA_AAAA_AAAA_AAAA_AAAA})
+        self.assertEqual(sim.inspect("even_bits"), 0)
+        self.assertEqual(sim.inspect("odd_bits"), 0xFFFF_FFFF_FFFF_FFFF_FFFF)
+        # 0x5 == 0b0101
+        sim.step({"bigger_input": 0x5555_5555_5555_5555_5555_5555_5555_5555_5555_5555})
+        self.assertEqual(sim.inspect("even_bits"), 0xFFFF_FFFF_FFFF_FFFF_FFFF)
+        self.assertEqual(sim.inspect("odd_bits"), 0)
+
     def test_reg_to_reg_simulation(self):
         self.r2 = pyrtl.Register(bitwidth=self.bitwidth, name="r2")
         self.r.next <<= self.r2