diff --git a/artiq/coredevice/core.py b/artiq/coredevice/core.py index 529c42f98..406a7c961 100644 --- a/artiq/coredevice/core.py +++ b/artiq/coredevice/core.py @@ -1,5 +1,6 @@ import os, sys import numpy +from functools import wraps from pythonparser import diagnostic @@ -120,7 +121,52 @@ class Core: except diagnostic.Error as error: raise CompileError(error.diagnostic) from error + def _run_compiled(self, kernel_library, embedding_map, symbolizer, demangler): + if self.first_run: + self.comm.check_system_info() + self.first_run = False + self.comm.load(kernel_library) + self.comm.run() + self.comm.serve(embedding_map, symbolizer, demangler) + def run(self, function, args, kwargs): + result = None + @rpc(flags={"async"}) + def set_result(new_result): + nonlocal result + result = new_result + embedding_map, kernel_library, symbolizer, demangler = \ + self.compile(function, args, kwargs, set_result) + self._run_compiled(kernel_library, embedding_map, symbolizer, demangler) + return result + + def precompile(self, function, *args, **kwargs): + """Precompile a kernel and return a callable that executes it on the core device + at a later time. + + Arguments to the kernel are set at compilation time and passed to this function, + as additional positional and keyword arguments. + The returned callable accepts no arguments. + + Precompiled kernels may use RPCs. + + Object attributes at the beginning of a precompiled kernel execution have the + values they had at precompilation time. If up-to-date values are required, + use RPC to read them. + Similarly, modified values are not written back, and explicit RPC should be used + to modify host objects. + Carefully review the source code of drivers calls used in precompiled kernels, as + they may rely on host object attributes being transfered between kernel calls. + Examples include code used to control DDS phase, and Urukul RF switch control + via the CPLD register. + + The return value of the callable is the return value of the kernel, if any. + + The callable may be called several times. + """ + if not hasattr(function, "artiq_embedded"): + raise ValueError("Argument is not a kernel") + result = None @rpc(flags={"async"}) def set_result(new_result): @@ -128,17 +174,15 @@ class Core: result = new_result embedding_map, kernel_library, symbolizer, demangler = \ - self.compile(function, args, kwargs, set_result) + self.compile(function, args, kwargs, set_result, attribute_writeback=False) - if self.first_run: - self.comm.check_system_info() - self.first_run = False + @wraps(function) + def run_precompiled(): + nonlocal result + self._run_compiled(kernel_library, embedding_map, symbolizer, demangler) + return result - self.comm.load(kernel_library) - self.comm.run() - self.comm.serve(embedding_map, symbolizer, demangler) - - return result + return run_precompiled @portable def seconds_to_mu(self, seconds): diff --git a/artiq/examples/kc705_nist_clock/repository/precompile.py b/artiq/examples/kc705_nist_clock/repository/precompile.py new file mode 100644 index 000000000..6bfaaf057 --- /dev/null +++ b/artiq/examples/kc705_nist_clock/repository/precompile.py @@ -0,0 +1,20 @@ +from artiq.experiment import * + + +class Precompile(EnvExperiment): + def build(self): + self.setattr_device("core") + self.hello_str = "hello ARTIQ" + + def prepare(self): + self.precompiled = self.core.precompile(self.hello, "world") + + @kernel + def hello(self, arg): + print(self.hello_str, arg) + self.hello_str = "nowriteback" + + def run(self): + self.precompiled() + self.hello_str = "noupdate" + self.precompiled() diff --git a/artiq/test/coredevice/test_compile.py b/artiq/test/coredevice/test_compile.py index 060cd1b8f..efb6073d5 100644 --- a/artiq/test/coredevice/test_compile.py +++ b/artiq/test/coredevice/test_compile.py @@ -20,6 +20,28 @@ class CheckLog(EnvExperiment): core_log("test_artiq_compile") + +class _Precompile(EnvExperiment): + def build(self): + self.setattr_device("core") + self.x = 1 + self.y = 2 + self.z = 3 + + def set_attr(self, value): + self.x = value + + @kernel + def the_kernel(self, arg): + self.set_attr(arg + self.y) + self.z = 23 + + def run(self): + precompiled = self.core.precompile(self.the_kernel, 40) + self.y = 0 + precompiled() + + class TestCompile(ExperimentCase): def test_compile(self): core_addr = self.device_mgr.get_desc("core")["arguments"]["host"] @@ -34,3 +56,9 @@ class TestCompile(ExperimentCase): log = mgmt.get_log() self.assertIn("test_artiq_compile", log) mgmt.close() + + def test_precompile(self): + exp = self.create(_Precompile) + exp.run() + self.assertEqual(exp.x, 42) + self.assertEqual(exp.z, 3)