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

[Ross Brunton via llvm-commits]

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

>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 1/2] [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

>From a114728fa0276de7391f3a92362fd7b859ff3aa1 Mon Sep 17 00:00:00 2001
From: Ross Brunton <ross at codeplay.com>
Date: Tue, 19 Aug 2025 10:05:50 +0100
Subject: [PATCH 2/2] Use function pointer type for callback function

---
 offload/plugins-nextgen/amdgpu/src/rtl.cpp | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index f7cf09da7a62d..da3ce2d1feae3 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -908,6 +908,10 @@ struct AMDGPUQueueTy {
 /// devices. This class relies on signals to implement streams and define the
 /// dependencies between asynchronous operations.
 struct AMDGPUStreamTy {
+public:
+  /// Function pointer type for `pushHostCallback`
+  using HostFnType = void (*)(void *);
+
 private:
   /// Utility struct holding arguments for async H2H memory copies.
   struct MemcpyArgsTy {
@@ -1064,7 +1068,7 @@ struct AMDGPUStreamTy {
   bool UseMultipleSdmaEngines;
 
   struct CallbackDataType {
-    void (*UserFn)(void *);
+    HostFnType UserFn;
     void *UserData;
     AMDGPUSignalTy *OutputSignal;
   };
@@ -1510,7 +1514,7 @@ struct AMDGPUStreamTy {
                                    OutputSignal->get());
   }
 
-  Error pushHostCallback(void (*Callback)(void *), void *UserData) {
+  Error pushHostCallback(HostFnType Callback, void *UserData) {
     // Retrieve an available signal for the operation's output.
     AMDGPUSignalTy *OutputSignal = nullptr;
     if (auto Err = SignalManager.getResource(OutputSignal))
@@ -2609,7 +2613,7 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
     return Plugin::success();
   }
 
-  Error enqueueHostCallImpl(void (*Callback)(void *), void *UserData,
+  Error enqueueHostCallImpl(AMDGPUStreamTy::HostFnType Callback, void *UserData,
                             AsyncInfoWrapperTy &AsyncInfo) override {
     AMDGPUStreamTy *Stream = nullptr;
     if (auto Err = getStream(AsyncInfo, Stream))