From a2ed01bc3ee4b8472768fc64777d160325bb0ee6 Mon Sep 17 00:00:00 2001
From: "Dunajski, Bartosz" <bartosz.dunajski@intel.com>
Date: Thu, 25 Apr 2019 09:59:08 +0200
Subject: [PATCH] Use waitForTaskCountWithKmdNotifyFallback in blitFromHostPtr
 method

Change-Id: Ib70df07e2a7cef514037c9327575ef8e867e6ca6
Related-To: NEO-3020
Signed-off-by: Dunajski, Bartosz <bartosz.dunajski@intel.com>
---
 .../command_stream_receiver_hw.inl            |  3 +-
 .../command_stream_receiver_hw_tests.cpp      | 39 +++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/runtime/command_stream/command_stream_receiver_hw.inl b/runtime/command_stream/command_stream_receiver_hw.inl
index 36cd90ccfb..b12fdfe886 100644
--- a/runtime/command_stream/command_stream_receiver_hw.inl
+++ b/runtime/command_stream/command_stream_receiver_hw.inl
@@ -820,9 +820,10 @@ void CommandStreamReceiverHw<GfxFamily>::blitFromHostPtr(MemObj &destinationMemO
 
     latestFlushedTaskCount = newTaskCount;
     taskCount = newTaskCount;
+    auto flushStampToWait = flushStamp->peekStamp();
 
     lock.unlock();
-    waitForCompletionWithTimeout(false, 0, newTaskCount);
+    waitForTaskCountWithKmdNotifyFallback(newTaskCount, flushStampToWait, false, false);
 }
 
 } // namespace NEO
diff --git a/unit_tests/command_stream/command_stream_receiver_hw_tests.cpp b/unit_tests/command_stream/command_stream_receiver_hw_tests.cpp
index 2d0b8d0fc9..480236a7c4 100644
--- a/unit_tests/command_stream/command_stream_receiver_hw_tests.cpp
+++ b/unit_tests/command_stream/command_stream_receiver_hw_tests.cpp
@@ -411,3 +411,42 @@ HWTEST_F(BcsTests, givenBufferWhenBlitCalledThenFlushCommandBuffer) {
 
     EXPECT_EQ(newTaskCount, csr.latestWaitForCompletionWithTimeoutTaskCount.load());
 }
+
+HWTEST_F(BcsTests, whenBlitFromHostPtrCalledThenCallWaitWithKmdFallback) {
+    class MyMockCsr : public UltCommandStreamReceiver<FamilyType> {
+      public:
+        using UltCommandStreamReceiver<FamilyType>::UltCommandStreamReceiver;
+
+        void waitForTaskCountWithKmdNotifyFallback(uint32_t taskCountToWait, FlushStamp flushStampToWait,
+                                                   bool useQuickKmdSleep, bool forcePowerSavingMode) override {
+            waitForTaskCountWithKmdNotifyFallbackCalled++;
+            taskCountToWaitPassed = taskCountToWait;
+            flushStampToWaitPassed = flushStampToWait;
+            useQuickKmdSleepPassed = useQuickKmdSleep;
+            forcePowerSavingModePassed = forcePowerSavingMode;
+        }
+
+        uint32_t taskCountToWaitPassed = 0;
+        FlushStamp flushStampToWaitPassed = 0;
+        bool useQuickKmdSleepPassed = false;
+        bool forcePowerSavingModePassed = false;
+        uint32_t waitForTaskCountWithKmdNotifyFallbackCalled = 0;
+    };
+
+    auto myMockCsr = std::make_unique<::testing::NiceMock<MyMockCsr>>(*pDevice->getExecutionEnvironment());
+    auto &bcsOsContext = pDevice->getUltCommandStreamReceiver<FamilyType>().getOsContext();
+    myMockCsr->initializeTagAllocation();
+    myMockCsr->setupContext(bcsOsContext);
+
+    cl_int retVal = CL_SUCCESS;
+    auto buffer = clUniquePtr<Buffer>(Buffer::create(context.get(), CL_MEM_READ_WRITE, 1, nullptr, retVal));
+    void *hostPtr = reinterpret_cast<void *>(0x12340000);
+
+    myMockCsr->blitFromHostPtr(*buffer, hostPtr, 1);
+
+    EXPECT_EQ(1u, myMockCsr->waitForTaskCountWithKmdNotifyFallbackCalled);
+    EXPECT_EQ(myMockCsr->taskCount, myMockCsr->taskCountToWaitPassed);
+    EXPECT_EQ(myMockCsr->flushStamp->peekStamp(), myMockCsr->flushStampToWaitPassed);
+    EXPECT_FALSE(myMockCsr->useQuickKmdSleepPassed);
+    EXPECT_FALSE(myMockCsr->forcePowerSavingModePassed);
+}