mwojcik/nac3
1
0
Fork 0
forked from M-Labs/nac3
Code Pull Requests Activity

artiq: support async rpcs

Browse Source
This commit is contained in:
mwojcik 2024-09-06 16:33:58 +08:00
parent 3e92c491f5
commit 2d1cd95c2c
4 changed files with 133 additions and 52 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
nac3artiq/demo/min_artiq.py
View File

@ -112,10 +112,15 @@ def extern(function):
register_function(function) register_function(function)
return 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): def kernel(function_or_method):
"""Decorates a function or method to be executed on the core device.""" """Decorates a function or method to be executed on the core device."""

83
nac3artiq/src/codegen.rs
View File

@ -824,6 +824,7 @@ fn rpc_codegen_callback_fn<'ctx>(
fun: (&FunSignature, DefinitionId), fun: (&FunSignature, DefinitionId),
args: Vec<(Option<StrRef>, ValueEnum<'ctx>)>, args: Vec<(Option<StrRef>, ValueEnum<'ctx>)>,
generator: &mut dyn CodeGenerator, generator: &mut dyn CodeGenerator,
is_async: bool,
) -> Result<Option<BasicValueEnum<'ctx>>, String> { ) -> Result<Option<BasicValueEnum<'ctx>>, String> {
let int8 = ctx.ctx.i8_type(); let int8 = ctx.ctx.i8_type();
let int32 = ctx.ctx.i32_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 ptr_type = int8.ptr_type(AddressSpace::default());
let tag_ptr_type = ctx.ctx.struct_type(&[ptr_type.into(), size_type.into()], false); 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); let service_id = int32.const_int(fun.1 .0 as u64, false);
// -- setup rpc tags // -- setup rpc tags
let mut tag = Vec::new(); let mut tag = Vec::new();
@ -932,35 +934,60 @@ fn rpc_codegen_callback_fn<'ctx>(
} }
// call // call
let rpc_send = ctx.module.get_function("rpc_send").unwrap_or_else(|| { if is_async {
ctx.module.add_function( let rpc_send_async = ctx.module.get_function("rpc_send_async").unwrap_or_else(|| {
"rpc_send", ctx.module.add_function(
ctx.ctx.void_type().fn_type( "rpc_send_async",
&[ ctx.ctx.void_type().fn_type(
int32.into(), &[
tag_ptr_type.ptr_type(AddressSpace::default()).into(), int32.into(),
ptr_type.ptr_type(AddressSpace::default()).into(), tag_ptr_type.ptr_type(AddressSpace::default()).into(),
], ptr_type.ptr_type(AddressSpace::default()).into(),
false, ],
), false,
None, ),
) None,
}); )
ctx.builder });
.build_call(rpc_send, &[service_id.into(), tag_ptr.into(), args_ptr.into()], "rpc.send") ctx.builder
.unwrap(); .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",
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, &[service_id.into(), tag_ptr.into(), args_ptr.into()], "rpc.send")
.unwrap();
}
// reclaim stack space used by arguments // reclaim stack space used by arguments
call_stackrestore(ctx, stackptr); call_stackrestore(ctx, stackptr);
let result = format_rpc_ret(generator, ctx, fun.0.ret); if is_async {
// async RPCs do not return any values
if !result.is_some_and(|res| res.get_type().is_pointer_type()) { Ok(None)
// An RPC returning an NDArray would not touch here. } else {
call_stackrestore(ctx, stackptr); let result = format_rpc_ret(generator, ctx, fun.0.ret);
if !result.is_some_and(|res| res.get_type().is_pointer_type()) {
// An RPC returning an NDArray would not touch here.
call_stackrestore(ctx, stackptr);
}
Ok(result)
} }
Ok(result)
} }
pub fn attributes_writeback( pub fn attributes_writeback(
@ -1055,7 +1082,7 @@ pub fn attributes_writeback(
let args: Vec<_> = let args: Vec<_> =
values.into_iter().map(|(_, val)| (None, ValueEnum::Dynamic(val))).collect(); values.into_iter().map(|(_, val)| (None, ValueEnum::Dynamic(val))).collect();
if let Err(e) = 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)); return Ok(Err(e));
} }
@ -1065,9 +1092,9 @@ pub fn attributes_writeback(
Ok(()) 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)
}))) })))
} }

76
nac3artiq/src/lib.rs
View File

@ -34,8 +34,12 @@ use nac3core::inkwell::{
targets::*, targets::*,
OptimizationLevel, 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, parser::parse_program,
}; };
use nac3core::toplevel::builtins::get_exn_constructor; use nac3core::toplevel::builtins::get_exn_constructor;
@ -135,7 +139,7 @@ struct Nac3 {
string_store: Arc<RwLock<HashMap<String, i32>>>, string_store: Arc<RwLock<HashMap<String, i32>>>,
exception_ids: Arc<RwLock<HashMap<usize, usize>>>, exception_ids: Arc<RwLock<HashMap<usize, usize>>>,
deferred_eval_store: DeferredEvaluationStore, deferred_eval_store: DeferredEvaluationStore,
/// LLVM-related options for code generation. /// LLVM-related options for code generzation.
llvm_options: CodeGenLLVMOptions, llvm_options: CodeGenLLVMOptions,
} }
@ -194,10 +198,8 @@ impl Nac3 {
body.retain(|stmt| { body.retain(|stmt| {
if let StmtKind::FunctionDef { ref decorator_list, .. } = stmt.node { if let StmtKind::FunctionDef { ref decorator_list, .. } = stmt.node {
decorator_list.iter().any(|decorator| { 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 { } else {
false false
} }
@ -210,9 +212,8 @@ impl Nac3 {
} }
StmtKind::FunctionDef { ref decorator_list, .. } => { StmtKind::FunctionDef { ref decorator_list, .. } => {
decorator_list.iter().any(|decorator| { 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 { } else {
false false
} }
@ -478,9 +479,13 @@ impl Nac3 {
match &stmt.node { match &stmt.node {
StmtKind::FunctionDef { decorator_list, .. } => { 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(); 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, .. } => { StmtKind::ClassDef { name, body, .. } => {
@ -488,14 +493,18 @@ impl Nac3 {
let class_obj = module.getattr(py, class_name.as_str()).unwrap(); let class_obj = module.getattr(py, class_name.as_str()).unwrap();
for stmt in body { for stmt in body {
if let StmtKind::FunctionDef { name, decorator_list, .. } = &stmt.node { 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() { if name == &"__init__".into() {
return Err(CompileError::new_err(format!( return Err(CompileError::new_err(format!(
"compilation failed\n----------\nThe constructor of class {} should not be decorated with rpc decorator (at {})", "compilation failed\n----------\nThe constructor of class {} should not be decorated with rpc decorator (at {})",
class_name, stmt.location 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 top_level = Arc::new(composer.make_top_level_context());
{ {
let rpc_codegen = rpc_codegen_callback();
let defs = top_level.definitions.read(); 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(); let mut def = defs[id.0].write();
match &mut *def { match &mut *def {
TopLevelDef::Function { codegen_callback, .. } => { TopLevelDef::Function { codegen_callback, .. } => {
*codegen_callback = Some(rpc_codegen.clone()); *codegen_callback = Some(rpc_codegen_callback(*is_async).clone());
} }
TopLevelDef::Class { methods, .. } => { TopLevelDef::Class { methods, .. } => {
let (class_def, method_name) = class_data.as_ref().unwrap(); let (class_def, method_name) = class_data.as_ref().unwrap();
@ -613,7 +621,7 @@ impl Nac3 {
if let TopLevelDef::Function { codegen_callback, .. } = if let TopLevelDef::Function { codegen_callback, .. } =
&mut *defs[id.0].write() &mut *defs[id.0].write()
{ {
*codegen_callback = Some(rpc_codegen.clone()); *codegen_callback = Some(rpc_codegen_callback(*is_async).clone());
store_fun store_fun
.call1( .call1(
py, py,
@ -844,6 +852,38 @@ impl Nac3 {
} }
} }
fn decorator_id_string(decorator: &Located<ExprKind>) -> Option<String> {
if let ExprKind::Name { id, .. } = decorator.node {
return Some(id.to_string());
} else if let ExprKind::Call { func, .. } = &decorator.node {
// decorators with flags (e.g. rpc async) have Call for the node,
// 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> {
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<()> { fn link_with_lld(elf_filename: String, obj_filename: String) -> PyResult<()> {
let linker_args = vec![ let linker_args = vec![
"-shared".to_string(), "-shared".to_string(),

13
nac3core/src/toplevel/composer.rs
View File

@ -1673,7 +1673,7 @@ impl TopLevelComposer {
// they may be changed with our use of placeholders // they may be changed with our use of placeholders
for (def, _) in definition_ast_list.iter().skip(self.builtin_num) { for (def, _) in definition_ast_list.iter().skip(self.builtin_num) {
if let TopLevelDef::Function { signature, var_id, .. } = &mut *def.write() { if let TopLevelDef::Function { signature, var_id, .. } = &mut *def.write() {
if let TypeEnum::TFunc(FunSignature { args, ret, vars }) = if let TypeEnum::TFunc(FunSignature { args, ret, vars}) =
unifier.get_ty(*signature).as_ref() unifier.get_ty(*signature).as_ref()
{ {
let new_var_ids = vars let new_var_ids = vars
@ -1894,7 +1894,7 @@ impl TopLevelComposer {
} = &mut *function_def } = &mut *function_def
{ {
let signature_ty_enum = unifier.get_ty(*signature); 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 { else {
unreachable!("must be typeenum::tfunc") unreachable!("must be typeenum::tfunc")
}; };
@ -2057,6 +2057,15 @@ impl TopLevelComposer {
instance_to_symbol.insert(String::new(), simple_name.to_string()); instance_to_symbol.insert(String::new(), simple_name.to_string());
continue; 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 let fun_body = body
.into_iter() .into_iter()