forked from M-Labs/artiq
support for pre-compiling subkernels
This commit is contained in:
parent
56418e342e
commit
1a28069aa2
|
@ -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()))
|
||||
|
|
|
@ -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):
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue