From a125f3b180e0bd4e6fccb47776d41c979b9f7169 Mon Sep 17 00:00:00 2001
From: Joseph Schuchart <joseph.schuchart@stonybrook.edu>
Date: Tue, 29 Oct 2024 16:16:01 -0400
Subject: [PATCH] Add Buffer::is_current_on()

This allows querying whether the buffer is current on the given device.
Also add a note that Buffer::current_device() cannot reflect the fact
that all devices may have the current data (e.g., after a pushout).

Also add static ttg::device::Device::host() that returns the host device.


Signed-off-by: Joseph Schuchart <joseph.schuchart@stonybrook.edu>
---
 ttg/ttg/device/device.h  |  4 ++++
 ttg/ttg/madness/buffer.h |  6 ++++++
 ttg/ttg/parsec/buffer.h  | 17 ++++++++++++++++-
 3 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/ttg/ttg/device/device.h b/ttg/ttg/device/device.h
index 6dcb3c722..d2e311ee9 100644
--- a/ttg/ttg/device/device.h
+++ b/ttg/ttg/device/device.h
@@ -51,6 +51,10 @@ namespace ttg::device {
     bool is_invalid() const {
       return (m_space == ttg::ExecutionSpace::Invalid);
     }
+
+    static Device host() {
+      return {};
+    }
   };
 } // namespace ttg::device
 
diff --git a/ttg/ttg/madness/buffer.h b/ttg/ttg/madness/buffer.h
index 4f8b0e0b2..27e29d05b 100644
--- a/ttg/ttg/madness/buffer.h
+++ b/ttg/ttg/madness/buffer.h
@@ -112,6 +112,12 @@ struct Buffer : private Allocator {
     /* no-op */
   }
 
+
+  bool is_current_on(ttg::device::Device dev) const {
+    assert(is_valid());
+    return true;
+  }
+
   /* Get the owner device ID, i.e., the last updated
    * device buffer. */
   ttg::device::Device get_owner_device() const {
diff --git a/ttg/ttg/parsec/buffer.h b/ttg/ttg/parsec/buffer.h
index 075493078..7bcccbfd5 100644
--- a/ttg/ttg/parsec/buffer.h
+++ b/ttg/ttg/parsec/buffer.h
@@ -191,8 +191,23 @@ struct Buffer : public detail::ttg_parsec_data_wrapper_t
     m_data->owner_device = parsec_id;
   }
 
+  bool is_current_on(ttg::device::Device dev) const {
+    if (empty()) return true; // empty is current everywhere
+    int parsec_id = detail::ttg_device_to_parsec_device(dev);
+    uint32_t max_version = 0;
+    for (int i = 0; i < parsec_nb_devices; ++i) {
+      if (nullptr == m_data->device_copies[i]) continue;
+      max_version = std::max(max_version, m_data->device_copies[i]->version);
+    }
+    return (m_data->device_copies[parsec_id] &&
+            m_data->device_copies[parsec_id]->version == max_version);
+  }
+
   /* Get the owner device ID, i.e., the last updated
-   * device buffer. */
+   * device buffer.
+   * NOTE: there may be more than one device with the current
+   *       data so the result may not always be what is expected.
+   *       Use is_current_on() to check for a specific device. */
   ttg::device::Device get_owner_device() const {
     assert(is_valid());
     if (empty()) return ttg::device::current_device(); // empty is always valid