From 325f95f61b29f93e2783599ee3526036c9dedeb3 Mon Sep 17 00:00:00 2001
From: Callum Fare <callum@codeplay.com>
Date: Tue, 20 Jun 2023 11:44:55 +0100
Subject: [PATCH] [UR] Add CUDA-specific queue flags

---
 include/ur.py               |  4 ++++
 include/ur_api.h            |  6 +++++-
 scripts/core/queue.yml      |  6 ++++++
 source/common/ur_params.hpp | 30 ++++++++++++++++++++++++++++++
 4 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/include/ur.py b/include/ur.py
index 21d467255e..a4b1ea622c 100644
--- a/include/ur.py
+++ b/include/ur.py
@@ -1632,6 +1632,10 @@ class ur_queue_flags_v(IntEnum):
                                                     ## semantics. Implementation chooses submission mode.
     SUBMISSION_IMMEDIATE = UR_BIT(8)                ## Hint: enqueue and submit immediately. No change in queue semantics.
                                                     ## Implementation chooses submission mode.
+    USE_DEFAULT_STREAM = UR_BIT(9)                  ## Use the default stream. Only meaningful for CUDA. Other platforms may
+                                                    ## ignore this flag.
+    SYNC_WITH_DEFAULT_STREAM = UR_BIT(10)           ## Synchronize with the default stream. Only meaningful for CUDA. Other
+                                                    ## platforms may ignore this flag.
 
 class ur_queue_flags_t(c_int):
     def __str__(self):
diff --git a/include/ur_api.h b/include/ur_api.h
index 0470bcccc9..96b9d4c91b 100644
--- a/include/ur_api.h
+++ b/include/ur_api.h
@@ -4355,13 +4355,17 @@ typedef enum ur_queue_flag_t {
                                                              ///< semantics. Implementation chooses submission mode.
     UR_QUEUE_FLAG_SUBMISSION_IMMEDIATE = UR_BIT(8),          ///< Hint: enqueue and submit immediately. No change in queue semantics.
                                                              ///< Implementation chooses submission mode.
+    UR_QUEUE_FLAG_USE_DEFAULT_STREAM = UR_BIT(9),            ///< Use the default stream. Only meaningful for CUDA. Other platforms may
+                                                             ///< ignore this flag.
+    UR_QUEUE_FLAG_SYNC_WITH_DEFAULT_STREAM = UR_BIT(10),     ///< Synchronize with the default stream. Only meaningful for CUDA. Other
+                                                             ///< platforms may ignore this flag.
     /// @cond
     UR_QUEUE_FLAG_FORCE_UINT32 = 0x7fffffff
     /// @endcond
 
 } ur_queue_flag_t;
 /// @brief Bit Mask for validating ur_queue_flags_t
-#define UR_QUEUE_FLAGS_MASK 0xfffffe00
+#define UR_QUEUE_FLAGS_MASK 0xfffff800
 
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Query information about a command queue
diff --git a/scripts/core/queue.yml b/scripts/core/queue.yml
index 5671bf1c9f..86d167b7b9 100644
--- a/scripts/core/queue.yml
+++ b/scripts/core/queue.yml
@@ -68,6 +68,12 @@ etors:
     - name: SUBMISSION_IMMEDIATE
       value: "$X_BIT(8)"
       desc: "Hint: enqueue and submit immediately. No change in queue semantics. Implementation chooses submission mode."
+    - name: USE_DEFAULT_STREAM
+      desc: "Use the default stream. Only meaningful for CUDA. Other platforms may ignore this flag."
+      value: "$X_BIT(9)"
+    - name: SYNC_WITH_DEFAULT_STREAM
+      desc: "Synchronize with the default stream. Only meaningful for CUDA. Other platforms may ignore this flag."
+      value: "$X_BIT(10)"
 --- #--------------------------------------------------------------------------
 type: function
 desc: "Query information about a command queue"
diff --git a/source/common/ur_params.hpp b/source/common/ur_params.hpp
index f076fa1575..32d4e867ee 100644
--- a/source/common/ur_params.hpp
+++ b/source/common/ur_params.hpp
@@ -7749,6 +7749,14 @@ inline std::ostream &operator<<(std::ostream &os, enum ur_queue_flag_t value) {
     case UR_QUEUE_FLAG_SUBMISSION_IMMEDIATE:
         os << "UR_QUEUE_FLAG_SUBMISSION_IMMEDIATE";
         break;
+
+    case UR_QUEUE_FLAG_USE_DEFAULT_STREAM:
+        os << "UR_QUEUE_FLAG_USE_DEFAULT_STREAM";
+        break;
+
+    case UR_QUEUE_FLAG_SYNC_WITH_DEFAULT_STREAM:
+        os << "UR_QUEUE_FLAG_SYNC_WITH_DEFAULT_STREAM";
+        break;
     default:
         os << "unknown enumerator";
         break;
@@ -7859,6 +7867,28 @@ inline void serializeFlag<ur_queue_flag_t>(std::ostream &os, uint32_t flag) {
         }
         os << UR_QUEUE_FLAG_SUBMISSION_IMMEDIATE;
     }
+
+    if ((val & UR_QUEUE_FLAG_USE_DEFAULT_STREAM) ==
+        (uint32_t)UR_QUEUE_FLAG_USE_DEFAULT_STREAM) {
+        val ^= (uint32_t)UR_QUEUE_FLAG_USE_DEFAULT_STREAM;
+        if (!first) {
+            os << " | ";
+        } else {
+            first = false;
+        }
+        os << UR_QUEUE_FLAG_USE_DEFAULT_STREAM;
+    }
+
+    if ((val & UR_QUEUE_FLAG_SYNC_WITH_DEFAULT_STREAM) ==
+        (uint32_t)UR_QUEUE_FLAG_SYNC_WITH_DEFAULT_STREAM) {
+        val ^= (uint32_t)UR_QUEUE_FLAG_SYNC_WITH_DEFAULT_STREAM;
+        if (!first) {
+            os << " | ";
+        } else {
+            first = false;
+        }
+        os << UR_QUEUE_FLAG_SYNC_WITH_DEFAULT_STREAM;
+    }
     if (val != 0) {
         std::bitset<32> bits(val);
         if (!first) {