transforms.artiq_ir_generator: emit ir.Parallel for with parallel:.

m-labs · Oct 9, 2015 · 32ce33a · 32ce33a
1 parent 48f1f48
commit 32ce33a
Showing 2 changed files with 50 additions and 1 deletion.
diff --git a/artiq/compiler/ir.py b/artiq/compiler/ir.py
@@ -1000,6 +1000,7 @@ def destinations(self):
         return self.operands[1:]
 
     def add_destination(self, destination):
+        destination.uses.add(self)
         self.operands.append(destination)
 
     def _operands_as_string(self):
@@ -1176,3 +1177,22 @@ def _operands_as_string(self):
             else:
                 table.append("{} => {}".format(types.TypePrinter().name(typ), target.as_operand()))
         return "cleanup {}, [{}]".format(self.cleanup().as_operand(), ", ".join(table))
+
+class Parallel(Terminator):
+    """
+    An instruction that schedules several threads of execution
+    in parallel.
+    """
+
+    def __init__(self, destinations=[], name=""):
+        super().__init__(destinations, builtins.TNone(), name)
+
+    def opcode(self):
+        return "parallel"
+
+    def destinations(self):
+        return self.operands
+
+    def add_destination(self, destination):
+        destination.uses.add(self)
+        self.operands.append(destination)
diff --git a/artiq/compiler/transforms/artiq_ir_generator.py b/artiq/compiler/transforms/artiq_ir_generator.py
@@ -669,7 +669,36 @@ def visit_Try(self, node):
                 if not post_handler.is_terminated():
                     post_handler.append(ir.Branch(tail))
 
-    # TODO: With
+    def visit_With(self, node):
+        if len(node.items) != 1:
+            diag = diagnostic.Diagnostic("fatal",
+                "only one expression per 'with' statement is supported",
+                {"type": types.TypePrinter().name(typ)},
+                node.context_expr.loc)
+            self.engine.process(diag)
+
+        context_expr_node  = node.items[0].context_expr
+        optional_vars_node = node.items[0].optional_vars
+
+        if types.is_builtin(context_expr_node.type, "sequential"):
+            self.visit(node.body)
+        elif types.is_builtin(context_expr_node.type, "parallel"):
+            parallel = self.append(ir.Parallel())
+
+            heads, tails = [], []
+            for stmt in node.body:
+                self.current_block = self.add_block()
+                heads.append(self.current_block)
+                self.visit(stmt)
+                tails.append(self.current_block)
+
+            for head in heads:
+                parallel.add_destination(head)
+
+            self.current_block = self.add_block()
+            for tail in tails:
+                if not tail.is_terminated():
+                    tail.append(ir.Branch(self.current_block))
 
     # Expression visitors
     # These visitors return a node in addition to mutating