fix: download staging buffer before reading from it

Related-To: NEO-14154

Signed-off-by: Szymon Morek <[email protected]>

Author: smorek-intel

shared/source/utilities/staging_buffer_manager.cpp

@@ -216,7 +216,7 @@ WaitStatus StagingBufferManager::copyStagingToHost(const std::pair<UserData, Sta
     if (status == WaitStatus::gpuHang) {
         return status;
     }
-
+    transfer.second.csr->downloadAllocations(true);
     auto &userData = transfer.first;
     tracker = transfer.second;
     auto stagingBuffer = addrToPtr(tracker.chunkAddress);

shared/test/unit_test/utilities/staging_buffer_manager_tests.cpp

@@ -521,6 +521,16 @@ TEST_F(StagingBufferManagerTest, givenStagingBufferWhenPerformImageReadThenRegio
     imageTransferThroughStagingBuffers(true, stagingBufferSize, globalOrigin, globalRegion, expectedChunks);
 }
 
+HWTEST_F(StagingBufferManagerTest, givenStagingBufferWhenPerformImageReadThenDownloadAllocationsCalledForAllReadChunks) {
+    size_t expectedChunks = 8;
+    const size_t globalOrigin[3] = {0, 0, 0};
+    const size_t globalRegion[3] = {4, expectedChunks, 1};
+    imageTransferThroughStagingBuffers(true, stagingBufferSize, globalOrigin, globalRegion, expectedChunks);
+    auto ultCsr = reinterpret_cast<UltCommandStreamReceiver<FamilyType> *>(csr);
+    EXPECT_EQ(expectedChunks, ultCsr->downloadAllocationsCalledCount);
+    EXPECT_TRUE(ultCsr->latestDownloadAllocationsBlocking);
+}
+
 TEST_F(StagingBufferManagerTest, givenStagingBufferWhenPerformImageReadThenWholeRegionCovered) {
     size_t expectedChunks = 8;
     const size_t globalOrigin[3] = {0, 0, 0};