fix: track shifted contextIds in bitset in bindlessHeapsHelper

- bitset is 64 bit in size, context ids may go beyond that limit
when multiple devices are available
- this change subtracts contextId of first context for a given root
device - tracked state dirty contexts ids are now zero-based

Resolves: GSD-10025

Signed-off-by: Mateusz Hoppe <[email protected]>
Source: 5ae2552

Author: HoppeMateusz

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

@@ -885,5 +885,69 @@ HWTEST_F(CommandQueueExecuteCommandListsSimpleTest, GivenRegisterInstructionCach
     commandQueue->destroy();
 }
 
+using CommandQueueExecuteCommandListsMultiDeviceTest = Test<MultiDeviceFixture>;
+
+HWTEST_F(CommandQueueExecuteCommandListsMultiDeviceTest, GivenDirtyFlagForContextInBindlessHelperWhenExecutingCmdListsThenStateCacheInvalidateIsSent) {
+    using PIPE_CONTROL = typename FamilyType::PIPE_CONTROL;
+    ze_command_queue_desc_t queueDesc = {};
+    ze_result_t returnValue;
+    auto neoDevice = driverHandle->devices[numRootDevices - 1]->getNEODevice();
+    auto device = driverHandle->devices[numRootDevices - 1];
+
+    auto bindlessHeapsHelper = std::make_unique<MockBindlesHeapsHelper>(neoDevice, neoDevice->getNumGenericSubDevices() > 1);
+    MockBindlesHeapsHelper *bindlessHeapsHelperPtr = bindlessHeapsHelper.get();
+    neoDevice->getExecutionEnvironment()->rootDeviceEnvironments[neoDevice->getRootDeviceIndex()]->bindlessHeapsHelper.reset(bindlessHeapsHelper.release());
+
+    queueDesc.mode = ZE_COMMAND_QUEUE_MODE_SYNCHRONOUS;
+    auto commandQueue = whiteboxCast(CommandQueue::create(productFamily, device, neoDevice->getDefaultEngine().commandStreamReceiver, &queueDesc, false, false, false, returnValue));
+    ASSERT_NE(nullptr, commandQueue);
+
+    uint32_t contextId = commandQueue->getCsr()->getOsContext().getContextId();
+    uint32_t contextIdFirst = neoDevice->getMemoryManager()->getFirstContextIdForRootDevice(numRootDevices - 1);
+    EXPECT_LE(contextIdFirst, contextId);
+
+    uint32_t firstDeviceFirstContextId = neoDevice->getMemoryManager()->getFirstContextIdForRootDevice(0);
+    EXPECT_EQ(0u, firstDeviceFirstContextId);
+
+    bindlessHeapsHelperPtr->stateCacheDirtyForContext.set();
+
+    auto usedSpaceBefore = commandQueue->commandStream.getUsed();
+
+    ze_command_list_handle_t commandLists[] = {
+        CommandList::create(productFamily, device, NEO::EngineGroupType::renderCompute, 0u, returnValue, false)->toHandle(),
+        CommandList::create(productFamily, device, NEO::EngineGroupType::renderCompute, 0u, returnValue, false)->toHandle()};
+    uint32_t numCommandLists = sizeof(commandLists) / sizeof(commandLists[0]);
+    CommandList::fromHandle(commandLists[0])->close();
+    CommandList::fromHandle(commandLists[1])->close();
+    auto result = commandQueue->executeCommandLists(numCommandLists, commandLists, nullptr, true, nullptr);
+
+    ASSERT_EQ(ZE_RESULT_SUCCESS, result);
+
+    auto usedSpaceAfter = commandQueue->commandStream.getUsed();
+    ASSERT_GT(usedSpaceAfter, usedSpaceBefore);
+
+    GenCmdList cmdList;
+    ASSERT_TRUE(FamilyType::Parse::parseCommandBuffer(
+        cmdList, ptrOffset(commandQueue->commandStream.getCpuBase(), 0), usedSpaceAfter));
+
+    auto pipeControls = findAll<PIPE_CONTROL *>(cmdList.begin(), cmdList.end());
+    ASSERT_NE(0u, pipeControls.size());
+
+    auto pipeControl = reinterpret_cast<PIPE_CONTROL *>(*pipeControls[0]);
+    EXPECT_TRUE(pipeControl->getCommandStreamerStallEnable());
+    EXPECT_TRUE(pipeControl->getStateCacheInvalidationEnable());
+    EXPECT_TRUE(pipeControl->getTextureCacheInvalidationEnable());
+    EXPECT_TRUE(pipeControl->getRenderTargetCacheFlushEnable());
+
+    EXPECT_FALSE(bindlessHeapsHelperPtr->getStateDirtyForContext(commandQueue->getCsr()->getOsContext().getContextId()));
+
+    for (auto i = 0u; i < numCommandLists; i++) {
+        auto commandList = CommandList::fromHandle(commandLists[i]);
+        commandList->destroy();
+    }
+
+    commandQueue->destroy();
+}
+
 } // namespace ult
 } // namespace L0

shared/source/helpers/bindless_heaps_helper.cpp

@@ -75,13 +75,19 @@ GraphicsAllocation *BindlessHeapsHelper::getHeapAllocation(size_t heapSize, size
 void BindlessHeapsHelper::clearStateDirtyForContext(uint32_t osContextId) {
     std::lock_guard<std::mutex> autolock(this->mtx);
 
-    stateCacheDirtyForContext.reset(osContextId);
+    uint32_t contextIdShifted = osContextId - memManager->getFirstContextIdForRootDevice(rootDeviceIndex);
+    DEBUG_BREAK_IF(contextIdShifted >= stateCacheDirtyForContext.size());
+
+    stateCacheDirtyForContext.reset(contextIdShifted);
 }
 
 bool BindlessHeapsHelper::getStateDirtyForContext(uint32_t osContextId) {
     std::lock_guard<std::mutex> autolock(this->mtx);
 
-    return stateCacheDirtyForContext.test(osContextId);
+    uint32_t contextIdShifted = osContextId - memManager->getFirstContextIdForRootDevice(rootDeviceIndex);
+    DEBUG_BREAK_IF(contextIdShifted >= stateCacheDirtyForContext.size());
+
+    return stateCacheDirtyForContext.test(contextIdShifted);
 }
 
 SurfaceStateInHeapInfo BindlessHeapsHelper::allocateSSInHeap(size_t ssSize, GraphicsAllocation *surfaceAllocation, BindlesHeapType heapType) {

shared/source/memory_manager/memory_manager.cpp

@@ -368,6 +368,14 @@ void MemoryManager::updateLatestContextIdForRootDevice(uint32_t rootDeviceIndex)
     }
 }
 
+uint32_t MemoryManager::getFirstContextIdForRootDevice(uint32_t rootDeviceIndex) {
+    auto entry = rootDeviceIndexToContextId.find(rootDeviceIndex);
+    if (entry != rootDeviceIndexToContextId.end()) {
+        return entry->second + 1;
+    }
+    return 0;
+}
+
 OsContext *MemoryManager::createAndRegisterOsContext(CommandStreamReceiver *commandStreamReceiver,
                                                      const EngineDescriptor &engineDescriptor) {
     auto rootDeviceIndex = commandStreamReceiver->getRootDeviceIndex();

shared/source/memory_manager/memory_manager.h

@@ -314,6 +314,7 @@ class MemoryManager {
 
     size_t getUsedLocalMemorySize(uint32_t rootDeviceIndex) const { return localMemAllocsSize[rootDeviceIndex]; }
     size_t getUsedSystemMemorySize() const { return sysMemAllocsSize; }
+    uint32_t getFirstContextIdForRootDevice(uint32_t rootDeviceIndex);
 
   protected:
     bool getAllocationData(AllocationData &allocationData, const AllocationProperties &properties, const void *hostPtr, const StorageInfo &storageInfo);