From b971cc8cdf55bc8e612bbebe7bbbff4d872cad1a Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Wed, 2 Sep 2015 17:46:09 -0600
Subject: [PATCH] compiler.{iodelay,transforms.iodelay_estimator}: implement.

---
 artiq/compiler/builtins.py                    |   6 +
 artiq/compiler/iodelay.py                     | 285 ++++++++++++++++++
 artiq/compiler/module.py                      |   2 +
 artiq/compiler/prelude.py                     |   4 +
 artiq/compiler/testbench/signature.py         |  17 +-
 artiq/compiler/transforms/__init__.py         |   1 +
 artiq/compiler/transforms/inferencer.py       |  10 +-
 .../compiler/transforms/iodelay_estimator.py  | 252 ++++++++++++++++
 artiq/compiler/types.py                       |  23 +-
 lit-test/test/inferencer/with.py              |   5 +
 lit-test/test/iodelay/argument.py             |   7 +
 lit-test/test/iodelay/arith.py                |   8 +
 lit-test/test/iodelay/call.py                 |  10 +
 lit-test/test/iodelay/call_subst.py           |  13 +
 lit-test/test/iodelay/error_arith.py          |  21 ++
 lit-test/test/iodelay/error_builtinfn.py      |  14 +
 lit-test/test/iodelay/error_call_nested.py    |  11 +
 lit-test/test/iodelay/error_call_subst.py     |  13 +
 lit-test/test/iodelay/error_control_flow.py   |  23 ++
 lit-test/test/iodelay/error_for.py            |   9 +
 lit-test/test/iodelay/error_goto.py           |  14 +
 lit-test/test/iodelay/error_return.py         |   8 +
 lit-test/test/iodelay/goto.py                 |  14 +
 lit-test/test/iodelay/linear.py               |  12 +
 lit-test/test/iodelay/loop.py                 |  13 +
 lit-test/test/iodelay/parallel.py             |  15 +
 lit-test/test/iodelay/range.py                |  21 ++
 lit-test/test/iodelay/return.py               |  16 +
 28 files changed, 839 insertions(+), 8 deletions(-)
 create mode 100644 artiq/compiler/iodelay.py
 create mode 100644 artiq/compiler/transforms/iodelay_estimator.py
 create mode 100644 lit-test/test/inferencer/with.py
 create mode 100644 lit-test/test/iodelay/argument.py
 create mode 100644 lit-test/test/iodelay/arith.py
 create mode 100644 lit-test/test/iodelay/call.py
 create mode 100644 lit-test/test/iodelay/call_subst.py
 create mode 100644 lit-test/test/iodelay/error_arith.py
 create mode 100644 lit-test/test/iodelay/error_builtinfn.py
 create mode 100644 lit-test/test/iodelay/error_call_nested.py
 create mode 100644 lit-test/test/iodelay/error_call_subst.py
 create mode 100644 lit-test/test/iodelay/error_control_flow.py
 create mode 100644 lit-test/test/iodelay/error_for.py
 create mode 100644 lit-test/test/iodelay/error_goto.py
 create mode 100644 lit-test/test/iodelay/error_return.py
 create mode 100644 lit-test/test/iodelay/goto.py
 create mode 100644 lit-test/test/iodelay/linear.py
 create mode 100644 lit-test/test/iodelay/loop.py
 create mode 100644 lit-test/test/iodelay/parallel.py
 create mode 100644 lit-test/test/iodelay/range.py
 create mode 100644 lit-test/test/iodelay/return.py

diff --git a/artiq/compiler/builtins.py b/artiq/compiler/builtins.py
index 4133015cc..e4a0c7257 100644
--- a/artiq/compiler/builtins.py
+++ b/artiq/compiler/builtins.py
@@ -138,6 +138,12 @@ def fn_print():
 def fn_kernel():
     return types.TBuiltinFunction("kernel")
 
+def fn_parallel():
+    return types.TBuiltinFunction("parallel")
+
+def fn_sequential():
+    return types.TBuiltinFunction("sequential")
+
 def fn_now():
     return types.TBuiltinFunction("now")
 
diff --git a/artiq/compiler/iodelay.py b/artiq/compiler/iodelay.py
new file mode 100644
index 000000000..5cbb3bb37
--- /dev/null
+++ b/artiq/compiler/iodelay.py
@@ -0,0 +1,285 @@
+"""
+The :mod:`iodelay` module contains the classes describing
+the statically inferred RTIO delay arising from executing
+a function.
+"""
+
+from pythonparser import diagnostic
+
+
+class Expr:
+    def __add__(lhs, rhs):
+        assert isinstance(rhs, Expr)
+        return Add(lhs, rhs)
+    __iadd__ = __add__
+
+    def __sub__(lhs, rhs):
+        assert isinstance(rhs, Expr)
+        return Sub(lhs, rhs)
+    __isub__ = __sub__
+
+    def __mul__(lhs, rhs):
+        assert isinstance(rhs, Expr)
+        return Mul(lhs, rhs)
+    __imul__ = __mul__
+
+    def __truediv__(lhs, rhs):
+        assert isinstance(rhs, Expr)
+        return TrueDiv(lhs, rhs)
+    __itruediv__ = __truediv__
+
+    def __floordiv__(lhs, rhs):
+        assert isinstance(rhs, Expr)
+        return FloorDiv(lhs, rhs)
+    __ifloordiv__ = __floordiv__
+
+    def __ne__(lhs, rhs):
+        return not (lhs == rhs)
+
+    def free_vars(self):
+        return set()
+
+    def fold(self, vars=None):
+        return self
+
+class Const(Expr):
+    _priority = 1
+
+    def __init__(self, value):
+        assert isinstance(value, (int, float))
+        self.value = value
+
+    def __str__(self):
+        return str(self.value)
+
+    def __eq__(lhs, rhs):
+        return rhs.__class__ == lhs.__class__ and lhs.value == rhs.value
+
+    def eval(self, env):
+        return self.value
+
+class Var(Expr):
+    _priority = 1
+
+    def __init__(self, name):
+        assert isinstance(name, str)
+        self.name = name
+
+    def __str__(self):
+        return self.name
+
+    def __eq__(lhs, rhs):
+        return rhs.__class__ == lhs.__class__ and lhs.value == rhs.value
+
+    def free_vars(self):
+        return {self.name}
+
+    def fold(self, vars=None):
+        if vars is not None and self.name in vars:
+            return vars[self.name]
+        else:
+            return self
+
+class Conv(Expr):
+    _priority = 1
+
+    def __init__(self, operand, ref_period):
+        assert isinstance(operand, Expr)
+        assert isinstance(ref_period, float)
+        self.operand, self.ref_period = operand, ref_period
+
+    def free_vars(self):
+        return self.operand.free_vars()
+
+    def fold(self, vars=None):
+        return self.__class__(self.operand.fold(vars),
+                              ref_period=self.ref_period)
+
+class MUToS(Conv):
+    def __str__(self):
+        return "mu->s({})".format(self.operand)
+
+    def eval(self, env):
+        return self.operand.eval(env) * self.ref_period
+
+class SToMU(Conv):
+    def __str__(self):
+        return "s->mu({})".format(self.operand)
+
+    def eval(self, env):
+        return self.operand.eval(env) / self.ref_period
+
+class BinOp(Expr):
+    def __init__(self, lhs, rhs):
+        self.lhs, self.rhs = lhs, rhs
+
+    def __str__(self):
+        lhs = "({})".format(self.lhs) if self.lhs._priority > self._priority else str(self.lhs)
+        rhs = "({})".format(self.rhs) if self.rhs._priority > self._priority else str(self.rhs)
+        return "{} {} {}".format(lhs, self._symbol, rhs)
+
+    def __eq__(lhs, rhs):
+        return rhs.__class__ == lhs.__class__ and lhs.lhs == rhs.lhs and rhs.lhs == rhs.rhs
+
+    def eval(self, env):
+        return self.__class__._op(self.lhs.eval(env), self.rhs.eval(env))
+
+    def free_vars(self):
+        return self.lhs.free_vars() | self.rhs.free_vars()
+
+    def _fold_binop(self, lhs, rhs):
+        if isinstance(lhs, Const) and lhs.__class__ == rhs.__class__:
+            return Const(self.__class__._op(lhs.value, rhs.value))
+        elif isinstance(lhs, (MUToS, SToMU)) and lhs.__class__ == rhs.__class__:
+            return lhs.__class__(self.__class__(lhs.operand, rhs.operand),
+                                 ref_period=lhs.ref_period).fold()
+        else:
+            return self.__class__(lhs, rhs)
+
+    def fold(self, vars=None):
+        return self._fold_binop(self.lhs.fold(vars), self.rhs.fold(vars))
+
+class BinOpFixpoint(BinOp):
+    def _fold_binop(self, lhs, rhs):
+        if isinstance(lhs, Const) and lhs.value == self._fixpoint:
+            return rhs
+        elif isinstance(rhs, Const) and rhs.value == self._fixpoint:
+            return lhs
+        else:
+            return super()._fold_binop(lhs, rhs)
+
+class Add(BinOpFixpoint):
+    _priority = 2
+    _symbol   = "+"
+    _op       = lambda a, b: a + b
+    _fixpoint = 0
+
+class Sub(BinOpFixpoint):
+    _priority = 2
+    _symbol   = "-"
+    _op       = lambda a, b: a - b
+    _fixpoint = 0
+
+class Mul(BinOpFixpoint):
+    _priority = 1
+    _symbol   = "*"
+    _op       = lambda a, b: a * b
+    _fixpoint = 1
+
+class Div(BinOp):
+    def _fold_binop(self, lhs, rhs):
+        if isinstance(rhs, Const) and rhs.value == 1:
+            return lhs
+        else:
+            return super()._fold_binop(lhs, rhs)
+
+class TrueDiv(Div):
+    _priority = 1
+    _symbol   = "/"
+    _op       = lambda a, b: a / b if b != 0 else 0
+
+class FloorDiv(Div):
+    _priority = 1
+    _symbol   = "//"
+    _op       = lambda a, b: a // b if b != 0 else 0
+
+class Max(Expr):
+    _priority = 1
+
+    def __init__(self, operands):
+        assert isinstance(operands, list)
+        assert all([isinstance(operand, Expr) for operand in operands])
+        self.operands = operands
+
+    def __str__(self):
+        return "max({})".format(", ".join([str(operand) for operand in self.operands]))
+
+    def __eq__(lhs, rhs):
+        return rhs.__class__ == lhs.__class__ and lhs.operands == rhs.operands
+
+    def free_vars(self):
+        return reduce(lambda a, b: a | b, [operand.free_vars() for operand in self.operands])
+
+    def eval(self, env):
+        return max([operand.eval() for operand in self.operands])
+
+    def fold(self, vars=None):
+        consts, exprs = [], []
+        for operand in self.operands:
+            operand = operand.fold(vars)
+            if isinstance(operand, Const):
+                consts.append(operand.value)
+            else:
+                exprs.append(operand)
+        if any(consts):
+            exprs.append(Const(max(consts)))
+        if len(exprs) == 1:
+            return exprs[0]
+        else:
+            return Max(exprs)
+
+def is_const(expr, value=None):
+    expr = expr.fold()
+    if value is None:
+        return isinstance(expr, Const)
+    else:
+        return isinstance(expr, Const) and expr.value == value
+
+def is_zero(expr):
+    return is_const(expr, 0)
+
+
+class Delay:
+    pass
+
+class Unknown(Delay):
+    """
+    Unknown delay, that is, IO delay that we have not
+    tried to compute yet.
+    """
+
+    def __repr__(self):
+        return "{}.Unknown()".format(__name__)
+
+class Indeterminate(Delay):
+    """
+    Indeterminate delay, that is, IO delay that can vary from
+    invocation to invocation.
+
+    :ivar cause: (:class:`pythonparser.diagnostic.Diagnostic`)
+        reason for the delay not being inferred
+    """
+
+    def __init__(self, cause):
+        assert isinstance(cause, diagnostic.Diagnostic)
+        self.cause = cause
+
+    def __repr__(self):
+        return "<{}.Indeterminate>".format(__name__)
+
+class Fixed(Delay):
+    """
+    Fixed delay, that is, IO delay that is always the same
+    for every invocation.
+
+    :ivar length: (int) delay in machine units
+    """
+
+    def __init__(self, length):
+        assert isinstance(length, Expr)
+        self.length = length
+
+    def __repr__(self):
+        return "{}.Fixed({})".format(__name__, self.length)
+
+def is_unknown(delay):
+    return isinstance(delay, Unknown)
+
+def is_indeterminate(delay):
+    return isinstance(delay, Indeterminate)
+
+def is_fixed(delay, length=None):
+    if length is None:
+        return isinstance(delay, Fixed)
+    else:
+        return isinstance(delay, Fixed) and delay.length == length
diff --git a/artiq/compiler/module.py b/artiq/compiler/module.py
index 1e4232694..8c320f8c7 100644
--- a/artiq/compiler/module.py
+++ b/artiq/compiler/module.py
@@ -50,6 +50,7 @@ class Module:
         inferencer = transforms.Inferencer(engine=self.engine)
         monomorphism_validator = validators.MonomorphismValidator(engine=self.engine)
         escape_validator = validators.EscapeValidator(engine=self.engine)
+        iodelay_estimator = transforms.IODelayEstimator(ref_period=ref_period)
         artiq_ir_generator = transforms.ARTIQIRGenerator(engine=self.engine,
                                                          module_name=src.name,
                                                          ref_period=ref_period)
@@ -62,6 +63,7 @@ class Module:
         inferencer.visit(src.typedtree)
         monomorphism_validator.visit(src.typedtree)
         escape_validator.visit(src.typedtree)
+        iodelay_estimator.visit_fixpoint(src.typedtree)
         self.artiq_ir = artiq_ir_generator.visit(src.typedtree)
         dead_code_eliminator.process(self.artiq_ir)
         local_access_validator.process(self.artiq_ir)
diff --git a/artiq/compiler/prelude.py b/artiq/compiler/prelude.py
index 2929c9c7a..26870cb52 100644
--- a/artiq/compiler/prelude.py
+++ b/artiq/compiler/prelude.py
@@ -28,6 +28,10 @@ def globals():
         # ARTIQ decorators
         "kernel":               builtins.fn_kernel(),
 
+        # ARTIQ context managers
+        "parallel":             builtins.fn_parallel(),
+        "sequential":           builtins.fn_sequential(),
+
         # ARTIQ time management functions
         "now":                  builtins.fn_now(),
         "delay":                builtins.fn_delay(),
diff --git a/artiq/compiler/testbench/signature.py b/artiq/compiler/testbench/signature.py
index 774259540..d0139a32f 100644
--- a/artiq/compiler/testbench/signature.py
+++ b/artiq/compiler/testbench/signature.py
@@ -1,6 +1,6 @@
 import sys, fileinput
 from pythonparser import diagnostic
-from .. import Module, Source
+from .. import types, iodelay, Module, Source
 
 def main():
     if len(sys.argv) > 1 and sys.argv[1] == "+diag":
@@ -13,15 +13,28 @@ def main():
     else:
         diag = False
         def process_diagnostic(diag):
-            print("\n".join(diag.render()))
+            print("\n".join(diag.render(colored=True)))
             if diag.level in ("fatal", "error"):
                 exit(1)
 
+    if len(sys.argv) > 1 and sys.argv[1] == "+delay":
+        del sys.argv[1]
+        force_delays = True
+    else:
+        force_delays = False
+
     engine = diagnostic.Engine()
     engine.process = process_diagnostic
 
     try:
         mod = Module(Source.from_string("".join(fileinput.input()).expandtabs(), engine=engine))
+
+        if force_delays:
+            for var in mod.globals:
+                typ = mod.globals[var].find()
+                if types.is_function(typ) and iodelay.is_indeterminate(typ.delay):
+                    process_diagnostic(typ.delay.cause)
+
         print(repr(mod))
     except:
         if not diag: raise
diff --git a/artiq/compiler/transforms/__init__.py b/artiq/compiler/transforms/__init__.py
index 576ec7e10..39831fea3 100644
--- a/artiq/compiler/transforms/__init__.py
+++ b/artiq/compiler/transforms/__init__.py
@@ -1,6 +1,7 @@
 from .asttyped_rewriter import ASTTypedRewriter
 from .inferencer import Inferencer
 from .int_monomorphizer import IntMonomorphizer
+from .iodelay_estimator import IODelayEstimator
 from .artiq_ir_generator import ARTIQIRGenerator
 from .dead_code_eliminator import DeadCodeEliminator
 from .llvm_ir_generator import LLVMIRGenerator
diff --git a/artiq/compiler/transforms/inferencer.py b/artiq/compiler/transforms/inferencer.py
index d82c32cd0..8cc6dc53a 100644
--- a/artiq/compiler/transforms/inferencer.py
+++ b/artiq/compiler/transforms/inferencer.py
@@ -945,13 +945,19 @@ class Inferencer(algorithm.Visitor):
 
     def visit_withitem(self, node):
         self.generic_visit(node)
-        if True: # none are supported yet
+
+        typ = node.context_expr.type
+        if not types.is_builtin(typ, "parallel") or types.is_builtin(typ, "sequential"):
             diag = diagnostic.Diagnostic("error",
                 "value of type {type} cannot act as a context manager",
-                {"type": types.TypePrinter().name(node.context_expr.type)},
+                {"type": types.TypePrinter().name(typ)},
                 node.context_expr.loc)
             self.engine.process(diag)
 
+        if node.optional_vars is not None:
+            self._unify(node.optional_vars.type, node.context_expr.type,
+                        node.optional_vars.loc, node.context_expr.loc)
+
     def visit_ExceptHandlerT(self, node):
         self.generic_visit(node)
 
diff --git a/artiq/compiler/transforms/iodelay_estimator.py b/artiq/compiler/transforms/iodelay_estimator.py
new file mode 100644
index 000000000..eb203c5dd
--- /dev/null
+++ b/artiq/compiler/transforms/iodelay_estimator.py
@@ -0,0 +1,252 @@
+"""
+:class:`IODelayEstimator` calculates the amount of time
+elapsed from the point of view of the RTIO core for
+every function.
+"""
+
+from pythonparser import ast, algorithm, diagnostic
+from .. import types, iodelay, builtins, asttyped
+
+class _UnknownDelay(Exception):
+    pass
+
+class IODelayEstimator(algorithm.Visitor):
+    def __init__(self, ref_period):
+        self.ref_period     = ref_period
+        self.changed        = False
+        self.current_delay  = iodelay.Const(0)
+        self.current_args   = None
+        self.current_goto   = None
+        self.current_return = None
+
+    def evaluate(self, node, abort):
+        if isinstance(node, asttyped.NumT):
+            return iodelay.Const(node.n)
+        elif isinstance(node, asttyped.CoerceT):
+            return self.evaluate(node.value, abort)
+        elif isinstance(node, asttyped.NameT):
+            if self.current_args is None:
+                note = diagnostic.Diagnostic("note",
+                    "this variable is not an argument", {},
+                    node.loc)
+                abort([note])
+            elif node.id in [arg.arg for arg in self.current_args.args]:
+                return iodelay.Var(node.id)
+            else:
+                notes = [
+                    diagnostic.Diagnostic("note",
+                        "this variable is not an argument of the innermost function", {},
+                        node.loc),
+                    diagnostic.Diagnostic("note",
+                        "only these arguments are in scope of analysis", {},
+                        self.current_args.loc)
+                ]
+                abort(notes)
+        elif isinstance(node, asttyped.BinOpT):
+            lhs = self.evaluate(node.left, abort)
+            rhs = self.evaluate(node.right, abort)
+            if isinstance(node.op, ast.Add):
+                return lhs + rhs
+            elif isinstance(node.op, ast.Sub):
+                return lhs - rhs
+            elif isinstance(node.op, ast.Mult):
+                return lhs * rhs
+            elif isinstance(node.op, ast.Div):
+                return lhs / rhs
+            elif isinstance(node.op, ast.FloorDiv):
+                return lhs // rhs
+            else:
+                note = diagnostic.Diagnostic("note",
+                    "this operator is not supported", {},
+                    node.op.loc)
+                abort([note])
+        else:
+            note = diagnostic.Diagnostic("note",
+                "this expression is not supported", {},
+                node.loc)
+            abort([note])
+
+    def abort(self, message, loc, notes=[]):
+        diag = diagnostic.Diagnostic("error", message, {}, loc, notes=notes)
+        raise diagnostic.Error(diag)
+
+    def visit_fixpoint(self, node):
+        while True:
+            self.changed = False
+            self.visit(node)
+            if not self.changed:
+                return
+
+    def visit_ModuleT(self, node):
+        try:
+            self.visit(node.body)
+        except (diagnostic.Error, _UnknownDelay):
+            pass # we don't care; module-level code is never interleaved
+
+    def visit_function(self, args, body, typ):
+        old_args, self.current_args = self.current_args, args
+        old_return, self.current_return = self.current_return, None
+        old_delay, self.current_delay = self.current_delay, iodelay.Const(0)
+        try:
+            self.visit(body)
+            if not iodelay.is_zero(self.current_delay) and self.current_return is not None:
+                self.abort("only return statement at the end of the function "
+                           "can be interleaved", self.current_return.loc)
+
+            delay = iodelay.Fixed(self.current_delay.fold())
+        except diagnostic.Error as error:
+            delay = iodelay.Indeterminate(error.diagnostic)
+        self.current_delay = old_delay
+        self.current_return = old_return
+        self.current_args = old_args
+
+        if iodelay.is_unknown(typ.delay) or iodelay.is_indeterminate(typ.delay):
+            typ.delay = delay
+        elif iodelay.is_fixed(typ.delay):
+            assert typ.delay.value == delay.value
+        else:
+            assert False
+
+    def visit_FunctionDefT(self, node):
+        self.visit(node.args.defaults)
+        self.visit(node.args.kw_defaults)
+
+        # We can only handle return in tail position.
+        if isinstance(node.body[-1], ast.Return):
+            body = node.body[:-1]
+        else:
+            body = node.body
+        self.visit_function(node.args, body, node.signature_type.find())
+
+    def visit_LambdaT(self, node):
+        self.visit_function(node.args, node.body, node.type.find())
+
+    def get_iterable_length(self, node):
+        def abort(notes):
+            self.abort("for statement cannot be interleaved because "
+                       "trip count is indeterminate",
+                       node.loc, notes)
+
+        def evaluate(node):
+            return self.evaluate(node, abort)
+
+        if isinstance(node, asttyped.CallT) and types.is_builtin(node.func.type, "range"):
+            range_min, range_max, range_step = iodelay.Const(0), None, iodelay.Const(1)
+            if len(node.args) == 3:
+                range_min, range_max, range_step = map(evaluate, node.args)
+            elif len(node.args) == 2:
+                range_min, range_max = map(evaluate, node.args)
+            elif len(node.args) == 1:
+                range_max, = map(evaluate, node.args)
+            return (range_max - range_min) // range_step
+        else:
+            note = diagnostic.Diagnostic("note",
+                "this value is not a constant range literal", {},
+                node.loc)
+            abort([note])
+
+    def visit_For(self, node):
+        self.visit(node.iter)
+
+        old_goto, self.current_goto = self.current_goto, None
+        old_delay, self.current_delay = self.current_delay, iodelay.Const(0)
+        self.visit(node.body)
+        if iodelay.is_zero(self.current_delay):
+            self.current_delay = old_delay
+        else:
+            if self.current_goto is not None:
+                self.abort("loop trip count is indeterminate because of control flow",
+                           self.current_goto.loc)
+
+            trip_count = self.get_iterable_length(node.iter)
+            self.current_delay = old_delay + self.current_delay * trip_count
+        self.current_goto = old_goto
+
+        self.visit(node.orelse)
+
+    def visit_goto(self, node):
+        self.current_goto = node
+
+    visit_Break    = visit_goto
+    visit_Continue = visit_goto
+
+    def visit_control_flow(self, kind, node):
+        old_delay, self.current_delay = self.current_delay, iodelay.Const(0)
+        self.generic_visit(node)
+        if not iodelay.is_zero(self.current_delay):
+            self.abort("{} cannot be interleaved".format(kind), node.loc)
+        self.current_delay = old_delay
+
+    visit_While  = lambda self, node: self.visit_control_flow("while statement", node)
+    visit_If     = lambda self, node: self.visit_control_flow("if statement",    node)
+    visit_IfExpT = lambda self, node: self.visit_control_flow("if expression",   node)
+    visit_Try    = lambda self, node: self.visit_control_flow("try statement",   node)
+
+    def visit_Return(self, node):
+        self.current_return = node
+
+    def visit_With(self, node):
+        self.visit(node.items)
+
+        context_expr = node.items[0].context_expr
+        if len(node.items) == 1 and types.is_builtin(context_expr.type, "parallel"):
+            delays = []
+            for stmt in node.body:
+                old_delay, self.current_delay = self.current_delay, iodelay.Const(0)
+                self.visit(stmt)
+                delays.append(self.current_delay)
+                self.current_delay = old_delay
+
+            self.current_delay += iodelay.Max(delays)
+        elif len(node.items) == 1 and types.is_builtin(context_expr.type, "serial"):
+            self.visit(node.body)
+        else:
+            self.abort("with statement cannot be interleaved", node.loc)
+
+    def visit_CallT(self, node):
+        typ = node.func.type.find()
+        def abort(notes):
+            self.abort("this call cannot be interleaved because "
+                       "an argument cannot be statically evaluated",
+                       node.loc, notes)
+
+        if types.is_builtin(typ, "delay"):
+            value = self.evaluate(node.args[0], abort=abort)
+            self.current_delay += iodelay.SToMU(value, ref_period=self.ref_period)
+        elif types.is_builtin(typ, "delay_mu"):
+            value = self.evaluate(node.args[0], abort=abort)
+            self.current_delay += value
+        elif not types.is_builtin(typ):
+            if types.is_function(typ):
+                offset = 0
+            elif types.is_method(typ):
+                offset = 1
+                typ = types.get_method_function(typ)
+            else:
+                assert False
+
+            delay = typ.find().delay
+            if iodelay.is_unknown(delay):
+                raise _UnknownDelay()
+            elif iodelay.is_indeterminate(delay):
+                cause = delay.cause
+                note = diagnostic.Diagnostic("note",
+                    "function called here", {},
+                    node.loc)
+                diag = diagnostic.Diagnostic(cause.level, cause.reason, cause.arguments,
+                                             cause.location, cause.highlights, cause.notes + [note])
+                raise diagnostic.Error(diag)
+            elif iodelay.is_fixed(delay):
+                args = {}
+                for kw_node in node.keywords:
+                    args[kw_node.arg] = kw_node.value
+                for arg_name, arg_node in zip(typ.args, node.args[offset:]):
+                    args[arg_name] = arg_node
+
+                print(args)
+
+                free_vars = delay.length.free_vars()
+                self.current_delay += delay.length.fold(
+                    { arg: self.evaluate(args[arg], abort=abort) for arg in free_vars })
+            else:
+                assert False
diff --git a/artiq/compiler/types.py b/artiq/compiler/types.py
index 04e736d65..ea2d1d875 100644
--- a/artiq/compiler/types.py
+++ b/artiq/compiler/types.py
@@ -5,6 +5,7 @@ in :mod:`asttyped`.
 
 import string
 from collections import OrderedDict
+from . import iodelay
 
 
 class UnificationError(Exception):
@@ -191,6 +192,8 @@ class TFunction(Type):
         optional arguments
     :ivar ret: (:class:`Type`)
         return type
+    :ivar delay: (:class:`iodelay.Delay`)
+        RTIO delay expression
     """
 
     attributes = OrderedDict([
@@ -198,11 +201,12 @@ class TFunction(Type):
         ('__closure__', _TPointer()),
     ])
 
-    def __init__(self, args, optargs, ret):
+    def __init__(self, args, optargs, ret, delay=iodelay.Unknown()):
         assert isinstance(args, OrderedDict)
         assert isinstance(optargs, OrderedDict)
         assert isinstance(ret, Type)
-        self.args, self.optargs, self.ret = args, optargs, ret
+        assert isinstance(delay, iodelay.Delay)
+        self.args, self.optargs, self.ret, self.delay = args, optargs, ret, delay
 
     def arity(self):
         return len(self.args) + len(self.optargs)
@@ -256,7 +260,8 @@ class TRPCFunction(TFunction):
     attributes = OrderedDict()
 
     def __init__(self, args, optargs, ret, service):
-        super().__init__(args, optargs, ret)
+        super().__init__(args, optargs, ret,
+                         delay=iodelay.Fixed(iodelay.Constant(0)))
         self.service = service
 
     def unify(self, other):
@@ -278,7 +283,8 @@ class TCFunction(TFunction):
     attributes = OrderedDict()
 
     def __init__(self, args, ret, name):
-        super().__init__(args, OrderedDict(), ret)
+        super().__init__(args, OrderedDict(), ret,
+                         delay=iodelay.Fixed(iodelay.Constant(0)))
         self.name = name
 
     def unify(self, other):
@@ -547,6 +553,15 @@ class TypePrinter(object):
             args += ["?%s:%s" % (arg, self.name(typ.optargs[arg])) for arg in typ.optargs]
             signature = "(%s)->%s" % (", ".join(args), self.name(typ.ret))
 
+            if iodelay.is_unknown(typ.delay) or iodelay.is_fixed(typ.delay, 0):
+                pass
+            elif iodelay.is_fixed(typ.delay):
+                signature += " delay({} mu)".format(typ.delay.length)
+            elif iodelay.is_indeterminate(typ.delay):
+                signature += " delay(?)"
+            else:
+                assert False
+
             if isinstance(typ, TRPCFunction):
                 return "rpc({}) {}".format(typ.service, signature)
             if isinstance(typ, TCFunction):
diff --git a/lit-test/test/inferencer/with.py b/lit-test/test/inferencer/with.py
new file mode 100644
index 000000000..b11f3da13
--- /dev/null
+++ b/lit-test/test/inferencer/with.py
@@ -0,0 +1,5 @@
+# RUN: %python -m artiq.compiler.testbench.inferencer %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+# CHECK-L: as x:<function parallel>
+with parallel as x: pass
diff --git a/lit-test/test/iodelay/argument.py b/lit-test/test/iodelay/argument.py
new file mode 100644
index 000000000..f026c78f1
--- /dev/null
+++ b/lit-test/test/iodelay/argument.py
@@ -0,0 +1,7 @@
+# RUN: %python -m artiq.compiler.testbench.signature %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+# CHECK-L: f: (a:float, b:int(width=64))->NoneType delay(s->mu(a) + b mu)
+def f(a, b):
+    delay(a)
+    delay_mu(b)
diff --git a/lit-test/test/iodelay/arith.py b/lit-test/test/iodelay/arith.py
new file mode 100644
index 000000000..0eba79715
--- /dev/null
+++ b/lit-test/test/iodelay/arith.py
@@ -0,0 +1,8 @@
+# RUN: %python -m artiq.compiler.testbench.signature %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+# CHECK-L: f: (a:int(width=32), b:int(width=32), c:int(width=32), d:int(width=32), e:int(width=32))->NoneType delay(s->mu(a * b // c + d - 10 / e) mu)
+def f(a, b, c, d, e):
+    delay(a * b // c + d - 10 / e)
+
+f(1,2,3,4,5)
diff --git a/lit-test/test/iodelay/call.py b/lit-test/test/iodelay/call.py
new file mode 100644
index 000000000..65580812a
--- /dev/null
+++ b/lit-test/test/iodelay/call.py
@@ -0,0 +1,10 @@
+# RUN: %python -m artiq.compiler.testbench.signature %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+def f():
+    delay(1.0)
+
+# CHECK-L: g: ()->NoneType delay(s->mu(2.0) mu)
+def g():
+    f()
+    f()
diff --git a/lit-test/test/iodelay/call_subst.py b/lit-test/test/iodelay/call_subst.py
new file mode 100644
index 000000000..fddc007c6
--- /dev/null
+++ b/lit-test/test/iodelay/call_subst.py
@@ -0,0 +1,13 @@
+# RUN: %python -m artiq.compiler.testbench.signature %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+def pulse(len):
+    # "on"
+    delay_mu(len)
+    # "off"
+    delay_mu(len)
+
+# CHECK-L: f: ()->NoneType delay(600 mu)
+def f():
+    pulse(100)
+    pulse(200)
diff --git a/lit-test/test/iodelay/error_arith.py b/lit-test/test/iodelay/error_arith.py
new file mode 100644
index 000000000..54e30aa24
--- /dev/null
+++ b/lit-test/test/iodelay/error_arith.py
@@ -0,0 +1,21 @@
+# RUN: %python -m artiq.compiler.testbench.signature +diag +delay %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+def f(a):
+    b = 1.0
+    # CHECK-L: ${LINE:+3}: error: this call cannot be interleaved
+    # CHECK-L: ${LINE:+2}: note: this variable is not an argument of the innermost function
+    # CHECK-L: ${LINE:-4}: note: only these arguments are in scope of analysis
+    delay(b)
+
+def g():
+    # CHECK-L: ${LINE:+2}: error: this call cannot be interleaved
+    # CHECK-L: ${LINE:+1}: note: this operator is not supported
+    delay(2.0**2)
+
+def h():
+    # CHECK-L: ${LINE:+2}: error: this call cannot be interleaved
+    # CHECK-L: ${LINE:+1}: note: this expression is not supported
+    delay_mu(1 if False else 2)
+
+f(1)
diff --git a/lit-test/test/iodelay/error_builtinfn.py b/lit-test/test/iodelay/error_builtinfn.py
new file mode 100644
index 000000000..ad44c972c
--- /dev/null
+++ b/lit-test/test/iodelay/error_builtinfn.py
@@ -0,0 +1,14 @@
+# RUN: %python -m artiq.compiler.testbench.signature +diag +delay %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+def f():
+    x = 1
+    # CHECK-L: ${LINE:+1}: error: this call cannot be interleaved because an argument cannot be statically evaluated
+    delay_mu(x)
+
+def g():
+    x = 1.0
+    # CHECK-L: ${LINE:+1}: error: this call cannot be interleaved
+    delay(x)
+
+
diff --git a/lit-test/test/iodelay/error_call_nested.py b/lit-test/test/iodelay/error_call_nested.py
new file mode 100644
index 000000000..b283c0917
--- /dev/null
+++ b/lit-test/test/iodelay/error_call_nested.py
@@ -0,0 +1,11 @@
+# RUN: %python -m artiq.compiler.testbench.signature +diag +delay %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+def f():
+    # CHECK-L: ${LINE:+1}: error: this call cannot be interleaved
+    delay(1.0**2)
+
+def g():
+    # CHECK-L: ${LINE:+1}: note: function called here
+    f()
+    f()
diff --git a/lit-test/test/iodelay/error_call_subst.py b/lit-test/test/iodelay/error_call_subst.py
new file mode 100644
index 000000000..62c6bb29a
--- /dev/null
+++ b/lit-test/test/iodelay/error_call_subst.py
@@ -0,0 +1,13 @@
+# RUN: %python -m artiq.compiler.testbench.signature +diag +delay %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+def pulse(len):
+    # "on"
+    delay_mu(len)
+    # "off"
+    delay_mu(len)
+
+def f():
+    a = 100
+    # CHECK-L: ${LINE:+1}: error: this call cannot be interleaved
+    pulse(a)
diff --git a/lit-test/test/iodelay/error_control_flow.py b/lit-test/test/iodelay/error_control_flow.py
new file mode 100644
index 000000000..c179c95b9
--- /dev/null
+++ b/lit-test/test/iodelay/error_control_flow.py
@@ -0,0 +1,23 @@
+# RUN: %python -m artiq.compiler.testbench.signature +diag +delay %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+def f():
+    # CHECK-L: ${LINE:+1}: error: while statement cannot be interleaved
+    while True:
+        delay_mu(1)
+
+def g():
+    # CHECK-L: ${LINE:+1}: error: if statement cannot be interleaved
+    if True:
+        delay_mu(1)
+
+def h():
+    # CHECK-L: ${LINE:+1}: error: if expression cannot be interleaved
+    delay_mu(1) if True else delay_mu(2)
+
+def i():
+    # CHECK-L: ${LINE:+1}: error: try statement cannot be interleaved
+    try:
+        delay_mu(1)
+    finally:
+        pass
diff --git a/lit-test/test/iodelay/error_for.py b/lit-test/test/iodelay/error_for.py
new file mode 100644
index 000000000..5c5b10ffc
--- /dev/null
+++ b/lit-test/test/iodelay/error_for.py
@@ -0,0 +1,9 @@
+# RUN: %python -m artiq.compiler.testbench.signature +diag +delay %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+def f():
+    r = range(10)
+    # CHECK-L: ${LINE:+2}: error: for statement cannot be interleaved because trip count is indeterminate
+    # CHECK-L: ${LINE:+1}: note: this value is not a constant range literal
+    for _ in r:
+        delay_mu(1)
diff --git a/lit-test/test/iodelay/error_goto.py b/lit-test/test/iodelay/error_goto.py
new file mode 100644
index 000000000..02686cd86
--- /dev/null
+++ b/lit-test/test/iodelay/error_goto.py
@@ -0,0 +1,14 @@
+# RUN: %python -m artiq.compiler.testbench.signature +diag +delay %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+def f():
+    for _ in range(10):
+        delay_mu(10)
+        # CHECK-L: ${LINE:+1}: error: loop trip count is indeterminate because of control flow
+        break
+
+def g():
+    for _ in range(10):
+        delay_mu(10)
+        # CHECK-L: ${LINE:+1}: error: loop trip count is indeterminate because of control flow
+        continue
diff --git a/lit-test/test/iodelay/error_return.py b/lit-test/test/iodelay/error_return.py
new file mode 100644
index 000000000..c047de2b0
--- /dev/null
+++ b/lit-test/test/iodelay/error_return.py
@@ -0,0 +1,8 @@
+# RUN: %python -m artiq.compiler.testbench.signature +diag +delay %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+def f():
+    if True:
+        # CHECK-L: ${LINE:+1}: error: only return statement at the end of the function can be interleaved
+        return 1
+    delay_mu(1)
diff --git a/lit-test/test/iodelay/goto.py b/lit-test/test/iodelay/goto.py
new file mode 100644
index 000000000..d80de43f9
--- /dev/null
+++ b/lit-test/test/iodelay/goto.py
@@ -0,0 +1,14 @@
+# RUN: %python -m artiq.compiler.testbench.signature %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+# CHECK-L: f: ()->NoneType delay(10 mu)
+def f():
+    delay_mu(10)
+    for _ in range(10):
+        break
+
+# CHECK-L: g: ()->NoneType delay(10 mu)
+def g():
+    delay_mu(10)
+    for _ in range(10):
+        continue
diff --git a/lit-test/test/iodelay/linear.py b/lit-test/test/iodelay/linear.py
new file mode 100644
index 000000000..7037516f6
--- /dev/null
+++ b/lit-test/test/iodelay/linear.py
@@ -0,0 +1,12 @@
+# RUN: %python -m artiq.compiler.testbench.signature %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+# CHECK-L: f: ()->NoneType delay(s->mu(1.0) + 1000 mu)
+def f():
+    delay(1.0)
+    delay_mu(1000)
+
+# CHECK-L: g: ()->NoneType delay(s->mu(5.0) mu)
+def g():
+    delay(1.0)
+    delay(2.0 * 2)
diff --git a/lit-test/test/iodelay/loop.py b/lit-test/test/iodelay/loop.py
new file mode 100644
index 000000000..228756224
--- /dev/null
+++ b/lit-test/test/iodelay/loop.py
@@ -0,0 +1,13 @@
+# RUN: %python -m artiq.compiler.testbench.signature %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+# CHECK-L: f: ()->NoneType delay(s->mu(1.5) * 10 mu)
+def f():
+    for _ in range(10):
+        delay(1.5)
+
+# CHECK-L: g: ()->NoneType delay(s->mu(1.5) * 2 * 10 mu)
+def g():
+    for _ in range(10):
+        for _ in range(2):
+            delay(1.5)
diff --git a/lit-test/test/iodelay/parallel.py b/lit-test/test/iodelay/parallel.py
new file mode 100644
index 000000000..0976f1780
--- /dev/null
+++ b/lit-test/test/iodelay/parallel.py
@@ -0,0 +1,15 @@
+# RUN: %python -m artiq.compiler.testbench.signature %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+# CHECK-L: f: (a:int(width=64), b:int(width=64))->NoneType delay(max(a, b) mu)
+def f(a, b):
+    with parallel:
+        delay_mu(a)
+        delay_mu(b)
+
+# CHECK-L: g: (a:int(width=64))->NoneType delay(max(a, 200) mu)
+def g(a):
+    with parallel:
+        delay_mu(100)
+        delay_mu(200)
+        delay_mu(a)
diff --git a/lit-test/test/iodelay/range.py b/lit-test/test/iodelay/range.py
new file mode 100644
index 000000000..498ee6eb5
--- /dev/null
+++ b/lit-test/test/iodelay/range.py
@@ -0,0 +1,21 @@
+# RUN: %python -m artiq.compiler.testbench.signature %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+# CHECK-L: f: (a:int(width=32))->NoneType delay(s->mu(1.5) * a mu)
+def f(a):
+    for _ in range(a):
+        delay(1.5)
+
+# CHECK-L: g: (a:int(width=32), b:int(width=32))->NoneType delay(s->mu(1.5) * (b - a) mu)
+def g(a, b):
+    for _ in range(a, b):
+        delay(1.5)
+
+# CHECK-L: h: (a:int(width=32), b:int(width=32), c:int(width=32))->NoneType delay(s->mu(1.5) * (b - a) // c mu)
+def h(a, b, c):
+    for _ in range(a, b, c):
+        delay(1.5)
+
+f(1)
+g(1,2)
+h(1,2,3)
diff --git a/lit-test/test/iodelay/return.py b/lit-test/test/iodelay/return.py
new file mode 100644
index 000000000..3a3dc09fa
--- /dev/null
+++ b/lit-test/test/iodelay/return.py
@@ -0,0 +1,16 @@
+# RUN: %python -m artiq.compiler.testbench.signature %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+# CHECK-L: f: ()->int(width=32) delay(s->mu(1.5) * 10 mu)
+def f():
+    for _ in range(10):
+        delay(1.5)
+    return 10
+
+# CHECK-L: g: (x:float)->int(width=32) delay(0 mu)
+def g(x):
+    if x > 1.0:
+        return 1
+    return 0
+
+g(1.0)