Poll csr completion at CmdQueue

Related-To: NEO-6090

Signed-off-by: Dominik Dabek <dominik.dabek@intel.com>

level_zero/core/source/cmdqueue/cmdqueue_hw.inl

@@ -474,6 +474,8 @@ ze_result_t CommandQueueHw<gfxCoreFamily>::executeCommandLists(
 
     this->heapContainer.clear();
 
+    csr->pollForCompletion();
+
     return ZE_RESULT_SUCCESS;
 }
 

level_zero/core/test/unit_tests/fixtures/device_fixture.h

@@ -15,12 +15,15 @@
 #include "shared/test/common/mocks/mock_memory_manager.h"
 #include "shared/test/unit_test/page_fault_manager/cpu_page_fault_manager_tests_fixture.h"
 
+#include "opencl/test/unit_test/libult/ult_aub_command_stream_receiver.h"
+
 #include "level_zero/core/source/context/context_imp.h"
 #include "level_zero/core/test/unit_tests/mocks/mock_driver_handle.h"
 
 namespace NEO {
 struct UltDeviceFactory;
-}
+extern CommandStreamReceiverCreateFunc commandStreamReceiverFactory[2 * IGFX_MAX_CORE];
+} // namespace NEO
 
 namespace L0 {
 struct Context;
@@ -74,6 +77,24 @@ struct ContextFixture : DeviceFixture {
     void TearDown() override;
 };
 
+struct AubCsrFixture : ContextFixture {
+    template <typename T>
+    void SetUpT() {
+        auto csrCreateFcn = &commandStreamReceiverFactory[IGFX_MAX_CORE + NEO::defaultHwInfo->platform.eRenderCoreFamily];
+        variableBackup = std::make_unique<VariableBackup<CommandStreamReceiverCreateFunc>>(csrCreateFcn);
+        *csrCreateFcn = UltAubCommandStreamReceiver<T>::create;
+        ContextFixture::SetUp();
+    }
+    template <typename T>
+    void TearDownT() {
+        ContextFixture::TearDown();
+    }
+
+    void SetUp() override{};
+    void TearDown() override{};
+    std::unique_ptr<VariableBackup<CommandStreamReceiverCreateFunc>> variableBackup;
+};
+
 struct MultipleDevicesWithCustomHwInfo {
     void SetUp();
     void TearDown() {}

level_zero/core/test/unit_tests/sources/cmdqueue/test_cmdqueue.cpp

@@ -20,6 +20,7 @@
 #include "shared/test/common/mocks/mock_memory_manager.h"
 #include "shared/test/common/mocks/ult_device_factory.h"
 
+#include "opencl/test/unit_test/libult/ult_aub_command_stream_receiver.h"
 #include "opencl/test/unit_test/libult/ult_command_stream_receiver.h"
 #include "test.h"
 
@@ -1518,6 +1519,31 @@ HWTEST2_F(ExecuteCommandListTests, givenTwoCommandQueuesHavingTwoB2BCommandLists
     commandQueue1->destroy();
 }
 
+using AubCsrTest = Test<AubCsrFixture>;
+
+HWTEST_TEMPLATED_F(AubCsrTest, givenAubCsrWhenCallingExecuteCommandListsThenPollForCompletionIsCalled) {
+    auto csr = neoDevice->getDefaultEngine().commandStreamReceiver;
+    ze_result_t returnValue;
+    ze_command_queue_desc_t desc = {};
+    ze_command_queue_handle_t commandQueue = {};
+    ze_result_t res = context->createCommandQueue(device, &desc, &commandQueue);
+    ASSERT_EQ(ZE_RESULT_SUCCESS, res);
+    ASSERT_NE(nullptr, commandQueue);
+
+    auto aub_csr = static_cast<NEO::UltAubCommandStreamReceiver<FamilyType> *>(csr);
+    CommandQueue *queue = static_cast<CommandQueue *>(L0::CommandQueue::fromHandle(commandQueue));
+    queue->setCommandQueuePreemptionMode(PreemptionMode::Disabled);
+    EXPECT_EQ(aub_csr->pollForCompletionCalled, 0u);
+
+    std::unique_ptr<L0::CommandList> commandList(L0::CommandList::create(productFamily, device, NEO::EngineGroupType::Compute, 0u, returnValue));
+    ASSERT_NE(nullptr, commandList);
+
+    auto commandListHandle = commandList->toHandle();
+    queue->executeCommandLists(1, &commandListHandle, nullptr, false);
+    EXPECT_EQ(aub_csr->pollForCompletionCalled, 1u);
+
+    L0::CommandQueue::fromHandle(commandQueue)->destroy();
+}
 using CommandQueueSynchronizeTest = Test<ContextFixture>;
 
 template <typename GfxFamily>

opencl/test/unit_test/command_stream/command_stream_receiver_with_aub_dump_tests.cpp

@@ -236,6 +236,31 @@ HWTEST_F(CommandStreamReceiverWithAubDumpSimpleTest, givenCsrWithAubDumpWhenWait
     csrWithAubDump.waitForTaskCountWithKmdNotifyFallback(1, 0, false, false);
 }
 
+HWTEST_F(CommandStreamReceiverWithAubDumpSimpleTest, givenCsrWithAubDumpWhenPollForCompletionCalledThenAubCsrPollForCompletionCalled) {
+    auto executionEnvironment = pDevice->getExecutionEnvironment();
+    executionEnvironment->initializeMemoryManager();
+
+    auto gmmHelper = executionEnvironment->rootDeviceEnvironments[0]->getGmmHelper();
+    MockAubCenter *mockAubCenter = new MockAubCenter(defaultHwInfo.get(), *gmmHelper, false, "file_name.aub", CommandStreamReceiverType::CSR_HW_WITH_AUB);
+    mockAubCenter->aubManager = std::unique_ptr<MockAubManager>(new MockAubManager());
+
+    executionEnvironment->rootDeviceEnvironments[0]->aubCenter = std::unique_ptr<MockAubCenter>(mockAubCenter);
+    DeviceBitfield deviceBitfield(1);
+    CommandStreamReceiverWithAUBDump<UltCommandStreamReceiver<FamilyType>> csrWithAubDump("file_name.aub", *executionEnvironment, 0, deviceBitfield);
+    csrWithAubDump.initializeTagAllocation();
+
+    csrWithAubDump.aubCSR.reset(nullptr);
+    csrWithAubDump.pollForCompletion();
+
+    auto mockAubCsr = new MockAubCsr<FamilyType>("file_name.aub", false, *executionEnvironment, 0, deviceBitfield);
+    mockAubCsr->initializeTagAllocation();
+    csrWithAubDump.aubCSR.reset(mockAubCsr);
+
+    EXPECT_FALSE(mockAubCsr->pollForCompletionCalled);
+    csrWithAubDump.pollForCompletion();
+    EXPECT_TRUE(mockAubCsr->pollForCompletionCalled);
+}
+
 HWTEST_F(CommandStreamReceiverWithAubDumpSimpleTest, givenCsrWithAubDumpWhenCreatingAubCsrThenInitializeTagAllocation) {
     auto executionEnvironment = pDevice->getExecutionEnvironment();
     executionEnvironment->initializeMemoryManager();

opencl/test/unit_test/libult/ult_aub_command_stream_receiver.h

@@ -54,6 +54,12 @@ class UltAubCommandStreamReceiver : public AUBCommandStreamReceiverHw<GfxFamily>
         return BaseClass::blitBuffer(blitPropertiesContainer, blocking, profilingEnabled, device);
     }
 
+    void pollForCompletion() override {
+        pollForCompletionCalled++;
+        BaseClass::pollForCompletion();
+    }
+
     uint32_t blitBufferCalled = 0;
+    uint32_t pollForCompletionCalled = 0;
 };
 } // namespace NEO

shared/source/command_stream/command_stream_receiver.h

@@ -85,6 +85,7 @@ class CommandStreamReceiver {
 
     virtual bool flushBatchedSubmissions() = 0;
     MOCKABLE_VIRTUAL bool submitBatchBuffer(BatchBuffer &batchBuffer, ResidencyContainer &allocationsForResidency);
+    virtual void pollForCompletion() {}
     virtual void programHardwareContext(LinearStream &cmdStream) = 0;
     virtual size_t getCmdsSizeForHardwareContext() const = 0;
 

shared/source/command_stream/command_stream_receiver_simulated_common_hw.h

@@ -52,7 +52,6 @@ class CommandStreamReceiverSimulatedCommonHw : public CommandStreamReceiverHw<Gf
     virtual bool expectMemoryEqual(void *gfxAddress, const void *srcAddress, size_t length);
     virtual bool expectMemoryNotEqual(void *gfxAddress, const void *srcAddress, size_t length);
     virtual bool expectMemoryCompressed(void *gfxAddress, const void *srcAddress, size_t length);
-    virtual void pollForCompletion() = 0;
     virtual void pollForCompletionImpl(){};
     virtual bool writeMemory(GraphicsAllocation &gfxAllocation) = 0;
     virtual void writeMemory(uint64_t gpuAddress, void *cpuAddress, size_t size, uint32_t memoryBank, uint64_t entryBits) = 0;

shared/source/command_stream/command_stream_receiver_with_aub_dump.h

@@ -46,6 +46,8 @@ class CommandStreamReceiverWithAUBDump : public BaseCSR {
 
     void addAubComment(const char *comment) override;
 
+    void pollForCompletion() override;
+
     std::unique_ptr<CommandStreamReceiver> aubCSR;
 };
 

shared/source/command_stream/command_stream_receiver_with_aub_dump.inl

@@ -85,4 +85,12 @@ void CommandStreamReceiverWithAUBDump<BaseCSR>::addAubComment(const char *commen
     }
     BaseCSR::addAubComment(comment);
 }
+
+template <typename BaseCSR>
+void CommandStreamReceiverWithAUBDump<BaseCSR>::pollForCompletion() {
+    if (aubCSR) {
+        aubCSR->pollForCompletion();
+    }
+    BaseCSR::pollForCompletion();
+}
 } // namespace NEO