initial subkernels commit

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

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."""

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
                         }

nac3artiq/src/subkernels.rs

@@ -0,0 +1,91 @@
+pub struct Subkernels {}
+
+impl Subkernels {
+    fn get_subkernel_builtins() -> Vec<Box<BuiltinFuncSpec>> {
+        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
+    }
+}