fix hipStreamAddCallback, block future work on stream (#1934)

hipamd/src/hip_hcc.cpp

@@ -1520,20 +1520,6 @@ hipError_t ihipStreamSynchronize(TlsData *tls, hipStream_t stream) {
     return e;
 }
 
-void ihipStreamCallbackHandler(ihipStreamCallback_t* cb) {
-    hipError_t e = hipSuccess;
-
-    // Synchronize stream
-    tprintf(DB_SYNC, "ihipStreamCallbackHandler wait on stream %s\n",
-            ToString(cb->_stream).c_str());
-    GET_TLS();
-    e = ihipStreamSynchronize(tls, cb->_stream);
-
-    // Call registered callback function
-    cb->_callback(cb->_stream, e, cb->_userData);
-    delete cb;
-}
-
 //---
 // Get the stream to use for a command submission.
 //

hipamd/src/hip_hcc_internal.h

@@ -654,19 +654,6 @@ class ihipStream_t {
 };
 
 
-//----
-// Internal structure for stream callback handler
-class ihipStreamCallback_t {
-   public:
-    ihipStreamCallback_t(hipStream_t stream, hipStreamCallback_t callback, void* userData)
-        : _stream(stream), _callback(callback), _userData(userData) {
-    };
-    hipStream_t _stream;
-    hipStreamCallback_t _callback;
-    void* _userData;
-};
-
-
 //----
 // Internal event structure:
 enum hipEventStatus_t {
@@ -980,7 +967,6 @@ hipError_t hipModuleGetFunctionEx(hipFunction_t* hfunc, hipModule_t hmod,
 
 hipStream_t ihipSyncAndResolveStream(hipStream_t, bool lockAcquired = 0);
 hipError_t ihipStreamSynchronize(TlsData *tls, hipStream_t stream);
-void ihipStreamCallbackHandler(ihipStreamCallback_t* cb);
 
 // Stream printf functions:
 inline std::ostream& operator<<(std::ostream& os, const ihipStream_t& s) {

hipamd/src/hip_stream.cpp

@@ -257,11 +257,39 @@ hipError_t hipStreamGetPriority(hipStream_t stream, int* priority) {
 hipError_t hipStreamAddCallback(hipStream_t stream, hipStreamCallback_t callback, void* userData,
                                 unsigned int flags) {
     HIP_INIT_API(hipStreamAddCallback, stream, callback, userData, flags);
-    hipError_t e = hipSuccess;
 
-    // Create a thread in detached mode to handle callback
-    ihipStreamCallback_t* cb = new ihipStreamCallback_t(stream, callback, userData);
-    std::thread(ihipStreamCallbackHandler, cb).detach();
+    auto stream_original{stream};
+    stream = ihipSyncAndResolveStream(stream);
 
-    return ihipLogStatus(e);
+    if (!stream) return hipErrorInvalidValue;
+
+    LockedAccessor_StreamCrit_t cs{stream->criticalData()};
+
+    // create first marker
+    auto cf = cs->_av.create_marker(hc::no_scope);
+    // get its signal
+    auto signal = *reinterpret_cast<hsa_signal_t*>(cf.get_native_handle());
+    // increment its signal value
+    hsa_signal_add_relaxed(signal, 1);
+
+    // create callback that can be passed to hsa_amd_signal_async_handler
+    // this function will call the user's callback, then sets first packet's signal to 0 to indicate completion
+    auto t{new std::function<void()>{[=]() {
+        callback(stream_original, hipSuccess, userData);
+        hsa_signal_store_relaxed(signal, 0);
+    }}};
+
+    // register above callback with HSA runtime to be called when first packet's signal
+    // is decremented from 2 to 1 by CP (or it is already at 1)
+    hsa_amd_signal_async_handler(signal, HSA_SIGNAL_CONDITION_EQ, 1,
+        [](hsa_signal_value_t x, void* p) {
+            (*static_cast<decltype(t)>(p))();
+            delete static_cast<decltype(t)>(p);
+            return false;
+        }, t);
+
+    // create additional marker that blocks on the first one
+    cs->_av.create_blocking_marker(cf, hc::no_scope);
+
+    return ihipLogStatus(hipSuccess);
 }