[llvm] [Offload] Check Expected from hasPendingWorkImpl in AMDGPU dataFill (PR #223329)

[via llvm-commits]

https://github.com/StevenYangCC updated https://github.com/llvm/llvm-project/pull/223329

>From 247e12749ac3a65862d90396ff55b3db5df7f124 Mon Sep 17 00:00:00 2001
From: "chengcang.yang" <yangchengcang at gmail.com>
Date: Mon, 14 Sep 2026 16:25:57 +0800
Subject: [PATCH] [Offload] Check Expected from hasPendingWorkImpl in AMDGPU
 dataFill

hasPendingWorkImpl returns Expected<bool>. Using that result directly as
a condition tests whether a value is present, not whether the queue has
pending work. The asynchronous fill path ran on every successful query,
and a failed query left an Error unconsumed.

Unwrap the Error first, then test the boolean. Add olMemFill tests for
the idle-queue and pending-work four-byte fill paths on AMDGPU.

Release ManuallyTriggeredTask from the destructor if a test asserts
before trigger(), so a waiting host callback is not destroyed mid-wait.
Set Enqueued only after olCreateEvent succeeds. If the host callback is
already waiting and event creation fails, unblock the wait without
syncing a missing completion event.
---
 offload/plugins-nextgen/amdgpu/src/rtl.cpp    |  6 +-
 .../unittests/OffloadAPI/common/Fixtures.hpp  | 28 ++++++++-
 .../unittests/OffloadAPI/memory/olMemFill.cpp | 57 +++++++++++++++++++
 3 files changed, 88 insertions(+), 3 deletions(-)

diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index ad9b98795c6edf..68e784d8063620 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -3050,7 +3050,11 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
         llvm_unreachable("Invalid pattern size");
       }
 
-      if (hasPendingWorkImpl(AsyncInfoWrapper)) {
+      auto Pending = hasPendingWorkImpl(AsyncInfoWrapper);
+      if (auto Err = Pending.takeError())
+        return Err;
+
+      if (*Pending) {
         AMDGPUStreamTy *Stream = nullptr;
         if (auto Err = getStream(AsyncInfoWrapper, Stream))
           return Err;
diff --git a/offload/unittests/OffloadAPI/common/Fixtures.hpp b/offload/unittests/OffloadAPI/common/Fixtures.hpp
index a05be01648ebc8..b88d3e3c04f3af 100644
--- a/offload/unittests/OffloadAPI/common/Fixtures.hpp
+++ b/offload/unittests/OffloadAPI/common/Fixtures.hpp
@@ -148,7 +148,9 @@ struct ManuallyTriggeredTask {
   std::mutex M;
   std::condition_variable CV;
   bool Flag = false;
-  ol_event_handle_t CompleteEvent;
+  bool Enqueued = false;
+  bool Triggered = false;
+  ol_event_handle_t CompleteEvent = nullptr;
 
   ol_result_t enqueue(ol_queue_handle_t Queue) {
     if (auto Err = olLaunchHostFunction(
@@ -159,7 +161,17 @@ struct ManuallyTriggeredTask {
             this))
       return Err;
 
-    return olCreateEvent(Queue, OL_EVENT_FLAGS_NONE, &CompleteEvent);
+    if (auto Err = olCreateEvent(Queue, OL_EVENT_FLAGS_NONE, &CompleteEvent)) {
+      // The host callback is already waiting. Unblock it without treating the
+      // task as fully enqueued, since there is no completion event to sync.
+      Flag = true;
+      CV.notify_one();
+      CompleteEvent = nullptr;
+      return Err;
+    }
+
+    Enqueued = true;
+    return nullptr;
   }
 
   void wait() {
@@ -169,11 +181,23 @@ struct ManuallyTriggeredTask {
   }
 
   ol_result_t trigger() {
+    if (Triggered)
+      return nullptr;
+    Triggered = true;
     Flag = true;
     CV.notify_one();
 
+    if (!CompleteEvent)
+      return nullptr;
     return olSyncEvent(CompleteEvent);
   }
+
+  /// ASSERT macros return from the test before later statements run. Release a
+  /// waiting host callback so teardown does not destroy it while still blocked.
+  ~ManuallyTriggeredTask() {
+    if (Enqueued && !Triggered)
+      (void)trigger();
+  }
 };
 
 struct OffloadTest : ::testing::Test {
diff --git a/offload/unittests/OffloadAPI/memory/olMemFill.cpp b/offload/unittests/OffloadAPI/memory/olMemFill.cpp
index 467a551c48c94c..2565c910f02d11 100644
--- a/offload/unittests/OffloadAPI/memory/olMemFill.cpp
+++ b/offload/unittests/OffloadAPI/memory/olMemFill.cpp
@@ -9,6 +9,7 @@
 #include "../common/Fixtures.hpp"
 #include <OffloadAPI.h>
 #include <array>
+#include <cstring>
 #include <gtest/gtest.h>
 #include <vector>
 
@@ -75,6 +76,62 @@ TEST_P(olMemFillTest, Success32Enqueue) {
   test_body<uint32_t, 0xDEADBEEF, 1024, true>();
 }
 
+// The AMDGPU plugin chooses a synchronous HSA fill when the queue is idle and
+// a host-callback fill when work is already pending. Those two paths used to
+// be collapsed because Expected<bool> was tested for "has a value" instead of
+// "has pending work".
+TEST_P(olMemFillTest, Success32IdleQueueCompletesWithoutSync) {
+  if (getPlatformBackend() != OL_PLATFORM_BACKEND_AMDGPU)
+    GTEST_SKIP() << "Idle-queue synchronous fill is AMDGPU-specific";
+
+  constexpr size_t Size = 1024;
+  void *Alloc;
+  ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED, Size, &Alloc));
+
+  uint32_t Pattern = 0xDEADBEEF;
+  ASSERT_SUCCESS(olMemFill(Queue, Alloc, sizeof(Pattern), &Pattern, Size));
+
+  bool IsQueueWorkCompleted = false;
+  ASSERT_SUCCESS(olQueryQueue(Queue, &IsQueueWorkCompleted));
+  ASSERT_TRUE(IsQueueWorkCompleted);
+
+  auto *AllocPtr = reinterpret_cast<uint32_t *>(Alloc);
+  for (size_t i = 0; i < Size / sizeof(Pattern); ++i)
+    ASSERT_EQ(AllocPtr[i], Pattern);
+
+  olMemFree(Alloc);
+}
+
+TEST_P(olMemFillTest, Success32EnqueueStaysPendingUntilHostTask) {
+  if (getPlatformBackend() != OL_PLATFORM_BACKEND_AMDGPU)
+    GTEST_SKIP() << "Pending-work asynchronous fill is AMDGPU-specific";
+
+  ManuallyTriggeredTask Manual;
+  ASSERT_SUCCESS(Manual.enqueue(Queue));
+
+  constexpr size_t Size = 1024;
+  void *Alloc;
+  ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED, Size, &Alloc));
+  std::memset(Alloc, 0, Size);
+
+  uint32_t Pattern = 0xDEADBEEF;
+  ASSERT_SUCCESS(olMemFill(Queue, Alloc, sizeof(Pattern), &Pattern, Size));
+
+  bool IsQueueWorkCompleted = false;
+  ASSERT_SUCCESS(olQueryQueue(Queue, &IsQueueWorkCompleted));
+  ASSERT_FALSE(IsQueueWorkCompleted);
+
+  auto *AllocPtr = reinterpret_cast<uint32_t *>(Alloc);
+  ASSERT_EQ(AllocPtr[0], 0u);
+
+  ASSERT_SUCCESS(Manual.trigger());
+  ASSERT_SUCCESS(olSyncQueue(Queue));
+  for (size_t i = 0; i < Size / sizeof(Pattern); ++i)
+    ASSERT_EQ(AllocPtr[i], Pattern);
+
+  olMemFree(Alloc);
+}
+
 TEST_P(olMemFillTest, SuccessLarge) {
   constexpr size_t Size = 1024;
   void *Alloc;