Merge pull request #346 from pllab/docs-and-tests

Adding documentation and some tests for the passes and transforms

Author: mdko

pyrtl/passes.py

@@ -216,8 +216,7 @@ def replace_net_with_wire(new_wire):
 
 
 def common_subexp_elimination(block=None, abs_thresh=1, percent_thresh=0):
-    """
-    Common Subexpression Elimination for PyRTL blocks
+    """ Common Subexpression Elimination for PyRTL blocks.
 
     :param block: the block to run the subexpression elimination on
     :param abs_thresh: absolute threshold for stopping optimization
@@ -235,6 +234,18 @@ def common_subexp_elimination(block=None, abs_thresh=1, percent_thresh=0):
 
 
 def _find_common_subexps(block):
+    """ Finds nets that can be considered the same based on op type, op param, and arguments.
+
+    :param block: Block to operate over
+    :return dict[LogicNet, [LogicNet]]: mapping from a logic net (with a placehold dest)
+        representing the common subexp, to a list of nets matching that common subexp that
+        can be replaced with the single common subexp.
+
+    Nets are the "same" if 1) their op types are the same, 2) their op_params are
+    the same (e.g. same memory if a memory-related op), and 3) their arguments are
+    the same (same constant value and bitwidth for const wires, otherwise same wire
+    object). The destination wire for a net is not considered.
+    """
     net_table = {}  # {net (without dest) : [net, ...]
     t = tuple()  # just a placeholder
     const_dict = {}
@@ -252,6 +263,11 @@ def _find_common_subexps(block):
 
 
 def _const_to_int(wire, const_dict):
+    """ Return a repr a Const (a tuple composed of width and value) for comparison with an 'is'.
+
+    If the wire is not a Const, just return the wire itself; comparison will be
+    done on the identity of the wire object instead.
+    """
     if isinstance(wire, Const):
         # a very bad hack to make sure two consts will compare
         # correctly with an 'is'
@@ -268,6 +284,12 @@ def _const_to_int(wire, const_dict):
 
 
 def _replace_subexps(block, net_table):
+    """ Removes unnecessary nets, connecting the common net's dest wire to unnecessary net's dest.
+
+    :param block: The block to operate over.
+    :param net_table: A mapping from common subexpression (a net) to a list of nets
+        that can be replaced with that common net.
+    """
     wire_map = {}
     unnecessary_nets = []
     for nets in net_table.values():
@@ -282,6 +304,16 @@ def _has_normal_dest_wire(net):
 
 
 def _process_nets_to_discard(nets, wire_map, unnecessary_nets):
+    """ Helper for tracking how a group of related nets should be replaced with a common one.
+
+    :param nets: List of nets that are considered equal and which should
+        be replaced by a single common net.
+    :param wire_map: Dict that will be updated with a mapping from every
+        old destination wire that needs to be removed, to the new destination
+        wire with which it should be replaced.
+    :param unnecessary_nets: List of nets that are to be discarded.
+
+    """
     if len(nets) == 1:
         return  # also deals with nets with no dest wires
     nets_to_consider = list(filter(_has_normal_dest_wire, nets))
@@ -297,7 +329,9 @@ def _process_nets_to_discard(nets, wire_map, unnecessary_nets):
 
 
 def _remove_unlistened_nets(block):
-    """ Removes all nets that are not connected to an output wirevector
+    """ Removes all nets that are not connected to an output wirevector.
+
+    :param block: The block to operate over.
     """
 
     listened_nets = set()
@@ -326,7 +360,12 @@ def add_to_listened(net):
 
 
 def _remove_unused_wires(block, keep_inputs=True):
-    """ Removes all unconnected wires from a block"""
+    """ Removes all unconnected wires from a block's wirevector_set.
+
+    :param block: The block to operate over.
+    :param keep_inputs: If True, retain any Input wires that are not connected
+        to any net.
+    """
     valid_wires = set()
     for logic_net in block.logic:
         valid_wires.update(logic_net.args, logic_net.dests)

pyrtl/transform.py

@@ -25,7 +25,7 @@
 
 
 def net_transform(transform_func, block=None, **kwargs):
-    """ Maps nets to new sets of nets according to a custom function
+    """ Maps nets to new sets of nets according to a custom function.
 
     :param transform_func:
         Function signature: func(orig_net (logicnet)) -> keep_orig_net (bool)
@@ -44,7 +44,7 @@ def net_transform(transform_func, block=None, **kwargs):
 
 
 def all_nets(transform_func):
-    """Decorator that wraps a net transform function"""
+    """ Decorator that wraps a net transform function. """
     @functools.wraps(transform_func)
     def t_res(**kwargs):
         net_transform(transform_func, **kwargs)
@@ -107,10 +107,10 @@ def replace_wire(orig_wire, new_src, new_dst, block=None):
 
 
 def replace_wires(wire_map, block=None):
-    """ Quickly replace all wires in a block
+    """ Quickly replace all wires in a block.
 
-    :param {old_wire: new_wire} wire_map: mapping of old wires to
-        new wires
+    :param {old_wire: new_wire} wire_map: mapping of old wires to new wires
+    :param block: block to operate over (defaults to working block)
     """
     block = working_block(block)
     src_nets, dst_nets = block.net_connections(include_virtual_nodes=False)
@@ -119,6 +119,80 @@ def replace_wires(wire_map, block=None):
 
 
 def replace_wire_fast(orig_wire, new_src, new_dst, src_nets, dst_nets, block=None):
+    """
+    Replace orig_wire with new_src and/or new_dst. The net that orig_wire originates from
+    (its source net) will now feed into new_src as its destination, and the nets that
+    orig_wire went to (its destination nets) will be fed from new_dst as
+    their respective arguments.
+
+    :param WireVector orig_wire: Wire to be replaced
+    :param WireVector new_src: Wire to replace orig_wire, anywhere orig_wire is the
+        destination of a net. Ignored if orig_wire equals new_src.
+    :param WireVector new_dst: Wire to replace orig_wire, anywhere orig_wire is an
+        argument of a net. Ignored if orig_wire equals new_dst.
+    :param {WireVector: LogicNet} src_nets: Maps a wire to the net where it is a dest
+    :param {WireVector: List[LogicNet]} dst_nets: Maps a wire to list of nets where it is an arg
+    :param Block block: The block on which to operate (defaults to working block)
+
+    new_src will now originate from orig_wire's source net (meaning new_src will be that net's
+    destination). new_dst will be now 
+
+    This *updates* the src_nets and dst_nets maps that are passed in, such that:
+
+    ```
+        old_src_net = src_nets[orig_wire]
+        src_nets[new_src] = old_src_net (where old_src_net.dests = (new_src,))
+    ``` 
+    and
+    ```
+        old_dst_nets = dst_nets[orig_wire]
+        dst_nets[new_dst] = [old_dst_net (where old_dst_net.args replaces orig_wire with new_dst) foreach old_dst_net]  # noqa
+    ```
+
+    This also removes and/or adds nets to the block's logic set.
+
+    For example, given the graph on left, `replace_wire_fast(w1, w4, w1, ...)` produces on right:
+
+    ```
+      a b c d                   a b    c d
+      | | | |                   | |    | |
+      net net                   net    net
+        | |                      |      |
+       w1 w2  ==> produces  ==>  w4 w1 w2
+        | |                          | |
+        net                          net
+         |                            |
+         w3                           w3
+    ```
+
+    And given the graph on the left, `replace_wire_fast(w1, w1, w4, ...)` produces on the right:
+    ```
+      a b c d                   a b    c d
+      | | | |                   | |    | |
+      net net                   net    net
+        | |                      |      |
+       w1 w2  ==> produces  ==>  w1 w4 w2
+        | |                          | |
+        net                          net
+         |                            |
+         w3                           w3
+    ```
+
+    Calling `replace_wire_fast(w1, w4, w4, ...)`, then, fully replaces w1 with w3 in both
+    its argument and dest positions:
+
+    ```
+      a b c d                   a b c d
+      | | | |                   | | | |
+      net net                   net net
+        | |                      |   |
+       w1 w2  ==> produces  ==>  w4 w2
+        | |                       | |
+        net                       net
+         |                         |
+         w3                        w3
+    ```
+    """
     def remove_net(net_):
         for arg in set(net_.args):
             dst_nets[arg].remove(net_)

tests/test_transform.py

@@ -197,6 +197,66 @@ def test_wire_used_in_multiple_places(self):
             self.assertNotIn(old_wire, block.wirevector_set)
         block.sanity_check()
 
+    def test_replace_only_src_wire(self):
+        a, b, c, d = pyrtl.input_list('a/1 b/1 c/1 d/1')
+        w1 = a & b
+        w1.name = 'w1'
+        w2 = c | d
+        w2.name = 'w2'
+        w3 = w1 ^ w2
+        w3.name = 'w3'
+        o = pyrtl.Output(1, 'o')
+        o <<= w3
+
+        w4 = pyrtl.WireVector(1, 'w4')
+        src_nets, dst_nets = pyrtl.working_block().net_connections()
+
+        w1_src_net = src_nets[w1]
+        w1_dst_net = dst_nets[w1][0]
+        self.assertEqual(w1_src_net.args, (a, b))
+        self.assertEqual(w1_src_net.dests, (w1,))
+        self.assertEqual(w1_dst_net.args, (w1, w2))
+        self.assertEqual(w1_dst_net.dests, (w3,))
+        self.assertNotIn(w4, src_nets)
+
+        pyrtl.transform.replace_wire_fast(w1, w4, w1, src_nets, dst_nets)
+
+        self.assertNotIn(w1, src_nets)  # The maps have been updated...
+        self.assertEqual(dst_nets[w1], [w1_dst_net])
+        w4_src_net = src_nets[w4]  # ...but the net can't be, so new updated versions were created
+        self.assertEqual(w4_src_net.args, w1_src_net.args)
+        self.assertEqual(w4_src_net.dests, (w4,))
+
+    def test_replace_only_dst_wire(self):
+        a, b, c, d = pyrtl.input_list('a/1 b/1 c/1 d/1')
+        w1 = a & b
+        w1.name = 'w1'
+        w2 = c | d
+        w2.name = 'w2'
+        w3 = w1 ^ w2
+        w3.name = 'w3'
+        o = pyrtl.Output(1, 'o')
+        o <<= w3
+
+        w4 = pyrtl.WireVector(1, 'w4')
+        src_nets, dst_nets = pyrtl.working_block().net_connections()
+
+        w1_src_net = src_nets[w1]
+        w1_dst_net = dst_nets[w1][0]
+        self.assertEqual(w1_src_net.args, (a, b))
+        self.assertEqual(w1_src_net.dests, (w1,))
+        self.assertEqual(w1_dst_net.args, (w1, w2))
+        self.assertEqual(w1_dst_net.dests, (w3,))
+        self.assertNotIn(w4, src_nets)
+
+        pyrtl.transform.replace_wire_fast(w1, w1, w4, src_nets, dst_nets)
+
+        self.assertNotIn(w1, dst_nets)  # The maps have been updated...
+        self.assertEqual(src_nets[w1], w1_src_net)
+        w4_dst_net = dst_nets[w4][0]  # ...but the net can't be, so new versions were created
+        self.assertEqual(w4_dst_net.args, (w4, w2))
+        self.assertEqual(w4_dst_net.dests, w1_dst_net.dests)
+
 
 # this code needs mocking from python 3's unittests to work
 """