transforms.interleaver: implement (without inlining).

artiq/compiler/ir.py

@@ -953,6 +953,9 @@ class Branch(Terminator):
     def target(self):
         return self.operands[0]
 
+    def set_target(self, new_target):
+        self.operands[0] = new_target
+
 class BranchIf(Terminator):
     """
     A conditional branch instruction.
@@ -1213,6 +1216,9 @@ class Delay(Terminator):
     def target(self):
         return self.operands[0]
 
+    def set_target(self, new_target):
+        self.operands[0] = new_target
+
     def substs(self):
         return zip(self.var_names, self.operands[1:])
 

artiq/compiler/transforms/interleaver.py

@@ -3,7 +3,7 @@
 the timestamp would always monotonically nondecrease.
 """
 
-from .. import ir
+from .. import ir, iodelay
 from ..analyses import domination
 
 class Interleaver:
@@ -15,8 +15,65 @@ class Interleaver:
             self.process_function(func)
 
     def process_function(self, func):
-        domtree = domination.PostDominatorTree(func)
+        for insn in func.instructions():
-        print(func)
+            if isinstance(insn, ir.Delay):
-        for block in func.basic_blocks:
+                if any(insn.expr.free_vars()):
-            idom = domtree.immediate_dominator(block)
+                    # If a function has free variables in delay expressions,
-            print(block.name, "->", idom.name if idom else "<exit>")
+                    # that means its IO delay depends on arguments.
+                    # Do not change such functions in any way so that it will
+                    # be successfully inlined and then removed by DCE.
+                    return
+
+        postdom_tree = None
+        for insn in func.instructions():
+            if isinstance(insn, ir.Parallel):
+                # Lazily compute dominators.
+                if postdom_tree is None:
+                    postdom_tree = domination.PostDominatorTree(func)
+
+                interleave_until = postdom_tree.immediate_dominator(insn.basic_block)
+                assert (interleave_until is not None) # no nonlocal flow in `with parallel`
+
+                target_block  = insn.basic_block
+                target_time   = 0
+                source_blocks = insn.basic_block.successors()
+                source_times  = [0 for _ in source_blocks]
+
+                while len(source_blocks) > 0:
+                    def iodelay_of_block(block):
+                        terminator = block.terminator()
+                        if isinstance(terminator, ir.Delay):
+                            # We should be able to fold everything without free variables.
+                            assert iodelay.is_const(terminator.expr)
+                            return terminator.expr.value
+                        else:
+                            return 0
+
+                    def time_after_block(pair):
+                        index, block = pair
+                        return source_times[index] + iodelay_of_block(block)
+
+                    index, source_block = min(enumerate(source_blocks), key=time_after_block)
+                    source_block_delay  = iodelay_of_block(source_block)
+
+                    target_terminator = target_block.terminator()
+                    if isinstance(target_terminator, (ir.Delay, ir.Branch)):
+                        target_terminator.set_target(source_block)
+                    elif isinstance(target_terminator, ir.Parallel):
+                        target_terminator.replace_with(ir.Branch(source_block))
+                    else:
+                        assert False
+
+                    new_source_block = postdom_tree.immediate_dominator(source_block)
+                    assert (new_source_block is not None)
+
+                    target_block          = source_block
+                    target_time          += source_block_delay
+
+                    if new_source_block == interleave_until:
+                        # We're finished with this branch.
+                        del source_blocks[index]
+                        del source_times[index]
+                    else:
+                        source_blocks[index] = new_source_block
+                        source_times[index]  = target_time

lit-test/test/interleaving/basic.py

@@ -0,0 +1,22 @@
+# RUN: %python -m artiq.compiler.testbench.jit %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+def g():
+    with parallel:
+        with sequential:
+            print("A", now_mu())
+            delay_mu(3)
+            print("B", now_mu())
+        with sequential:
+            print("C", now_mu())
+            delay_mu(2)
+            print("D", now_mu())
+            delay_mu(2)
+            print("E", now_mu())
+
+# CHECK-L: C 0
+# CHECK-L: A 2
+# CHECK-L: D 5
+# CHECK-L: B 7
+# CHECK-L: E 7
+g()