From cbc660e740453a8b43f7c13ea0358caed2bc9f57 Mon Sep 17 00:00:00 2001
From: mwojcik <mw@m-labs.hk>
Date: Mon, 17 Apr 2023 12:47:17 +0800
Subject: [PATCH] ddma: pass "uses_ddma" flag

---
 src/runtime/src/comms.rs      |  2 +-
 src/runtime/src/kernel/dma.rs | 29 ++++++++++++++++++-----------
 src/runtime/src/kernel/mod.rs |  2 +-
 src/runtime/src/rtio_dma.rs   | 19 ++++++++++++++++---
 4 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/src/runtime/src/comms.rs b/src/runtime/src/comms.rs
index da67ba1b..b5db79d3 100644
--- a/src/runtime/src/comms.rs
+++ b/src/runtime/src/comms.rs
@@ -330,7 +330,7 @@ async fn handle_run_kernel(
                 rtio_dma::erase(name, aux_mutex, routing_table, timer).await;
             }
             kernel::Message::DmaGetRequest(name) => {
-                let result = rtio_dma::retrieve(name);
+                let result = rtio_dma::retrieve(name).await;
                 control
                     .borrow_mut()
                     .tx
diff --git a/src/runtime/src/kernel/dma.rs b/src/runtime/src/kernel/dma.rs
index 997625d0..c77a6aa6 100644
--- a/src/runtime/src/kernel/dma.rs
+++ b/src/runtime/src/kernel/dma.rs
@@ -10,6 +10,7 @@ use crate::{artiq_raise, pl::csr, rtio};
 pub struct DmaTrace {
     duration: i64,
     address: i32,
+    uses_ddma: bool,
 }
 
 #[derive(Clone, Debug)]
@@ -151,8 +152,12 @@ pub extern "C" fn dma_retrieve(name: CSlice<u8>) -> DmaTrace {
     }
     match unsafe { KERNEL_CHANNEL_0TO1.as_mut().unwrap() }.recv() {
         Message::DmaGetReply(None) => (),
-        Message::DmaGetReply(Some((address, duration))) => {
-            return DmaTrace { address, duration };
+        Message::DmaGetReply(Some((address, duration, uses_ddma))) => {
+            return DmaTrace {
+                address,
+                duration,
+                uses_ddma,
+            };
         }
         _ => panic!("Expected DmaGetReply after DmaGetRequest!"),
     }
@@ -160,7 +165,7 @@ pub extern "C" fn dma_retrieve(name: CSlice<u8>) -> DmaTrace {
     artiq_raise!("DMAError", "DMA trace not found");
 }
 
-pub extern "C" fn dma_playback(timestamp: i64, ptr: i32) {
+pub extern "C" fn dma_playback(timestamp: i64, ptr: i32, _uses_ddma: bool) {
     unsafe {
         csr::rtio_dma::base_address_write(ptr as u32);
         csr::rtio_dma::time_offset_write(timestamp as u64);
@@ -168,13 +173,15 @@ pub extern "C" fn dma_playback(timestamp: i64, ptr: i32) {
         csr::cri_con::selected_write(1);
         csr::rtio_dma::enable_write(1);
         #[cfg(has_drtio)]
-        KERNEL_CHANNEL_1TO0
-            .as_mut()
-            .unwrap()
-            .send(Message::DmaStartRemoteRequest {
-                id: ptr,
-                timestamp: timestamp,
-            });
+        if _uses_ddma {
+            KERNEL_CHANNEL_1TO0
+                .as_mut()
+                .unwrap()
+                .send(Message::DmaStartRemoteRequest {
+                    id: ptr,
+                    timestamp: timestamp,
+                });
+        }
         while csr::rtio_dma::enable_read() != 0 {}
         csr::cri_con::selected_write(0);
 
@@ -203,7 +210,7 @@ pub extern "C" fn dma_playback(timestamp: i64, ptr: i32) {
             }
         }
         #[cfg(has_drtio)]
-        {
+        if _uses_ddma {
             KERNEL_CHANNEL_1TO0
                 .as_mut()
                 .unwrap()
diff --git a/src/runtime/src/kernel/mod.rs b/src/runtime/src/kernel/mod.rs
index 308e4fe3..864ac017 100644
--- a/src/runtime/src/kernel/mod.rs
+++ b/src/runtime/src/kernel/mod.rs
@@ -52,7 +52,7 @@ pub enum Message {
     DmaPutRequest(DmaRecorder),
     DmaEraseRequest(String),
     DmaGetRequest(String),
-    DmaGetReply(Option<(i32, i64)>),
+    DmaGetReply(Option<(i32, i64, bool)>),
     #[cfg(has_drtio)]
     DmaStartRemoteRequest {
         id: i32,
diff --git a/src/runtime/src/rtio_dma.rs b/src/runtime/src/rtio_dma.rs
index d24c1796..d39e7de9 100644
--- a/src/runtime/src/rtio_dma.rs
+++ b/src/runtime/src/rtio_dma.rs
@@ -193,7 +193,7 @@ pub mod remote_dma {
             up: bool,
         ) {
             // update state of the destination, resend traces if it's up
-            if let Some(trace) = self.traces.lock().get_mut(&destination) {
+            if let Some(trace) = self.traces.async_lock().await.get_mut(&destination) {
                 if up {
                     match drtio::ddma_upload_trace(
                         aux_mutex,
@@ -213,6 +213,10 @@ pub mod remote_dma {
                 }
             }
         }
+
+        pub async fn is_empty(&self) -> bool {
+            self.traces.async_lock().await.is_empty()
+        }
     }
 
     static mut TRACES: BTreeMap<u32, TraceSet> = BTreeMap::new();
@@ -269,6 +273,11 @@ pub mod remote_dma {
                 .await;
         }
     }
+
+    pub async fn has_remote_traces(id: u32) -> bool {
+        let trace_set = unsafe { TRACES.get_mut(&id).unwrap() };
+        !(trace_set.is_empty().await)
+    }
 }
 
 pub async fn put_record(
@@ -343,7 +352,11 @@ pub async fn erase(name: String, _aux_mutex: &Rc<Mutex<bool>>, _routing_table: &
     }
 }
 
-pub fn retrieve(name: String) -> Option<(i32, i64)> {
+pub async fn retrieve(name: String) -> Option<(i32, i64, bool)> {
     let (ptr, _v, duration) = DMA_RECORD_STORE.lock().get(&name)?.clone();
-    Some((ptr as i32, duration))
+    #[cfg(has_drtio)]
+    let uses_ddma = remote_dma::has_remote_traces(ptr).await;
+    #[cfg(not(has_drtio))]
+    let uses_ddma = false;
+    Some((ptr as i32, duration, uses_ddma))
 }