diff --git a/artiq/coredevice/core.py b/artiq/coredevice/core.py index 022b3d968..5d4ae8725 100644 --- a/artiq/coredevice/core.py +++ b/artiq/coredevice/core.py @@ -35,6 +35,10 @@ def _make_debug_unparse(final): return _filtered_unparse +def _no_debug_unparse(label, node): + pass + + class Core: def __init__(self, core_com, runtime_env=None): if runtime_env is None: @@ -42,51 +46,55 @@ class Core: self.runtime_env = runtime_env self.core_com = core_com - def run(self, k_function, k_args, k_kwargs): - # transform/simplify AST - _debug_unparse = _make_debug_unparse("remove_dead_code_2") - - func_def, rpc_map, exception_map = inline( - self, k_function, k_args, k_kwargs) - _debug_unparse("inline", func_def) - + def transform_stack(self, func_def, rpc_map, exception_map, + debug_unparse=_no_debug_unparse): lower_units(func_def, rpc_map) - _debug_unparse("lower_units", func_def) + debug_unparse("lower_units", func_def) remove_inter_assigns(func_def) - _debug_unparse("remove_inter_assigns_1", func_def) + debug_unparse("remove_inter_assigns_1", func_def) fold_constants(func_def) - _debug_unparse("fold_constants_1", func_def) + debug_unparse("fold_constants_1", func_def) unroll_loops(func_def, 500) - _debug_unparse("unroll_loops", func_def) + debug_unparse("unroll_loops", func_def) interleave(func_def) - _debug_unparse("interleave", func_def) + debug_unparse("interleave", func_def) lower_time(func_def, getattr(self.runtime_env, "initial_time", 0), self.runtime_env.ref_period) - _debug_unparse("lower_time", func_def) + debug_unparse("lower_time", func_def) remove_inter_assigns(func_def) - _debug_unparse("remove_inter_assigns_2", func_def) + debug_unparse("remove_inter_assigns_2", func_def) fold_constants(func_def) - _debug_unparse("fold_constants_2", func_def) + debug_unparse("fold_constants_2", func_def) remove_dead_code(func_def) - _debug_unparse("remove_dead_code_1", func_def) + debug_unparse("remove_dead_code_1", func_def) remove_inter_assigns(func_def) - _debug_unparse("remove_inter_assigns_3", func_def) + debug_unparse("remove_inter_assigns_3", func_def) fold_constants(func_def) - _debug_unparse("fold_constants_3", func_def) + debug_unparse("fold_constants_3", func_def) remove_dead_code(func_def) - _debug_unparse("remove_dead_code_2", func_def) + debug_unparse("remove_dead_code_2", func_def) + + def run(self, k_function, k_args, k_kwargs): + # transform/simplify AST + debug_unparse = _make_debug_unparse("remove_dead_code_2") + + func_def, rpc_map, exception_map = inline( + self, k_function, k_args, k_kwargs) + debug_unparse("inline", func_def) + + self.transform_stack(func_def, rpc_map, exception_map, debug_unparse) # compile to machine code and run binary = get_runtime_binary(self.runtime_env, func_def) diff --git a/test/transforms.py b/test/transforms.py new file mode 100644 index 000000000..ea89a7322 --- /dev/null +++ b/test/transforms.py @@ -0,0 +1,45 @@ +import unittest +import ast + +from artiq.coredevice import comm_dummy, core +from artiq.transforms.unparse import unparse + + +# Original code before inline: +# +# n = time_to_cycles(1.2345*ns) +# ftw = self.dds.frequency_to_ftw(345*MHz) +# f = self.dds.ftw_to_frequency(ftw) +# phi = 1000*cycles_to_time(n)*f +# do_someting(int(phi)) +# +optimize_in = """ + +def run(): + dds_sysclk = Quantity(Fraction(1000000000, 1), 'Hz') + n = time_to_cycles((1.2345 * Quantity(Fraction(1, 1000000000), 's'))) + with sequential: + frequency = (345 * Quantity(Fraction(1000000, 1), 'Hz')) + frequency_to_ftw_return = int((((2 ** 32) * frequency) / dds_sysclk)) + ftw = frequency_to_ftw_return + with sequential: + ftw2 = ftw + ftw_to_frequency_return = ((ftw2 * dds_sysclk) / (2 ** 32)) + f = ftw_to_frequency_return + phi = ((1000 * cycles_to_time(n)) * f) + do_something(int(phi)) +""" + +optimize_out = """ + +def run(): + do_something(344) +""" + + +class OptimizeCase(unittest.TestCase): + def test_optimize(self): + coredev = core.Core(comm_dummy.Comm()) + func_def = ast.parse(optimize_in).body[0] + coredev.transform_stack(func_def, dict(), dict()) + self.assertEqual(unparse(func_def), optimize_out)