subkernel: separate tags and data

artiq/compiler/ir.py

@@ -764,6 +764,23 @@ class GetOptArgFromRemote(GetArgFromRemote):
     def opcode(self):
         return "getoptargfromremote({})".format(repr(self.arg_name))
 
+class SubkernelAwaitArgs(Instruction):
+    """
+    A builtin instruction that takes min and max received messages as operands,
+    and a list of received types.
+
+    :ivar arg_types: (list of types) types of passed arguments (including optional)
+    """
+
+    """
+    :param arg_types: (list of types) types of passed arguments (including optional)
+    """
+
+    def __init__(self, operands, arg_types, name=None):
+        assert isinstance(arg_types, list)
+        self.arg_types = arg_types
+        super().__init__(operands, builtins.TNone(), name)
+
 class GetAttr(Instruction):
     """
     An intruction that loads an attribute from an object,

artiq/compiler/transforms/artiq_ir_generator.py

@@ -2614,8 +2614,9 @@ class ARTIQIRGenerator(algorithm.Visitor):
                 min_args = ir.Constant(len(fn_typ.args)-offset, builtins.TInt8())
                 max_args = ir.Constant(fn_typ.arity()-offset, builtins.TInt8())
 
-                rcvd_count = self.append(ir.Builtin("subkernel_await_args", [min_args, max_args], builtins.TNone()))
+                arg_types = list(fn_typ.args.items())[offset:] 
-                arg_types = list(fn_typ.args.items())[offset:]
+                arg_type_list = [a[1] for a in arg_types] + [a[1] for a in fn_typ.optargs.items()]
+                rcvd_count = self.append(ir.SubkernelAwaitArgs([min_args, max_args], arg_type_list))
                 # obligatory arguments
                 for arg_name, arg_type in arg_types:
                     args[index] = self.append(ir.GetArgFromRemote(arg_name, arg_type,

artiq/compiler/transforms/llvm_ir_generator.py

@@ -405,7 +405,7 @@ class LLVMIRGenerator:
         elif name == "subkernel_await_finish":
             llty = ll.FunctionType(llvoid, [lli32, lli64])
         elif name == "subkernel_await_message":
-            llty = ll.FunctionType(lli8, [lli32, lli64, lli8, lli8])
+            llty = ll.FunctionType(lli8, [lli32, lli64, llsliceptr, lli8, lli8])
 
         # with now-pinning
         elif name == "now":
@@ -1400,12 +1400,6 @@ class LLVMIRGenerator:
                 return self.llbuilder.call(self.llbuiltin("delay_mu"), [llinterval])
         elif insn.op == "end_catch":
             return self.llbuilder.call(self.llbuiltin("__artiq_end_catch"), [])
-        elif insn.op == "subkernel_await_args":
-            llmin = self.map(insn.operands[0])
-            llmax = self.map(insn.operands[1])
-            return self.llbuilder.call(self.llbuiltin("subkernel_await_message"), 
-                                       [ll.Constant(lli32, 0), ll.Constant(lli64, 10_000), llmin, llmax],
-                                       name="subkernel.await.args")
         elif insn.op == "subkernel_await_finish":
             llsid = self.map(insn.operands[0])
             lltimeout = self.map(insn.operands[1])
@@ -1414,7 +1408,9 @@ class LLVMIRGenerator:
         elif insn.op == "subkernel_retrieve_return":
             llsid = self.map(insn.operands[0])
             lltimeout = self.map(insn.operands[1])
-            self.llbuilder.call(self.llbuiltin("subkernel_await_message"), [llsid, lltimeout, ll.Constant(lli8, 1), ll.Constant(lli8, 1)],
+            lltagptr = self._build_subkernel_tags([insn.type])
+            self.llbuilder.call(self.llbuiltin("subkernel_await_message"), 
+                                [llsid, lltimeout, lltagptr, ll.Constant(lli8, 1), ll.Constant(lli8, 1)],
                                 name="subkernel.await.message")
             llstackptr = self.llbuilder.call(self.llbuiltin("llvm.stacksave"), [],
                                              name="subkernel.arg.stack")
@@ -1426,6 +1422,14 @@ class LLVMIRGenerator:
         else:
             assert False
 
+    def process_SubkernelAwaitArgs(self, insn):
+        llmin = self.map(insn.operands[0])
+        llmax = self.map(insn.operands[1])
+        lltagptr = self._build_subkernel_tags(insn.arg_types)
+        return self.llbuilder.call(self.llbuiltin("subkernel_await_message"), 
+                                    [ll.Constant(lli32, 0), ll.Constant(lli64, 10_000), lltagptr, llmin, llmax],
+                                    name="subkernel.await.args")
+
     def process_Closure(self, insn):
         llenv = self.map(insn.environment())
         llenv = self.llbuilder.bitcast(llenv, llptr)
@@ -1505,6 +1509,24 @@ class LLVMIRGenerator:
 
         return llfun, list(llargs), llarg_attrs, llcallstackptr
 
+    def _build_subkernel_tags(self, tag_list):
+        def ret_error_handler(typ):
+            printer = types.TypePrinter()
+            note = diagnostic.Diagnostic("note",
+                "value of type {type}",
+                {"type": printer.name(typ)},
+                fun_loc)
+            diag = diagnostic.Diagnostic("error",
+                "type {type} is not supported in subkernels",
+                {"type": printer.name(fun_type.ret)},
+                fun_loc, notes=[note])
+            self.engine.process(diag)
+        tag = b"".join([ir.rpc_tag(arg_type, ret_error_handler) for arg_type in tag_list])
+        lltag = self.llconst_of_const(ir.Constant(tag, builtins.TStr()))
+        lltagptr = self.llbuilder.alloca(lltag.type)
+        self.llbuilder.store(lltag, lltagptr)
+        return lltagptr
+
     def _build_rpc_recv(self, ret, llstackptr, llnormalblock=None, llunwindblock=None):
         # T result = {
         #   void *ret_ptr = alloca(sizeof(T));

artiq/firmware/ksupport/lib.rs

@@ -138,7 +138,7 @@ extern fn rpc_send_async(service: u32, tag: &CSlice<u8>, data: *const *const ())
     rpc_queue::enqueue(|mut slice| {
         let length = {
             let mut writer = Cursor::new(&mut slice[4..]);
-            rpc_proto::send_args(&mut writer, service, tag.as_ref(), data)?;
+            rpc_proto::send_args(&mut writer, service, tag.as_ref(), data, true)?;
             writer.position()
         };
         io::ProtoWrite::write_u32(&mut slice, length as u32)
@@ -499,8 +499,8 @@ extern fn subkernel_send_message(id: u32, count: u8, tag: &CSlice<u8>, data: *co
 }
 
 #[unwind(allowed)]
-extern fn subkernel_await_message(id: u32, timeout: u64, min: u8, max: u8) -> u8 {
+extern fn subkernel_await_message(id: u32, timeout: u64, tags: &CSlice<u8>, min: u8, max: u8) -> u8 {
-    send(&SubkernelMsgRecvRequest { id: id, timeout: timeout });
+    send(&SubkernelMsgRecvRequest { id: id, timeout: timeout, tags: tags.as_ref() });
     recv!(SubkernelMsgRecvReply { status, count } => {
         match status {
             SubkernelStatus::NoError => {

artiq/firmware/libproto_artiq/kernel_proto.rs

@@ -108,7 +108,7 @@ pub enum Message<'a> {
     SubkernelAwaitFinishRequest { id: u32, timeout: u64 },
     SubkernelAwaitFinishReply { status: SubkernelStatus },
     SubkernelMsgSend { id: u32, count: u8, tag: &'a [u8], data: *const *const () },
-    SubkernelMsgRecvRequest { id: u32, timeout: u64 },
+    SubkernelMsgRecvRequest { id: u32, timeout: u64, tags: &'a [u8] },
     SubkernelMsgRecvReply { status: SubkernelStatus, count: u8 },
 
     Log(fmt::Arguments<'a>),

artiq/firmware/libproto_artiq/rpc_proto.rs

@@ -190,9 +190,9 @@ unsafe fn recv_value<R, E>(reader: &mut R, tag: Tag, data: &mut *mut (),
     }
 }
 
-pub fn recv_return<R, E>(reader: &mut R, tag_bytes: &[u8], data: *mut (),
+pub fn recv_return<'a, R, E>(reader: &mut R, tag_bytes: &'a [u8], data: *mut (),
                          alloc: &dyn Fn(usize) -> Result<*mut (), E>)
-                        -> Result<(), E>
+                        -> Result<&'a [u8], E>
     where R: Read + ?Sized,
           E: From<Error<R::ReadError>>
 {
@@ -204,14 +204,16 @@ pub fn recv_return<R, E>(reader: &mut R, tag_bytes: &[u8], data: *mut (),
     let mut data = data;
     unsafe { recv_value(reader, tag, &mut data, alloc)? };
 
-    Ok(())
+    Ok(it.data)
 }
 
-unsafe fn send_elements<W>(writer: &mut W, elt_tag: Tag, length: usize, data: *const ())
+unsafe fn send_elements<W>(writer: &mut W, elt_tag: Tag, length: usize, data: *const (), write_tags: bool)
                           -> Result<(), Error<W::WriteError>>
     where W: Write + ?Sized
 {
-    writer.write_u8(elt_tag.as_u8())?;
+    if write_tags {
+        writer.write_u8(elt_tag.as_u8())?;
+    }
     match elt_tag {
         // we cannot use NativeEndian::from_slice_i32 as the data is not mutable,
         // and that is not needed as the data is already in native endian
@@ -230,14 +232,14 @@ unsafe fn send_elements<W>(writer: &mut W, elt_tag: Tag, length: usize, data: *c
         _ => {
             let mut data = data;
             for _ in 0..length {
-                send_value(writer, elt_tag, &mut data)?;
+                send_value(writer, elt_tag, &mut data, write_tags)?;
             }
         }
     }
     Ok(())
 }
 
-unsafe fn send_value<W>(writer: &mut W, tag: Tag, data: &mut *const ())
+unsafe fn send_value<W>(writer: &mut W, tag: Tag, data: &mut *const (), write_tags: bool)
                        -> Result<(), Error<W::WriteError>>
     where W: Write + ?Sized
 {
@@ -248,8 +250,9 @@ unsafe fn send_value<W>(writer: &mut W, tag: Tag, data: &mut *const ())
             $map
         })
     }
-
+    if write_tags {
-    writer.write_u8(tag.as_u8())?;
+        writer.write_u8(tag.as_u8())?;
+    }
     match tag {
         Tag::None => Ok(()),
         Tag::Bool =>
@@ -269,12 +272,14 @@ unsafe fn send_value<W>(writer: &mut W, tag: Tag, data: &mut *const ())
                 writer.write_bytes((*ptr).as_ref())),
         Tag::Tuple(it, arity) => {
             let mut it = it.clone();
-            writer.write_u8(arity)?;
+            if write_tags {
+                writer.write_u8(arity)?;
+            }
             let mut max_alignment = 0;
             for _ in 0..arity {
                 let tag = it.next().expect("truncated tag");
                 max_alignment = core::cmp::max(max_alignment, tag.alignment());
-                send_value(writer, tag, data)?
+                send_value(writer, tag, data, write_tags)?
             }
             *data = round_up_const(*data, max_alignment);
             Ok(())
@@ -286,11 +291,13 @@ unsafe fn send_value<W>(writer: &mut W, tag: Tag, data: &mut *const ())
                 let length = (**ptr).length as usize;
                 writer.write_u32((**ptr).length)?;
                 let tag = it.clone().next().expect("truncated tag");
-                send_elements(writer, tag, length, (**ptr).elements)
+                send_elements(writer, tag, length, (**ptr).elements, write_tags)
             })
         }
         Tag::Array(it, num_dims) => {
-            writer.write_u8(num_dims)?;
+            if write_tags {
+                writer.write_u8(num_dims)?;
+            }
             consume_value!(*const(), |buffer| {
                 let elt_tag = it.clone().next().expect("truncated tag");
 
@@ -302,14 +309,14 @@ unsafe fn send_value<W>(writer: &mut W, tag: Tag, data: &mut *const ())
                     })
                 }
                 let length = total_len as usize;
-                send_elements(writer, elt_tag, length, *buffer)
+                send_elements(writer, elt_tag, length, *buffer, write_tags)
             })
         }
         Tag::Range(it) => {
             let tag = it.clone().next().expect("truncated tag");
-            send_value(writer, tag, data)?;
+            send_value(writer, tag, data, write_tags)?;
-            send_value(writer, tag, data)?;
+            send_value(writer, tag, data, write_tags)?;
-            send_value(writer, tag, data)?;
+            send_value(writer, tag, data, write_tags)?;
             Ok(())
         }
         Tag::Keyword(it) => {
@@ -319,7 +326,7 @@ unsafe fn send_value<W>(writer: &mut W, tag: Tag, data: &mut *const ())
                 writer.write_string(str::from_utf8((*ptr).name.as_ref()).unwrap())?;
                 let tag = it.clone().next().expect("truncated tag");
                 let mut data = ptr.offset(1) as *const ();
-                send_value(writer, tag, &mut data)
+                send_value(writer, tag, &mut data, write_tags)
             })
             // Tag::Keyword never appears in composite types, so we don't have
             // to accurately advance data.
@@ -333,7 +340,7 @@ unsafe fn send_value<W>(writer: &mut W, tag: Tag, data: &mut *const ())
     }
 }
 
-pub fn send_args<W>(writer: &mut W, service: u32, tag_bytes: &[u8], data: *const *const ())
+pub fn send_args<W>(writer: &mut W, service: u32, tag_bytes: &[u8], data: *const *const (), write_tags: bool)
                    -> Result<(), Error<W::WriteError>>
     where W: Write + ?Sized
 {
@@ -350,7 +357,7 @@ pub fn send_args<W>(writer: &mut W, service: u32, tag_bytes: &[u8], data: *const
     for index in 0.. {
         if let Some(arg_tag) = args_it.next() {
             let mut data = unsafe { *data.offset(index) };
-            unsafe { send_value(writer, arg_tag, &mut data)? };
+            unsafe { send_value(writer, arg_tag, &mut data, write_tags)? };
         } else {
             break
         }
@@ -482,7 +489,7 @@ mod tag {
 
     #[derive(Debug, Clone, Copy)]
     pub struct TagIterator<'a> {
-        data: &'a [u8]
+        pub data: &'a [u8]
     }
 
     impl<'a> TagIterator<'a> {

artiq/firmware/runtime/kernel.rs

@@ -307,8 +307,7 @@ pub mod subkernel {
 
     pub struct Message {
         from_id: u32,
-        pub tag_count: u8,
+        pub count: u8,
-        pub tag: u8,
         pub data: Vec<u8>
     }
 
@@ -334,9 +333,8 @@ pub mod subkernel {
             None => unsafe {
                 CURRENT_MESSAGES.insert(id, Message {
                     from_id: id,
-                    tag_count: data[0],
+                    count: data[0],
-                    tag: data[1],
+                    data: data[1..length].to_vec()
-                    data: data[2..length].to_vec()
                 });
             }
         };
@@ -404,7 +402,7 @@ pub mod subkernel {
         let destination = unsafe { SUBKERNELS.get(&id).unwrap().destination };
 
         // reuse rpc code for sending arbitrary data
-        rpc::send_args(&mut writer, 0, tag, message)?;
+        rpc::send_args(&mut writer, 0, tag, message, false)?;
         // skip service tag, but overwrite first byte with tag count
         let data = &mut writer.into_inner()[3..];
         data[0] = count;

artiq/firmware/runtime/session.rs

@@ -502,7 +502,7 @@ fn process_kern_message(io: &Io, aux_mutex: &Mutex,
                     None => unexpected!("unexpected RPC in flash kernel"),
                     Some(ref mut stream) => {
                         host_write(stream, host::Reply::RpcRequest { async: async })?;
-                        rpc::send_args(stream, service, tag, data)?;
+                        rpc::send_args(stream, service, tag, data, true)?;
                         if !async {
                             session.kernel_state = KernelState::RpcWait
                         }
@@ -608,10 +608,10 @@ fn process_kern_message(io: &Io, aux_mutex: &Mutex,
                 kern_acknowledge()
             }
             #[cfg(has_drtio)]
-            &kern::SubkernelMsgRecvRequest { id, timeout } => {
+            &kern::SubkernelMsgRecvRequest { id, timeout, tags } => {
                 let message_received = subkernel::message_await(io, _subkernel_mutex, id, timeout);
                 let (status, count) = match message_received {
-                    Ok(ref message) => (kern::SubkernelStatus::NoError, message.tag_count),
+                    Ok(ref message) => (kern::SubkernelStatus::NoError, message.count),
                     Err(SubkernelError::Timeout) => (kern::SubkernelStatus::Timeout, 0),
                     Err(SubkernelError::IncorrectState) => (kern::SubkernelStatus::IncorrectState, 0),
                     Err(SubkernelError::SubkernelFinished) => {
@@ -628,11 +628,11 @@ fn process_kern_message(io: &Io, aux_mutex: &Mutex,
                     }
                     Err(_) => (kern::SubkernelStatus::OtherError, 0)
                 };
-                kern_send(io, &kern::SubkernelMsgRecvReply { status: status, count: count })?;
+                kern_send(io, &kern::SubkernelMsgRecvReply { status: status, count: count})?;
                 if let Ok(message) = message_received {
                     // receive code almost identical to RPC recv, except we are not reading from a stream
                     let mut reader = Cursor::new(message.data);
-                    let mut tag: [u8; 1] = [message.tag];
+                    let mut current_tags = tags;
                     let mut i = 0;
                     loop {
                         // kernel has to consume all arguments in the whole message
@@ -643,7 +643,7 @@ fn process_kern_message(io: &Io, aux_mutex: &Mutex,
                                     "expected root value slot from kernel CPU, not {:?}", other)
                             }
                         })?;
-                        let res = rpc::recv_return(&mut reader, &tag, slot, &|size| -> Result<_, Error<SchedError>> {
+                        let res = rpc::recv_return(&mut reader, current_tags, slot, &|size| -> Result<_, Error<SchedError>> {
                             if size == 0 {
                                 return Ok(0 as *mut ())
                             }
@@ -657,17 +657,19 @@ fn process_kern_message(io: &Io, aux_mutex: &Mutex,
                             })?)
                         });
                         match res {
-                            Ok(_) => kern_send(io, &kern::RpcRecvReply(Ok(0)))?,
+                            Ok(new_tags) => {
+                                kern_send(io, &kern::RpcRecvReply(Ok(0)))?;
+                                i += 1;
+                                if i < message.count {
+                                    // update the tag for next read
+                                    current_tags = new_tags;
+                                } else {
+                                    // should be done by then
+                                    break;
+                                }
+                            },
                             Err(_) => unexpected!("expected valid subkernel message data")
                         };
-                        i += 1;
-                        if i < message.tag_count {
-                            // update the tag for next read
-                            tag[0] = reader.read_u8()?;
-                        } else {
-                            // should be done by then
-                            break;
-                        }
                     }
                     Ok(())
                 } else {

artiq/firmware/satman/kernel.rs

@@ -6,7 +6,7 @@ use board_artiq::{mailbox, spi};
 use board_misoc::{csr, clock, i2c};
 use proto_artiq::{kernel_proto as kern, session_proto::Reply::KernelException as HostKernelException, rpc_proto as rpc};
 use eh::eh_artiq;
-use io::{Cursor, ProtoRead};
+use io::Cursor;
 use kernel::eh_artiq::StackPointerBacktrace;
 
 use ::{cricon_select, RtioMaster};
@@ -52,12 +52,12 @@ mod kernel_cpu {
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 enum KernelState {
     Absent,
     Loaded,
     Running,
-    MsgAwait { max_time: u64 },
+    MsgAwait { max_time: u64, tags: Vec<u8> },
     MsgSending
 }
 
@@ -99,7 +99,6 @@ pub struct Sliceable {
 /* represents interkernel messages */
 struct Message {
     count: u8,
-    tag: u8,
     data: Vec<u8>
 }
 
@@ -201,8 +200,7 @@ impl MessageManager {
             None => {
                 self.in_buffer = Some(Message {
                     count: data[0],
-                    tag: data[1],
+                    data: data[1..length].to_vec()
-                    data: data[2..length].to_vec()
                 });
             }
         };
@@ -264,7 +262,7 @@ impl MessageManager {
 
     pub fn accept_outgoing(&mut self, count: u8, tag: &[u8], data: *const *const ()) -> Result<(), Error>  {
         let mut writer = Cursor::new(Vec::new());
-        rpc::send_args(&mut writer, 0, tag, data)?;
+        rpc::send_args(&mut writer, 0, tag, data, false)?;
         // skip service tag, but write the count
         let mut data = writer.into_inner().split_off(3);
         data[0] = count;
@@ -507,17 +505,18 @@ impl Manager {
     }
 
     fn process_external_messages(&mut self) -> Result<(), Error> {
-        match self.session.kernel_state {
+        match &self.session.kernel_state {
-            KernelState::MsgAwait { max_time } => {
+            KernelState::MsgAwait { max_time, tags } => {
-                if clock::get_ms() > max_time {
+                if clock::get_ms() > *max_time {
                     kern_send(&kern::SubkernelMsgRecvReply { status: kern::SubkernelStatus::Timeout, count: 0 })?;
                     self.session.kernel_state = KernelState::Running;
                     return Ok(())
                 }
                 if let Some(message) = self.session.messages.get_incoming() {
                     kern_send(&kern::SubkernelMsgRecvReply { status: kern::SubkernelStatus::NoError, count: message.count })?;
+                    let tags = tags.clone();
                     self.session.kernel_state = KernelState::Running;
-                    pass_message_to_kernel(&message)
+                    pass_message_to_kernel(&message, &tags)
                 } else {
                     Err(Error::AwaitingMessage)
                 }
@@ -538,7 +537,7 @@ impl Manager {
         // returns Ok(with_exception) on finish
         // None if the kernel is still running
         kern_recv(|request| {
-            match (request, self.session.kernel_state) {
+            match (request, &self.session.kernel_state) {
                 (&kern::LoadReply(_), KernelState::Loaded) => {
                     // We're standing by; ignore the message.
                     return Ok(None)
@@ -611,9 +610,9 @@ impl Manager {
                     Ok(())
                 }
 
-                &kern::SubkernelMsgRecvRequest { id: _, timeout } => {
+                &kern::SubkernelMsgRecvRequest { id: _, timeout, tags } => {
                     let max_time = clock::get_ms() + timeout as u64;
-                    self.session.kernel_state = KernelState::MsgAwait { max_time: max_time };
+                    self.session.kernel_state = KernelState::MsgAwait { max_time: max_time, tags: tags.to_vec() };
                     Ok(())
                 },
 
@@ -695,10 +694,9 @@ fn slice_kernel_exception(exceptions: &[Option<eh_artiq::Exception>],
     }
 }
 
-fn pass_message_to_kernel(message: &Message) -> Result<(), Error> {
+fn pass_message_to_kernel(message: &Message, tags: &[u8]) -> Result<(), Error> {
     let mut reader = Cursor::new(&message.data);
-    let mut tag: [u8; 1] = [message.tag];
+    let mut current_tags = tags;
-    let count = message.count;
     let mut i = 0;
     loop {
         let slot = kern_recv_w_timeout(100, |reply| {
@@ -712,8 +710,7 @@ fn pass_message_to_kernel(message: &Message) -> Result<(), Error> {
                     "expected root value slot from kernel CPU, not {:?}", other)
             }
         })?;
-
+        let res = rpc::recv_return(&mut reader, current_tags, slot, &|size| -> Result<_, Error> {
-        let res = rpc::recv_return(&mut reader, &tag, slot, &|size| -> Result<_, Error> {
             if size == 0 {
                 return Ok(0 as *mut ())
             }
@@ -735,17 +732,19 @@ fn pass_message_to_kernel(message: &Message) -> Result<(), Error> {
             })?)
         });
         match res {
-            Ok(_) => kern_send(&kern::RpcRecvReply(Ok(0)))?,
+            Ok(new_tags) => {
+                kern_send(&kern::RpcRecvReply(Ok(0)))?;
+                i += 1;
+                if i < message.count {
+                    // update the tag for next read
+                    current_tags = new_tags;
+                } else {
+                    // should be done by then
+                    break;
+                }
+            },
             Err(_) => unexpected!("expected valid subkernel message data")
         };
-        i += 1;
-        if i < count {
-            // update the tag for next read
-            tag[0] = reader.read_u8()?;
-        } else {
-            // should be done by then
-            break;
-        }
     }
     Ok(())
 }

artiq/test/lit/embedding/subkernel_return.py

@@ -9,13 +9,13 @@ def entrypoint():
     # CHECK: call void @subkernel_load_run\(i32 1, i1 true\), !dbg !.
     # CHECK-NOT: call void @subkernel_send_message\(.*\), !dbg !.
     returning()
-    # CHECK: call i8 @subkernel_await_message\(i32 1, i64 10000, i8 1, i8 1\), !dbg !.
+    # CHECK: call i8 @subkernel_await_message\(i32 1, i64 10000, { i8\*, i32 }\* nonnull .*, i8 1, i8 1\), !dbg !.
     # CHECK: call void @subkernel_await_finish\(i32 1, i64 10000\), !dbg !.
     subkernel_await(returning)
 
 # CHECK-L: declare void @subkernel_load_run(i32, i1) local_unnamed_addr
 # CHECK-NOT-L: declare void @subkernel_send_message(i32, i8, { i8*, i32 }*, i8**) local_unnamed_addr
-# CHECK-L: declare i8 @subkernel_await_message(i32, i64, i8, i8) local_unnamed_addr
+# CHECK-L: declare i8 @subkernel_await_message(i32, i64, { i8*, i32 }*, i8, i8) local_unnamed_addr
 # CHECK-L: declare void @subkernel_await_finish(i32, i64) local_unnamed_addr
 @subkernel(destination=1)
 def returning() -> TInt32:

artiq/test/lit/embedding/subkernel_return_none.py

@@ -10,13 +10,13 @@ def entrypoint():
     # CHECK-NOT: call void @subkernel_send_message\(.*\), !dbg !.
     returning_none()
     # CHECK: call void @subkernel_await_finish\(i32 1, i64 10000\), !dbg !.
-    # CHECK-NOT: call void @subkernel_await_message\(i32 1, i64 10000\), !dbg !.
+    # CHECK-NOT: call i8 @subkernel_await_message\(i32 1, i64 10000\, .*\), !dbg !.
     subkernel_await(returning_none)
 
 # CHECK-L: declare void @subkernel_load_run(i32, i1) local_unnamed_addr
 # CHECK-NOT-L: declare void @subkernel_send_message(i32, { i8*, i32 }*, i8**) local_unnamed_addr
 # CHECK-L: declare void @subkernel_await_finish(i32, i64) local_unnamed_addr
-# CHECK-NOT-L: declare void @subkernel_await_message(i32, i64) local_unnamed_addr
+# CHECK-NOT-L: declare i8 @subkernel_await_message(i32, i64, { i8*, i32 }*, i8, i8) local_unnamed_addr
 @subkernel(destination=1)
 def returning_none() -> TNone:
     pass