From 1a28069aa2d4b5395f0f34b9e9368b48e409d7d8 Mon Sep 17 00:00:00 2001 From: mwojcik Date: Thu, 19 Oct 2023 17:16:26 +0800 Subject: [PATCH] support for pre-compiling subkernels --- artiq/compiler/embedding.py | 7 +++++ artiq/coredevice/core.py | 42 ++++++++++++++++------------- artiq/frontend/artiq_compile.py | 48 ++++++++++++++++++++++++++------- artiq/frontend/artiq_run.py | 20 +++++++++++++- 4 files changed, 89 insertions(+), 28 deletions(-) diff --git a/artiq/compiler/embedding.py b/artiq/compiler/embedding.py index 9c2f270d8..3b6f6ae6c 100644 --- a/artiq/compiler/embedding.py +++ b/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())) diff --git a/artiq/coredevice/core.py b/artiq/coredevice/core.py index ca9123bee..4ad85b780 100644 --- a/artiq/coredevice/core.py +++ b/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): diff --git a/artiq/frontend/artiq_compile.py b/artiq/frontend/artiq_compile.py index 9aeceb6d9..04a466563 100755 --- a/artiq/frontend/artiq_compile.py +++ b/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 " .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() diff --git a/artiq/frontend/artiq_run.py b/artiq/frontend/artiq_run.py index 948c34475..f0a28e91c 100755 --- a/artiq/frontend/artiq_run.py +++ b/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)