artiq: support async rpcs

nac3artiq/demo/min_artiq.py

@@ -112,10 +112,15 @@ def extern(function):
     register_function(function)
     return function
 
-def rpc(function):
+
-    """Decorates a function declaration defined by the core device runtime."""
+def rpc(arg=None, flags={}):
-    register_function(function)
+    """Decorates a function or method to be executed on the host interpreter."""
-    return function
+    if arg is None:
+        def inner_decorator(function):
+            return rpc(function, flags)
+        return inner_decorator
+    register_function(arg)
+    return arg
 
 def kernel(function_or_method):
     """Decorates a function or method to be executed on the core device."""

nac3artiq/src/codegen.rs

@@ -824,6 +824,7 @@ fn rpc_codegen_callback_fn<'ctx>(
     fun: (&FunSignature, DefinitionId),
     args: Vec<(Option<StrRef>, ValueEnum<'ctx>)>,
     generator: &mut dyn CodeGenerator,
+    is_async: bool,
 ) -> Result<Option<BasicValueEnum<'ctx>>, String> {
     let int8 = ctx.ctx.i8_type();
     let int32 = ctx.ctx.i32_type();
@@ -831,6 +832,7 @@ fn rpc_codegen_callback_fn<'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 service_id = int32.const_int(fun.1 .0 as u64, false);
     // -- setup rpc tags
     let mut tag = Vec::new();
@@ -932,6 +934,25 @@ fn rpc_codegen_callback_fn<'ctx>(
     }
 
     // call
+    if is_async {
+        let rpc_send_async = ctx.module.get_function("rpc_send_async").unwrap_or_else(|| {
+            ctx.module.add_function(
+                "rpc_send_async",
+                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(rpc_send_async, &[service_id.into(), tag_ptr.into(), args_ptr.into()], "rpc.send")
+            .unwrap();
+    } else {
         let rpc_send = ctx.module.get_function("rpc_send").unwrap_or_else(|| {
             ctx.module.add_function(
                 "rpc_send",
@@ -949,10 +970,15 @@ fn rpc_codegen_callback_fn<'ctx>(
         ctx.builder
             .build_call(rpc_send, &[service_id.into(), tag_ptr.into(), args_ptr.into()], "rpc.send")
             .unwrap();
+    }
 
     // reclaim stack space used by arguments
     call_stackrestore(ctx, stackptr);
 
+    if is_async {
+        // async RPCs do not return any values
+        Ok(None)
+    } else {
         let result = format_rpc_ret(generator, ctx, fun.0.ret);
     
         if !result.is_some_and(|res| res.get_type().is_pointer_type()) {
@@ -962,6 +988,7 @@ fn rpc_codegen_callback_fn<'ctx>(
     
         Ok(result)
     }
+}
 
 pub fn attributes_writeback(
     ctx: &mut CodeGenContext<'_, '_>,
@@ -1055,7 +1082,7 @@ pub fn attributes_writeback(
         let args: Vec<_> =
             values.into_iter().map(|(_, val)| (None, ValueEnum::Dynamic(val))).collect();
         if let Err(e) =
-            rpc_codegen_callback_fn(ctx, None, (&fun, PrimDef::Int32.id()), args, generator)
+            rpc_codegen_callback_fn(ctx, None, (&fun, PrimDef::Int32.id()), args, generator, false)
         {
             return Ok(Err(e));
         }
@@ -1065,9 +1092,9 @@ pub fn attributes_writeback(
     Ok(())
 }
 
-pub fn rpc_codegen_callback() -> Arc<GenCall> {
+pub fn rpc_codegen_callback(is_async: bool) -> Arc<GenCall> {
-    Arc::new(GenCall::new(Box::new(|ctx, obj, fun, args, generator| {
+    Arc::new(GenCall::new(Box::new(move |ctx, obj, fun, args, generator| {
-        rpc_codegen_callback_fn(ctx, obj, fun, args, generator)
+        rpc_codegen_callback_fn(ctx, obj, fun, args, generator, is_async)
     })))
 }
 

nac3artiq/src/lib.rs

@@ -34,8 +34,12 @@ use nac3core::inkwell::{
     targets::*,
     OptimizationLevel,
 };
-use nac3core::nac3parser::{
+use itertools::Itertools;
-    ast::{ExprKind, Stmt, StmtKind, StrRef},
+use nac3core::codegen::{gen_func_impl, CodeGenLLVMOptions, CodeGenTargetMachineOptions};
+use nac3core::toplevel::builtins::get_exn_constructor;
+use nac3core::typecheck::typedef::{into_var_map, TypeEnum, Unifier, VarMap};
+use nac3parser::{
+    ast::{Constant, ExprKind, Located, Stmt, StmtKind, StrRef},
     parser::parse_program,
 };
 use nac3core::toplevel::builtins::get_exn_constructor;
@@ -135,7 +139,7 @@ struct Nac3 {
     string_store: Arc<RwLock<HashMap<String, i32>>>,
     exception_ids: Arc<RwLock<HashMap<usize, usize>>>,
     deferred_eval_store: DeferredEvaluationStore,
-    /// LLVM-related options for code generation.
+    /// LLVM-related options for code generzation.
     llvm_options: CodeGenLLVMOptions,
 }
 
@@ -194,10 +198,8 @@ impl Nac3 {
                     body.retain(|stmt| {
                         if let StmtKind::FunctionDef { ref decorator_list, .. } = stmt.node {
                             decorator_list.iter().any(|decorator| {
-                                if let ExprKind::Name { id, .. } = decorator.node {
+                                if let Some(id) = decorator_id_string(decorator) {
-                                    id.to_string() == "kernel"
+                                    id == "kernel" || id == "portable" || id == "rpc"
-                                        || id.to_string() == "portable"
-                                        || id.to_string() == "rpc"
                                 } else {
                                     false
                                 }
@@ -210,9 +212,8 @@ impl Nac3 {
                 }
                 StmtKind::FunctionDef { ref decorator_list, .. } => {
                     decorator_list.iter().any(|decorator| {
-                        if let ExprKind::Name { id, .. } = decorator.node {
+                        if let Some(id) = decorator_id_string(decorator) {
-                            let id = id.to_string();
+                            id == "extern" || id == "kernel" || id == "portable" || id == "rpc"
-                            id == "extern" || id == "portable" || id == "kernel" || id == "rpc"
                         } else {
                             false
                         }
@@ -478,9 +479,13 @@ impl Nac3 {
 
             match &stmt.node {
                 StmtKind::FunctionDef { decorator_list, .. } => {
-                    if decorator_list.iter().any(|decorator| matches!(decorator.node, ExprKind::Name { id, .. } if id == "rpc".into())) {
+                    if decorator_list.iter().any(|decorator| decorator_id_string(decorator) == Some("rpc".to_string())) {
                         store_fun.call1(py, (def_id.0.into_py(py), module.getattr(py, name.to_string().as_str()).unwrap())).unwrap();
-                        rpc_ids.push((None, def_id));
+                        let is_async = decorator_list.iter().any(
+                            |decorator| decorator_get_flags(decorator).iter().any(
+                                |constant| *constant == Constant::Str("async".into())
+                            ));
+                        rpc_ids.push((None, def_id, is_async));
                     }
                 }
                 StmtKind::ClassDef { name, body, .. } => {
@@ -488,14 +493,18 @@ impl Nac3 {
                     let class_obj = module.getattr(py, class_name.as_str()).unwrap();
                     for stmt in body {
                         if let StmtKind::FunctionDef { name, decorator_list, .. } = &stmt.node {
-                            if decorator_list.iter().any(|decorator| matches!(decorator.node, ExprKind::Name { id, .. } if id == "rpc".into())) {
+                            if decorator_list.iter().any(|decorator| decorator_id_string(decorator) == Some("rpc".to_string())) {
+                                let is_async = decorator_list.iter().any(
+                                    |decorator| decorator_get_flags(decorator).iter().any(
+                                        |constant| *constant == Constant::Str("async".into())
+                                    ));
                                 if name == &"__init__".into() {
                                     return Err(CompileError::new_err(format!(
                                         "compilation failed\n----------\nThe constructor of class {} should not be decorated with rpc decorator (at {})",
                                         class_name, stmt.location
                                     )));
                                 }
-                                rpc_ids.push((Some((class_obj.clone(), *name)), def_id));
+                                rpc_ids.push((Some((class_obj.clone(), *name)), def_id, is_async));
                             }
                         }
                     }
@@ -596,13 +605,12 @@ impl Nac3 {
         let top_level = Arc::new(composer.make_top_level_context());
 
         {
-            let rpc_codegen = rpc_codegen_callback();
             let defs = top_level.definitions.read();
-            for (class_data, id) in &rpc_ids {
+            for (class_data, id, is_async) in &rpc_ids {
                 let mut def = defs[id.0].write();
                 match &mut *def {
                     TopLevelDef::Function { codegen_callback, .. } => {
-                        *codegen_callback = Some(rpc_codegen.clone());
+                        *codegen_callback = Some(rpc_codegen_callback(*is_async));
                     }
                     TopLevelDef::Class { methods, .. } => {
                         let (class_def, method_name) = class_data.as_ref().unwrap();
@@ -613,7 +621,7 @@ impl Nac3 {
                             if let TopLevelDef::Function { codegen_callback, .. } =
                                 &mut *defs[id.0].write()
                             {
-                                *codegen_callback = Some(rpc_codegen.clone());
+                                *codegen_callback = Some(rpc_codegen_callback(*is_async));
                                 store_fun
                                     .call1(
                                         py,
@@ -844,6 +852,41 @@ impl Nac3 {
     }
 }
 
+fn decorator_id_string(decorator: &Located<ExprKind>) -> Option<String> {
+    /// Retrieves the Name.id from a decorator, supports decorators with arguments.
+    if let ExprKind::Name { id, .. } = decorator.node {
+        // Bare decorator
+        return Some(id.to_string());
+    } else if let ExprKind::Call { func, .. } = &decorator.node {
+        // Decorators that are calls (e.g. "@rpc()") have Call for the node,
+        // need to extract the id from within.
+        if let ExprKind::Name { id, .. } = func.node {
+            return Some(id.to_string());
+        }
+    }
+    None
+}
+
+fn decorator_get_flags(decorator: &Located<ExprKind>) -> Vec<Constant> {
+    /// Retrieves flags from a decorator, if any.
+    let mut flags = vec![];
+    if let ExprKind::Call { keywords, .. } = &decorator.node {
+        for keyword in keywords.iter() {
+            if keyword.node.arg != Some("flags".into()) {
+                continue;
+            }
+            if let ExprKind::Set { elts } = &keyword.node.value.node {
+                for elt in elts {
+                    if let ExprKind::Constant { value, .. } = &elt.node {
+                        flags.push(value.clone());
+                    }
+                }
+            }
+        }
+    }
+    flags
+}
+
 fn link_with_lld(elf_filename: String, obj_filename: String) -> PyResult<()> {
     let linker_args = vec![
         "-shared".to_string(),

nac3core/src/toplevel/composer.rs

@@ -1894,7 +1894,7 @@ impl TopLevelComposer {
             } = &mut *function_def
             {
                 let signature_ty_enum = unifier.get_ty(*signature);
-                let TypeEnum::TFunc(FunSignature { args, ret, vars }) = signature_ty_enum.as_ref()
+                let TypeEnum::TFunc(FunSignature { args, ret, vars, .. }) = signature_ty_enum.as_ref()
                 else {
                     unreachable!("must be typeenum::tfunc")
                 };
@@ -2057,6 +2057,15 @@ impl TopLevelComposer {
                         instance_to_symbol.insert(String::new(), simple_name.to_string());
                         continue;
                     }
+                    if !decorator_list.is_empty() {
+                        if let ast::ExprKind::Call { func, .. } = &decorator_list[0].node {
+                            if matches!(&func.node, 
+                                ast::ExprKind::Name{ id, .. } if id == &"rpc".into()) {
+                                    instance_to_symbol.insert(String::new(), simple_name.to_string());
+                                    continue;
+                                }
+                        }
+                    }
 
                     let fun_body = body
                         .into_iter()