From 6ef28226438ded0684482da8478fd79d5ff33fdf Mon Sep 17 00:00:00 2001
From: "Milczarek, Slawomir" <slawomir.milczarek@intel.com>
Date: Thu, 31 Jan 2019 16:18:18 +0100
Subject: [PATCH] TBX CSR to call writeMemory on aubManager

Change-Id: Ie10a30944e0c3a8e136816fa91ff501887a3d0d2
---
 .../tbx_command_stream_receiver_hw.inl           |  4 ++--
 .../command_stream/tbx_command_stream_tests.cpp  | 16 +++++++++++++---
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/runtime/command_stream/tbx_command_stream_receiver_hw.inl b/runtime/command_stream/tbx_command_stream_receiver_hw.inl
index 9519ce4ded..d84bacd058 100644
--- a/runtime/command_stream/tbx_command_stream_receiver_hw.inl
+++ b/runtime/command_stream/tbx_command_stream_receiver_hw.inl
@@ -361,9 +361,9 @@ void TbxCommandStreamReceiverHw<GfxFamily>::pollForCompletion() {
 
 template <typename GfxFamily>
 void TbxCommandStreamReceiverHw<GfxFamily>::writeMemory(uint64_t gpuAddress, void *cpuAddress, size_t size, uint32_t memoryBank, uint64_t entryBits, DevicesBitfield devicesBitfield) {
-    if (hardwareContext) {
+    if (aubManager) {
         int hint = AubMemDump::DataTypeHintValues::TraceNotype;
-        hardwareContext->writeMemory(gpuAddress, cpuAddress, size, memoryBank, hint, MemoryConstants::pageSize64k);
+        aubManager->writeMemory(gpuAddress, cpuAddress, size, memoryBank, hint, MemoryConstants::pageSize64k);
         return;
     }
 
diff --git a/unit_tests/command_stream/tbx_command_stream_tests.cpp b/unit_tests/command_stream/tbx_command_stream_tests.cpp
index d4eda6680f..5ca97763c7 100644
--- a/unit_tests/command_stream/tbx_command_stream_tests.cpp
+++ b/unit_tests/command_stream/tbx_command_stream_tests.cpp
@@ -349,6 +349,7 @@ HWTEST_F(TbxCommandStreamTests, givenTbxCommandStreamReceiverWhenFlushIsCalledTh
     auto tbxExecutionEnvironment = getEnvironment<MockTbxCsr<FamilyType>>(true, true);
     auto tbxCsr = tbxExecutionEnvironment->template getCsr<MockTbxCsr<FamilyType>>();
     tbxCsr->hardwareContext = std::unique_ptr<MockHardwareContext>(mockHardwareContext);
+    tbxCsr->aubManager = mockManager.get();
 
     LinearStream cs(tbxExecutionEnvironment->commandBuffer);
     BatchBuffer batchBuffer{cs.getGraphicsAllocation(), 1, 0, nullptr, false, false, QueueThrottle::MEDIUM, cs.getUsed(), &cs};
@@ -358,9 +359,12 @@ HWTEST_F(TbxCommandStreamTests, givenTbxCommandStreamReceiverWhenFlushIsCalledTh
     tbxCsr->flush(batchBuffer, allocationsForResidency);
 
     EXPECT_TRUE(mockHardwareContext->initializeCalled);
-    EXPECT_TRUE(mockHardwareContext->writeMemoryCalled);
     EXPECT_TRUE(mockHardwareContext->submitCalled);
     EXPECT_TRUE(mockHardwareContext->pollForCompletionCalled);
+
+    //call writeMemory on aubManager to clone page tables
+    EXPECT_FALSE(mockHardwareContext->writeMemoryCalled);
+    EXPECT_TRUE(mockManager->writeMemoryCalled);
 }
 
 HWTEST_F(TbxCommandStreamTests, givenTbxCommandStreamReceiverInBatchedModeWhenFlushIsCalledThenItShouldMakeCommandBufferResident) {
@@ -373,6 +377,7 @@ HWTEST_F(TbxCommandStreamTests, givenTbxCommandStreamReceiverInBatchedModeWhenFl
     auto tbxExecutionEnvironment = getEnvironment<MockTbxCsr<FamilyType>>(true, true);
     auto tbxCsr = tbxExecutionEnvironment->template getCsr<MockTbxCsr<FamilyType>>();
     tbxCsr->hardwareContext = std::unique_ptr<MockHardwareContext>(mockHardwareContext);
+    tbxCsr->aubManager = mockManager.get();
 
     LinearStream cs(tbxExecutionEnvironment->commandBuffer);
     BatchBuffer batchBuffer{cs.getGraphicsAllocation(), 1, 0, nullptr, false, false, QueueThrottle::MEDIUM, cs.getUsed(), &cs};
@@ -380,7 +385,9 @@ HWTEST_F(TbxCommandStreamTests, givenTbxCommandStreamReceiverInBatchedModeWhenFl
 
     tbxCsr->flush(batchBuffer, allocationsForResidency);
 
-    EXPECT_TRUE(mockHardwareContext->writeMemoryCalled);
+    //call writeMemory on aubManager to clone page tables
+    EXPECT_FALSE(mockHardwareContext->writeMemoryCalled);
+    EXPECT_TRUE(mockManager->writeMemoryCalled);
     EXPECT_EQ(1u, batchBuffer.commandBufferAllocation->getResidencyTaskCount(tbxCsr->getOsContext().getContextId()));
 }
 
@@ -408,12 +415,15 @@ HWTEST_F(TbxCommandStreamTests, givenTbxCommandStreamReceiverWhenMakeResidentIsC
     auto tbxExecutionEnvironment = getEnvironment<MockTbxCsr<FamilyType>>(true, true);
     auto tbxCsr = tbxExecutionEnvironment->template getCsr<MockTbxCsr<FamilyType>>();
     tbxCsr->hardwareContext = std::unique_ptr<MockHardwareContext>(mockHardwareContext);
+    tbxCsr->aubManager = mockManager.get();
 
     MockGraphicsAllocation allocation(reinterpret_cast<void *>(0x1000), 0x1000);
     ResidencyContainer allocationsForResidency = {&allocation};
     tbxCsr->processResidency(allocationsForResidency);
 
-    EXPECT_TRUE(mockHardwareContext->writeMemoryCalled);
+    //call writeMemory on aubManager to clone page tables
+    EXPECT_FALSE(mockHardwareContext->writeMemoryCalled);
+    EXPECT_TRUE(mockManager->writeMemoryCalled);
 }
 
 HWTEST_F(TbxCommandStreamTests, givenTbxCommandStreamReceiverWhenMakeCoherentIsCalledThenItShouldCallTheExpectedHwContextFunctions) {