From 00ec574d73e0957a2c96a3a3c137ced8e072b18b Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Fri, 20 Nov 2015 00:03:26 +0800
Subject: [PATCH] transforms.interleaver: implement (without inlining).

---
 artiq/compiler/ir.py                     |  6 +++
 artiq/compiler/transforms/interleaver.py | 69 +++++++++++++++++++++---
 lit-test/test/interleaving/basic.py      | 22 ++++++++
 3 files changed, 91 insertions(+), 6 deletions(-)
 create mode 100644 lit-test/test/interleaving/basic.py

diff --git a/artiq/compiler/ir.py b/artiq/compiler/ir.py
index 60126b534..adae44a89 100644
--- a/artiq/compiler/ir.py
+++ b/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:])
 
diff --git a/artiq/compiler/transforms/interleaver.py b/artiq/compiler/transforms/interleaver.py
index 256c5508d..44dd4fc8c 100644
--- a/artiq/compiler/transforms/interleaver.py
+++ b/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)
-        print(func)
-        for block in func.basic_blocks:
-            idom = domtree.immediate_dominator(block)
-            print(block.name, "->", idom.name if idom else "<exit>")
+        for insn in func.instructions():
+            if isinstance(insn, ir.Delay):
+                if any(insn.expr.free_vars()):
+                    # If a function has free variables in delay expressions,
+                    # 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
diff --git a/lit-test/test/interleaving/basic.py b/lit-test/test/interleaving/basic.py
new file mode 100644
index 000000000..fbc9f2091
--- /dev/null
+++ b/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()