m-labs · Aug 18, 2014 · Aug 18, 2014 · Aug 18, 2014 · Aug 18, 2014 · Aug 18, 2014
diff --git a/artiq/compiler/fold_constants.py b/artiq/compiler/fold_constants.py
@@ -1,6 +1,7 @@
 import ast, operator
 
 from artiq.compiler.tools import *
+from artiq.language.core import int64, round64
 
 _ast_unops = {
 	ast.Invert: operator.inv,
@@ -57,4 +58,18 @@ def visit_BinOp(self, node):
 			return node
 		return ast.copy_location(result, node)
 
-fold_constants = make_stmt_transformer(_ConstantFolder)
+	def visit_Call(self, node):
+		self.generic_visit(node)
+		fn = node.func.id
+		if fn in ("int", "int64", "round", "round64"):
+			try:
+				arg = eval_constant(node.args[0])
+			except NotConstant:
+				return node
+			result = value_to_ast(globals()[fn](arg))
+			return ast.copy_location(result, node)
+		else:
+			return node
+
+def fold_constants(node):
+	_ConstantFolder().visit(node)
diff --git a/artiq/compiler/inline.py b/artiq/compiler/inline.py
@@ -26,11 +26,10 @@ def __init__(self):
 		# reserved names
 		for kg in core_language.kernel_globals:
 			self.use_count[kg] = 1
-		self.use_count["range"] = 1
-		self.use_count["Fraction"] = 1
-		self.use_count["Quantity"] = 1
-		self.use_count["s_unit"] = 1
-		self.use_count["Hz_unit"] = 1
+		for name in "int", "round", "int64", "round64", \
+		  "range", "Fraction", "Quantity", \
+		  "s_unit", "Hz_unit", "microcycle_unit":
+			self.use_count[name] = 1
 
 	def new_name(self, base_name):
 		if base_name[-1].isdigit():
@@ -102,7 +101,7 @@ def get_constants(self, r_obj, r_funcname):
 _embeddable_calls = {
 	core_language.delay, core_language.at, core_language.now,
 	core_language.syscall,
-	range,
+	range, int, round, core_language.int64, core_language.round64,
 	Fraction, units.Quantity
 }
 
@@ -147,7 +146,7 @@ def visit_Call(self, node):
 		elif hasattr(func, "k_function_info") and getattr(func.__self__, func.k_function_info.core_name) is self.core:
 			args = [func.__self__] + new_args
 			inlined, _ = inline(self.core, func.k_function_info.k_function, args, dict(), self.rm)
-			return inlined
+			return inlined.body
 		else:
 			args = [ast.Str("rpc"), value_to_ast(self.rm.rpc_map[func])]
 			args += new_args
@@ -169,6 +168,7 @@ def visit_Expr(self, node):
 			return node
 
 	def visit_FunctionDef(self, node):
+		node.args = ast.arguments(args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[])
 		node.decorator_list = []
 		self.generic_visit(node)
 		return node
@@ -226,4 +226,4 @@ def inline(core, k_function, k_args, k_kwargs, rm=None):
 		funcdef.body[0:0] = rm.kernel_attr_init
 
 	r_rpc_map = dict((rpc_num, rpc_fun) for rpc_fun, rpc_num in rm.rpc_map.items())
-	return funcdef.body, r_rpc_map
+	return funcdef, r_rpc_map
diff --git a/artiq/compiler/interleave.py b/artiq/compiler/interleave.py
@@ -1,12 +1,11 @@
 import ast, types
 
-from artiq.language import units
 from artiq.compiler.tools import *
 
 # -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 (in seconds)
+#     by the given amount (in microcycles)
 def _get_duration(stmt):
 	if isinstance(stmt, (ast.Expr, ast.Assign)):
 		return _get_duration(stmt.value)
@@ -21,10 +20,8 @@ def _get_duration(stmt):
 			try:
 				da = eval_constant(stmt.args[0])
 			except NotConstant:
-				return -1
-			if da.unit != units.s_unit:
-				raise units.DimensionError("Delay not expressed in seconds")
-			return da.amount
+				da = -1
+			return da
 		else:
 			return 0
 	else:
@@ -54,14 +51,9 @@ def _interleave_timelines(timelines):
 				stmt.delay -= dt
 				if stmt.delay == 0:
 					ref_stmt = stmt.stmt
-			da_expr = ast.copy_location(
-					ast.Call(func=ast.Name("Quantity", ast.Load()),
-					args=[value_to_ast(dt), ast.Name("s_unit", ast.Load())],
-					keywords=[], starargs=[], kwargs=[]),
-				ref_stmt)
 			delay_stmt = ast.copy_location(
 				ast.Expr(ast.Call(func=ast.Name("delay", ast.Load()),
-					args=[da_expr],
+					args=[value_to_ast(dt)],
 					keywords=[], starargs=[], kwargs=[])),
 				ref_stmt)
 			r.append(delay_stmt)
@@ -84,21 +76,21 @@ def _interleave_timelines(timelines):
 
 	return r
 
-def interleave(stmts):
+def _interleave_stmts(stmts):
 	replacements = []
 	for stmt_i, stmt in enumerate(stmts):
 		if isinstance(stmt, (ast.For, ast.While, ast.If)):
-			interleave(stmt.body)
-			interleave(stmt.orelse)
+			_interleave_stmts(stmt.body)
+			_interleave_stmts(stmt.orelse)
 		elif isinstance(stmt, ast.With):
 			btype = stmt.items[0].context_expr.id
 			if btype == "sequential":
-				interleave(stmt.body)
+				_interleave_stmts(stmt.body)
 				replacements.append((stmt_i, stmt.body))
 			elif btype == "parallel":
 				timelines = [[s] for s in stmt.body]
 				for timeline in timelines:
-					interleave(timeline)
+					_interleave_stmts(timeline)
 				merged = _interleave_timelines(timelines)
 				if merged is not None:
 					replacements.append((stmt_i, merged))
@@ -108,3 +100,6 @@ def interleave(stmts):
 	for location, new_stmts in replacements:
 		stmts[offset+location:offset+location+1] = new_stmts
 		offset += len(new_stmts) - 1
+
+def interleave(funcdef):
+	_interleave_stmts(funcdef.body)
diff --git a/artiq/compiler/ir.py b/artiq/compiler/ir.py
@@ -1,157 +1,26 @@
-import ast
-
 from llvm import core as lc
 from llvm import passes as lp
 
-class _Namespace:
-	def __init__(self, function):
-		self.function = function
-		self.bindings = dict()
-
-	def store(self, builder, val, name):
-		try:
-			allocated = self.bindings[name]
-		except KeyError:
-			entry = self.function.get_entry_basic_block()
-			builder_e = lc.Builder.new(entry)
-			builder_e.position_at_beginning(entry)
-			allocated = builder_e.alloca(lc.Type.int(), name=name)
-			self.bindings[name] = allocated
-		builder.store(val, allocated)
-
-	def load(self, builder, name):
-		return builder.load(self.bindings[name])
-
-_binop_to_builder = {
-	ast.Add: "add",
-	ast.Sub: "sub",
-	ast.Mult: "mul",
-	ast.FloorDiv: "sdiv",
-	ast.Mod: "srem",
-
-	ast.LShift: "shl",
-	ast.RShift: "ashr",
-
-	ast.BitAnd: "and_",
-	ast.BitOr: "or_",
-}
-
-def _emit_expr(env, builder, ns, node):
-	if isinstance(node, ast.Name):
-		return ns.load(builder, node.id)
-	elif isinstance(node, ast.Num):
-		return lc.Constant.int(lc.Type.int(), node.n)
-	elif isinstance(node, ast.BinOp):
-		left = _emit_expr(env, builder, ns, node.left)
-		right = _emit_expr(env, builder, ns, node.right)
-		bf = getattr(builder, _binop_to_builder[type(node.op)])
-		return bf(left, right)
-	elif isinstance(node, ast.Compare):
-		comparisons = []
-		old_comparator = _emit_expr(env, builder, ns, node.left)
-		for op, comparator_a in zip(node.ops, node.comparators):
-			comparator = _emit_expr(env, builder, ns, comparator_a)
-			mapping = {
-				ast.Eq: lc.ICMP_EQ,
-				ast.NotEq: lc.ICMP_NE,
-				ast.Lt: lc.ICMP_SLT,
-				ast.LtE: lc.ICMP_SLE,
-				ast.Gt: lc.ICMP_SGT,
-				ast.GtE: lc.ICMP_SGE
-			}
-			comparison = builder.icmp(mapping[type(op)], old_comparator, comparator)
-			comparisons.append(comparison)
-			old_comparator = comparator
-		r = comparisons[0]
-		for comparison in comparisons[1:]:
-			r = builder.band(r, comparison)
-		return r
-	elif isinstance(node, ast.Call):
-		if node.func.id == "syscall":
-			return env.emit_syscall(builder, node.args[0].s,
-				[_emit_expr(env, builder, ns, expr) for expr in node.args[1:]])
-		elif node.func.id == "Quantity":
-			return _emit_expr(env, builder, ns, node.args[0])
-		elif node.func.id == "Fraction":
-			if not isinstance(node.args[1], ast.Num) or node.args[1].n != 1:
-				raise NotImplementedError
-			return _emit_expr(env, builder, ns, node.args[0])
-		else:
-			raise NotImplementedError
-	else:
-		raise NotImplementedError
-
-def _emit_statements(env, builder, ns, stmts):
-	for stmt in stmts:
-		if isinstance(stmt, ast.Return):
-			val = _emit_expr(env, builder, ns, stmt.value)
-			builder.ret(val)
-		elif isinstance(stmt, ast.Assign):
-			val = _emit_expr(env, builder, ns, stmt.value)
-			for target in stmt.targets:
-				if isinstance(target, ast.Name):
-					ns.store(builder, val, target.id)
-				else:
-					raise NotImplementedError
-		elif isinstance(stmt, ast.AugAssign):
-			left = _emit_expr(env, builder, ns, stmt.target)
-			right = _emit_expr(env, builder, ns, stmt.value)
-			bf = getattr(builder, _binop_to_builder[type(stmt.op)])
-			result = bf(left, right)
-			ns.store(builder, result, stmt.target.id)
-		elif isinstance(stmt, ast.If):
-			function = builder.basic_block.function
-			then_block = function.append_basic_block("i_then")
-			else_block = function.append_basic_block("i_else")
-			merge_block = function.append_basic_block("i_merge")
-
-			condition = _emit_expr(env, builder, ns, stmt.test)
-			builder.cbranch(condition, then_block, else_block)
-
-			builder.position_at_end(then_block)
-			_emit_statements(env, builder, ns, stmt.body)
-			builder.branch(merge_block)
-
-			builder.position_at_end(else_block)
-			_emit_statements(env, builder, ns, stmt.orelse)
-			builder.branch(merge_block)
+from artiq.compiler import ir_infer_types, ir_ast_body
 
-			builder.position_at_end(merge_block)
-		elif isinstance(stmt, ast.While):
-			function = builder.basic_block.function
-			body_block = function.append_basic_block("w_body")
-			else_block = function.append_basic_block("w_else")
-			merge_block = function.append_basic_block("w_merge")
-
-			condition = _emit_expr(env, builder, ns, stmt.test)
-			builder.cbranch(condition, body_block, else_block)
-
-			builder.position_at_end(body_block)
-			_emit_statements(env, builder, ns, stmt.body)
-			condition = _emit_expr(env, builder, ns, stmt.test)
-			builder.cbranch(condition, body_block, merge_block)
-
-			builder.position_at_end(else_block)
-			_emit_statements(env, builder, ns, stmt.orelse)
-			builder.branch(merge_block)
+def compile_function(module, env, funcdef):
+	function_type = lc.Type.function(lc.Type.void(), [])
+	function = module.add_function(function_type, funcdef.name)
+	bb = function.append_basic_block("entry")
+	builder = lc.Builder.new(bb)
 
-			builder.position_at_end(merge_block)
-		elif isinstance(stmt, ast.Expr):
-			_emit_expr(env, builder, ns, stmt.value)
-		else:
-			raise NotImplementedError
+	ns = ir_infer_types.infer_types(env, funcdef)
+	for k, v in ns.items():
+		v.create_alloca(builder, k)
+	visitor = ir_ast_body.Visitor(env, ns, builder)
+	visitor.visit_statements(funcdef.body)
+	builder.ret_void()
 
-def get_runtime_binary(env, stmts):
+def get_runtime_binary(env, funcdef):
 	module = lc.Module.new("main")
 	env.set_module(module)
 
-	function_type = lc.Type.function(lc.Type.void(), [])
-	function = module.add_function(function_type, "run")
-	bb = function.append_basic_block("entry")
-	builder = lc.Builder.new(bb)
-	ns = _Namespace(function)
-	_emit_statements(env, builder, ns, stmts)
-	builder.ret_void()
+	compile_function(module, env, funcdef)
 
 	pass_manager = lp.PassManager.new()
 	pass_manager.add(lp.PASS_MEM2REG)

diff --git a/artiq/compiler/ir_ast_body.py b/artiq/compiler/ir_ast_body.py
@@ -6,7 +6,8 @@
 from artiq.compiler import ir_values
 
 class Visitor:
-	def __init__(self, ns, builder=None):
+	def __init__(self, env, ns, builder=None):
+		self.env = env
 		self.ns = ns
 		self.builder = builder
 
@@ -96,7 +97,7 @@ def _visit_expr_Compare(self, node):
 		old_comparator = self.visit_expression(node.left)
 		for op, comparator_a in zip(node.ops, node.comparators):
 			comparator = self.visit_expression(comparator_a)
-			comparison = ast_cmps[type(op)](old_comparator, comparator)
+			comparison = ast_cmps[type(op)](old_comparator, comparator, self.builder)
 			comparisons.append(comparison)
 			old_comparator = comparator
 		r = comparisons[0]
@@ -112,7 +113,15 @@ def _visit_expr_Call(self, node):
 			"round": ir_values.operators.round,
 			"round64": ir_values.operators.round64,
 		}
-		return ast_unfuns[node.func.id](self.visit_expression(node.args[0]), self.builder)
+		fn = node.func.id
+		if fn in ast_unfuns:
+			return ast_unfuns[fn](self.visit_expression(node.args[0]), self.builder)
+		elif fn == "syscall":
+			return self.env.syscall(node.args[0].s,
+				[self.visit_expression(expr) for expr in node.args[1:]],
+				self.builder)
+		else:
+			raise NameError("Function '{}' is not defined".format(fn))
 
 	def visit_statements(self, stmts):
 		for node in stmts:
@@ -127,14 +136,14 @@ def _visit_stmt_Assign(self, node):
 		val = self.visit_expression(node.value)
 		for target in node.targets:
 			if isinstance(target, ast.Name):
-				self.builder.store(val, self.ns[target.id])
+				self.builder.store(val.llvm_value, self.ns[target.id].llvm_value)
 			else:
 				raise NotImplementedError
 
 	def _visit_stmt_AugAssign(self, node):
 		val = self.visit_expression(ast.BinOp(op=node.op, left=node.target, right=node.value))
 		if isinstance(node.target, ast.Name):
-			self.builder.store(val, self.ns[node.target.id])
+			self.builder.store(val.llvm_value, self.ns[node.target.id].llvm_value)
 		else:
 			raise NotImplementedError
 
@@ -166,8 +175,8 @@ def _visit_stmt_While(self, node):
 		else_block = function.append_basic_block("w_else")
 		merge_block = function.append_basic_block("w_merge")
 
-		condition = self.visit_expression(node.test)
-		self.builder.cbranch(condition, body_block, else_block)
+		condition = ir_values.operators.bool(self.visit_expression(node.test), self.builder)
+		self.builder.cbranch(condition.llvm_value, body_block, else_block)
 
 		self.builder.position_at_end(body_block)
 		self.visit_statements(node.body)