fix: do not unregister shared allocation

Allocations imported from shared handle are not using registerAlloc in drm
manager but on free, unregister was called.

This could lead to problems with allocation size tracking.

This change will skip the unregisterAllocation call if allocation is
imported from shared handle.

Signed-off-by: Dominik Dabek <[email protected]>

Author: ddabek-i

shared/source/os_interface/linux/drm_memory_manager.cpp

@@ -1229,7 +1229,9 @@ void DrmMemoryManager::freeGraphicsMemoryImpl(GraphicsAllocation *gfxAllocation,
         return;
     }
     DrmAllocation *drmAlloc = static_cast<DrmAllocation *>(gfxAllocation);
-    this->unregisterAllocation(gfxAllocation);
+    if (Sharing::nonSharedResource == gfxAllocation->peekSharedHandle()) {
+        this->unregisterAllocation(gfxAllocation);
+    }
     auto rootDeviceIndex = gfxAllocation->getRootDeviceIndex();
     for (auto &engine : getRegisteredEngines(rootDeviceIndex)) {
         auto memoryOperationsInterface = static_cast<DrmMemoryOperationsHandler *>(executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface.get());

shared/source/os_interface/linux/drm_memory_manager.h

@@ -93,7 +93,7 @@ class DrmMemoryManager : public MemoryManager {
     std::vector<GraphicsAllocation *> &getLocalMemAllocs(uint32_t rootDeviceIndex);
     AllocationStatus registerSysMemAlloc(GraphicsAllocation *allocation) override;
     AllocationStatus registerLocalMemAlloc(GraphicsAllocation *allocation, uint32_t rootDeviceIndex) override;
-    void unregisterAllocation(GraphicsAllocation *allocation);
+    MOCKABLE_VIRTUAL void unregisterAllocation(GraphicsAllocation *allocation);
 
     static std::unique_ptr<MemoryManager> create(ExecutionEnvironment &executionEnvironment);
 

shared/test/common/mocks/linux/mock_drm_memory_manager.h

@@ -175,10 +175,28 @@ class TestedDrmMemoryManager : public MemoryManagerCreate<DrmMemoryManager> {
         return MemoryManager::computeStorageInfoMemoryBanks(properties, preferredBank, allBanks);
     }
 
+    AllocationStatus registerSysMemAlloc(GraphicsAllocation *allocation) override {
+        ++registerSysMemAllocCalled;
+        return DrmMemoryManager::registerSysMemAlloc(allocation);
+    }
+
+    AllocationStatus registerLocalMemAlloc(GraphicsAllocation *allocation, uint32_t rootDeviceIndex) override {
+        ++registerLocalMemAllocCalled;
+        return DrmMemoryManager::registerLocalMemAlloc(allocation, rootDeviceIndex);
+    }
+
+    void unregisterAllocation(GraphicsAllocation *allocation) override {
+        ++unregisterAllocationCalled;
+        DrmMemoryManager::unregisterAllocation(allocation);
+    }
+
     uint32_t acquireGpuRangeCalledTimes = 0u;
     uint32_t acquireGpuRangeWithCustomAlignmenCalledTimes = 0u;
     size_t acquireGpuRangeWithCustomAlignmenPassedAlignment = 0u;
     size_t computeStorageInfoMemoryBanksCalled = 0u;
+    size_t registerSysMemAllocCalled = 0u;
+    size_t registerLocalMemAllocCalled = 0u;
+    size_t unregisterAllocationCalled = 0u;
     ExecutionEnvironment *executionEnvironment = nullptr;
 
   protected:

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

@@ -3055,6 +3055,22 @@ TEST_F(DrmMemoryManagerBasic, givenMemoryManagerWhenCreateAllocationFromHandleIs
     memoryManager->freeGraphicsMemory(allocation);
 }
 
+TEST_F(DrmMemoryManagerBasic, givenMemoryManagerWhenCreateAllocationFromHandleIsCalledThenAllocationIsNotRegisteredNorUnregistered) {
+    std::unique_ptr<TestedDrmMemoryManager> memoryManager(new (std::nothrow) TestedDrmMemoryManager(false,
+                                                                                                    false,
+                                                                                                    true,
+                                                                                                    executionEnvironment));
+    TestedDrmMemoryManager::OsHandleData osHandleData{1u};
+    AllocationProperties properties(rootDeviceIndex, false, MemoryConstants::pageSize, AllocationType::sharedBuffer, false, {});
+
+    auto allocation = memoryManager->createGraphicsAllocationFromSharedHandle(osHandleData, properties, false, false, true, nullptr);
+    EXPECT_NE(nullptr, allocation);
+    memoryManager->freeGraphicsMemory(allocation);
+    EXPECT_EQ(0u, memoryManager->registerSysMemAllocCalled);
+    EXPECT_EQ(0u, memoryManager->registerLocalMemAllocCalled);
+    EXPECT_EQ(0u, memoryManager->unregisterAllocationCalled);
+}
+
 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDisabledForcePinAndEnabledValidateHostMemoryWhenPinBBAllocationFailsThenUnrecoverableIsCalled) {
     this->mock = static_cast<DrmMockCustom *>(executionEnvironment->rootDeviceEnvironments[0]->osInterface->getDriverModel()->as<Drm>());
     this->mock->reset();