From 4ec730f1dc7851e6d1ab16e3535d9236c42cd548 Mon Sep 17 00:00:00 2001
From: David Yat Sin <David.YatSin@amd.com>
Date: Tue, 5 Nov 2024 10:14:20 +0000
Subject: [PATCH] rocr: Add HSA_SIGNAL_WAIT_ABORT_TIMEOUT

Add support for abort timeout when hsa_signal_wait_relaxed is called and
signal does not clear within timeout.
timeout is in seconds

Change-Id: If1db5a8af33c82ddc4b48968c3d8eceb97d0ea6d
---
 .../core/runtime/default_signal.cpp           | 12 +++++++++++
 .../core/runtime/interrupt_signal.cpp         | 20 +++++++++++++++++--
 runtime/hsa-runtime/core/util/flag.h          |  8 ++++++++
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/runtime/hsa-runtime/core/runtime/default_signal.cpp b/runtime/hsa-runtime/core/runtime/default_signal.cpp
index 0c5e7ca596..5881f0e03f 100644
--- a/runtime/hsa-runtime/core/runtime/default_signal.cpp
+++ b/runtime/hsa-runtime/core/runtime/default_signal.cpp
@@ -86,6 +86,9 @@ hsa_signal_value_t BusyWaitSignal::WaitRelaxed(hsa_signal_condition_t condition,
   bool condition_met = false;
   int64_t value;
 
+  const uint32_t &signal_abort_timeout =
+    core::Runtime::runtime_singleton_->flag().signal_abort_timeout();
+
   debug_warning_n((!g_use_interrupt_wait || isIPC()) &&
                   "Use of non-host signal in host signal wait API.", 10);
 
@@ -139,6 +142,15 @@ hsa_signal_value_t BusyWaitSignal::WaitRelaxed(hsa_signal_condition_t condition,
       return hsa_signal_value_t(value);
     }
 
+    if (signal_abort_timeout) {
+      const timer::fast_clock::duration abort_timeout =
+          std::chrono::seconds(signal_abort_timeout);
+
+      if(time - start_time > abort_timeout)
+        throw AMD::hsa_exception(HSA_STATUS_ERROR_FATAL,
+                                 "Signal wait abort timeout.\n");
+    }
+
     if (time - start_time > kMaxElapsed) {
       os::uSleep(20);
 #if defined(__i386__) || defined(__x86_64__)
diff --git a/runtime/hsa-runtime/core/runtime/interrupt_signal.cpp b/runtime/hsa-runtime/core/runtime/interrupt_signal.cpp
index 6acb75c2e4..d8f46a8d0b 100644
--- a/runtime/hsa-runtime/core/runtime/interrupt_signal.cpp
+++ b/runtime/hsa-runtime/core/runtime/interrupt_signal.cpp
@@ -152,6 +152,9 @@ hsa_signal_value_t InterruptSignal::WaitRelaxed(
 
   uint64_t event_age = 1;
 
+  const uint32_t &signal_abort_timeout =
+    core::Runtime::runtime_singleton_->flag().signal_abort_timeout();
+
   if (!core::Runtime::runtime_singleton_->KfdVersion().supports_event_age) {
       event_age = 0;
       // Allow only the first waiter to sleep. Without event age tracking,
@@ -212,6 +215,15 @@ hsa_signal_value_t InterruptSignal::WaitRelaxed(
       return hsa_signal_value_t(value);
     }
 
+    if (signal_abort_timeout) {
+      const timer::fast_clock::duration abort_timeout =
+          std::chrono::seconds(signal_abort_timeout);
+
+      if(time - start_time > abort_timeout)
+        throw AMD::hsa_exception(HSA_STATUS_ERROR_FATAL,
+                                 "Signal wait abort timeout.\n");
+    }
+
     if (wait_hint == HSA_WAIT_STATE_ACTIVE) {
 #if defined(__i386__) || defined(__x86_64__)
       if (g_use_mwaitx) {
@@ -235,9 +247,13 @@ hsa_signal_value_t InterruptSignal::WaitRelaxed(
 
     uint32_t wait_ms;
     auto time_remaining = fast_timeout - (time - start_time);
-    uint64_t ct=timer::duration_cast<std::chrono::milliseconds>(
+    uint64_t ct = timer::duration_cast<std::chrono::milliseconds>(
       time_remaining).count();
-    wait_ms = (ct>0xFFFFFFFEu) ? 0xFFFFFFFEu : ct;
+
+    wait_ms = static_cast<uint32_t>(std::min(ct, 0xFFFFFFFEUL));
+    if (signal_abort_timeout)
+      wait_ms = std::min(wait_ms, signal_abort_timeout * 1000);
+
     hsaKmtWaitOnEvent_Ext(event_, wait_ms, &event_age);
   }
 }
diff --git a/runtime/hsa-runtime/core/util/flag.h b/runtime/hsa-runtime/core/util/flag.h
index e4783c3bd3..4a23e7a256 100644
--- a/runtime/hsa-runtime/core/util/flag.h
+++ b/runtime/hsa-runtime/core/util/flag.h
@@ -253,6 +253,10 @@ class Flag {
 
     var = os::GetEnvVar("HSA_WAIT_ANY_DEBUG");
     wait_any_ = (var == "1") ? true : false;
+
+    /* hsa_signal_wait_relaxed abort timeout  */
+    var = os::GetEnvVar("HSA_SIGNAL_WAIT_ABORT_TIMEOUT");
+    signal_abort_timeout_ = var.empty() ? 0 : atoi(var.c_str());
   }
 
   void parse_masks(uint32_t maxGpu, uint32_t maxCU) {
@@ -366,6 +370,9 @@ class Flag {
   size_t pc_sampling_max_device_buffer_size() const { return pc_sampling_max_device_buffer_size_; }
 
   bool dev_mem_queue() const { return dev_mem_queue_; }
+
+  uint32_t signal_abort_timeout() const { return signal_abort_timeout_; }
+
  private:
   bool check_flat_scratch_;
   bool enable_vm_fault_message_;
@@ -396,6 +403,7 @@ class Flag {
   bool enable_ipc_mode_legacy_;
   bool wait_any_;
   bool dev_mem_queue_;
+  uint32_t signal_abort_timeout_;
 
   SDMA_OVERRIDE enable_sdma_;
   SDMA_OVERRIDE enable_peer_sdma_;