Fix the ordering of the external stream. (#23598)

There was an issue with the previous fix that we would not retain the
data properly.

Signed-off-by: Andrew Woloszyn <andrew.woloszyn@gmail.com>

Author: AWoloszyn

runtime/src/iree/hal/drivers/hip/hip_device.c

@@ -1566,10 +1566,10 @@ static iree_status_t iree_hal_hip_device_perform_buffer_operation_now(
     }
   }
   IREE_TRACE_ZONE_END(z3);
-
+  iree_hal_hip_dispatch_completed_data_t* external_stream_dispatch_data = NULL;
   if (device->uses_external_stream) {
-    iree_hal_hip_set_external_stream_data_completed(
-        data->base.external_stream_dispatch_data);
+    external_stream_dispatch_data = data->base.external_stream_dispatch_data;
+    iree_hal_resource_retain(external_stream_dispatch_data);
   }
 
   const iree_hal_hip_dynamic_symbols_t* symbols = device->hip_symbols;
@@ -1589,6 +1589,12 @@ static iree_status_t iree_hal_hip_device_perform_buffer_operation_now(
     iree_hal_hip_device_destroy_buffer_callback_data(data);
   }
 
+  if (external_stream_dispatch_data) {
+    iree_hal_hip_set_external_stream_data_completed(
+        external_stream_dispatch_data);
+    iree_hal_resource_release(external_stream_dispatch_data);
+  }
+
   IREE_TRACE_ZONE_END(z0);
   return iree_status_join(
       status, IREE_HIP_CALL_TO_STATUS(symbols, hipCtxPopCurrent(NULL)));
@@ -2061,6 +2067,11 @@ static iree_status_t iree_hal_hip_device_perform_queue_read_now(
         device, device->devices[device_ordinal].hip_dispatch_stream,
         data->base.wait_semaphore_list, device_ordinal);
   }
+  iree_hal_hip_dispatch_completed_data_t* external_stream_dispatch_data = NULL;
+  if (device->uses_external_stream) {
+    external_stream_dispatch_data = data->base.external_stream_dispatch_data;
+    iree_hal_resource_retain(external_stream_dispatch_data);
+  }
 
   const iree_hal_hip_dynamic_symbols_t* symbols = device->hip_symbols;
   iree_device_size_t amount_left = data->length;
@@ -2152,20 +2163,22 @@ static iree_status_t iree_hal_hip_device_perform_queue_read_now(
     }
   }
 
-  if (device->uses_external_stream) {
-    iree_hal_hip_set_external_stream_data_completed(
-        data->base.external_stream_dispatch_data);
-  }
-
   if (!iree_status_is_ok(status)) {
     for (iree_host_size_t i = 0; i < data->base.signal_semaphore_list.count;
          ++i) {
       iree_hal_semaphore_fail(data->base.signal_semaphore_list.semaphores[i],
                               iree_status_clone(status));
     }
+
     iree_hal_hip_device_destroy_queue_read_callback_data(data);
   }
 
+  if (external_stream_dispatch_data) {
+    iree_hal_hip_set_external_stream_data_completed(
+        external_stream_dispatch_data);
+    iree_hal_resource_release(external_stream_dispatch_data);
+  }
+
   IREE_TRACE_ZONE_END(z0);
   return iree_status_join(
       status, IREE_HIP_CALL_TO_STATUS(symbols, hipCtxPopCurrent(NULL)));
@@ -2494,10 +2507,10 @@ static iree_status_t iree_hal_hip_device_execute_now(void* user_data,
   }
 
   IREE_TRACE_ZONE_END(z1);
-
+  iree_hal_hip_dispatch_completed_data_t* external_stream_dispatch_data = NULL;
   if (device->uses_external_stream) {
-    iree_hal_hip_set_external_stream_data_completed(
-        data->base.external_stream_dispatch_data);
+    external_stream_dispatch_data = data->base.external_stream_dispatch_data;
+    iree_hal_resource_retain(external_stream_dispatch_data);
   }
 
   // Store symbols, because the cleanup may trigger off-thread
@@ -2520,6 +2533,12 @@ static iree_status_t iree_hal_hip_device_execute_now(void* user_data,
     iree_hal_hip_device_destroy_callback_data(data);
   }
 
+  if (external_stream_dispatch_data) {
+    iree_hal_hip_set_external_stream_data_completed(
+        external_stream_dispatch_data);
+    iree_hal_resource_release(external_stream_dispatch_data);
+  }
+
   IREE_TRACE_ZONE_END(z0);
   return iree_status_join(
       status, IREE_HIP_CALL_TO_STATUS(symbols, hipCtxPopCurrent(NULL)));