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
|
subkernels[k] = v
|
||||||
return subkernels
|
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):
|
def has_rpc_or_subkernel(self):
|
||||||
return any(filter(lambda x: inspect.isfunction(x) or inspect.ismethod(x),
|
return any(filter(lambda x: inspect.isfunction(x) or inspect.ismethod(x),
|
||||||
self.object_forward_map.values()))
|
self.object_forward_map.values()))
|
||||||
|
|
|
@ -147,12 +147,11 @@ class Core:
|
||||||
result = new_result
|
result = new_result
|
||||||
embedding_map, kernel_library, symbolizer, demangler, subkernel_arg_types = \
|
embedding_map, kernel_library, symbolizer, demangler, subkernel_arg_types = \
|
||||||
self.compile(function, args, kwargs, set_result)
|
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)
|
self._run_compiled(kernel_library, embedding_map, symbolizer, demangler)
|
||||||
return result
|
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):
|
||||||
for sid, subkernel_fn in embedding_map.subkernels().items():
|
|
||||||
# pass self to subkernels (if applicable)
|
# pass self to subkernels (if applicable)
|
||||||
# assuming the first argument is self
|
# assuming the first argument is self
|
||||||
subkernel_args = getfullargspec(subkernel_fn.artiq_embedded.function)
|
subkernel_args = getfullargspec(subkernel_fn.artiq_embedded.function)
|
||||||
|
@ -169,6 +168,13 @@ class Core:
|
||||||
subkernel_arg_types=subkernel_arg_types.get(sid, []))
|
subkernel_arg_types=subkernel_arg_types.get(sid, []))
|
||||||
if object_map.has_rpc_or_subkernel():
|
if object_map.has_rpc_or_subkernel():
|
||||||
raise ValueError("Subkernel must not use RPC or subkernels in other destinations")
|
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)
|
self.comm.upload_subkernel(kernel_library, sid, destination)
|
||||||
|
|
||||||
def precompile(self, function, *args, **kwargs):
|
def precompile(self, function, *args, **kwargs):
|
||||||
|
|
|
@ -1,6 +1,6 @@
|
||||||
#!/usr/bin/env python3
|
#!/usr/bin/env python3
|
||||||
|
|
||||||
import os, sys, logging, argparse
|
import os, sys, io, tarfile, logging, argparse
|
||||||
|
|
||||||
from sipyco import common_args
|
from sipyco import common_args
|
||||||
|
|
||||||
|
@ -63,9 +63,16 @@ def main():
|
||||||
core_name = exp.run.artiq_embedded.core_name
|
core_name = exp.run.artiq_embedded.core_name
|
||||||
core = getattr(exp_inst, 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], {},
|
core.compile(exp.run, [exp_inst], {},
|
||||||
attribute_writeback=False, print_as_rpc=False)
|
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:
|
except CompileError as error:
|
||||||
return
|
return
|
||||||
finally:
|
finally:
|
||||||
|
@ -73,16 +80,39 @@ def main():
|
||||||
finally:
|
finally:
|
||||||
dataset_db.close_db()
|
dataset_db.close_db()
|
||||||
|
|
||||||
if object_map.has_rpc_or_subkernel():
|
if object_map.has_rpc():
|
||||||
raise ValueError("Experiment must not use RPC or subkernels")
|
raise ValueError("Experiment must not use RPC")
|
||||||
|
|
||||||
output = args.output
|
output = args.output
|
||||||
|
|
||||||
|
if not subkernels:
|
||||||
|
# just write the ELF file
|
||||||
if output is None:
|
if output is None:
|
||||||
basename, ext = os.path.splitext(args.file)
|
basename, ext = os.path.splitext(args.file)
|
||||||
output = "{}.elf".format(basename)
|
output = "{}.elf".format(basename)
|
||||||
|
|
||||||
with open(output, "wb") as f:
|
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__":
|
if __name__ == "__main__":
|
||||||
main()
|
main()
|
||||||
|
|
|
@ -4,6 +4,7 @@
|
||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
import sys
|
import sys
|
||||||
|
import tarfile
|
||||||
from operator import itemgetter
|
from operator import itemgetter
|
||||||
import logging
|
import logging
|
||||||
from collections import defaultdict
|
from collections import defaultdict
|
||||||
|
@ -86,6 +87,20 @@ class LLVMBitcodeRunner(FileRunner):
|
||||||
return self.target.link([self.target.assemble(llmodule)])
|
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:
|
class DummyScheduler:
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
self.rid = 0
|
self.rid = 0
|
||||||
|
@ -156,6 +171,7 @@ def _build_experiment(device_mgr, dataset_mgr, args):
|
||||||
argument_mgr = ProcessArgumentManager(arguments)
|
argument_mgr = ProcessArgumentManager(arguments)
|
||||||
managers = (device_mgr, dataset_mgr, argument_mgr, {})
|
managers = (device_mgr, dataset_mgr, argument_mgr, {})
|
||||||
if hasattr(args, "file"):
|
if hasattr(args, "file"):
|
||||||
|
is_tar = tarfile.is_tarfile(args.file)
|
||||||
is_elf = args.file.endswith(".elf")
|
is_elf = args.file.endswith(".elf")
|
||||||
is_ll = args.file.endswith(".ll")
|
is_ll = args.file.endswith(".ll")
|
||||||
is_bc = args.file.endswith(".bc")
|
is_bc = args.file.endswith(".bc")
|
||||||
|
@ -165,7 +181,9 @@ def _build_experiment(device_mgr, dataset_mgr, args):
|
||||||
if args.class_name:
|
if args.class_name:
|
||||||
raise ValueError("class-name not supported "
|
raise ValueError("class-name not supported "
|
||||||
"for precompiled kernels")
|
"for precompiled kernels")
|
||||||
if is_elf:
|
if is_tar:
|
||||||
|
return TARRunner(managers, file=args.file)
|
||||||
|
elif is_elf:
|
||||||
return ELFRunner(managers, file=args.file)
|
return ELFRunner(managers, file=args.file)
|
||||||
elif is_ll:
|
elif is_ll:
|
||||||
return LLVMIRRunner(managers, file=args.file)
|
return LLVMIRRunner(managers, file=args.file)
|
||||||
|
|
Loading…
Reference in New Issue