From 3a2b43a8a64f76aa003748ceecbbfe2a6de285b3 Mon Sep 17 00:00:00 2001
From: mwojcik <mw@m-labs.hk>
Date: Thu, 31 Oct 2024 15:18:37 +0800
Subject: [PATCH] nac3artiq: add kernel destination to compilation

---
 nac3artiq/demo/embedding_map.py | 11 ++++++-----
 nac3artiq/demo/min_artiq.py     |  4 +++-
 nac3artiq/src/lib.rs            | 11 +++++++----
 3 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/nac3artiq/demo/embedding_map.py b/nac3artiq/demo/embedding_map.py
index cecdd16a..69089acc 100644
--- a/nac3artiq/demo/embedding_map.py
+++ b/nac3artiq/demo/embedding_map.py
@@ -6,6 +6,7 @@ class EmbeddingMap:
         self.string_reverse_map = {}
         self.function_map = {}
         self.attributes_writeback = []
+        self.subkernel_map = {}
         self.subkernel_message_map = {}
 
         if not old_embedding_map is None:
@@ -45,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:
@@ -72,11 +77,7 @@ class EmbeddingMap:
         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
+        return self.subkernel_map
     
     def subkernel_messages(self):
         messages = {}
diff --git a/nac3artiq/demo/min_artiq.py b/nac3artiq/demo/min_artiq.py
index 4ed458ee..a70a6956 100644
--- a/nac3artiq/demo/min_artiq.py
+++ b/nac3artiq/demo/min_artiq.py
@@ -226,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):
diff --git a/nac3artiq/src/lib.rs b/nac3artiq/src/lib.rs
index 86b24a96..a52640d3 100644
--- a/nac3artiq/src/lib.rs
+++ b/nac3artiq/src/lib.rs
@@ -361,6 +361,7 @@ impl Nac3 {
         &self,
         obj: &PyAny,
         method_name: &str,
+        destination: u8,
         args: Vec<&PyAny>,
         embedding_map: &PyAny,
         py: Python,
@@ -1175,6 +1176,7 @@ impl Nac3 {
         &mut self,
         obj: &PyAny,
         method_name: &str,
+        destination: u8,
         args: Vec<&PyAny>,
         filename: &str,
         embedding_map: &PyAny,
@@ -1195,7 +1197,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
@@ -1213,7 +1215,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)
         }
     }
 
@@ -1221,6 +1223,7 @@ impl Nac3 {
         &mut self,
         obj: &PyAny,
         method_name: &str,
+        destination: u8,
         args: Vec<&PyAny>,
         embedding_map: &PyAny,
         py: Python,
@@ -1244,7 +1247,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
@@ -1257,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)
         }
     }
 }