occheung
/
artiq
forked from M-Labs/artiq
1
0
Fork 0
Code Pull Requests Activity

support for pre-compiling subkernels

This commit is contained in:
mwojcik 2023-10-19 17:16:26 +08:00 committed by Sébastien Bourdeauducq
parent 56418e342e
commit 1a28069aa2
4 changed files with 89 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
artiq/compiler/embedding.py
View File

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

12
artiq/coredevice/core.py
View File

@ -147,12 +147,11 @@ 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):
for sid, subkernel_fn in embedding_map.subkernels().items():
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)
@ -169,6 +168,13 @@ class Core:
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():
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):

40
artiq/frontend/artiq_compile.py
View File

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

20
artiq/frontend/artiq_run.py
View File

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