compiler.analyses.domination: fix PostDominatorTree.

artiq/compiler/analyses/domination.py

@@ -10,63 +10,21 @@ class GenericDominatorTree:
         self._assign_names()
         self._compute()
 
-    def _start_blocks(self):
-        """
-        Returns a starting collection of basic blocks (entry block
-        for dominator tree and exit blocks for postdominator tree).
-        """
+    def _traverse_in_postorder(self):
         raise NotImplementedError
 
-    def _next_blocks(self, block):
-        """
-        Returns the collection of blocks to be traversed after `block`
-        (successors for dominator tree and predecessors for postdominator
-        tree).
-        """
-        raise NotImplementedError
-
-    def _prev_blocks(self, block):
-        """
-        Returns the collection of blocks to be traversed before `block`
-        (predecessors for dominator tree and successors for postdominator
-        tree).
-        """
+    def _prev_block_names(self, block):
         raise NotImplementedError
 
     def _assign_names(self):
-        """Assigns names to basic blocks in postorder."""
-        visited = set()
-        postorder = []
+        postorder = self._traverse_in_postorder()
 
-        def visit(block):
-            visited.add(block)
-            for next_block in self._next_blocks(block):
-                if next_block not in visited:
-                    visit(next_block)
-            postorder.append(block)
-
-        for block in self._start_blocks():
-            visit(block)
-
-        self._last_name     = len(postorder)
+        self._start_name    = len(postorder) - 1
         self._block_of_name = postorder
         self._name_of_block = {}
         for block_name, block in enumerate(postorder):
-            # print("name", block_name + 1, block.name)
             self._name_of_block[block] = block_name
 
-    def _start_block_names(self):
-        for block in self._start_blocks():
-            yield self._name_of_block[block]
-
-    def _next_block_names(self, block_name):
-        for block in self._next_blocks(self._block_of_name[block_name]):
-            yield self._name_of_block[block]
-
-    def _prev_block_names(self, block_name):
-        for block in self._prev_blocks(self._block_of_name[block_name]):
-            yield self._name_of_block[block]
-
     def _intersect(self, block_name_1, block_name_2):
         finger_1, finger_2 = block_name_1, block_name_2
         while finger_1 != finger_2:
@@ -79,23 +37,23 @@ class GenericDominatorTree:
     def _compute(self):
         self._doms = {}
 
-        for block_name in range(self._last_name):
-            self._doms[block_name] = None
+        # Start block dominates itself.
+        self._doms[self._start_name] = self._start_name
 
-        start_block_names = set()
-        for block_name in self._start_block_names():
-            self._doms[block_name] = block_name
-            start_block_names.add(block_name)
+        # We don't yet know what blocks dominate all other blocks.
+        for block_name in range(self._start_name):
+            self._doms[block_name] = None
 
         changed = True
         while changed:
-            # print("doms", {k+1: self._doms[k]+1 if self._doms[k] is not None else None for k in self._doms})
-
             changed = False
-            for block_name in reversed(range(self._last_name)):
-                if block_name in start_block_names:
-                    continue
 
+            # For all blocks except start block, in reverse postorder...
+            for block_name in reversed(range(self._start_name)):
+                # Select a new immediate dominator from the blocks we have
+                # already processed, and remember all others.
+                # We've already processed at least one previous block because
+                # of the graph traverse order.
                 new_idom, prev_block_names = None, []
                 for prev_block_name in self._prev_block_names(block_name):
                     if new_idom is None and self._doms[prev_block_name] is not None:
@@ -103,12 +61,10 @@ class GenericDominatorTree:
                     else:
                         prev_block_names.append(prev_block_name)
 
-                # print("block_name", block_name + 1, "new_idom", new_idom + 1)
+                # Find a common previous block
                 for prev_block_name in prev_block_names:
-                    # print("prev_block_name", prev_block_name + 1)
                     if self._doms[prev_block_name] is not None:
                         new_idom = self._intersect(prev_block_name, new_idom)
-                        # print("new_idom+", new_idom + 1)
 
                 if self._doms[block_name] != new_idom:
                     self._doms[block_name] = new_idom
@@ -130,26 +86,52 @@ class DominatorTree(GenericDominatorTree):
         self.function = function
         super().__init__()
 
-    def _start_blocks(self):
-        return [self.function.entry()]
+    def _traverse_in_postorder(self):
+        postorder = []
 
-    def _next_blocks(self, block):
-        return block.successors()
+        visited = set()
+        def visit(block):
+            visited.add(block)
+            for next_block in block.successors():
+                if next_block not in visited:
+                    visit(next_block)
+            postorder.append(block)
 
-    def _prev_blocks(self, block):
-        return block.predecessors()
+        visit(self.function.entry())
+
+        return postorder
+
+    def _prev_block_names(self, block_name):
+        for block in self._block_of_name[block_name].predecessors():
+            yield self._name_of_block[block]
 
 class PostDominatorTree(GenericDominatorTree):
     def __init__(self, function):
         self.function = function
         super().__init__()
 
-    def _start_blocks(self):
-        return [block for block in self.function.basic_blocks
-                       if none(block.successors())]
+    def _traverse_in_postorder(self):
+        postorder = []
 
-    def _next_blocks(self, block):
-        return block.predecessors()
+        visited = set()
+        def visit(block):
+            visited.add(block)
+            for next_block in block.predecessors():
+                if next_block not in visited:
+                    visit(next_block)
+            postorder.append(block)
 
-    def _prev_blocks(self, block):
-        return block.successors()
+        for block in self.function.basic_blocks:
+            if not any(block.successors()):
+                visit(block)
+
+        postorder.append(None) # virtual exit block
+        return postorder
+
+    def _prev_block_names(self, block_name):
+        succ_blocks = self._block_of_name[block_name].successors()
+        if len(succ_blocks) > 0:
+            for block in succ_blocks:
+                yield self._name_of_block[block]
+        else:
+            yield self._start_name

artiq/test/compiler/domination.py

@@ -44,7 +44,8 @@ def dom(function, domtree):
 def idom(function, domtree):
     idom = {}
     for block in function.basic_blocks:
-        idom[block.name] = domtree.immediate_dominator(block).name
+        idom_block = domtree.immediate_dominator(block)
+        idom[block.name] = idom_block.name if idom_block else None
     return idom
 
 class TestDominatorTree(unittest.TestCase):
@@ -113,3 +114,53 @@ class TestDominatorTree(unittest.TestCase):
         self.assertEqual({
             1: 6, 2: 6, 3: 6, 4: 6, 5: 6, 6: 6
         }, idom(func, domtree))
+
+class TestPostDominatorTree(unittest.TestCase):
+    def test_linear(self):
+        func = makefn('A', {
+            'A': ['B'],
+            'B': ['C'],
+            'C': []
+        })
+        domtree = PostDominatorTree(func)
+        self.assertEqual({
+            'A': 'B', 'B': 'C', 'C': None
+        }, idom(func, domtree))
+
+    def test_diamond(self):
+        func = makefn('A', {
+            'A': ['B', 'D'],
+            'B': ['C'],
+            'C': ['E'],
+            'D': ['E'],
+            'E': []
+        })
+        domtree = PostDominatorTree(func)
+        self.assertEqual({
+            'E': None, 'D': 'E', 'C': 'E', 'B': 'C', 'A': 'E'
+        }, idom(func, domtree))
+
+    def test_multi_exit(self):
+        func = makefn('A', {
+            'A': ['B', 'C'],
+            'B': [],
+            'C': []
+        })
+        domtree = PostDominatorTree(func)
+        self.assertEqual({
+            'A': None, 'B': None, 'C': None
+        }, idom(func, domtree))
+
+    def test_multi_exit_diamond(self):
+        func = makefn('A', {
+            'A': ['B', 'C'],
+            'B': ['D'],
+            'C': ['D'],
+            'D': ['E', 'F'],
+            'E': [],
+            'F': []
+        })
+        domtree = PostDominatorTree(func)
+        self.assertEqual({
+            'A': 'D', 'B': 'D', 'C': 'D', 'D': None, 'E': None, 'F': None
+        }, idom(func, domtree))