forked from M-Labs/artiq
1
0
Fork 0

core: support precompilation of kernels

This commit is contained in:
Sebastien Bourdeauducq 2022-03-06 18:24:57 +08:00
parent 232f28c0e8
commit ac55da81d8
3 changed files with 101 additions and 9 deletions

View File

@ -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):

View File

@ -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()

View File

@ -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)