transform: collapse timeline

examples/collapse_test.py

@@ -2,12 +2,8 @@ def collapse_test():
 	for i in range(3):
 		with parallel:
 			with sequential:
-				pulse("a", 100*MHz, 10*us)
-				delay(10*us)
+				pulse("a", 100*MHz, 20*us)
 				pulse("b", 100*MHz, 10*us)
-				delay(20*us)
 			with sequential:
-				pulse("a", 100*MHz, 10*us)
-				delay(10*us)
-				pulse("b", 100*MHz, 10*us)
-				delay(10*us)
+				pulse("A", 100*MHz, 10*us)
+				pulse("B", 100*MHz, 10*us)

examples/transform.py

@@ -1,4 +1,4 @@
-import inspect, textwrap, ast
+import inspect, textwrap, ast, types
 
 from artiq import units, unparse
 
@@ -102,6 +102,100 @@ def unroll_loops(stmts, limit):
 		stmts[offset+location:offset+location+1] = new_stmts
 		offset += len(new_stmts) - 1
 
+# -1 statement duration could not be pre-determined
+#  0 statement has no effect on timeline
+# >0 statement is a static delay that advances the timeline by the given amount
+def _get_duration(stmt):
+	if isinstance(stmt, (ast.Expr, ast.Assign)):
+		return _get_duration(stmt.value)
+	elif isinstance(stmt, ast.Call) and isinstance(stmt.func, ast.Name):
+		name = stmt.func.id
+		if name == "delay":
+			if isinstance(stmt.args[0], ast.Num):
+				return stmt.args[0].n
+			else:
+				return -1
+		elif name == "wait_edge":
+			return -1
+		else:
+			return 0
+	else:
+		return -1
+
+def _merge_timelines(timelines):
+	r = []
+
+	current_stmts = []
+	for stmts in timelines:
+		it = iter(stmts)
+		try:
+			stmt = next(it)
+		except StopIteration:
+			pass
+		else:
+			current_stmts.append(types.SimpleNamespace(delay=_get_duration(stmt), stmt=stmt, it=it))
+
+	while current_stmts:
+		dt = min(stmt.delay for stmt in current_stmts)
+		if dt < 0:
+			# contains statement(s) with indeterminate duration
+			return None
+		if dt > 0:
+			# advance timeline by dt
+			for stmt in current_stmts:
+				stmt.delay -= dt
+				if stmt.delay == 0:
+					ref_stmt = stmt.stmt
+			delay_stmt = ast.copy_location(
+				ast.Expr(ast.Call(func=ast.Name(id="delay", ctx=ast.Load()),
+					args=[ast.Num(n=dt)],
+					keywords=[], starargs=[], kwargs=[])),
+				ref_stmt)
+			r.append(delay_stmt)
+		else:
+			for stmt in current_stmts:
+				if stmt.delay == 0:
+					r.append(stmt.stmt)
+		# discard executed statements
+		exhausted_list = []
+		for stmt_i, stmt in enumerate(current_stmts):
+			if stmt.delay == 0:
+				try:
+					stmt.stmt = next(stmt.it)
+				except StopIteration:
+					exhausted_list.append(stmt_i)
+				else:
+					stmt.delay = _get_duration(stmt.stmt)
+		for offset, i in enumerate(exhausted_list):
+			current_stmts.pop(i-offset)
+
+	return r
+
+def collapse(stmts):
+	replacements = []
+	for stmt_i, stmt in enumerate(stmts):
+		if isinstance(stmt, (ast.For, ast.While, ast.If)):
+			collapse(stmt.body)
+			collapse(stmt.orelse)
+		elif isinstance(stmt, ast.With):
+			btype = stmt.items[0].context_expr.id
+			if btype == "sequential":
+				collapse(stmt.body)
+				replacements.append((stmt_i, stmt.body))
+			elif btype == "parallel":
+				timelines = [[s] for s in stmt.body]
+				for timeline in timelines:
+					collapse(timeline)
+				merged = _merge_timelines(timelines)
+				if merged is not None:
+					replacements.append((stmt_i, merged))
+			else:
+				raise ValueError("Unknown block type: " + btype)
+	offset = 0
+	for location, new_stmts in replacements:
+		stmts[offset+location:offset+location+1] = new_stmts
+		offset += len(new_stmts) - 1
+
 if __name__ == "__main__":
 	import collapse_test
 	kernel = collapse_test.collapse_test
@@ -111,5 +205,6 @@ if __name__ == "__main__":
 
 	explicit_delays(node)
 	unroll_loops(node, 50)
+	collapse(node)
 
 	unparse.Unparser(node)