Register Allocations

Related-To: NEO-4964

Change-Id: I792dd7f7d6d594f51701ec7a40b2c0d36531b02b
Signed-off-by: Mateusz Hoppe <mateusz.hoppe@intel.com>

opencl/test/unit_test/mocks/linux/mock_drm_allocation.h

@@ -26,6 +26,7 @@ class MockDrmAllocation : public DrmAllocation {
   public:
     using DrmAllocation::bufferObjects;
     using DrmAllocation::memoryPool;
+    using DrmAllocation::registeredBoBindHandles;
 
     MockDrmAllocation(AllocationType allocationType, MemoryPool::Type pool) : DrmAllocation(0, allocationType, nullptr, nullptr, 0, static_cast<size_t>(0), pool) {
     }

opencl/test/unit_test/mocks/linux/mock_drm_memory_manager.h

@@ -68,6 +68,7 @@ class TestedDrmMemoryManager : public MemoryManagerCreate<DrmMemoryManager> {
     using DrmMemoryManager::pinBBs;
     using DrmMemoryManager::pinThreshold;
     using DrmMemoryManager::pushSharedBufferObject;
+    using DrmMemoryManager::registerAllocation;
     using DrmMemoryManager::releaseGpuRange;
     using DrmMemoryManager::setDomainCpu;
     using DrmMemoryManager::sharingBufferObjects;

opencl/test/unit_test/os_interface/linux/drm_buffer_object_tests.cpp

@@ -341,18 +341,3 @@ TEST(DrmBufferObject, whenBindExtHandleAddedThenItIsStored) {
     EXPECT_EQ(1u, bo.bindExtHandles.size());
     EXPECT_EQ(4u, bo.bindExtHandles[0]);
 }
-
-TEST(DrmBufferObject, givenBoWithBindExtHandlesWhenBoIsDestructedThenHandlesAreUnregistered) {
-    auto executionEnvironment = std::make_unique<ExecutionEnvironment>();
-    executionEnvironment->prepareRootDeviceEnvironments(1);
-    DrmMockResources drm(*executionEnvironment->rootDeviceEnvironments[0]);
-
-    {
-        MockBufferObject bo(&drm, 0, 0, 1);
-        bo.addBindExtHandle(4);
-        bo.addBindExtHandle(5);
-        bo.addBindExtHandle(6);
-    }
-    EXPECT_EQ(6u, drm.unregisteredHandle);
-    EXPECT_EQ(3u, drm.unregisterCalledCount);
-}

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

@@ -3870,4 +3870,85 @@ TEST(DrmAllocationTest, givenResourceRegistrationNotEnabledWhenRegisteringBindEx
     EXPECT_EQ(Drm::ResourceClass::MaxSize, drm.registeredClass);
 }
 
+TEST(DrmMemoryManager, givenTrackedAllocationTypeWhenAllocatingThenAllocationIsRegistered) {
+    auto executionEnvironment = std::make_unique<MockExecutionEnvironment>();
+    executionEnvironment->prepareRootDeviceEnvironments(1u);
+    executionEnvironment->rootDeviceEnvironments[0]->setHwInfo(defaultHwInfo.get());
+    auto memoryManager = std::make_unique<TestedDrmMemoryManager>(false, false, false, *executionEnvironment);
+    auto mockDrm = new DrmMockResources(*executionEnvironment->rootDeviceEnvironments[0]);
+    executionEnvironment->rootDeviceEnvironments[0]->osInterface = std::make_unique<OSInterface>();
+    executionEnvironment->rootDeviceEnvironments[0]->osInterface->get()->setDrm(mockDrm);
+
+    for (uint32_t i = 3; i < 3 + static_cast<uint32_t>(Drm::ResourceClass::MaxSize); i++) {
+        mockDrm->classHandles.push_back(i);
+    }
+
+    EXPECT_TRUE(mockDrm->resourceRegistrationEnabled());
+
+    NEO::AllocationProperties properties{0, true, MemoryConstants::pageSize,
+                                         NEO::GraphicsAllocation::AllocationType::DEBUG_SBA_TRACKING_BUFFER,
+                                         false,
+                                         CommonConstants::allDevicesBitfield};
+
+    properties.gpuAddress = 0x20000;
+    auto sbaAllocation = memoryManager->allocateGraphicsMemoryWithProperties(properties);
+    EXPECT_EQ(Drm::ResourceClass::SbaTrackingBuffer, mockDrm->registeredClass);
+
+    EXPECT_EQ(sizeof(uint64_t), mockDrm->registeredDataSize);
+    uint64_t *data = reinterpret_cast<uint64_t *>(mockDrm->registeredData);
+    EXPECT_EQ(properties.gpuAddress, *data);
+
+    memoryManager->freeGraphicsMemory(sbaAllocation);
+}
+
+TEST(DrmMemoryManager, givenTrackedAllocationTypeWhenFreeingThenRegisteredHandlesAreUnregistered) {
+    auto executionEnvironment = std::make_unique<MockExecutionEnvironment>();
+    executionEnvironment->prepareRootDeviceEnvironments(1u);
+    executionEnvironment->rootDeviceEnvironments[0]->setHwInfo(defaultHwInfo.get());
+    auto memoryManager = std::make_unique<TestedDrmMemoryManager>(false, false, false, *executionEnvironment);
+    auto mockDrm = new DrmMockResources(*executionEnvironment->rootDeviceEnvironments[0]);
+    executionEnvironment->rootDeviceEnvironments[0]->osInterface = std::make_unique<OSInterface>();
+    executionEnvironment->rootDeviceEnvironments[0]->osInterface->get()->setDrm(mockDrm);
+
+    for (uint32_t i = 3; i < 3 + static_cast<uint32_t>(Drm::ResourceClass::MaxSize); i++) {
+        mockDrm->classHandles.push_back(i);
+    }
+
+    EXPECT_TRUE(mockDrm->resourceRegistrationEnabled());
+
+    NEO::AllocationProperties properties{0, true, MemoryConstants::pageSize,
+                                         NEO::GraphicsAllocation::AllocationType::DEBUG_SBA_TRACKING_BUFFER,
+                                         false,
+                                         CommonConstants::allDevicesBitfield};
+
+    properties.gpuAddress = 0x20000;
+    auto sbaAllocation = memoryManager->allocateGraphicsMemoryWithProperties(properties);
+
+    EXPECT_EQ(0u, mockDrm->unregisterCalledCount);
+
+    memoryManager->freeGraphicsMemory(sbaAllocation);
+
+    EXPECT_EQ(DrmMockResources::registerResourceReturnHandle, mockDrm->unregisteredHandle);
+    EXPECT_EQ(1u, mockDrm->unregisterCalledCount);
+}
+
+TEST(DrmMemoryManager, givenNullBoWhenRegisteringBindExtHandleThenEarlyReturn) {
+    auto executionEnvironment = std::make_unique<MockExecutionEnvironment>();
+    executionEnvironment->prepareRootDeviceEnvironments(1u);
+    executionEnvironment->rootDeviceEnvironments[0]->setHwInfo(defaultHwInfo.get());
+    auto mockDrm = std::make_unique<DrmMockResources>(*executionEnvironment->rootDeviceEnvironments[0]);
+
+    for (uint32_t i = 3; i < 3 + static_cast<uint32_t>(Drm::ResourceClass::MaxSize); i++) {
+        mockDrm->classHandles.push_back(i);
+    }
+
+    EXPECT_TRUE(mockDrm->resourceRegistrationEnabled());
+
+    MockDrmAllocation gfxAllocation(GraphicsAllocation::AllocationType::DEBUG_SBA_TRACKING_BUFFER, MemoryPool::MemoryNull);
+
+    gfxAllocation.registerBOBindExtHandle(mockDrm.get());
+    EXPECT_EQ(1u, gfxAllocation.registeredBoBindHandles.size());
+    gfxAllocation.freeRegisteredBOBindExtHandles(mockDrm.get());
+}
+
 } // namespace NEO

opencl/test/unit_test/os_interface/linux/drm_mock.h

@@ -11,6 +11,7 @@
 #include "shared/source/execution_environment/root_device_environment.h"
 #include "shared/source/helpers/constants.h"
 #include "shared/source/helpers/hw_helper.h"
+#include "shared/source/helpers/string.h"
 #include "shared/source/os_interface/linux/drm_neo.h"
 
 #include "opencl/source/platform/platform.h"
@@ -220,6 +221,8 @@ class DrmMockResources : public DrmMock {
 
     uint32_t registerResource(ResourceClass classType, void *data, size_t size) override {
         registeredClass = classType;
+        memcpy_s(registeredData, sizeof(registeredData), data, size);
+        registeredDataSize = size;
         return registerResourceReturnHandle;
     }
 
@@ -234,4 +237,6 @@ class DrmMockResources : public DrmMock {
     uint32_t unregisterCalledCount = 0;
     ResourceClass registeredClass = ResourceClass::MaxSize;
     bool registerClassesCalled = false;
+    uint64_t registeredData[128];
+    size_t registeredDataSize;
 };

shared/source/memory_manager/memory_manager.cpp

@@ -411,6 +411,7 @@ GraphicsAllocation *MemoryManager::allocateGraphicsMemoryInPreferredPool(const A
         this->registerSysMemAlloc(allocation);
     }
     FileLoggerInstance().logAllocation(allocation);
+    registerAllocation(allocation);
     return allocation;
 }
 

shared/source/memory_manager/memory_manager.h

@@ -212,6 +212,7 @@ class MemoryManager {
     virtual void *lockResourceImpl(GraphicsAllocation &graphicsAllocation) = 0;
     virtual void unlockResourceImpl(GraphicsAllocation &graphicsAllocation) = 0;
     virtual void freeAssociatedResourceImpl(GraphicsAllocation &graphicsAllocation) { return unlockResourceImpl(graphicsAllocation); };
+    virtual void registerAllocation(GraphicsAllocation *allocation) {}
 
     bool forceNonSvmForExternalHostPtr = false;
     bool force32bitAllocations = false;

shared/source/os_interface/linux/drm_allocation.cpp

@@ -89,12 +89,20 @@ void DrmAllocation::registerBOBindExtHandle(Drm *drm) {
     if (resourceClass != Drm::ResourceClass::MaxSize) {
         uint64_t gpuAddress = getGpuAddress();
         auto handle = drm->registerResource(resourceClass, &gpuAddress, sizeof(gpuAddress));
+        registeredBoBindHandles.push_back(handle);
         auto &bos = getBOs();
 
         for (auto bo : bos) {
-            bo->addBindExtHandle(handle);
+            if (bo) {
+                bo->addBindExtHandle(handle);
+            }
         }
     }
 }
 
+void DrmAllocation::freeRegisteredBOBindExtHandles(Drm *drm) {
+    for (auto &i : registeredBoBindHandles) {
+        drm->unregisterResource(i);
+    }
+}
 } // namespace NEO

shared/source/os_interface/linux/drm_allocation.h

@@ -67,8 +67,10 @@ class DrmAllocation : public GraphicsAllocation {
     void bindBO(BufferObject *bo, OsContext *osContext, uint32_t vmHandleId, std::vector<BufferObject *> *bufferObjects, bool bind);
     void bindBOs(OsContext *osContext, uint32_t vmHandleId, std::vector<BufferObject *> *bufferObjects, bool bind);
     void registerBOBindExtHandle(Drm *drm);
+    void freeRegisteredBOBindExtHandles(Drm *drm);
 
   protected:
     BufferObjects bufferObjects{};
+    StackVec<uint32_t, 1> registeredBoBindHandles;
 };
 } // namespace NEO

shared/source/os_interface/linux/drm_buffer_object.cpp

@@ -47,12 +47,6 @@ BufferObject::BufferObject(Drm *drm, int handle, size_t size, size_t maxOsContex
     }
 }
 
-BufferObject::~BufferObject() {
-    for (auto &i : bindExtHandles) {
-        drm->unregisterResource(i);
-    }
-};
-
 uint32_t BufferObject::getRefCount() const {
     return this->refCount.load();
 }

shared/source/os_interface/linux/drm_buffer_object.h

@@ -32,7 +32,7 @@ class BufferObject {
 
   public:
     BufferObject(Drm *drm, int handle, size_t size, size_t maxOsContextCount);
-    MOCKABLE_VIRTUAL ~BufferObject();
+    MOCKABLE_VIRTUAL ~BufferObject() = default;
 
     struct Deleter {
         void operator()(BufferObject *bo) {

shared/source/os_interface/linux/drm_memory_manager.cpp

@@ -650,6 +650,7 @@ void DrmMemoryManager::removeAllocationFromHostPtrManager(GraphicsAllocation *gf
 }
 
 void DrmMemoryManager::freeGraphicsMemoryImpl(GraphicsAllocation *gfxAllocation) {
+    DrmAllocation *drmAlloc = static_cast<DrmAllocation *>(gfxAllocation);
     this->unregisterAllocation(gfxAllocation);
 
     for (auto &engine : this->registeredEngines) {
@@ -676,6 +677,8 @@ void DrmMemoryManager::freeGraphicsMemoryImpl(GraphicsAllocation *gfxAllocation)
     releaseGpuRange(gfxAllocation->getReservedAddressPtr(), gfxAllocation->getReservedAddressSize(), gfxAllocation->getRootDeviceIndex());
     alignedFreeWrapper(gfxAllocation->getDriverAllocatedCpuPtr());
 
+    drmAlloc->freeRegisteredBOBindExtHandles(&getDrm(drmAlloc->getRootDeviceIndex()));
+
     delete gfxAllocation;
 }
 
@@ -907,4 +910,11 @@ void DrmMemoryManager::unregisterAllocation(GraphicsAllocation *allocation) {
                                                                        allocation),
                                                            localMemAllocs[allocation->getRootDeviceIndex()].end());
 }
+
+void DrmMemoryManager::registerAllocation(GraphicsAllocation *allocation) {
+    if (allocation) {
+        auto drmAllocation = static_cast<DrmAllocation *>(allocation);
+        drmAllocation->registerBOBindExtHandle(&getDrm(drmAllocation->getRootDeviceIndex()));
+    }
+}
 } // namespace NEO

shared/source/os_interface/linux/drm_memory_manager.h

@@ -97,6 +97,7 @@ class DrmMemoryManager : public MemoryManager {
     DrmAllocation *allocate32BitGraphicsMemoryImpl(const AllocationData &allocationData, bool useLocalMemory) override;
     GraphicsAllocation *allocateGraphicsMemoryInDevicePool(const AllocationData &allocationData, AllocationStatus &status) override;
     bool createDrmAllocation(Drm *drm, DrmAllocation *allocation, uint64_t gpuAddress, size_t maxOsContextCount);
+    void registerAllocation(GraphicsAllocation *allocation) override;
 
     Drm &getDrm(uint32_t rootDeviceIndex) const;
     uint32_t getRootDeviceIndex(const Drm *drm);