diff --git a/artiq/compiler/lower_time.py b/artiq/compiler/lower_time.py new file mode 100644 index 000000000..00039dfad --- /dev/null +++ b/artiq/compiler/lower_time.py @@ -0,0 +1,42 @@ +import ast + +class _TimeLowerer(ast.NodeTransformer): + def __init__(self, ref_period): + self.ref_period = ref_period + + def visit_Call(self, node): + if isinstance(node.func, ast.Name) \ + and node.func.id == "Quantity" and node.args[1].id == "base_s_unit": + return ast.copy_location( + ast.BinOp(left=node.args[0], op=ast.FloorDiv(), right=ast.Num(self.ref_period.amount)), + node) + elif isinstance(node.func, ast.Name) and node.func.id == "now": + return ast.copy_location(ast.Name("now", ast.Load()), node) + else: + self.generic_visit(node) + return node + + def visit_Expr(self, node): + self.generic_visit(node) + if isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Name): + funcname = node.value.func.id + if funcname == "delay": + return ast.copy_location( + ast.AugAssign(target=ast.Name("now", ast.Store()), op=ast.Add(), value=node.value.args[0]), + node) + elif funcname == "at": + return ast.copy_location( + ast.Assign(targets=[ast.Name("now", ast.Store())], value=node.value.args[0]), + node) + else: + return node + else: + return node + +def lower_time(stmts, ref_period): + transformer = _TimeLowerer(ref_period) + new_stmts = [transformer.visit(stmt) for stmt in stmts] + new_stmts.insert(0, ast.copy_location( + ast.Assign(targets=[ast.Name("now", ast.Store())], value=ast.Num(0)), + stmts[0])) + stmts[:] = new_stmts diff --git a/artiq/devices/core.py b/artiq/devices/core.py index 5bfa7bfbc..bc6eb4dff 100644 --- a/artiq/devices/core.py +++ b/artiq/devices/core.py @@ -2,6 +2,7 @@ from artiq.compiler.inline import inline from artiq.compiler.fold_constants import fold_constants from artiq.compiler.unroll_loops import unroll_loops from artiq.compiler.interleave import interleave +from artiq.compiler.lower_time import lower_time from artiq.compiler.ir import get_runtime_binary class Core: @@ -16,6 +17,8 @@ class Core: fold_constants(stmts) unroll_loops(stmts, 50) interleave(stmts) + lower_time(stmts, self.runtime_env.ref_period) + fold_constants(stmts) binary = get_runtime_binary(self.runtime_env, stmts) self.core_com.run(binary)