await sk message before receiving it

nac3artiq/demo/embedding_map.py

@@ -1,11 +1,19 @@
 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_map = {}
+        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",
@@ -38,6 +46,10 @@ class EmbeddingMap:
         self.function_map[key] = fun
         return key
 
+    def store_subkernel(self, key, fun):
+        self.subkernel_map[key] = fun
+        return key
+
     def store_object(self, obj):
         obj_id = id(obj)
         if obj_id in self.object_inverse_map:
@@ -64,3 +76,19 @@ class EmbeddingMap:
     def retrieve_str(self, key):
         return self.string_map[key]
 
+    def subkernels(self):
+        return self.subkernel_map
+    
+    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."""
@@ -215,8 +226,10 @@ class Core:
         else:
             obj = method
             name = ""
+        
+        dest = getattr(method, "__artiq_destination__", 0)
 
-        compiler.compile_method_to_file(obj, name, args, "module.elf", embedding)
+        compiler.compile_method_to_file(obj, name, dest, args, "module.elf", embedding)
 
     @kernel
     def reset(self):

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
                         }
@@ -361,6 +361,7 @@ impl Nac3 {
         &self,
         obj: &PyAny,
         method_name: &str,
+        destination: u8,
         args: Vec<&PyAny>,
         embedding_map: &PyAny,
         py: Python,
@@ -382,6 +383,7 @@ impl Nac3 {
         let store_obj = embedding_map.getattr("store_object").unwrap().to_object(py);
         let store_str = embedding_map.getattr("store_str").unwrap().to_object(py);
         let store_fun = embedding_map.getattr("store_function").unwrap().to_object(py);
+        let store_subkernel = embedding_map.getattr("store_subkernel").unwrap().to_object(py);
         let host_attributes = embedding_map.getattr("attributes_writeback").unwrap().to_object(py);
         let global_value_ids: Arc<RwLock<HashMap<_, _>>> = Arc::new(RwLock::new(HashMap::new()));
         let helper = PythonHelper {
@@ -412,6 +414,7 @@ impl Nac3 {
         let mut module_to_resolver_cache: HashMap<u64, _> = HashMap::new();
 
         let mut rpc_ids = vec![];
+        let mut subkernel_ids = vec![];
         for (stmt, path, module) in &self.top_levels {
             let py_module: &PyAny = module.extract(py)?;
             let module_id: u64 = id_fn.call1((py_module,))?.extract()?;
@@ -495,6 +498,32 @@ impl Nac3 {
                                 .any(|constant| *constant == Constant::Str("async".into()))
                         });
                         rpc_ids.push((None, def_id, is_async));
+                    } 
+                    else if decorator_list
+                        .iter()
+                        .any(|decorator| decorator_id_string(decorator) == Some("subkernel".to_string()))
+                    {
+                        if let Constant::Int(sk_dest) = decorator_get_destination(decorator) {
+                            if sk_dest < 0 || sk_dest > 255 {
+                                return Err(CompileError::new_err(format!(
+                                    "compilation failed\n----------\nSubkernel destination must be between 0 and 255 (at {})",
+                                    stmt.location
+                                )));
+                            }
+                            if sk_dest != destination {
+                                // subkernels with the same destination as currently compiled kernel
+                                // are treated as normal kernels
+                                store_subkernel.call1(py, (def_id.0.into_py(py), module.getattr(py, name.to_string().as_str()).unwrap())).unwrap();
+                                subkernel_ids.push((None, def_id, destination));
+                            }
+                        } else {
+                            return Err(CompileError::new_err(format!(
+                                "compilation failed\n----------\nDestination must be provided for subkernels (at {})",
+                                stmt.location
+                            )));
+                        }
+                        
+                        
                     }
                 }
                 StmtKind::ClassDef { name, body, .. } => {
@@ -517,6 +546,31 @@ impl Nac3 {
                                     )));
                                 }
                                 rpc_ids.push((Some((class_obj.clone(), *name)), def_id, is_async));
+                            } else if decorator_list.iter().any(|decorator| matches!(decorator.node, ExprKind::Name { id, .. } if id == "subkernel".into())) {
+                                if let Constant::Int(sk_dest) = decorator_get_destination(decorator) {
+                                    if name == &"__init__".into() {
+                                        return Err(CompileError::new_err(format!(
+                                            "compilation failed\n----------\nThe constructor of class {} should not be decorated with subkernel decorator (at {})",
+                                            class_name, stmt.location
+                                        )));
+                                    }
+                                    if sk_dest < 0 || sk_dest > 255 {
+                                        return Err(CompileError::new_err(format!(
+                                            "compilation failed\n----------\nSubkernel destination must be between 0 and 255 (at {})",
+                                            stmt.location
+                                        )));
+                                    }
+                                    if sk_dest != destination {
+                                        subkernel_ids.push((Some((class_obj.clone(), *name)), def_id, sk_dest));
+                                    }
+                                } else {
+                                    return Err(CompileError::new_err(format!(
+                                        "compilation failed\n----------\nDestination must be provided for subkernels (at {})",
+                                        stmt.location
+                                    )));
+                                }
+
+                                
                             }
                         }
                     }
@@ -592,7 +646,7 @@ impl Nac3 {
         );
         let signature = store.add_cty(signature);
 
-        if let Err(e) = composer.start_analysis(true) {
+        if let Err(e) = composer.start_analysis(true, Some(destination)) {
             // report error of __modinit__ separately
             return if e.iter().any(|err| err.contains("<nac3_synthesized_modinit>")) {
                 let msg = Self::report_modinit(
@@ -655,6 +709,43 @@ impl Nac3 {
                     }
                 }
             }
+            for (class_data, id, sk_dest) in &subkernel_ids {
+                let mut def = defs[id.0].write();
+                match &mut *def {
+                    TopLevelDef::Function { codegen_callback, .. } => {
+                        *codegen_callback = Some(subkernel_codegen_callback());
+                    }
+                    TopLevelDef::Class { methods, .. } => {
+                        let (class_def, method_name) = class_data.as_ref().unwrap();
+                        for (name, _, id) in &*methods {
+                            if name != method_name {
+                                continue;
+                            }
+                            if let TopLevelDef::Function { codegen_callback, .. } =
+                                &mut *defs[id.0].write()
+                            {
+                                *codegen_callback = Some(subkernel_codegen_callback());
+                                store_subkernel
+                                    .call1(
+                                        py,
+                                        (
+                                            id.0.into_py(py),
+                                            class_def
+                                                .getattr(py, name.to_string().as_str())
+                                                .unwrap(),
+                                        ),
+                                    )
+                                    .unwrap();
+                            }
+                        }
+                    }
+                    TopLevelDef::Variable { .. } => {
+                        return Err(CompileError::new_err(String::from(
+                            "Unsupported @subkernel annotation on global variable",
+                        )))
+                    }
+                }
+            }
         }
 
         let instance = {
@@ -904,6 +995,20 @@ fn decorator_get_flags(decorator: &Located<ExprKind>) -> Vec<Constant> {
     flags
 }
 
+fn decorator_get_destination(decorator: &Located<ExprKind>) -> Option<Constant> {
+    if let ExprKind::Call { keywords, .. } = &decorator.node {
+        for keyword in keywords {
+            if keyword.node.arg != Some("destination".into()) {
+                continue;
+            }
+            if let ExprKind::Constant { value, .. } = &keyword.node.value.node {
+                return Some(value.clone());
+            }
+        }
+    }
+    None
+}
+
 fn link_with_lld(elf_filename: String, obj_filename: String) -> PyResult<()> {
     let linker_args = vec![
         "-shared".to_string(),
@@ -1116,6 +1221,7 @@ impl Nac3 {
         &mut self,
         obj: &PyAny,
         method_name: &str,
+        destination: u8,
         args: Vec<&PyAny>,
         filename: &str,
         embedding_map: &PyAny,
@@ -1136,7 +1242,7 @@ impl Nac3 {
                 Ok(())
             };
 
-            self.compile_method(obj, method_name, args, embedding_map, py, &link_fn)
+            self.compile_method(obj, method_name, destination, args, embedding_map, py, &link_fn)
         } else {
             let link_fn = |module: &Module| {
                 let object_mem = target_machine
@@ -1154,7 +1260,7 @@ impl Nac3 {
                 }
             };
 
-            self.compile_method(obj, method_name, args, embedding_map, py, &link_fn)
+            self.compile_method(obj, method_name, destination, args, embedding_map, py, &link_fn)
         }
     }
 
@@ -1162,6 +1268,7 @@ impl Nac3 {
         &mut self,
         obj: &PyAny,
         method_name: &str,
+        destination: u8,
         args: Vec<&PyAny>,
         embedding_map: &PyAny,
         py: Python,
@@ -1185,7 +1292,7 @@ impl Nac3 {
                 Ok(PyBytes::new(py, &fs::read(filename).unwrap()).into())
             };
 
-            self.compile_method(obj, method_name, args, embedding_map, py, &link_fn)
+            self.compile_method(obj, method_name, destination, args, embedding_map, py, &link_fn)
         } else {
             let link_fn = |module: &Module| {
                 let object_mem = target_machine
@@ -1198,7 +1305,7 @@ impl Nac3 {
                 }
             };
 
-            self.compile_method(obj, method_name, args, embedding_map, py, &link_fn)
+            self.compile_method(obj, method_name, destination, args, embedding_map, py, &link_fn)
         }
     }
 }

nac3artiq/src/subkernels.rs

@@ -0,0 +1,365 @@
+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, '_>,
+        obj: &Option<(Type, ValueEnum<'ctx>)>,
+        fun: (&FunSignature, DefinitionId),
+        args: &[(Option<StrRef>, ValueEnum<'ctx>)],
+        generator: &mut dyn CodeGenerator,
+    ) {
+        let sid_type = ctx.ctx.i32_type();
+        let timeout_type = ctx.ctx.i64_type();
+        assert!(matches!(args.len(), 1..=2));
+        let timeout = if args.len() == 1 {
+            timeout_type.const_zero().to_basic_value_enum(context, generator, obj_ty)?; // ?
+        } else { args[0].1.clone().to_basic_value_enum(context, generator, obj_ty)?; } // ?
+
+        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, '_>,
+        obj: &Option<(Type, ValueEnum<'ctx>)>,
+        fun: (&FunSignature, DefinitionId),
+        args: &[(Option<StrRef>, ValueEnum<'ctx>)],
+        generator: &mut dyn CodeGenerator,
+    ) {
+        assert_eq!(args.len(), 1);
+        let sid_type = ctx.ctx.i32_type();
+        let dest_type = ctx.ctx.i8_type();
+        let run_type = ctx.ctx.bool_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.into(), run_type.into()], false),
+                None,
+            )
+        });
+
+        let subkernel_id = int32.const_int(fun.1 .0 as u64, false);
+        let destination = int32.const_int(fun.? as u64, false); // TODO
+        ctx.builder
+            .build_call_or_invoke(subkernel_load_run, &[subkernel_id.into(), destination.into(), run_type.const_zero()], "subkernel.preload")
+            .unwrap();
+    }
+
+    fn subkernel_callback_fn<'ctx>(
+        ctx: &mut CodeGenContext<'ctx, '_>,
+        obj: Option<(Type, ValueEnum<'ctx>)>,
+        fun: (&FunSignature, DefinitionId),
+        args: Vec<(Option<StrRef>, ValueEnum<'ctx>)>,
+        generator: &mut dyn CodeGenerator,
+    ) -> Result<Option<BasicValueEnum<'ctx>>, String> {
+        let int8 = ctx.ctx.i8_type();
+        let int32 = ctx.ctx.i32_type();
+        let size_type = generator.get_size_type(ctx.ctx);
+        let ptr_type = int8.ptr_type(AddressSpace::default());
+        let tag_ptr_type = ctx.ctx.struct_type(&[ptr_type.into(), size_type.into()], false);
+    
+        let subkernel_id = int32.const_int(fun.1 .0 as u64, false);
+        let destination = int32.const_int(fun.? as u64, false); // TODO
+
+        // -- start the subkernel
+        let sid_type = ctx.ctx.i32_type();
+        let dest_type = ctx.ctx.i8_type();
+        let run_type = ctx.ctx.i1_type();
+        let subkernel_start = ctx.module.get_function("subkernel_load_run").unwrap_or_else(|| {
+            ctx.module.add_function(
+                "subkernel_load_run",
+                ctx.ctx.void_type().fn_type(
+                    &[
+                        ctx.ctx.void_type().fn_type(&[sid_type.into(), dest_type.into(), run_type.into()], false),
+                    ],
+                    false,
+                ),
+                None,
+            )
+        });
+        ctx.builder
+            .build_call_or_invoke(subkernel_start, &[subkernel_id.into(), destination.into(), run_type.const_int(1, false)], "subkernel.run")
+            .unwrap();
+        // -- setup rpc tags
+        let mut tag = Vec::new();
+        if obj.is_some() {
+            tag.push(b'O');
+        }
+        for arg in &fun.0.args {
+            gen_rpc_tag(ctx, arg.ty, &mut tag)?;
+        }
+        tag.push(b':');
+        gen_rpc_tag(ctx, fun.0.ret, &mut tag)?;
+    
+        let mut hasher = DefaultHasher::new();
+        tag.hash(&mut hasher);
+        let hash = format!("{}", hasher.finish());
+    
+        let tag_ptr = ctx
+            .module
+            .get_global(hash.as_str())
+            .unwrap_or_else(|| {
+                let tag_arr_ptr = ctx.module.add_global(
+                    int8.array_type(tag.len() as u32),
+                    None,
+                    format!("tagptr{}", fun.1 .0).as_str(),
+                );
+                tag_arr_ptr.set_initializer(&int8.const_array(
+                    &tag.iter().map(|v| int8.const_int(u64::from(*v), false)).collect::<Vec<_>>(),
+                ));
+                tag_arr_ptr.set_linkage(Linkage::Private);
+                let tag_ptr = ctx.module.add_global(tag_ptr_type, None, &hash);
+                tag_ptr.set_linkage(Linkage::Private);
+                tag_ptr.set_initializer(&ctx.ctx.const_struct(
+                    &[
+                        tag_arr_ptr.as_pointer_value().const_cast(ptr_type).into(),
+                        size_type.const_int(tag.len() as u64, false).into(),
+                    ],
+                    false,
+                ));
+                tag_ptr
+            })
+            .as_pointer_value();
+    
+        let arg_length = args.len() + usize::from(obj.is_some());
+    
+        let stackptr = call_stacksave(ctx, Some("subkernel.stack"));
+        let args_ptr = ctx
+            .builder
+            .build_array_alloca(
+                ptr_type,
+                ctx.ctx.i32_type().const_int(arg_length as u64, false),
+                "argptr",
+            )
+            .unwrap();
+    
+        // -- rpc args handling
+        let mut keys = fun.0.args.clone();
+        let mut mapping = HashMap::new();
+        for (key, value) in args {
+            mapping.insert(key.unwrap_or_else(|| keys.remove(0).name), value);
+        }
+        // default value handling
+        for k in keys {
+            mapping
+                .insert(k.name, ctx.gen_symbol_val(generator, &k.default_value.unwrap(), k.ty).into());
+        }
+        // reorder the parameters
+        let mut real_params = fun
+            .0
+            .args
+            .iter()
+            .map(|arg| {
+                mapping
+                    .remove(&arg.name)
+                    .unwrap()
+                    .to_basic_value_enum(ctx, generator, arg.ty)
+                    .map(|llvm_val| (llvm_val, arg.ty))
+            })
+            .collect::<Result<Vec<(_, _)>, _>>()?;
+        if let Some(obj) = obj {
+            if let ValueEnum::Static(obj_val) = obj.1 {
+                real_params.insert(0, (obj_val.get_const_obj(ctx, generator), obj.0));
+            } else {
+                // should be an error here...
+                panic!("only host object is allowed");
+            }
+        }
+    
+        for (i, (arg, arg_ty)) in real_params.iter().enumerate() {
+            let arg_slot = format_rpc_arg(generator, ctx, (*arg, *arg_ty, i));
+            let arg_ptr = unsafe {
+                ctx.builder.build_gep(
+                    args_ptr,
+                    &[int32.const_int(i as u64, false)],
+                    &format!("subkernel.arg{i}"),
+                )
+            }
+            .unwrap();
+            ctx.builder.build_store(arg_ptr, arg_slot).unwrap();
+        }
+    
+        // send the message
+        let subkernel_send = ctx.module.get_function("subkernel_send_message").unwrap_or_else(|| {
+            ctx.module.add_function(
+                "subkernel_send_message",
+                ctx.ctx.void_type().fn_type(
+                    &[
+                        int32.into(),
+                        tag_ptr_type.ptr_type(AddressSpace::default()).into(),
+                        ptr_type.ptr_type(AddressSpace::default()).into(),
+                    ],
+                    false,
+                ),
+                None,
+            )
+        });
+        ctx.builder
+            .build_call_or_invoke(subkernel_send, &[service_id.into(), tag_ptr.into(), args_ptr.into()], "subkernel.send")
+            .unwrap();
+        // reclaim stack space used by arguments
+        call_stackrestore(ctx, stackptr);
+    }
+
+    pub fn subkernel_codegen_callback() -> Arc<GenCall> {
+        Arc::new(GenCall::new(Box::new(|ctx, obj, fun, args, generator| {
+            subkernel_codegen_callback_fn(ctx, obj, fun, args, generator)
+        })))
+    }
+
+    fn subkernel_recv_message<'ctx>(
+        ctx: &mut CodeGenContext<'ctx, '_>,
+        obj: Option<(Type, ValueEnum<'ctx>)>,
+        fun: (&FunSignature, DefinitionId),
+        args: Vec<(Option<StrRef>, ValueEnum<'ctx>)>,
+        generator: &mut dyn CodeGenerator,
+    ) -> Result<Option<BasicValueEnum<'ctx>>, String>  {
+
+        // await for the message first
+        let int8 = ctx.ctx.i8_type();
+        let int64 = ctx.ctx.i64_type();
+        let size_type = generator.get_size_type(ctx.ctx);
+        let ptr_type = int8.ptr_type(AddressSpace::default());
+        let tag_ptr_type = ctx.ctx.struct_type(&[ptr_type.into(), size_type.into()], false);
+    
+        let subkernel_await_message = ctx.module.get_function("subkernel_await_message").unwrap_or_else(|| {
+            ctx.module.add_function(
+                "subkernel_await_message",
+                ctx.ctx.void_type().fn_type(
+                    &[
+                        ctx.ctx.void_type().fn_type(&[int8.into(), int64.into(), tag_ptr_type.into(), int8.into(), int8.into()], false),
+                    ],
+                    false,
+                ),
+                None,
+            )
+        });
+
+        ctx.builder
+            .build_call_or_invoke(subkernel_await_message, &[subkernel_id.into(), destination.into(), run_type.const_int(1, false)], "subkernel.run")
+            .unwrap();
+        
+        // -- receive value:
+        // T result = {
+        //   void *ret_ptr = alloca(sizeof(T));
+        //   void *ptr = ret_ptr;
+        //   loop: int size = rpc_recv(ptr);
+        //   // Non-zero: Provide `size` bytes of extra storage for variable-length data.
+        //   if(size) { ptr = alloca(size); goto loop; }
+        //   else *(T*)ret_ptr
+        // }
+        let rpc_recv = ctx.module.get_function("rpc_recv").unwrap_or_else(|| {
+            ctx.module.add_function("rpc_recv", int32.fn_type(&[ptr_type.into()], false), None)
+        });
+    
+        if ctx.unifier.unioned(fun.0.ret, ctx.primitives.none) {
+            ctx.build_call_or_invoke(rpc_recv, &[ptr_type.const_null().into()], "rpc_recv");
+            return Ok(None);
+        }
+    
+        let prehead_bb = ctx.builder.get_insert_block().unwrap();
+        let current_function = prehead_bb.get_parent().unwrap();
+        let head_bb = ctx.ctx.append_basic_block(current_function, "rpc.head");
+        let alloc_bb = ctx.ctx.append_basic_block(current_function, "rpc.continue");
+        let tail_bb = ctx.ctx.append_basic_block(current_function, "rpc.tail");
+    
+        let ret_ty = ctx.get_llvm_abi_type(generator, fun.0.ret);
+        let need_load = !ret_ty.is_pointer_type();
+        let slot = ctx.builder.build_alloca(ret_ty, "rpc.ret.slot").unwrap();
+        let slotgen = ctx.builder.build_bitcast(slot, ptr_type, "rpc.ret.ptr").unwrap();
+        ctx.builder.build_unconditional_branch(head_bb).unwrap();
+        ctx.builder.position_at_end(head_bb);
+    
+        let phi = ctx.builder.build_phi(ptr_type, "rpc.ptr").unwrap();
+        phi.add_incoming(&[(&slotgen, prehead_bb)]);
+        let alloc_size = ctx
+            .build_call_or_invoke(rpc_recv, &[phi.as_basic_value()], "rpc.size.next")
+            .unwrap()
+            .into_int_value();
+        let is_done = ctx
+            .builder
+            .build_int_compare(inkwell::IntPredicate::EQ, int32.const_zero(), alloc_size, "rpc.done")
+            .unwrap();
+    
+        ctx.builder.build_conditional_branch(is_done, tail_bb, alloc_bb).unwrap();
+        ctx.builder.position_at_end(alloc_bb);
+    
+        let alloc_ptr = ctx.builder.build_array_alloca(ptr_type, alloc_size, "rpc.alloc").unwrap();
+        let alloc_ptr = ctx.builder.build_bitcast(alloc_ptr, ptr_type, "rpc.alloc.ptr").unwrap();
+        phi.add_incoming(&[(&alloc_ptr, alloc_bb)]);
+        ctx.builder.build_unconditional_branch(head_bb).unwrap();
+    
+        ctx.builder.position_at_end(tail_bb);
+    
+        let result = ctx.builder.build_load(slot, "rpc.result").unwrap();
+        if need_load {
+            call_stackrestore(ctx, stackptr);
+        }
+        Ok(Some(result))
+    }
+}

nac3core/src/toplevel/composer.rs

@@ -463,13 +463,13 @@ impl TopLevelComposer {
         Ok((name, DefinitionId(self.definition_ast_list.len() - 1), Some(ty_to_be_unified)))
     }
 
-    pub fn start_analysis(&mut self, inference: bool) -> Result<(), HashSet<String>> {
+    pub fn start_analysis(&mut self, inference: bool, destination: Option<u8>) -> Result<(), HashSet<String>> {
         self.analyze_top_level_class_type_var()?;
         self.analyze_top_level_class_bases()?;
         self.analyze_top_level_class_fields_methods()?;
         self.analyze_top_level_function()?;
         if inference {
-            self.analyze_function_instance()?;
+            self.analyze_function_instance(destination)?;
         }
         self.analyze_top_level_variables()?;
         Ok(())
@@ -1736,7 +1736,8 @@ impl TopLevelComposer {
 
     /// step 5, analyze and call type inferencer to fill the `instance_to_stmt` of
     /// [`TopLevelDef::Function`]
-    fn analyze_function_instance(&mut self) -> Result<(), HashSet<String>> {
+    /// destination is an ARTIQ-only argument necessary for proper support of subkernels
+    fn analyze_function_instance(&mut self, destination: Option<u8>) -> Result<(), HashSet<String>> {
         // first get the class constructor type correct for the following type check in function body
         // also do class field instantiation check
         let init_str_id = "__init__".into();
@@ -2134,6 +2135,28 @@ impl TopLevelComposer {
                         instance_to_symbol.insert(String::new(), simple_name.to_string());
                         continue;
                     }
+                    if let ExprKind::Call { func, keywords, .. } = &decorator_list[0].node {
+                        if matches!(&func.node, ExprKind::Name { id, .. } if id == &"subkernel".into())
+                    {
+                        let subkernel_dest = keywords.iter().find(|keyword| {
+                                if keyword.node.arg == Some("destination".into()) {
+                                    match &keyword.node.value.node {
+                                        ExprKind::Constant { .. } => true,
+                                        _ => false
+                                    }
+                                } else {
+                                    false
+                                }
+                            }
+                        ).node.value.node.value;
+                        // treat subkernel as kernel if the destination is the same as currently compiled
+                        if let Some(Constant::Int(sk_dest)) = subkernel_dest {
+                            if sk_dest != destination.unwrap_or(0) {
+                                instance_to_symbol.insert(String::new(), simple_name.to_string());
+                            }
+                        }
+                        continue;
+                    }
 
                     if matches!(&decorator_list[0].node, ExprKind::Name { id, .. } if id == &"rpc".into())
                     {

nac3core/src/toplevel/test.rs

@@ -199,7 +199,7 @@ fn test_simple_function_analyze(source: &[&str], tys: &[&str], names: &[&str]) {
         }
     }
 
-    composer.start_analysis(true).unwrap();
+    composer.start_analysis(true, None).unwrap();
 
     for (i, (def, _)) in composer.definition_ast_list.iter().skip(composer.builtin_num).enumerate()
     {
@@ -563,7 +563,7 @@ fn test_analyze(source: &[&str], res: &[&str]) {
         }
     }
 
-    if let Err(msg) = composer.start_analysis(false) {
+    if let Err(msg) = composer.start_analysis(false, None) {
         if print {
             println!("{}", msg.iter().sorted().join("\n----------\n"));
         } else {
@@ -748,7 +748,7 @@ fn test_inference(source: Vec<&str>, res: &[&str]) {
         }
     }
 
-    if let Err(msg) = composer.start_analysis(true) {
+    if let Err(msg) = composer.start_analysis(true, None) {
         if print {
             println!("{}", msg.iter().sorted().join("\n----------\n"));
         } else {

nac3standalone/src/main.rs

@@ -408,7 +408,7 @@ fn main() {
     let signature = store.from_signature(&mut composer.unifier, &primitive, &signature, &mut cache);
     let signature = store.add_cty(signature);
 
-    if let Err(errors) = composer.start_analysis(true) {
+    if let Err(errors) = composer.start_analysis(true, None) {
         let error_count = errors.len();
         eprintln!("{error_count} error(s) occurred during top level analysis.");