Correct memory banks passed to hardware contexts

Related-To: NEO-2881

Change-Id: I60d4a649c1ab3db634c56792e407bb318831de1b
Signed-off-by: Mateusz Jablonski <[email protected]>

Author: JablonskiMateusz

runtime/helpers/hardware_context_controller.cpp

@@ -48,8 +48,15 @@ void HardwareContextController::submit(uint64_t batchBufferGpuAddress, const voi
 }
 
 void HardwareContextController::writeMemory(uint64_t gfxAddress, const void *memory, size_t size, uint32_t memoryBanks, int hint, size_t pageSize) {
-    for (auto &hardwareContext : hardwareContexts) {
-        hardwareContext->writeMemory(gfxAddress, memory, size, memoryBanks, hint, pageSize);
+    if (hardwareContexts.size() == 1u) {
+        hardwareContexts.at(0)->writeMemory(gfxAddress, memory, size, memoryBanks, hint, pageSize);
+        return;
+    }
+    for (auto bankId = 0u; bankId < hardwareContexts.size(); bankId++) {
+        auto &hardwareContext = hardwareContexts.at(bankId);
+        auto selectedBank = memoryBanks & (1 << bankId);
+        UNRECOVERABLE_IF(selectedBank == 0);
+        hardwareContext->writeMemory(gfxAddress, memory, size, memoryBanks & (1 << bankId), hint, pageSize);
     }
 }
 

unit_tests/command_stream/aub_command_stream_receiver_1_tests.cpp

@@ -1177,6 +1177,20 @@ TEST_F(HardwareContextContainerTests, givenSingleHwContextWhenSubmitMethodIsCall
     EXPECT_FALSE(mockHwContext0->writeMemoryCalled);
 }
 
+TEST_F(HardwareContextContainerTests, givenSingleHwContextWhenWriteMemoryIsCalledThenWholeMemoryBanksArePassed) {
+    MockAubManager aubManager;
+    MockOsContext osContext(1, getDeviceBitfieldForNDevices(1), aub_stream::ENGINE_RCS, PreemptionMode::Disabled, false);
+    HardwareContextController hwContextContainer(aubManager, osContext, 0);
+    EXPECT_EQ(1u, hwContextContainer.hardwareContexts.size());
+
+    auto mockHwContext0 = static_cast<MockHardwareContext *>(hwContextContainer.hardwareContexts[0].get());
+
+    hwContextContainer.writeMemory(1, reinterpret_cast<const void *>(0x123), 2, 3u, 4, 5);
+
+    EXPECT_TRUE(mockHwContext0->writeMemoryCalled);
+    EXPECT_EQ(3u, mockHwContext0->memoryBanksPassed);
+}
+
 TEST_F(HardwareContextContainerTests, givenMultipleHwContextWhenSingleMethodIsCalledThenUseAllContexts) {
     MockAubManager aubManager;
     MockOsContext osContext(1, getDeviceBitfieldForNDevices(2), aub_stream::ENGINE_RCS, PreemptionMode::Disabled, false);
@@ -1201,7 +1215,7 @@ TEST_F(HardwareContextContainerTests, givenMultipleHwContextWhenSingleMethodIsCa
     hwContextContainer.pollForCompletion();
     hwContextContainer.expectMemory(1, reinterpret_cast<const void *>(0x123), 2, 0);
     hwContextContainer.submit(1, reinterpret_cast<const void *>(0x123), 2, 0, 1);
-    hwContextContainer.writeMemory(1, reinterpret_cast<const void *>(0x123), 2, 3, 4, 5);
+    hwContextContainer.writeMemory(1, reinterpret_cast<const void *>(0x123), 2, 3u, 4, 5);
 
     EXPECT_TRUE(mockHwContext0->initializeCalled);
     EXPECT_TRUE(mockHwContext1->initializeCalled);
@@ -1213,6 +1227,8 @@ TEST_F(HardwareContextContainerTests, givenMultipleHwContextWhenSingleMethodIsCa
     EXPECT_TRUE(mockHwContext1->submitCalled);
     EXPECT_TRUE(mockHwContext0->writeMemoryCalled);
     EXPECT_TRUE(mockHwContext1->writeMemoryCalled);
+    EXPECT_EQ(1u, mockHwContext0->memoryBanksPassed);
+    EXPECT_EQ(2u, mockHwContext1->memoryBanksPassed);
 }
 
 TEST_F(HardwareContextContainerTests, givenMultipleHwContextWhenSingleMethodIsCalledThenUseFirstContext) {

unit_tests/mocks/mock_aub_manager.h

@@ -24,6 +24,7 @@ struct MockHardwareContext : public aub_stream::HardwareContext {
     void writeMemory(uint64_t gfxAddress, const void *memory, size_t size, uint32_t memoryBanks, int hint, size_t pageSize = 65536) override {
         writeMemoryCalled = true;
         writeMemoryPageSizePassed = pageSize;
+        memoryBanksPassed = memoryBanks;
     }
     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; }
@@ -43,6 +44,7 @@ struct MockHardwareContext : public aub_stream::HardwareContext {
     bool dumpSurfaceCalled = false;
 
     size_t writeMemoryPageSizePassed = 0;
+    uint32_t memoryBanksPassed = 0;
 
     const uint32_t deviceIndex;
 };