feature(sysman): reinitialize gfxPartition on reset

Related-To: NEO-13203

Signed-off-by: Mateusz Hoppe <mateusz.hoppe@intel.com>
Source: 0589a70

Author: HoppeMateusz

level_zero/tools/source/sysman/linux/os_sysman_imp.cpp

@@ -341,6 +341,7 @@ void LinuxSysmanImp::releaseDeviceResources() {
     executionEnvironment->memoryManager->releaseDeviceSpecificMemResources(rootDeviceIndex);
     executionEnvironment->releaseRootDeviceEnvironmentResources(executionEnvironment->rootDeviceEnvironments[rootDeviceIndex].get());
     executionEnvironment->rootDeviceEnvironments[rootDeviceIndex].reset();
+    executionEnvironment->memoryManager->releaseDeviceSpecificGfxPartition(rootDeviceIndex);
 }
 
 void LinuxSysmanImp::reInitSysmanDeviceResources() {

level_zero/tools/test/unit_tests/sources/sysman/diagnostics/linux/test_zes_sysman_diagnostics.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020-2023 Intel Corporation
+ * Copyright (C) 2020-2024 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -764,6 +764,19 @@ TEST_F(ZesDiagnosticsFixture, GivenValidSysmanImpPointerWhenCallingReleaseResour
     EXPECT_EQ(ZE_RESULT_SUCCESS, pLinuxSysmanImp->initDevice());
 }
 
+TEST_F(ZesDiagnosticsFixture, GivenSysmanImpPointerWhenCallingReleaseResourcesThenGfxPartitionIsRemovedForRootDevice) {
+    pLinuxSysmanImp->diagnosticsReset = true;
+    auto devicePtr = static_cast<DeviceImp *>(pLinuxSysmanImp->pDevice);
+    auto executionEnvironment = devicePtr->getNEODevice()->getExecutionEnvironment();
+    auto rootIndex = devicePtr->getNEODevice()->getRootDeviceIndex();
+
+    pLinuxSysmanImp->releaseDeviceResources();
+
+    EXPECT_EQ(nullptr, executionEnvironment->memoryManager->getGfxPartition(rootIndex));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, pLinuxSysmanImp->initDevice());
+    EXPECT_NE(nullptr, executionEnvironment->memoryManager->getGfxPartition(rootIndex));
+}
+
 HWTEST2_F(ZesDiagnosticsFixture, GivenValidDiagnosticsHandleAndHandleCountZeroWhenCallingReInitThenValidCountIsReturnedAndVerifyzesDeviceEnumDiagnosticTestSuitesSucceeds, IsPVC) {
     uint32_t count = 0;
     ze_result_t result = zesDeviceEnumDiagnosticTestSuites(device->toHandle(), &count, nullptr);

shared/source/memory_manager/memory_manager.h

@@ -279,6 +279,9 @@ class MemoryManager {
 
     virtual void releaseDeviceSpecificMemResources(uint32_t rootDeviceIndex){};
     virtual void createDeviceSpecificMemResources(uint32_t rootDeviceIndex){};
+    virtual void releaseDeviceSpecificGfxPartition(uint32_t rootDeviceIndex){};
+    virtual bool reInitDeviceSpecificGfxPartition(uint32_t rootDeviceIndex) { return true; };
+
     void reInitLatestContextId() {
         latestContextId = std::numeric_limits<uint32_t>::max();
     }

shared/source/os_interface/device_factory.cpp

@@ -249,6 +249,7 @@ std::unique_ptr<Device> DeviceFactory::createDevice(ExecutionEnvironment &execut
         return device;
     }
 
+    executionEnvironment.memoryManager->reInitDeviceSpecificGfxPartition(rootDeviceIndex);
     executionEnvironment.memoryManager->createDeviceSpecificMemResources(rootDeviceIndex);
     executionEnvironment.memoryManager->reInitLatestContextId();
     device = createRootDeviceFunc(executionEnvironment, rootDeviceIndex);

shared/source/os_interface/linux/drm_memory_manager.cpp

@@ -2948,4 +2948,23 @@ void DrmMemoryManager::getExtraDeviceProperties(uint32_t rootDeviceIndex, uint32
     getDrm(rootDeviceIndex).getIoctlHelper()->queryDeviceParams(moduleId, serverType);
 }
 
+bool DrmMemoryManager::reInitDeviceSpecificGfxPartition(uint32_t rootDeviceIndex) {
+    if (gfxPartitions.at(rootDeviceIndex) == nullptr) {
+        auto gpuAddressSpace = executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->getHardwareInfo()->capabilityTable.gpuAddressSpace;
+
+        gfxPartitions.at(rootDeviceIndex) = std::make_unique<GfxPartition>(reservedCpuAddressRange);
+
+        uint64_t gfxTop{};
+        getDrm(rootDeviceIndex).queryGttSize(gfxTop, false);
+
+        if (getGfxPartition(rootDeviceIndex)->init(gpuAddressSpace, getSizeToReserve(), rootDeviceIndex, gfxPartitions.size(), heapAssigners[rootDeviceIndex]->apiAllowExternalHeapForSshAndDsh, DrmMemoryManager::getSystemSharedMemory(rootDeviceIndex), gfxTop)) {
+            return true;
+        }
+    }
+    return false;
+}
+
+void DrmMemoryManager::releaseDeviceSpecificGfxPartition(uint32_t rootDeviceIndex) {
+    gfxPartitions.at(rootDeviceIndex).reset();
+}
 } // namespace NEO

shared/source/os_interface/linux/drm_memory_manager.h

@@ -100,6 +100,8 @@ class DrmMemoryManager : public MemoryManager {
     DrmAllocation *createUSMHostAllocationFromSharedHandle(osHandle handle, const AllocationProperties &properties, void *mappedPtr, bool reuseSharedAllocation);
     void releaseDeviceSpecificMemResources(uint32_t rootDeviceIndex) override;
     void createDeviceSpecificMemResources(uint32_t rootDeviceIndex) override;
+    void releaseDeviceSpecificGfxPartition(uint32_t rootDeviceIndex) override;
+    bool reInitDeviceSpecificGfxPartition(uint32_t rootDeviceIndex) override;
     bool allowIndirectAllocationsAsPack(uint32_t rootDeviceIndex) override;
     Drm &getDrm(uint32_t rootDeviceIndex) const;
     size_t getSizeOfChunk(size_t allocSize);

shared/test/common/mocks/mock_memory_manager.cpp

@@ -10,6 +10,7 @@
 #include "shared/source/command_stream/command_stream_receiver.h"
 #include "shared/source/gmm_helper/gmm.h"
 #include "shared/source/helpers/aligned_memory.h"
+#include "shared/source/helpers/hw_info.h"
 #include "shared/source/helpers/surface_format_info.h"
 #include "shared/source/memory_manager/deferred_deleter.h"
 #include "shared/source/memory_manager/gfx_partition.h"
@@ -248,6 +249,25 @@ bool MockMemoryManager::copyMemoryToAllocationBanks(GraphicsAllocation *graphics
     return OsAgnosticMemoryManager::copyMemoryToAllocationBanks(graphicsAllocation, destinationOffset, memoryToCopy, sizeToCopy, handleMask);
 };
 
+bool MockMemoryManager::reInitDeviceSpecificGfxPartition(uint32_t rootDeviceIndex) {
+    if (gfxPartitions.at(rootDeviceIndex) == nullptr) {
+        // 4 x sizeof(Heap32) + 2 x sizeof(Standard/Standard64k)
+        size_t reservedCpuAddressRangeSize = static_cast<size_t>((4 * 4 + 2 * 4)) * static_cast<size_t>(MemoryConstants::gigaByte);
+        gfxPartitions.at(rootDeviceIndex) = std::make_unique<GfxPartition>(reservedCpuAddressRange);
+
+        auto gpuAddressSpace = executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->getHardwareInfo()->capabilityTable.gpuAddressSpace;
+        auto gfxTop = gpuAddressSpace + 1;
+        if (getGfxPartition(rootDeviceIndex)->init(gpuAddressSpace, reservedCpuAddressRangeSize, rootDeviceIndex, gfxPartitions.size(), heapAssigners[rootDeviceIndex]->apiAllowExternalHeapForSshAndDsh, OsAgnosticMemoryManager::getSystemSharedMemory(rootDeviceIndex), gfxTop)) {
+            return true;
+        }
+    }
+    return false;
+}
+
+void MockMemoryManager::releaseDeviceSpecificGfxPartition(uint32_t rootDeviceIndex) {
+    gfxPartitions.at(rootDeviceIndex).reset();
+}
+
 void *MockAllocSysMemAgnosticMemoryManager::allocateSystemMemory(size_t size, size_t alignment) {
     constexpr size_t minAlignment = 16;
     alignment = std::max(alignment, minAlignment);

shared/test/common/mocks/mock_memory_manager.h

@@ -251,6 +251,9 @@ class MockMemoryManager : public MemoryManagerCreate<OsAgnosticMemoryManager> {
         *serverType = MockMemoryManager::serverType;
     }
 
+    bool reInitDeviceSpecificGfxPartition(uint32_t rootDeviceIndex) override;
+    void releaseDeviceSpecificGfxPartition(uint32_t rootDeviceIndex) override;
+
     MockGraphicsAllocation *mockGa;
     size_t ipcAllocationSize = 4096u;
     uint32_t copyMemoryToAllocationBanksCalled = 0u;

shared/test/unit_test/os_interface/linux/drm_memory_manager_tests.cpp

@@ -8476,3 +8476,46 @@ TEST_F(DrmMemoryManagerTest, givenVmAdviseAtomicAttributeNotPresentWhenCreateSha
 
     memoryManager->freeGraphicsMemory(sharedUSM);
 }
+
+TEST_F(DrmMemoryManagerTest, givenGfxPartitionWhenReleasedAndReinitializedThenNewGfxPartitionIsCorrect) {
+
+    auto gfxPartition = memoryManager->getGfxPartition(0);
+
+    auto heapExternal = gfxPartition->getHeapBase(HeapIndex::heapExternal);
+    auto heapStandard = gfxPartition->getHeapBase(HeapIndex::heapStandard);
+    auto heapStandard64KB = gfxPartition->getHeapBase(HeapIndex::heapStandard64KB);
+    auto heapSvm = gfxPartition->getHeapBase(HeapIndex::heapSvm);
+    auto heapExtended = gfxPartition->getHeapBase(HeapIndex::heapExtended);
+    auto heapExternalFrontWindow = gfxPartition->getHeapBase(HeapIndex::heapExternalFrontWindow);
+    auto heapExternalDeviceFrontWindow = gfxPartition->getHeapBase(HeapIndex::heapExternalDeviceFrontWindow);
+    auto heapInternalFrontWindow = gfxPartition->getHeapBase(HeapIndex::heapInternalFrontWindow);
+    auto heapInternalDeviceFrontWindow = gfxPartition->getHeapBase(HeapIndex::heapInternalDeviceFrontWindow);
+
+    memoryManager->releaseDeviceSpecificGfxPartition(0);
+    EXPECT_EQ(nullptr, memoryManager->getGfxPartition(0));
+    memoryManager->reInitDeviceSpecificGfxPartition(0);
+
+    EXPECT_NE(nullptr, memoryManager->getGfxPartition(0));
+
+    gfxPartition = memoryManager->getGfxPartition(0);
+
+    auto heapExternal2 = gfxPartition->getHeapBase(HeapIndex::heapExternal);
+    auto heapStandard2 = gfxPartition->getHeapBase(HeapIndex::heapStandard);
+    auto heapStandard64KB2 = gfxPartition->getHeapBase(HeapIndex::heapStandard64KB);
+    auto heapSvm2 = gfxPartition->getHeapBase(HeapIndex::heapSvm);
+    auto heapExtended2 = gfxPartition->getHeapBase(HeapIndex::heapExtended);
+    auto heapExternalFrontWindow2 = gfxPartition->getHeapBase(HeapIndex::heapExternalFrontWindow);
+    auto heapExternalDeviceFrontWindow2 = gfxPartition->getHeapBase(HeapIndex::heapExternalDeviceFrontWindow);
+    auto heapInternalFrontWindow2 = gfxPartition->getHeapBase(HeapIndex::heapInternalFrontWindow);
+    auto heapInternalDeviceFrontWindow2 = gfxPartition->getHeapBase(HeapIndex::heapInternalDeviceFrontWindow);
+
+    EXPECT_EQ(heapExternal, heapExternal2);
+    EXPECT_EQ(heapStandard, heapStandard2);
+    EXPECT_EQ(heapStandard64KB, heapStandard64KB2);
+    EXPECT_EQ(heapSvm, heapSvm2);
+    EXPECT_EQ(heapExtended, heapExtended2);
+    EXPECT_EQ(heapExternalFrontWindow, heapExternalFrontWindow2);
+    EXPECT_EQ(heapExternalDeviceFrontWindow, heapExternalDeviceFrontWindow2);
+    EXPECT_EQ(heapInternalFrontWindow, heapInternalFrontWindow2);
+    EXPECT_EQ(heapInternalDeviceFrontWindow, heapInternalDeviceFrontWindow2);
+}