Mark Scratch and Private allocations for capture

Related-To: NEO-5237

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

Author: HoppeMateusz

opencl/test/unit_test/memory_manager/memory_manager_tests.cpp

@@ -2568,3 +2568,10 @@ TEST(MemoryManagerTest, givenDebugModuleAreaAllocationTypeWhenCallingGetAllocati
     mockMemoryManager.getAllocationData(allocData, properties, nullptr, mockMemoryManager.createStorageInfoFromProperties(properties));
     EXPECT_EQ(1u, allocData.flags.use32BitFrontWindow);
 }
+
+TEST(MemoryManagerTest, WhenCallingIsAllocationTypeToCaptureThenScratchAndPrivateTypesReturnTrue) {
+    MockMemoryManager mockMemoryManager;
+
+    EXPECT_TRUE(mockMemoryManager.isAllocationTypeToCapture(GraphicsAllocation::AllocationType::SCRATCH_SURFACE));
+    EXPECT_TRUE(mockMemoryManager.isAllocationTypeToCapture(GraphicsAllocation::AllocationType::PRIVATE_SURFACE));
+}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019-2020 Intel Corporation
+ * Copyright (C) 2019-2021 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -36,7 +36,13 @@ class MockDrmAllocation : public DrmAllocation {
         DrmAllocation::registerBOBindExtHandle(drm);
     }
 
+    void markForCapture() override {
+        markedForCapture = true;
+        DrmAllocation::markForCapture();
+    }
+
     bool registerBOBindExtHandleCalled = false;
+    bool markedForCapture = false;
 };
 
 } // namespace NEO

opencl/test/unit_test/mocks/mock_memory_manager.h

@@ -49,6 +49,7 @@ class MockMemoryManager : public MemoryManagerCreate<OsAgnosticMemoryManager> {
     using OsAgnosticMemoryManager::allocateGraphicsMemoryForImageFromHostPtr;
     using MemoryManagerCreate<OsAgnosticMemoryManager>::MemoryManagerCreate;
     using MemoryManager::enable64kbpages;
+    using MemoryManager::isAllocationTypeToCapture;
     using MemoryManager::isCopyRequired;
     using MemoryManager::localMemorySupported;
     using MemoryManager::reservedMemory;

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

@@ -4081,6 +4081,32 @@ TEST(DrmMemoryManager, givenTrackedAllocationTypeAndDisabledRegistrationInDrmWhe
     EXPECT_EQ(Drm::ResourceClass::MaxSize, mockDrm->registeredClass);
 }
 
+TEST(DrmMemoryManager, givenResourceRegistrationEnabledAndAllocTypeToCaptureWhenRegisteringAllocationInOsThenItIsMarkedForCapture) {
+    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);
+
+    // mock resource registration enabling by storing class handles
+    mockDrm->classHandles.push_back(1);
+
+    MockBufferObject bo(mockDrm, 0, 0, 1);
+    MockDrmAllocation allocation(GraphicsAllocation::AllocationType::SCRATCH_SURFACE, MemoryPool::System4KBPages);
+    allocation.bufferObjects[0] = &bo;
+    memoryManager->registerAllocationInOs(&allocation);
+
+    EXPECT_TRUE(allocation.markedForCapture);
+
+    MockDrmAllocation allocation2(GraphicsAllocation::AllocationType::BUFFER, MemoryPool::System4KBPages);
+    allocation2.bufferObjects[0] = &bo;
+    memoryManager->registerAllocationInOs(&allocation2);
+
+    EXPECT_FALSE(allocation2.markedForCapture);
+}
+
 TEST(DrmMemoryManager, givenTrackedAllocationTypeWhenAllocatingThenAllocationIsRegistered) {
     auto executionEnvironment = std::make_unique<MockExecutionEnvironment>();
     executionEnvironment->prepareRootDeviceEnvironments(1u);
@@ -4279,6 +4305,22 @@ TEST(DrmAllocationTest, givenDrmAllocationWhenCacheRegionIsSetSuccessfullyThenSe
     }
 }
 
+TEST(DrmAllocationTest, givenBoWhenMarkingForCaptureThenBosAreMarked) {
+    auto executionEnvironment = std::make_unique<ExecutionEnvironment>();
+    executionEnvironment->prepareRootDeviceEnvironments(1);
+
+    DrmMock drm(*executionEnvironment->rootDeviceEnvironments[0]);
+
+    MockBufferObject bo(&drm, 0, 0, 1);
+    MockDrmAllocation allocation(GraphicsAllocation::AllocationType::SCRATCH_SURFACE, MemoryPool::System4KBPages);
+    allocation.markForCapture();
+
+    allocation.bufferObjects[0] = &bo;
+    allocation.markForCapture();
+
+    EXPECT_TRUE(bo.isMarkedForCapture());
+}
+
 TEST_F(DrmMemoryManagerTest, givenDrmAllocationWithHostPtrWhenItIsCreatedWithCacheRegionThenSetRegionInBufferObject) {
     mock->ioctl_expected.total = -1;
     auto drm = static_cast<DrmMockCustom *>(executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface->get()->getDrm());

shared/source/memory_manager/memory_manager.cpp

@@ -748,6 +748,17 @@ void MemoryManager::overrideAllocationData(AllocationData &allocationData, const
     }
 }
 
+bool MemoryManager::isAllocationTypeToCapture(GraphicsAllocation::AllocationType type) const {
+    switch (type) {
+    case GraphicsAllocation::AllocationType::SCRATCH_SURFACE:
+    case GraphicsAllocation::AllocationType::PRIVATE_SURFACE:
+        return true;
+    default:
+        break;
+    }
+    return false;
+}
+
 bool MemoryTransferHelper::transferMemoryToAllocation(bool useBlitter, const Device &device, GraphicsAllocation *dstAllocation, size_t dstOffset, const void *srcMemory, size_t srcSize) {
     if (useBlitter) {
         return (BlitHelperFunctions::blitMemoryToAllocation(device, dstAllocation, dstOffset, srcMemory, {srcSize, 1, 1}) == BlitOperationResult::Success);

shared/source/memory_manager/memory_manager.h

@@ -223,6 +223,7 @@ class MemoryManager {
     virtual void unlockResourceImpl(GraphicsAllocation &graphicsAllocation) = 0;
     virtual void freeAssociatedResourceImpl(GraphicsAllocation &graphicsAllocation) { return unlockResourceImpl(graphicsAllocation); };
     virtual void registerAllocationInOs(GraphicsAllocation *allocation) {}
+    bool isAllocationTypeToCapture(GraphicsAllocation::AllocationType type) const;
 
     bool initialized = false;
     bool forceNonSvmForExternalHostPtr = false;

shared/source/os_interface/linux/drm_allocation.cpp

@@ -147,4 +147,13 @@ void DrmAllocation::freeRegisteredBOBindExtHandles(Drm *drm) {
         drm->unregisterResource(i);
     }
 }
+
+void DrmAllocation::markForCapture() {
+    auto &bos = getBOs();
+    for (auto bo : bos) {
+        if (bo) {
+            bo->markForCapture();
+        }
+    }
+}
 } // namespace NEO

shared/source/os_interface/linux/drm_allocation.h

@@ -78,6 +78,7 @@ class DrmAllocation : public GraphicsAllocation {
     MOCKABLE_VIRTUAL void registerBOBindExtHandle(Drm *drm);
     void freeRegisteredBOBindExtHandles(Drm *drm);
     void linkWithRegisteredHandle(uint32_t handle);
+    MOCKABLE_VIRTUAL void markForCapture();
 
   protected:
     BufferObjects bufferObjects{};

shared/source/os_interface/linux/drm_memory_manager.cpp

@@ -995,6 +995,10 @@ void DrmMemoryManager::registerAllocationInOs(GraphicsAllocation *allocation) {
     if (allocation && getDrm(allocation->getRootDeviceIndex()).resourceRegistrationEnabled()) {
         auto drmAllocation = static_cast<DrmAllocation *>(allocation);
         drmAllocation->registerBOBindExtHandle(&getDrm(drmAllocation->getRootDeviceIndex()));
+
+        if (isAllocationTypeToCapture(drmAllocation->getAllocationType())) {
+            drmAllocation->markForCapture();
+        }
     }
 }
 } // namespace NEO