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

initial subkernels commit

Browse Source
This commit is contained in:
mwojcik 2024-09-03 16:58:57 +08:00
parent 979209a526
commit a6dcd8dca3
4 changed files with 132 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
nac3artiq/demo/embedding_map.py
View File

@ -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

11
nac3artiq/demo/min_artiq.py
View File

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

4
nac3artiq/src/lib.rs
View File

@ -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
}

91
nac3artiq/src/subkernels.rs Normal file
View File

@ -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
}
}