Pass additional allocation parameters to aub_stream

BartoszDunajski · Compute-Runtime-Automation · commit 15f3353351d0 · 2021-02-17T14:20:08.000+01:00
Signed-off-by: Bartosz Dunajski &lt;bartosz.dunajski@intel.com&gt;
diff --git a/opencl/source/command_stream/definitions/command_stream_receiver_simulated_hw.h b/opencl/source/command_stream/definitions/command_stream_receiver_simulated_hw.h
@@ -60,6 +60,10 @@ class CommandStreamReceiverSimulatedHw : public CommandStreamReceiverSimulatedCo
         aub_stream::AllocationParams allocationParams(gpuAddress, cpuAddress, size, this->getMemoryBank(&graphicsAllocation),
                                                       hint, graphicsAllocation.getUsedPageSize());
 
+        auto gmm = graphicsAllocation.getDefaultGmm();
+
+        allocationParams.additionalParams.compressionEnabled = gmm ? gmm->isRenderCompressed : false;
+
         this->aubManager->writeMemory2(allocationParams);
     }
 
diff --git a/opencl/test/unit_test/command_stream/aub_command_stream_receiver_1_tests.cpp b/opencl/test/unit_test/command_stream/aub_command_stream_receiver_1_tests.cpp
@@ -1272,7 +1272,7 @@ HWTEST_F(AubCsrTest, WhenWriteWithAubManagerIsCalledThenAubManagerIsInvokedWithC
     EXPECT_TRUE(aubManager.writeMemory2Called);
     EXPECT_EQ(AubMemDump::DataTypeHintValues::TraceBatchBuffer, aubManager.hintToWriteMemory);
 
-    aubManager.writeMemoryCalled = false;
+    aubManager.writeMemory2Called = false;
 
     auto allocation2 = executionEnvironment->memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{rootDeviceIndex, true, MemoryConstants::pageSize, GraphicsAllocation::AllocationType::LINEAR_STREAM});
 
@@ -1353,9 +1353,7 @@ TEST_F(HardwareContextContainerTests, givenMultipleHwContextWhenSingleMethodIsCa
     EXPECT_FALSE(mockHwContext1->expectMemoryCalled);
     EXPECT_FALSE(mockHwContext0->submitCalled);
     EXPECT_FALSE(mockHwContext1->submitCalled);
-    EXPECT_FALSE(mockHwContext0->writeMemoryCalled);
     EXPECT_FALSE(mockHwContext0->writeMemory2Called);
-    EXPECT_FALSE(mockHwContext1->writeMemoryCalled);
     EXPECT_FALSE(mockHwContext1->writeMemory2Called);
 
     aub_stream::AllocationParams params(1, reinterpret_cast<const void *>(0x123), 2, 3u, 4, 5);
@@ -1374,9 +1372,7 @@ TEST_F(HardwareContextContainerTests, givenMultipleHwContextWhenSingleMethodIsCa
     EXPECT_TRUE(mockHwContext1->expectMemoryCalled);
     EXPECT_TRUE(mockHwContext0->submitCalled);
     EXPECT_TRUE(mockHwContext1->submitCalled);
-    EXPECT_FALSE(mockHwContext0->writeMemoryCalled);
     EXPECT_TRUE(mockHwContext0->writeMemory2Called);
-    EXPECT_FALSE(mockHwContext1->writeMemoryCalled);
     EXPECT_TRUE(mockHwContext1->writeMemory2Called);
     EXPECT_EQ(1u, mockHwContext0->memoryBanksPassed);
     EXPECT_EQ(2u, mockHwContext1->memoryBanksPassed);
diff --git a/opencl/test/unit_test/mocks/mock_aub_manager.h b/opencl/test/unit_test/mocks/mock_aub_manager.h
@@ -7,6 +7,8 @@
 
 #pragma once
 
+#include "shared/source/helpers/debug_helpers.h"
+
 #include "third_party/aub_stream/headers/allocation_params.h"
 #include "third_party/aub_stream/headers/aub_manager.h"
 #include "third_party/aub_stream/headers/aubstream.h"
@@ -23,26 +25,29 @@ struct MockHardwareContext : public aub_stream::HardwareContext {
     void writeAndSubmitBatchBuffer(uint64_t gfxAddress, const void *batchBuffer, size_t size, uint32_t memoryBank, size_t pageSize) override { writeAndSubmitCalled = true; }
     void submitBatchBuffer(uint64_t gfxAddress, bool overrideRingHead) override { submitCalled = true; }
     void writeMemory(uint64_t gfxAddress, const void *memory, size_t size, uint32_t memoryBanks, int hint, size_t pageSize) override {
-        writeMemoryCalled = true;
-        writeMemoryPageSizePassed = pageSize;
-        memoryBanksPassed = memoryBanks;
+        UNRECOVERABLE_IF(true); // shouldnt be used
     }
     void writeMemory2(aub_stream::AllocationParams allocationParams) override {
         writeMemory2Called = true;
         writeMemoryPageSizePassed = allocationParams.pageSize;
         memoryBanksPassed = allocationParams.memoryBanks;
+
+        if (storeAllocationParams) {
+            storedAllocationParams.push_back(allocationParams);
+        }
     }
     void freeMemory(uint64_t gfxAddress, size_t size) override { freeMemoryCalled = true; }
     void expectMemory(uint64_t gfxAddress, const void *memory, size_t size, uint32_t compareOperation) override { expectMemoryCalled = true; }
     void readMemory(uint64_t gfxAddress, void *memory, size_t size, uint32_t memoryBank, size_t pageSize) override { readMemoryCalled = true; }
     void dumpBufferBIN(uint64_t gfxAddress, size_t size) override { dumpBufferBINCalled = true; }
     void dumpSurface(const SurfaceInfo &surfaceInfo) override { dumpSurfaceCalled = true; }
 
+    std::vector<aub_stream::AllocationParams> storedAllocationParams;
+    bool storeAllocationParams = false;
     bool initializeCalled = false;
     bool pollForCompletionCalled = false;
     bool writeAndSubmitCalled = false;
     bool submitCalled = false;
-    bool writeMemoryCalled = false;
     bool writeMemory2Called = false;
     bool freeMemoryCalled = false;
     bool expectMemoryCalled = false;
@@ -105,15 +110,17 @@ class MockAubManager : public aub_stream::AubManager {
     }
 
     void writeMemory(uint64_t gfxAddress, const void *memory, size_t size, uint32_t memoryBanks, int hint, size_t pageSize) override {
-        writeMemoryCalled = true;
-        hintToWriteMemory = hint;
-        writeMemoryPageSizePassed = pageSize;
+        UNRECOVERABLE_IF(true); // shouldnt be used
     }
 
     void writeMemory2(aub_stream::AllocationParams allocationParams) override {
         writeMemory2Called = true;
         hintToWriteMemory = allocationParams.hint;
         writeMemoryPageSizePassed = allocationParams.pageSize;
+
+        if (storeAllocationParams) {
+            storedAllocationParams.push_back(allocationParams);
+        }
     }
 
     void writePageTableEntries(uint64_t gfxAddress, size_t size, uint32_t memoryBanks, int hint,
@@ -129,6 +136,7 @@ class MockAubManager : public aub_stream::AubManager {
         freeMemoryCalled = true;
     }
 
+    std::vector<aub_stream::AllocationParams> storedAllocationParams;
     uint32_t openCalledCnt = 0;
     std::string fileName = "";
     bool closeCalled = false;
@@ -137,11 +145,11 @@ class MockAubManager : public aub_stream::AubManager {
     bool isPaused = false;
     bool addCommentCalled = false;
     std::string receivedComment = "";
-    bool writeMemoryCalled = false;
     bool writeMemory2Called = false;
     bool writePageTableEntriesCalled = false;
     bool writePhysicalMemoryPagesCalled = false;
     bool freeMemoryCalled = false;
+    bool storeAllocationParams = false;
     uint32_t contextFlags = 0;
     int hintToWriteMemory = 0;
     size_t writeMemoryPageSizePassed = 0;
diff --git a/shared/source/os_interface/aub_memory_operations_handler.cpp b/shared/source/os_interface/aub_memory_operations_handler.cpp
@@ -8,6 +8,7 @@
 #include "shared/source/os_interface/aub_memory_operations_handler.h"
 
 #include "shared/source/aub_mem_dump/aub_mem_dump.h"
+#include "shared/source/gmm_helper/gmm.h"
 #include "shared/source/memory_manager/graphics_allocation.h"
 
 #include "third_party/aub_stream/headers/allocation_params.h"
@@ -27,12 +28,18 @@ MemoryOperationsStatus AubMemoryOperationsHandler::makeResident(Device *device,
     auto lock = acquireLock(resourcesLock);
     int hint = AubMemDump::DataTypeHintValues::TraceNotype;
     for (const auto &allocation : gfxAllocations) {
-        aubManager->writeMemory2({allocation->getGpuAddress(),
-                                  allocation->getUnderlyingBuffer(),
-                                  allocation->getUnderlyingBufferSize(),
-                                  allocation->storageInfo.getMemoryBanks(),
-                                  hint,
-                                  allocation->getUsedPageSize()});
+        aub_stream::AllocationParams params(allocation->getGpuAddress(),
+                                            allocation->getUnderlyingBuffer(),
+                                            allocation->getUnderlyingBufferSize(),
+                                            allocation->storageInfo.getMemoryBanks(),
+                                            hint,
+                                            allocation->getUsedPageSize());
+
+        auto gmm = allocation->getDefaultGmm();
+
+        params.additionalParams.compressionEnabled = gmm ? gmm->isRenderCompressed : false;
+
+        aubManager->writeMemory2(params);
         residentAllocations.push_back(allocation);
     }
     return MemoryOperationsStatus::SUCCESS;
diff --git a/shared/test/unit_test/os_interface/aub_memory_operations_handler_tests.cpp b/shared/test/unit_test/os_interface/aub_memory_operations_handler_tests.cpp
@@ -10,6 +10,7 @@
 #include "shared/source/aub_mem_dump/aub_mem_dump.h"
 
 #include "opencl/test/unit_test/mocks/mock_aub_manager.h"
+#include "opencl/test/unit_test/mocks/mock_gmm.h"
 
 TEST_F(AubMemoryOperationsHandlerTests, givenNullPtrAsAubManagerWhenMakeResidentCalledThenFalseReturned) {
     getMemoryOperationsHandler()->setAubManager(nullptr);
@@ -27,6 +28,25 @@ TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerWhenMakeResidentCalledThe
     EXPECT_TRUE(aubManager.writeMemory2Called);
 }
 
+TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerWhenMakeResidentCalledOnCompressedAllocationThenPassCorrectParams) {
+    MockAubManager aubManager;
+    aubManager.storeAllocationParams = true;
+
+    getMemoryOperationsHandler()->setAubManager(&aubManager);
+    auto memoryOperationsInterface = getMemoryOperationsHandler();
+
+    MockGmm gmm;
+    gmm.isRenderCompressed = true;
+    allocPtr->setDefaultGmm(&gmm);
+
+    auto result = memoryOperationsInterface->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
+    EXPECT_EQ(result, MemoryOperationsStatus::SUCCESS);
+
+    EXPECT_TRUE(aubManager.writeMemory2Called);
+    EXPECT_EQ(1u, aubManager.storedAllocationParams.size());
+    EXPECT_TRUE(aubManager.storedAllocationParams[0].additionalParams.compressionEnabled);
+}
+
 TEST_F(AubMemoryOperationsHandlerTests, givenAllocationWhenMakeResidentCalledThenTraceNotypeHintReturned) {
     MockAubManager aubManager;
     getMemoryOperationsHandler()->setAubManager(&aubManager);

Original file line number	Diff line number	Diff line change
`@@ -60,6 +60,10 @@ class CommandStreamReceiverSimulatedHw : public CommandStreamReceiverSimulatedCo`
`60`	`60`	`aub_stream::AllocationParams allocationParams(gpuAddress, cpuAddress, size, this->getMemoryBank(&graphicsAllocation),`
`61`	`61`	`hint, graphicsAllocation.getUsedPageSize());`
`62`	`62`
	`63`	`+ auto gmm = graphicsAllocation.getDefaultGmm();`
	`64`	`+`
	`65`	`+ allocationParams.additionalParams.compressionEnabled = gmm ? gmm->isRenderCompressed : false;`
	`66`	`+`
`63`	`67`	`this->aubManager->writeMemory2(allocationParams);`
`64`	`68`	`}`
`65`	`69`