From 61000c0dd48eb825c909bc8b6101c50de56839c5 Mon Sep 17 00:00:00 2001
From: "Milczarek, Slawomir" <slawomir.milczarek@intel.com>
Date: Fri, 24 Aug 2018 00:16:32 +0200
Subject: [PATCH] CSR AUB + HW mode - flush to poll MMIO for completion

Ensures that each submit of LRCA be serialized through the simulator
like it is for AUBs captured in the standalone mode.

Change-Id: I1e3ad500012dce960d0e64b56af1cb60142772da
---
 .../command_stream/aub_command_stream_receiver_hw.h |  2 +-
 .../aub_command_stream_receiver_hw.inl              |  2 +-
 .../aub_command_stream_receiver_tests.cpp           |  1 +
 unit_tests/command_stream/aub_file_stream_tests.cpp | 13 +++++++++++++
 unit_tests/mocks/mock_aub_csr.h                     |  4 ++++
 5 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/runtime/command_stream/aub_command_stream_receiver_hw.h b/runtime/command_stream/aub_command_stream_receiver_hw.h
index 0b0a65e57e..05ead52133 100644
--- a/runtime/command_stream/aub_command_stream_receiver_hw.h
+++ b/runtime/command_stream/aub_command_stream_receiver_hw.h
@@ -44,7 +44,7 @@ class AUBCommandStreamReceiverHw : public CommandStreamReceiverSimulatedHw<GfxFa
 
     // Family specific version
     void submitLRCA(EngineType engineType, const MiContextDescriptorReg &contextDescriptor);
-    void pollForCompletion(EngineType engineType);
+    MOCKABLE_VIRTUAL void pollForCompletion(EngineType engineType);
     void initGlobalMMIO();
     void initEngineMMIO(EngineType engineType);
 
diff --git a/runtime/command_stream/aub_command_stream_receiver_hw.inl b/runtime/command_stream/aub_command_stream_receiver_hw.inl
index 78dc147d3b..5b67d32840 100644
--- a/runtime/command_stream/aub_command_stream_receiver_hw.inl
+++ b/runtime/command_stream/aub_command_stream_receiver_hw.inl
@@ -477,8 +477,8 @@ FlushStamp AUBCommandStreamReceiverHw<GfxFamily>::flush(BatchBuffer &batchBuffer
         submitLRCA(engineType, contextDescriptor);
     }
 
+    pollForCompletion(engineType);
     if (this->standalone) {
-        pollForCompletion(engineType);
         *this->tagAddress = this->peekLatestSentTaskCount();
     }
 
diff --git a/unit_tests/command_stream/aub_command_stream_receiver_tests.cpp b/unit_tests/command_stream/aub_command_stream_receiver_tests.cpp
index f637a9e15d..dd0c2eb30a 100644
--- a/unit_tests/command_stream/aub_command_stream_receiver_tests.cpp
+++ b/unit_tests/command_stream/aub_command_stream_receiver_tests.cpp
@@ -1303,6 +1303,7 @@ HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverWhenAddPatc
     ResidencyContainer allocationsForResidency;
 
     EXPECT_CALL(*aubCsr, addPatchInfoComments()).Times(0);
+
     aubCsr->flush(batchBuffer, engineType, allocationsForResidency, *pDevice->getOsContext());
 }
 
diff --git a/unit_tests/command_stream/aub_file_stream_tests.cpp b/unit_tests/command_stream/aub_file_stream_tests.cpp
index dd38293eb9..726de507ec 100644
--- a/unit_tests/command_stream/aub_file_stream_tests.cpp
+++ b/unit_tests/command_stream/aub_file_stream_tests.cpp
@@ -117,6 +117,19 @@ HWTEST_F(AubFileStreamTests, givenAubCommandStreamReceiverWhenFlushIsCalledThenF
     EXPECT_TRUE(mockAubFileStreamPtr->lockStreamCalled);
 }
 
+HWTEST_F(AubFileStreamTests, givenAubCommandStreamReceiverWhenFlushIsCalledThenItShouldCallPollForCompletion) {
+    auto aubExecutionEnvironment = getEnvironment<MockAubCsr<FamilyType>>(false, true, false);
+    auto aubCsr = aubExecutionEnvironment->template getCsr<MockAubCsr<FamilyType>>();
+    LinearStream cs(aubExecutionEnvironment->commandBuffer);
+
+    BatchBuffer batchBuffer{cs.getGraphicsAllocation(), 0, 0, nullptr, false, false, QueueThrottle::MEDIUM, cs.getUsed(), &cs};
+    auto engineType = OCLRT::ENGINE_RCS;
+    ResidencyContainer allocationsForResidency = {};
+
+    aubCsr->flush(batchBuffer, engineType, allocationsForResidency, *pDevice->getOsContext());
+    EXPECT_TRUE(aubCsr->pollForCompletionCalled);
+}
+
 HWTEST_F(AubFileStreamTests, givenAubCommandStreamReceiverWhenFlushIsCalledThenFileStreamShouldBeFlushed) {
     auto aubExecutionEnvironment = getEnvironment<AUBCommandStreamReceiverHw<FamilyType>>(true, true, true);
     auto aubCsr = aubExecutionEnvironment->template getCsr<AUBCommandStreamReceiverHw<FamilyType>>();
diff --git a/unit_tests/mocks/mock_aub_csr.h b/unit_tests/mocks/mock_aub_csr.h
index 21ea8f37c6..d328a510bd 100644
--- a/unit_tests/mocks/mock_aub_csr.h
+++ b/unit_tests/mocks/mock_aub_csr.h
@@ -43,8 +43,12 @@ struct MockAubCsr : public AUBCommandStreamReceiverHw<GfxFamily> {
     void initProgrammingFlags() override {
         initProgrammingFlagsCalled = true;
     }
+    void pollForCompletion(EngineType engineType) override {
+        pollForCompletionCalled = true;
+    }
     bool flushBatchedSubmissionsCalled = false;
     bool initProgrammingFlagsCalled = false;
+    bool pollForCompletionCalled = false;
 
     void initFile(const std::string &fileName) override {
         fileIsOpen = true;