From a6dcd8dca3cc6517702043c2c755095ba77d9b7c Mon Sep 17 00:00:00 2001 From: mwojcik Date: Tue, 3 Sep 2024 16:58:57 +0800 Subject: [PATCH] initial subkernels commit --- nac3artiq/demo/embedding_map.py | 29 ++++++++++- nac3artiq/demo/min_artiq.py | 11 ++++ nac3artiq/src/lib.rs | 4 +- nac3artiq/src/subkernels.rs | 91 +++++++++++++++++++++++++++++++++ 4 files changed, 132 insertions(+), 3 deletions(-) create mode 100644 nac3artiq/src/subkernels.rs diff --git a/nac3artiq/demo/embedding_map.py b/nac3artiq/demo/embedding_map.py index 2116c69e..cecdd16a 100644 --- a/nac3artiq/demo/embedding_map.py +++ b/nac3artiq/demo/embedding_map.py @@ -1,11 +1,18 @@ class EmbeddingMap: - def __init__(self): + def __init__(self, old_embedding_map=None): self.object_inverse_map = {} self.object_map = {} self.string_map = {} self.string_reverse_map = {} self.function_map = {} self.attributes_writeback = [] + self.subkernel_message_map = {} + + if not old_embedding_map is None: + for key, obj_ref in old_embedding_map.subkernels().items(): + self.object_map[key] = obj_ref + obj_id = id(obj_ref) + self.object_inverse_map[obj_id] = key # preallocate exception names self.preallocate_runtime_exception_names(["RuntimeError", @@ -64,3 +71,23 @@ class EmbeddingMap: def retrieve_str(self, key): return self.string_map[key] + def subkernels(self): + subkernels = {} + for k, v in self.object_map.items(): + if getattr(v, "__artiq_destination__") is not None: + subkernels[k] = v + return subkernels + + def subkernel_messages(self): + messages = {} + for msg_id in self.subkernel_message_map.values(): + messages[msg_id] = self.retrieve_object(msg_id) + return messages + + def subkernel_messages_unpaired(self): + unpaired = [] + for msg_id in self.subkernel_message_map.values(): + msg_obj = self.retrieve_object(msg_id) + if msg_obj.send_loc is None or msg_obj.recv_loc is None: + unpaired.append(msg_obj) + return unpaired \ No newline at end of file diff --git a/nac3artiq/demo/min_artiq.py b/nac3artiq/demo/min_artiq.py index 3840a57a..4ed458ee 100644 --- a/nac3artiq/demo/min_artiq.py +++ b/nac3artiq/demo/min_artiq.py @@ -135,8 +135,19 @@ def kernel(function_or_method): @wraps(function_or_method) def run_on_core(*args, **kwargs): raise RuntimeError("Kernel functions need explicit core.run()") + run_on_core.__artiq_kernel__ = True + run_on_core.__artiq_destination__ = None return run_on_core +def subkernel(function_or_method, destination=0): + """Decorates a function or method to be executed on a satellite core device.""" + _register_function(function_or_method) + @wraps(function_or_method) + def run_on_core(*args, **kwargs): + raise RuntimeError("Subkernels cannot be called by the host") + run_on_core.__artiq_kernel__ = True + run_on_core.__artiq_destination__ = destination + return function_or_method def portable(function): """Decorates a function or method to be executed on the same device (host/core device) as the caller.""" diff --git a/nac3artiq/src/lib.rs b/nac3artiq/src/lib.rs index 6e80fd03..c280f74c 100644 --- a/nac3artiq/src/lib.rs +++ b/nac3artiq/src/lib.rs @@ -196,7 +196,7 @@ impl Nac3 { if let StmtKind::FunctionDef { ref decorator_list, .. } = stmt.node { decorator_list.iter().any(|decorator| { if let Some(id) = decorator_id_string(decorator) { - id == "kernel" || id == "portable" || id == "rpc" + id == "kernel" || id == "portable" || id == "rpc" || id == "subkernel" } else { false } @@ -210,7 +210,7 @@ impl Nac3 { StmtKind::FunctionDef { ref decorator_list, .. } => { decorator_list.iter().any(|decorator| { if let Some(id) = decorator_id_string(decorator) { - id == "extern" || id == "kernel" || id == "portable" || id == "rpc" + id == "extern" || id == "kernel" || id == "portable" || id == "rpc" || id == "subkernel" } else { false } diff --git a/nac3artiq/src/subkernels.rs b/nac3artiq/src/subkernels.rs new file mode 100644 index 00000000..3bbd1725 --- /dev/null +++ b/nac3artiq/src/subkernels.rs @@ -0,0 +1,91 @@ +pub struct Subkernels {} + +impl Subkernels { + fn get_subkernel_builtins() -> Vec> { + vec![ + Box::new(|primitives, unifier| { + let arg_ty = unifier.get_fresh_var(Some("T".into()), None); + ( + "subkernel_await".into(), + FunSignature { + args: vec![FuncArg { + name: "subkernel".into(), + ty: arg_ty.ty, + default_value: None, + is_vararg: false, + }, + FuncArg { + name: "timeout".into(), + ty: primitives.int32, + default_value: 0, + is_vararg: false + }], + // todo: figure out how to return the type function returns + // or catch it in the type inferencer + ret: primitives.none, + vars: into_var_map([arg_ty]), + }, + Arc::new(GenCall::new(Box::new(move |ctx, obj, fun, args, generator| { + gen_subkernel_await(ctx, &obj, fun, &args, generator)?; + + Ok(None) + }))), + ) + }), + Box::new(|primitives, unifier| { + let arg_ty = unifier.get_fresh_var(Some("T".into()), None); + ( + "subkernel_preload".into(), + FunSignature { + args: vec![ + FuncArg { + name: "subkernel".into(), + ty: arg_ty.ty, + default_value: None, + is_vararg: false, + } + ], + ret: primitives.none, + vars: into_var_map([arg_ty]), + }, + Arc::new(GenCall::new(Box::new(move |ctx, obj, fun, args, generator| { + gen_subkernel_preload(ctx, &obj, fun, &args, generator)?; + + Ok(None) + }))), + ) + }), + ] + } + + fn gen_subkernel_await<'ctx>(ctx: &mut CodeGenContext<'ctx, '_>, awaited: BasicValueEnum<'ctx>, timeout: BasicValueEnum<'ctx>) { + let sid_type = ctx.ctx.i32_type(); + // how to deal with optional arguments? + let timeout_type = ctx.ctx.i64_type(); + let subkernel_await_finish = ctx.module.get_function("subkernel_await_finish").unwrap_or_else(|| { + ctx.module.add_function( + "subkernel_await_finish", + ctx.ctx.void_type().fn_type(&[sid_type.into(), timeout_type.into()], false), + None, + ) + }); + + // call or invoke + // generate RPC for receiving return value depending on fun ret + } + + fn gen_subkernel_preload<'ctx>(ctx: &mut CodeGenContext<'ctx, '_>, preloaded: BasicValueEnum<'ctx>) { + let sid_type = ctx.ctx.i32_type(); + let dest_type = ctx.ctx.i8_type(); + let run_type = ctx.ctx.i1_type(); + let subkernel_load_run = ctx.module.get_function("subkernel_load_run").unwrap_or_else(|| { + ctx.module.add_function( + "subkernel_load_run", + ctx.ctx.void_type().fn_type(&[sid_type.into(), dest_type, run_type], false), + None, + ) + }); + // retrieve destination and sid from the fn (?) + // call or invoke + } +}