m-labs · May 9, 2015 · May 9, 2015
Showing with 83 additions and 11 deletions.

+16 −0 artiq/coredevice/dds.py

+60 −3 artiq/transforms/inline.py

+4 −4 examples/master/repository/dds_test.py

+3 −4 examples/master/repository/photon_histogram.py
diff --git a/artiq/coredevice/dds.py b/artiq/coredevice/dds.py
@@ -10,13 +10,29 @@
 PHASE_MODE_TRACKING = 2
 
 
+class _BatchContextManager:
+    def __init__(self, dds_bus):
+        self.dds_bus = dds_bus
+
+    @kernel
+    def __enter__(self):
+        self.dds_bus.batch_enter()
+
+    @kernel
+    def __exit__(self, type, value, traceback):
+        self.dds_bus.batch_exit()
+
+
 class DDSBus(AutoDB):
     """Core device Direct Digital Synthesis (DDS) bus batching driver.
 
     Manages batching of DDS commands on a DDS shared bus."""
     class DBKeys:
         core = Device()
 
+    def build(self):
+        self.batch = _BatchContextManager(self)
+
     @kernel
     def batch_enter(self):
         """Starts a DDS command batch. All DDS commands are buffered

diff --git a/artiq/transforms/inline.py b/artiq/transforms/inline.py
@@ -3,7 +3,6 @@
 import ast
 import types
 import builtins
-from copy import copy
 from fractions import Fraction
 from collections import OrderedDict
 from functools import partial
@@ -199,8 +198,10 @@ def code_visit(self, node):
 
     # This is ast.NodeTransformer.generic_visit from CPython, modified
     # to update self._insertion_point.
-    def code_generic_visit(self, node):
+    def code_generic_visit(self, node, exclude_fields=set()):
         for field, old_value in ast.iter_fields(node):
+            if field in exclude_fields:
+                continue
             old_value = getattr(node, field, None)
             if isinstance(old_value, list):
                 prev_insertion_point = self._insertion_point
@@ -378,6 +379,60 @@ def code_visit_ExceptHandler(self, node):
         self.code_generic_visit(node)
         return node
 
+    def get_user_ctxm(self, context_expr):
+        try:
+            ctxm = self.static_visit(context_expr)
+        except:
+            # this also catches watchdog()
+            return None
+        else:
+            if (ctxm is core_language.sequential
+                    or ctxm is core_language.parallel):
+                return None
+            return ctxm
+
+    def code_visit_With(self, node):
+        if len(node.items) != 1:
+            raise NotImplementedError
+        item = node.items[0]
+        if item.optional_vars is not None:
+            raise NotImplementedError
+        ctxm = self.get_user_ctxm(item.context_expr)
+        if ctxm is None:
+            self.code_generic_visit(node)
+            return node
+
+        # user context manager
+        self.code_generic_visit(node, {"items"})
+        if (not hasattr(ctxm, "__enter__")
+                or not hasattr(ctxm.__enter__, "k_function_info")):
+            raise NotImplementedError
+        enter = get_inline(self.core,
+                           self.attribute_namespace, self.in_use_names,
+                           None, self.mappers,
+                           ctxm.__enter__.k_function_info.k_function,
+                           [ctxm], dict())
+        if (not hasattr(ctxm, "__exit__")
+                or not hasattr(ctxm.__exit__, "k_function_info")):
+            raise NotImplementedError
+        exit = get_inline(self.core,
+                          self.attribute_namespace, self.in_use_names,
+                          None, self.mappers,
+                          ctxm.__exit__.k_function_info.k_function,
+                          [ctxm, None, None, None], dict())
+        try_stmt = ast.copy_location(
+            ast.Try(body=node.body,
+                    handlers=[],
+                    orelse=[],
+                    finalbody=exit.body), node)
+        return ast.copy_location(
+            ast.With(
+                items=[ast.withitem(context_expr=ast.Name(id="sequential",
+                                                          ctx=ast.Load()),
+                                    optional_vars=None)],
+                body=enter.body + [try_stmt]),
+            node)
+
     def code_visit_FunctionDef(self, node):
         node.args = ast.arguments(args=[], vararg=None, kwonlyargs=[],
                                   kw_defaults=[], kwarg=None, defaults=[])
@@ -470,7 +525,9 @@ def get_attr_writeback(attribute_namespace, rpc_mapper, loc_node):
 def inline(core, k_function, k_args, k_kwargs, with_attr_writeback):
     # OrderedDict prevents non-determinism in attribute init
     attribute_namespace = OrderedDict()
-    in_use_names = copy(embeddable_func_names)
+    # NOTE: in_use_names will be mutated. Do not mutate embeddable_func_names!
+    in_use_names = embeddable_func_names | {"sequential", "parallel",
+                                            "watchdog"}
     mappers = types.SimpleNamespace(
         rpc=HostObjectMapper(),
         exception=HostObjectMapper(core_language.first_user_eid)

diff --git a/examples/master/repository/dds_test.py b/examples/master/repository/dds_test.py
@@ -17,10 +17,10 @@ class DBKeys:
 
     @kernel
     def run(self):
-        self.dds_bus.batch_enter()
-        self.dds1.set(120*MHz)
-        self.dds2.set(200*MHz)
-        self.dds_bus.batch_exit()
+        with self.dds_bus.batch:
+            self.dds1.set(120*MHz)
+            self.dds2.set(200*MHz)
+        delay(1*us)
 
         for i in range(10000):
             if i & 0x200:

diff --git a/examples/master/repository/photon_histogram.py b/examples/master/repository/photon_histogram.py
@@ -27,10 +27,9 @@ class DBKeys:
 
     @kernel
     def program_cooling(self):
-        self.dds_bus.batch_enter()
-        self.bd_dds.set(200*MHz)
-        self.bdd_dds.set(300*MHz)
-        self.dds_bus.batch_exit()
+        with self.dds_bus.batch:
+            self.bd_dds.set(200*MHz)
+            self.bdd_dds.set(300*MHz)
 
     @kernel
     def cool_detect(self):