Implement dumb 'with parallel' (#265).

m-labs · Feb 22, 2016 · bc81be1 · bc81be1
1 parent 51a5910
commit bc81be1
Showing 10 changed files with 52 additions and 2 deletions.
diff --git a/artiq/compiler/builtins.py b/artiq/compiler/builtins.py
@@ -144,6 +144,9 @@ def fn_print():
 def fn_kernel():
     return types.TBuiltinFunction("kernel")
 
+def obj_parallel():
+    return types.TBuiltin("parallel")
+
 def obj_interleave():
     return types.TBuiltin("interleave")
 

diff --git a/artiq/compiler/prelude.py b/artiq/compiler/prelude.py
@@ -30,6 +30,7 @@ def globals():
         "portable":             builtins.fn_kernel(),
 
         # ARTIQ context managers
+        "parallel":             builtins.obj_parallel(),
         "interleave":           builtins.obj_interleave(),
         "sequential":           builtins.obj_sequential(),
         "watchdog":             builtins.fn_watchdog(),

diff --git a/artiq/compiler/transforms/artiq_ir_generator.py b/artiq/compiler/transforms/artiq_ir_generator.py
@@ -765,6 +765,28 @@ def visit_With(self, node):
                 if not tail.is_terminated():
                     tail.append(ir.Branch(self.current_block))
             return
+        elif types.is_builtin(context_expr_node.type, "parallel"):
+            start_mu = self.append(ir.Builtin("now_mu", [], builtins.TInt64()))
+            end_mu   = start_mu
+
+            for stmt in node.body:
+                self.append(ir.Builtin("at_mu", [start_mu], builtins.TNone()))
+
+                block = self.add_block()
+                if self.current_block.is_terminated():
+                    self.warn_unreachable(stmt[0])
+                else:
+                    self.append(ir.Branch(block))
+                self.current_block = block
+
+                self.visit(stmt)
+
+                mid_mu = self.append(ir.Builtin("now_mu", [], builtins.TInt64()))
+                gt_mu  = self.append(ir.Compare(ast.Gt(loc=None), mid_mu, end_mu))
+                end_mu = self.append(ir.Select(gt_mu, mid_mu, end_mu))
+
+            self.append(ir.Builtin("at_mu", [end_mu], builtins.TNone()))
+            return
 
         cleanup = []
         for item_node in node.items:

diff --git a/artiq/compiler/transforms/inferencer.py b/artiq/compiler/transforms/inferencer.py
@@ -994,6 +994,7 @@ def visit_withitemT(self, node):
 
         typ = node.context_expr.type
         if (types.is_builtin(typ, "interleave") or types.is_builtin(typ, "sequential") or
+            types.is_builtin(typ, "parallel") or
                 (isinstance(node.context_expr, asttyped.CallT) and
                  types.is_builtin(node.context_expr.func.type, "watchdog"))):
             # builtin context managers
@@ -1092,8 +1093,8 @@ def visit_With(self, node):
 
         for item_node in node.items:
             typ = item_node.context_expr.type.find()
-            if (types.is_builtin(typ, "interleave") or types.is_builtin(typ, "sequential")) and \
-                    len(node.items) != 1:
+            if (types.is_builtin(typ, "parallel") or types.is_builtin(typ, "interleave") or
+                types.is_builtin(typ, "sequential")) and len(node.items) != 1:
                 diag = diagnostic.Diagnostic("error",
                     "the '{kind}' context manager must be the only one in a 'with' statement",
                     {"kind": typ.name},

diff --git a/artiq/test/lit/lit.cfg b/artiq/test/lit/lit.cfg
@@ -34,3 +34,4 @@ if os.name == "posix":
     config.environment["LIBARTIQ_SUPPORT"] = support_lib
 
     config.available_features.add("exceptions")
+    config.available_features.add("time")
diff --git a/artiq/test/lit/time/advance.py b/artiq/test/lit/time/advance.py
@@ -1,4 +1,5 @@
 # RUN: %python -m artiq.compiler.testbench.jit %s
+# REQUIRES: time
 
 assert now() == 0.0
 delay(100.0)

diff --git a/artiq/test/lit/time/advance_mu.py b/artiq/test/lit/time/advance_mu.py
@@ -1,4 +1,5 @@
 # RUN: %python -m artiq.compiler.testbench.jit %s
+# REQUIRES: time
 
 assert now_mu() == 0
 delay_mu(100)

diff --git a/artiq/test/lit/time/conversion.py b/artiq/test/lit/time/conversion.py
@@ -1,4 +1,5 @@
 # RUN: %python -m artiq.compiler.testbench.jit %s
+# REQUIRES: time
 
 assert seconds_to_mu(2.0) == 2000000
 assert mu_to_seconds(1500000) == 1.5
diff --git a/artiq/test/lit/time/parallel.py b/artiq/test/lit/time/parallel.py
@@ -0,0 +1,18 @@
+# RUN: %python -m artiq.compiler.testbench.jit %s
+# REQUIRES: time
+
+with parallel:
+    with sequential:
+        assert now_mu() == 0
+        delay_mu(10)
+        assert now_mu() == 10
+    with sequential:
+        assert now_mu() == 0
+        delay_mu(20)
+        assert now_mu() == 20
+    with sequential:
+        assert now_mu() == 0
+        delay_mu(15)
+        assert now_mu() == 15
+    assert now_mu() == 0
+assert now_mu() == 20
diff --git a/artiq/test/lit/time/watchdog.py b/artiq/test/lit/time/watchdog.py
@@ -1,5 +1,6 @@
 # RUN: %python -m artiq.compiler.testbench.jit %s >%t
 # RUN: OutputCheck %s --file-to-check=%t
+# REQUIRES: time
 
 def f():
     with watchdog(1.0):
Original file line number	Diff line number	Diff line change
		@@ -34,3 +34,4 @@ if os.name == "posix":
		config.environment["LIBARTIQ_SUPPORT"] = support_lib

		config.available_features.add("exceptions")
		config.available_features.add("time")