From ac55da81d85d21b8f47e85ace068a815a1e55137 Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Sun, 6 Mar 2022 18:24:57 +0800
Subject: [PATCH] core: support precompilation of kernels

---
 artiq/coredevice/core.py                      | 62 ++++++++++++++++---
 .../kc705_nist_clock/repository/precompile.py | 20 ++++++
 artiq/test/coredevice/test_compile.py         | 28 +++++++++
 3 files changed, 101 insertions(+), 9 deletions(-)
 create mode 100644 artiq/examples/kc705_nist_clock/repository/precompile.py

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)