[llvm] [Offload] Use `amd_signal_async_handler` for host function calls (PR #154131)

[Ross Brunton via llvm-commits]

https://github.com/RossBrunton created https://github.com/llvm/llvm-project/pull/154131

None

>From b02f33ae520e0936ec64bfcc772df8a2846258df Mon Sep 17 00:00:00 2001
From: Ross Brunton <ross at codeplay.com>
Date: Mon, 18 Aug 2025 16:06:52 +0100
Subject: [PATCH] [Offload] Use `amd_signal_async_handler` for host function
 calls

---
 offload/plugins-nextgen/amdgpu/src/rtl.cpp | 44 +++++++++++++---------
 1 file changed, 27 insertions(+), 17 deletions(-)

diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index 83280fe0a49c9..f7cf09da7a62d 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -1063,18 +1063,19 @@ struct AMDGPUStreamTy {
   /// Indicate to spread data transfers across all available SDMAs
   bool UseMultipleSdmaEngines;
 
+  struct CallbackDataType {
+    void (*UserFn)(void *);
+    void *UserData;
+    AMDGPUSignalTy *OutputSignal;
+  };
   /// Wrapper function for implementing host callbacks
-  static void CallbackWrapper(AMDGPUSignalTy *InputSignal,
-                              AMDGPUSignalTy *OutputSignal,
-                              void (*Callback)(void *), void *UserData) {
-    // The wait call will not error in this context.
-    if (InputSignal)
-      if (auto Err = InputSignal->wait())
-        reportFatalInternalError(std::move(Err));
-
-    Callback(UserData);
-
-    OutputSignal->signal();
+  static bool callbackWrapper([[maybe_unused]] hsa_signal_value_t Signal,
+                              void *UserData) {
+    auto CallbackData = reinterpret_cast<CallbackDataType *>(UserData);
+    CallbackData->UserFn(CallbackData->UserData);
+    CallbackData->OutputSignal->signal();
+    delete CallbackData;
+    return false;
   }
 
   /// Return the current number of asynchronous operations on the stream.
@@ -1525,13 +1526,22 @@ struct AMDGPUStreamTy {
       InputSignal = consume(OutputSignal).second;
     }
 
-    // "Leaking" the thread here is consistent with other work added to the
-    // queue. The input and output signals will remain valid until the output is
-    // signaled.
-    std::thread(CallbackWrapper, InputSignal, OutputSignal, Callback, UserData)
-        .detach();
+    auto *CallbackData = new CallbackDataType{Callback, UserData, OutputSignal};
+    if (InputSignal && InputSignal->load()) {
+      hsa_status_t Status = hsa_amd_signal_async_handler(
+          InputSignal->get(), HSA_SIGNAL_CONDITION_EQ, 0, callbackWrapper,
+          CallbackData);
 
-    return Plugin::success();
+      return Plugin::check(Status, "error in hsa_amd_signal_async_handler: %s");
+    } else {
+      // No dependencies - schedule it now.
+      // Using a seperate thread because this function should run asynchronously
+      // and not block the main thread.
+      std::thread([](void *CallbackData) { callbackWrapper(0, CallbackData); },
+                  CallbackData)
+          .detach();
+      return Plugin::success();
+    }
   }
 
   /// Synchronize with the stream. The current thread waits until all operations