support for pre-compiling subkernels

artiq/compiler/embedding.py

@@ -193,6 +193,13 @@ class EmbeddingMap:
                     subkernels[k] = v
         return subkernels
 
+    def has_rpc(self):
+        return any(filter(
+                lambda x: (inspect.isfunction(x) or inspect.ismethod(x)) and \
+                    (not hasattr(x, "artiq_embedded") or x.artiq_embedded.destination is None),
+                self.object_forward_map.values()
+            ))
+
     def has_rpc_or_subkernel(self):
         return any(filter(lambda x: inspect.isfunction(x) or inspect.ismethod(x),
                           self.object_forward_map.values()))

artiq/coredevice/core.py

@@ -147,28 +147,34 @@ class Core:
             result = new_result
         embedding_map, kernel_library, symbolizer, demangler, subkernel_arg_types = \
             self.compile(function, args, kwargs, set_result)
-        self.compile_subkernels(embedding_map, args, subkernel_arg_types)
+        self.compile_and_upload_subkernels(embedding_map, args, subkernel_arg_types)
         self._run_compiled(kernel_library, embedding_map, symbolizer, demangler)
         return result
 
-    def compile_subkernels(self, embedding_map, args, subkernel_arg_types):
+    def compile_subkernel(self, sid, subkernel_fn, embedding_map, args, subkernel_arg_types):
+        # pass self to subkernels (if applicable)
+        # assuming the first argument is self
+        subkernel_args = getfullargspec(subkernel_fn.artiq_embedded.function)
+        self_arg = []
+        if len(subkernel_args[0]) > 0:
+            if subkernel_args[0][0] == 'self':
+                self_arg = args[:1]
+        destination = subkernel_fn.artiq_embedded.destination
+        destination_tgt = self.satellite_cpu_targets[destination]
+        target = get_target_cls(destination_tgt)(subkernel_id=sid)
+        object_map, kernel_library, _, _, _ = \
+            self.compile(subkernel_fn, self_arg, {}, attribute_writeback=False,
+                        print_as_rpc=False, target=target, destination=destination, 
+                        subkernel_arg_types=subkernel_arg_types.get(sid, []))
+        if object_map.has_rpc_or_subkernel():
+            raise ValueError("Subkernel must not use RPC or subkernels in other destinations")
+        return destination, kernel_library
+
+    def compile_and_upload_subkernels(self, embedding_map, args, subkernel_arg_types):
         for sid, subkernel_fn in embedding_map.subkernels().items():
-            # pass self to subkernels (if applicable)
-            # assuming the first argument is self
-            subkernel_args = getfullargspec(subkernel_fn.artiq_embedded.function)
-            self_arg = []
-            if len(subkernel_args[0]) > 0:
-                if subkernel_args[0][0] == 'self':
-                    self_arg = args[:1]
-            destination = subkernel_fn.artiq_embedded.destination
-            destination_tgt = self.satellite_cpu_targets[destination]
-            target = get_target_cls(destination_tgt)(subkernel_id=sid)
-            object_map, kernel_library, _, _, _ = \
-                self.compile(subkernel_fn, self_arg, {}, attribute_writeback=False,
-                            print_as_rpc=False, target=target, destination=destination, 
-                            subkernel_arg_types=subkernel_arg_types.get(sid, []))
-            if object_map.has_rpc_or_subkernel():
-                raise ValueError("Subkernel must not use RPC or subkernels in other destinations")
+            destination, kernel_library = \
+                self.compile_subkernel(sid, subkernel_fn, embedding_map,
+                                       args, subkernel_arg_types)
             self.comm.upload_subkernel(kernel_library, sid, destination)
 
     def precompile(self, function, *args, **kwargs):

artiq/frontend/artiq_compile.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-import os, sys, logging, argparse
+import os, sys, io, tarfile, logging, argparse
 
 from sipyco import common_args
 
@@ -63,9 +63,16 @@ def main():
             core_name = exp.run.artiq_embedded.core_name
             core = getattr(exp_inst, core_name)
 
-            object_map, kernel_library, _, _, _ = \
+            object_map, main_kernel_library, _, _, subkernel_arg_types = \
                 core.compile(exp.run, [exp_inst], {},
                              attribute_writeback=False, print_as_rpc=False)
+
+            subkernels = {}
+            for sid, subkernel_fn in object_map.subkernels().items():
+                destination, subkernel_library = core.compile_subkernel(
+                    sid, subkernel_fn, object_map, 
+                    [exp_inst], subkernel_arg_types)
+                subkernels[sid] = (destination, subkernel_library)
         except CompileError as error:
             return
         finally:
@@ -73,16 +80,39 @@ def main():
     finally:
         dataset_db.close_db()
 
-    if object_map.has_rpc_or_subkernel():
-        raise ValueError("Experiment must not use RPC or subkernels")
+    if object_map.has_rpc():
+        raise ValueError("Experiment must not use RPC")
 
     output = args.output
-    if output is None:
-        basename, ext = os.path.splitext(args.file)
-        output = "{}.elf".format(basename)
 
-    with open(output, "wb") as f:
-        f.write(kernel_library)
+    if not subkernels:
+        # just write the ELF file
+        if output is None:
+            basename, ext = os.path.splitext(args.file)
+            output = "{}.elf".format(basename)
+
+        with open(output, "wb") as f:
+            f.write(main_kernel_library)
+    else:
+        # combine them in a tar archive
+        if output is None:
+            basename, ext = os.path.splitext(args.file)
+            output = "{}.tar".format(basename)
+
+        with tarfile.open(output, "w:") as tar:
+            # write the main lib as "main.elf"
+            main_kernel_fileobj = io.BytesIO(main_kernel_library)
+            main_kernel_info = tarfile.TarInfo(name="main.elf")
+            main_kernel_info.size = len(main_kernel_library)
+            tar.addfile(main_kernel_info, fileobj=main_kernel_fileobj)
+
+            # subkernels as "<sid> <destination>.elf"
+            for sid, (destination, subkernel_library) in subkernels.items():
+                subkernel_fileobj = io.BytesIO(subkernel_library)
+                subkernel_info = tarfile.TarInfo(name="{} {}.elf".format(sid, destination))
+                subkernel_info.size = len(subkernel_library)
+                tar.addfile(subkernel_info, fileobj=subkernel_fileobj)
+
 
 if __name__ == "__main__":
     main()

artiq/frontend/artiq_run.py

@@ -4,6 +4,7 @@
 
 import argparse
 import sys
+import tarfile
 from operator import itemgetter
 import logging
 from collections import defaultdict
@@ -86,6 +87,20 @@ class LLVMBitcodeRunner(FileRunner):
         return self.target.link([self.target.assemble(llmodule)])
 
 
+class TARRunner(FileRunner):
+    def compile(self):
+        with tarfile.open(self.file, "r:") as tar:
+            for entry in tar:
+                if entry.name == 'main.elf':
+                    main_lib = tar.extractfile(entry).read()
+                else:
+                    subkernel_name = entry.name.removesuffix(".elf")
+                    sid, dest = tuple(map(lambda x: int(x), subkernel_name.split(" ")))
+                    subkernel_lib = tar.extractfile(entry).read()
+                    self.core.comm.upload_subkernel(subkernel_lib, sid, dest)
+        return main_lib
+
+
 class DummyScheduler:
     def __init__(self):
         self.rid = 0
@@ -156,6 +171,7 @@ def _build_experiment(device_mgr, dataset_mgr, args):
     argument_mgr = ProcessArgumentManager(arguments)
     managers = (device_mgr, dataset_mgr, argument_mgr, {})
     if hasattr(args, "file"):
+        is_tar = tarfile.is_tarfile(args.file)
         is_elf = args.file.endswith(".elf")
         is_ll  = args.file.endswith(".ll")
         is_bc  = args.file.endswith(".bc")
@@ -165,7 +181,9 @@ def _build_experiment(device_mgr, dataset_mgr, args):
             if args.class_name:
                 raise ValueError("class-name not supported "
                                  "for precompiled kernels")
-        if is_elf:
+        if is_tar:
+            return TARRunner(managers, file=args.file)
+        elif is_elf:
             return ELFRunner(managers, file=args.file)
         elif is_ll:
             return LLVMIRRunner(managers, file=args.file)