mirror of https://github.com/m-labs/artiq.git
core: support precompilation of kernels
This commit is contained in:
parent
232f28c0e8
commit
ac55da81d8
|
@ -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):
|
||||
|
|
|
@ -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()
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue