Change pollForCompletion() insertion locations

Poll is done on:
- Aub CSR destruction
- expectMemory
- blocking calls

Poll is not done on flush

Change-Id: I1a776a932cb608c01f0de249e7cef26b00147f31
Signed-off-by: Maciej Dziuban <[email protected]>

Author: DziubanMaciejIntel

runtime/command_stream/aub_command_stream_receiver_hw.h

@@ -15,6 +15,7 @@
 #include "runtime/memory_manager/page_table.h"
 #include "runtime/memory_manager/physical_address_allocator.h"
 #include "runtime/memory_manager/os_agnostic_memory_manager.h"
+#include "runtime/utilities/spinlock.h"
 
 namespace OCLRT {
 
@@ -62,6 +63,8 @@ class AUBCommandStreamReceiverHw : public CommandStreamReceiverSimulatedHw<GfxFa
     // Family specific version
     MOCKABLE_VIRTUAL void submitBatchBuffer(uint64_t batchBufferGpuAddress, const void *batchBuffer, size_t batchBufferSize, uint32_t memoryBank, uint64_t entryBits);
     MOCKABLE_VIRTUAL void pollForCompletion();
+    void pollForCompletionImpl();
+    void waitForTaskCountWithKmdNotifyFallback(uint32_t taskCountToWait, FlushStamp flushStampToWait, bool useQuickKmdSleep, bool forcePowerSavingMode) override;
 
     uint32_t getDumpHandle();
     MOCKABLE_VIRTUAL void addContextToken(uint32_t dumpHandle);
@@ -115,5 +118,8 @@ class AUBCommandStreamReceiverHw : public CommandStreamReceiverSimulatedHw<GfxFa
 
     bool dumpAubNonWritable = false;
     ExternalAllocationsContainer externalAllocations;
+
+    uint32_t pollForCompletionTaskCount = 0u;
+    SpinLock pollForCompletionLock;
 };
 } // namespace OCLRT

runtime/command_stream/aub_command_stream_receiver_hw.inl

@@ -32,6 +32,7 @@
 
 #include <algorithm>
 #include <cstring>
+#include "aub_command_stream_receiver_hw.h"
 
 namespace OCLRT {
 
@@ -89,6 +90,9 @@ AUBCommandStreamReceiverHw<GfxFamily>::AUBCommandStreamReceiverHw(const Hardware
 
 template <typename GfxFamily>
 AUBCommandStreamReceiverHw<GfxFamily>::~AUBCommandStreamReceiverHw() {
+    if (osContext) {
+        pollForCompletion();
+    }
     freeEngineInfoTable();
 }
 
@@ -342,10 +346,6 @@ FlushStamp AUBCommandStreamReceiverHw<GfxFamily>::flush(BatchBuffer &batchBuffer
 
     submitBatchBuffer(batchBufferGpuAddress, pBatchBuffer, sizeBatchBuffer, this->getMemoryBank(batchBuffer.commandBufferAllocation), this->getPPGTTAdditionalBits(batchBuffer.commandBufferAllocation));
 
-    if (!DebugManager.flags.AUBDumpConcurrentCS.get()) {
-        pollForCompletion();
-    }
-
     if (this->standalone) {
         *this->tagAddress = this->peekLatestSentTaskCount();
     }
@@ -561,6 +561,17 @@ void AUBCommandStreamReceiverHw<GfxFamily>::submitBatchBuffer(uint64_t batchBuff
 
 template <typename GfxFamily>
 void AUBCommandStreamReceiverHw<GfxFamily>::pollForCompletion() {
+    const auto lock = std::unique_lock<decltype(pollForCompletionLock)>{pollForCompletionLock};
+    if (this->pollForCompletionTaskCount == this->taskCount) {
+        return;
+    }
+    pollForCompletionImpl();
+}
+
+template <typename GfxFamily>
+void AUBCommandStreamReceiverHw<GfxFamily>::pollForCompletionImpl() {
+    this->pollForCompletionTaskCount = this->taskCount;
+
     if (hardwareContext) {
         hardwareContext->pollForCompletion();
         return;
@@ -578,6 +589,12 @@ void AUBCommandStreamReceiverHw<GfxFamily>::pollForCompletion() {
         AubMemDump::CmdServicesMemTraceRegisterPoll::TimeoutActionValues::Abort);
 }
 
+template <typename GfxFamily>
+inline void AUBCommandStreamReceiverHw<GfxFamily>::waitForTaskCountWithKmdNotifyFallback(uint32_t taskCountToWait, FlushStamp flushStampToWait, bool useQuickKmdSleep, bool forcePowerSavingMode) {
+    CommandStreamReceiverSimulatedHw<GfxFamily>::waitForTaskCountWithKmdNotifyFallback(taskCountToWait, flushStampToWait, useQuickKmdSleep, forcePowerSavingMode);
+    pollForCompletion();
+}
+
 template <typename GfxFamily>
 constexpr uint32_t AUBCommandStreamReceiverHw<GfxFamily>::getMaskAndValueForPollForCompletion() {
     return 0x100;
@@ -691,6 +708,8 @@ void AUBCommandStreamReceiverHw<GfxFamily>::expectMemoryNotEqual(void *gfxAddres
 template <typename GfxFamily>
 void AUBCommandStreamReceiverHw<GfxFamily>::expectMemory(const void *gfxAddress, const void *srcAddress,
                                                          size_t length, uint32_t compareOperation) {
+    pollForCompletion();
+
     if (hardwareContext) {
         hardwareContext->expectMemory(reinterpret_cast<uint64_t>(gfxAddress), srcAddress, length, compareOperation);
     }
@@ -846,4 +865,5 @@ template <typename GfxFamily>
 int AUBCommandStreamReceiverHw<GfxFamily>::getAddressSpaceFromPTEBits(uint64_t entryBits) const {
     return AubMemDump::AddressSpaceValues::TraceNonlocal;
 }
+
 } // namespace OCLRT

unit_tests/aub_tests/command_stream/aub_command_stream_fixture.h

@@ -77,14 +77,17 @@ class AUBCommandStreamFixture : public CommandStreamFixture {
     }
 
     template <typename FamilyType>
-    void pollForCompletion() {
+    AUBCommandStreamReceiverHw<FamilyType> *getAubCsr() {
         CommandStreamReceiver *csr = pCommandStreamReceiver;
         if (testMode == TestMode::AubTestsWithTbx) {
             csr = reinterpret_cast<CommandStreamReceiverWithAUBDump<TbxCommandStreamReceiverHw<FamilyType>> *>(pCommandStreamReceiver)->aubCSR;
         }
+        return reinterpret_cast<AUBCommandStreamReceiverHw<FamilyType> *>(csr);
+    }
 
-        auto aubCsr = reinterpret_cast<AUBCommandStreamReceiverHw<FamilyType> *>(csr);
-        aubCsr->pollForCompletion();
+    template <typename FamilyType>
+    void pollForCompletion() {
+        getAubCsr<FamilyType>()->pollForCompletion();
     }
 
     GraphicsAllocation *createResidentAllocationAndStoreItInCsr(const void *address, size_t size) {

unit_tests/aub_tests/command_stream/aub_command_stream_tests.cpp

@@ -58,6 +58,7 @@ struct AUBFixture : public AUBCommandStreamFixture,
         ResidencyContainer allocationsForResidency;
         pCommandStreamReceiver->flush(batchBuffer, allocationsForResidency);
 
+        AUBCommandStreamFixture::getAubCsr<FamilyType>()->pollForCompletionImpl();
         auto engineType = pCommandStreamReceiver->getOsContext().getEngineType();
         auto mmioBase = CommandStreamReceiverSimulatedCommonHw<FamilyType>::getCsTraits(engineType.type).mmioBase;
         AUBCommandStreamFixture::expectMMIO<FamilyType>(AubMemDump::computeRegisterOffset(mmioBase, 0x2094), noopId);
@@ -75,6 +76,7 @@ HWTEST_F(AUBcommandstreamTests, testFlushTwice) {
     BatchBuffer batchBuffer2{pCS->getGraphicsAllocation(), 0, 0, nullptr, false, false, QueueThrottle::MEDIUM, pCS->getUsed(), pCS};
     ResidencyContainer allocationsForResidency2;
     pCommandStreamReceiver->flush(batchBuffer2, allocationsForResidency);
+    AUBCommandStreamFixture::getAubCsr<FamilyType>()->pollForCompletionImpl();
 }
 
 HWTEST_F(AUBcommandstreamTests, testNoopIdRcs) {

unit_tests/command_stream/aub_command_stream_receiver_2_tests.cpp

@@ -851,8 +851,9 @@ HWTEST_F(AubCommandStreamReceiverTests, givenAubCsrWhenAskedForMemoryExpectation
     uint32_t compareNotEqual = AubMemDump::CmdServicesMemTraceMemoryCompare::CompareOperationValues::CompareNotEqual;
     uint32_t compareEqual = AubMemDump::CmdServicesMemTraceMemoryCompare::CompareOperationValues::CompareEqual;
 
-    MyMockAubCsr myMockCsr(**platformDevices, std::string(), true, *pDevice->getExecutionEnvironment());
     auto mockStream = std::make_unique<MockAubFileStream>();
+    MyMockAubCsr myMockCsr(**platformDevices, std::string(), true, *pDevice->getExecutionEnvironment());
+    myMockCsr.setupContext(pDevice->getExecutionEnvironment()->commandStreamReceivers[0][0]->getOsContext());
     myMockCsr.stream = mockStream.get();
 
     myMockCsr.expectMemoryNotEqual(mockAddress, mockAddress, 1);