feature: support explicit memory locking

Added lockMemory in context to explicitly locking memory,
Added a boolean flag in graphics_allocation to indicate the allocation
is locked, and modified memory_operations_handler to add lock().

Related-To: NEO-8277
Signed-off-by: Young Jin Yoon <young.jin.yoon@intel.com>

level_zero/core/source/context/CMakeLists.txt

@@ -7,8 +7,9 @@
 target_sources(${L0_STATIC_LIB_NAME}
                PRIVATE
                ${CMAKE_CURRENT_SOURCE_DIR}/CMakeLists.txt
+               ${CMAKE_CURRENT_SOURCE_DIR}/${BRANCH_DIR_SUFFIX}context_imp_helper.cpp
                ${CMAKE_CURRENT_SOURCE_DIR}/context_imp_${DRIVER_MODEL}/context_imp_${DRIVER_MODEL}.cpp
-               ${CMAKE_CURRENT_SOURCE_DIR}/context_imp_${DRIVER_MODEL}${BRANCH_DIR_SUFFIX}context_imp_helper.cpp
+               ${CMAKE_CURRENT_SOURCE_DIR}/context_imp_${DRIVER_MODEL}${BRANCH_DIR_SUFFIX}context_imp_${DRIVER_MODEL}_helper.cpp
                ${CMAKE_CURRENT_SOURCE_DIR}/context_imp.cpp
                ${CMAKE_CURRENT_SOURCE_DIR}/context_imp.h
                ${CMAKE_CURRENT_SOURCE_DIR}/context.h

level_zero/core/source/context/context.h

@@ -162,7 +162,7 @@ struct Context : _ze_context_handle_t {
     virtual ze_result_t getVirtualAddressSpaceIpcHandle(ze_device_handle_t hDevice,
                                                         ze_ipc_mem_handle_t *pIpcHandle) = 0;
     virtual ze_result_t putVirtualAddressSpaceIpcHandle(ze_ipc_mem_handle_t ipcHandle) = 0;
-
+    virtual ze_result_t lockMemory(ze_device_handle_t hDevice, void *ptr, size_t size) = 0;
     virtual bool isShareableMemory(const void *exportDesc, bool exportableMemory, NEO::Device *neoDevice) = 0;
     virtual void *getMemHandlePtr(ze_device_handle_t hDevice, uint64_t handle, NEO::AllocationType allocationType, ze_ipc_memory_flags_t flags) = 0;
 

level_zero/core/source/context/context_imp.cpp

@@ -471,6 +471,9 @@ ze_result_t ContextImp::makeMemoryResident(ze_device_handle_t hDevice, void *ptr
             return ZE_RESULT_ERROR_INVALID_ARGUMENT;
         }
     }
+    if (allocation->isLockedMemory()) {
+        return ZE_RESULT_ERROR_INVALID_ARGUMENT;
+    }
 
     NEO::MemoryOperationsHandler *memoryOperationsIface = neoDevice->getRootDeviceEnvironment().memoryOperationsInterface.get();
     auto success = memoryOperationsIface->makeResident(neoDevice, ArrayRef<NEO::GraphicsAllocation *>(&allocation, 1));

level_zero/core/source/context/context_imp.h

@@ -66,6 +66,7 @@ struct ContextImp : Context {
                                  void **ptr) override;
     ze_result_t getIpcHandleFromFd(uint64_t handle, ze_ipc_mem_handle_t *pIpcHandle) override;
     ze_result_t getFdFromIpcHandle(ze_ipc_mem_handle_t ipcHandle, uint64_t *pHandle) override;
+    ze_result_t lockMemory(ze_device_handle_t hDevice, void *ptr, size_t size) override;
 
     ze_result_t
     getIpcMemHandles(

level_zero/core/source/context/context_imp_helper.cpp

@@ -0,0 +1,18 @@
+/*
+ * Copyright (C) 2024 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "level_zero/core/source/context/context_imp.h"
+#include "level_zero/core/source/device/device.h"
+#include "level_zero/core/source/driver/driver_handle_imp.h"
+
+namespace L0 {
+
+ze_result_t ContextImp::lockMemory(ze_device_handle_t device, void *ptr, size_t size) {
+    return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE;
+}
+
+} // namespace L0

level_zero/core/source/context/context_imp_wddm/context_imp_wddm_helper.cpp

@@ -18,4 +18,5 @@ ze_result_t ContextImp::getVirtualAddressSpaceIpcHandle(ze_device_handle_t hDevi
 ze_result_t ContextImp::putVirtualAddressSpaceIpcHandle(ze_ipc_mem_handle_t ipcHandle) {
     return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }
+
 } // namespace L0

level_zero/core/test/unit_tests/sources/context/CMakeLists.txt

@@ -8,6 +8,7 @@ target_sources(${TARGET_NAME} PRIVATE
                ${CMAKE_CURRENT_SOURCE_DIR}/CMakeLists.txt
                ${CMAKE_CURRENT_SOURCE_DIR}/test_context.cpp
                ${CMAKE_CURRENT_SOURCE_DIR}/test_context_${DRIVER_MODEL}.cpp
+               ${CMAKE_CURRENT_SOURCE_DIR}${BRANCH_DIR_SUFFIX}test_context_helper.cpp
 )
 
 add_subdirectories()

level_zero/core/test/unit_tests/sources/context/linux/CMakeLists.txt

@@ -7,7 +7,7 @@
 if(UNIX)
   target_sources(${TARGET_NAME} PRIVATE
                  ${CMAKE_CURRENT_SOURCE_DIR}/CMakeLists.txt
-                 ${CMAKE_CURRENT_SOURCE_DIR}${BRANCH_DIR_SUFFIX}test_context_helper.cpp
+                 ${CMAKE_CURRENT_SOURCE_DIR}${BRANCH_DIR_SUFFIX}test_context_helper_linux.cpp
   )
 endif()
 

level_zero/core/test/unit_tests/sources/context/linux/test_context_helper_linux.cpp

@@ -10,12 +10,14 @@
 #include "shared/test/common/mocks/mock_builtins.h"
 #include "shared/test/common/mocks/mock_device.h"
 #include "shared/test/common/mocks/mock_driver_model.h"
+#include "shared/test/common/mocks/mock_memory_operations_handler.h"
 #include "shared/test/common/test_macros/test.h"
 
 #include "level_zero/core/source/context/context_imp.h"
 #include "level_zero/core/source/driver/driver_handle_imp.h"
 #include "level_zero/core/source/driver/driver_imp.h"
 #include "level_zero/core/test/unit_tests/fixtures/device_fixture.h"
+#include "level_zero/core/test/unit_tests/fixtures/host_pointer_manager_fixture.h"
 #include "level_zero/core/test/unit_tests/mocks/mock_driver_handle.h"
 
 #include "gtest/gtest.h"
@@ -41,5 +43,6 @@ TEST_F(ContextGetVirtualAddressSpaceTests, givenDrmDriverModelWhenCallingGetVirt
     res = contextImp->destroy();
     EXPECT_EQ(ZE_RESULT_SUCCESS, res);
 }
+
 } // namespace ult
 } // namespace L0

level_zero/core/test/unit_tests/sources/context/test_context.cpp

@@ -580,6 +580,31 @@ TEST_F(ContextMakeMemoryResidentTests,
     context->freeMem(ptr);
 }
 
+TEST_F(ContextMakeMemoryResidentTests,
+       givenValidAllocationwithLockedWhenCallingMakeMemoryResidentThenInvalidArgumentIsReturned) {
+    const size_t size = 4096;
+    void *ptr = nullptr;
+    ze_host_mem_alloc_desc_t hostDesc = {};
+    ze_device_mem_alloc_desc_t deviceDesc = {};
+    ze_result_t res = context->allocSharedMem(device->toHandle(),
+                                              &deviceDesc,
+                                              &hostDesc,
+                                              size,
+                                              0,
+                                              &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    DriverHandleImp *driverHandleImp = static_cast<DriverHandleImp *>(hostDriverHandle.get());
+    auto allocation = driverHandleImp->getDriverSystemMemoryAllocation(ptr, size, neoDevice->getRootDeviceIndex(), nullptr);
+    allocation->setLockedMemory(true);
+
+    mockMemoryInterface->makeResidentResult = NEO::MemoryOperationsStatus::success;
+    res = context->makeMemoryResident(device, ptr, size);
+    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);
+
+    context->freeMem(ptr);
+}
+
 TEST_F(ContextMakeMemoryResidentTests,
        whenMakingASharedMemoryResidentThenIsAddedToVectorOfResidentAllocations) {
     const size_t size = 4096;

level_zero/core/test/unit_tests/sources/context/test_context_helper.cpp

@@ -0,0 +1,48 @@
+/*
+ * Copyright (C) 2023-2024 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "shared/source/execution_environment/root_device_environment.h"
+#include "shared/source/os_interface/os_interface.h"
+#include "shared/test/common/mocks/mock_builtins.h"
+#include "shared/test/common/mocks/mock_device.h"
+#include "shared/test/common/mocks/mock_driver_model.h"
+#include "shared/test/common/mocks/mock_memory_operations_handler.h"
+#include "shared/test/common/test_macros/test.h"
+
+#include "level_zero/core/source/context/context_imp.h"
+#include "level_zero/core/source/driver/driver_handle_imp.h"
+#include "level_zero/core/source/driver/driver_imp.h"
+#include "level_zero/core/test/unit_tests/fixtures/device_fixture.h"
+#include "level_zero/core/test/unit_tests/fixtures/host_pointer_manager_fixture.h"
+#include "level_zero/core/test/unit_tests/mocks/mock_driver_handle.h"
+
+#include "gtest/gtest.h"
+namespace L0 {
+namespace ult {
+
+using ContextLockMemoryTests = Test<HostPointerManagerFixure>;
+
+TEST_F(ContextLockMemoryTests, givenValidPointerWhenCallingLockMemoryThenUnsupportedErrorIsReturned) {
+    const size_t size = 4096;
+    void *ptr = nullptr;
+    ze_device_mem_alloc_desc_t deviceDesc = {};
+    ze_result_t res = context->allocDeviceMem(device->toHandle(),
+                                              &deviceDesc,
+                                              size,
+                                              0,
+                                              &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    mockMemoryInterface->lockResult = NEO::MemoryOperationsStatus::success;
+    res = context->lockMemory(device, ptr, size);
+    EXPECT_EQ(ZE_RESULT_ERROR_UNSUPPORTED_FEATURE, res);
+
+    context->freeMem(ptr);
+}
+
+} // namespace ult
+} // namespace L0

level_zero/core/test/unit_tests/sources/context/windows/CMakeLists.txt

@@ -7,6 +7,6 @@
 if(WIN32)
   target_sources(${TARGET_NAME} PRIVATE
                  ${CMAKE_CURRENT_SOURCE_DIR}/CMakeLists.txt
-                 ${CMAKE_CURRENT_SOURCE_DIR}${BRANCH_DIR_SUFFIX}test_context_helper.cpp
+                 ${CMAKE_CURRENT_SOURCE_DIR}${BRANCH_DIR_SUFFIX}test_context_helper_windows.cpp
   )
 endif()

level_zero/core/test/unit_tests/sources/context/windows/test_context_helper_windows.cpp

@@ -10,12 +10,14 @@
 #include "shared/test/common/mocks/mock_builtins.h"
 #include "shared/test/common/mocks/mock_device.h"
 #include "shared/test/common/mocks/mock_driver_model.h"
+#include "shared/test/common/mocks/mock_memory_operations_handler.h"
 #include "shared/test/common/test_macros/test.h"
 
 #include "level_zero/core/source/context/context_imp.h"
 #include "level_zero/core/source/driver/driver_handle_imp.h"
 #include "level_zero/core/source/driver/driver_imp.h"
 #include "level_zero/core/test/unit_tests/fixtures/device_fixture.h"
+#include "level_zero/core/test/unit_tests/fixtures/host_pointer_manager_fixture.h"
 #include "level_zero/core/test/unit_tests/mocks/mock_driver_handle.h"
 
 #include "gtest/gtest.h"

shared/source/debug_settings/debug_variables_base.inl

@@ -612,6 +612,7 @@ DECLARE_DEBUG_VARIABLE(int32_t, EnableSetPair, -1, "Use SET_PAIR to pair two buf
 DECLARE_DEBUG_VARIABLE(int32_t, ForcePreferredAllocationMethod, -1, "Sets preferred allocation method for Wddm paths; values = -1: driver default, 0: UseUmdSystemPtr, 1: AllocateByKmd")
 DECLARE_DEBUG_VARIABLE(int32_t, EventTimestampRefreshIntervalInMilliSec, -1, "-1: use driver default, This value sets the refresh interval for getting synchronized GPU and CPU timestamp")
 DECLARE_DEBUG_VARIABLE(int64_t, ReadOnlyAllocationsTypeMask, 0, "0: default,  >0: (bitmask) for given Graphics Allocation Type, set as read only resource.")
+
 /* Binary Cache */
 DECLARE_DEBUG_VARIABLE(bool, BinaryCacheTrace, false, "enable cl_cache to produce .trace files with information about hash computation")
 

shared/source/memory_manager/graphics_allocation.h

@@ -124,6 +124,8 @@ class GraphicsAllocation : public IDNode<GraphicsAllocation> {
     bool peekEvictable() const { return allocationInfo.flags.evictable; }
     bool isFlushL3Required() const { return allocationInfo.flags.flushL3Required; }
     void setFlushL3Required(bool flushL3Required) { allocationInfo.flags.flushL3Required = flushL3Required; }
+    bool isLockedMemory() const { return allocationInfo.flags.lockedMemory; }
+    void setLockedMemory(bool locked) { allocationInfo.flags.lockedMemory = locked; }
 
     bool isUncacheable() const { return allocationInfo.flags.uncacheable; }
     void setUncacheable(bool uncacheable) { allocationInfo.flags.uncacheable = uncacheable; }
@@ -349,7 +351,8 @@ class GraphicsAllocation : public IDNode<GraphicsAllocation> {
                 uint32_t flushL3Required : 1;
                 uint32_t uncacheable : 1;
                 uint32_t is32BitAllocation : 1;
-                uint32_t reserved : 27;
+                uint32_t lockedMemory : 1;
+                uint32_t reserved : 26;
             } flags;
             uint32_t allFlags = 0u;
         };
@@ -359,6 +362,7 @@ class GraphicsAllocation : public IDNode<GraphicsAllocation> {
             flags.evictable = true;
             flags.flushL3Required = true;
             flags.is32BitAllocation = false;
+            flags.lockedMemory = false;
         }
     };
 

shared/source/memory_manager/memory_operations_handler.h

@@ -20,6 +20,7 @@ class MemoryOperationsHandler {
     virtual ~MemoryOperationsHandler() = default;
 
     virtual MemoryOperationsStatus makeResident(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) = 0;
+    virtual MemoryOperationsStatus lock(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) = 0;
     virtual MemoryOperationsStatus evict(Device *device, GraphicsAllocation &gfxAllocation) = 0;
     virtual MemoryOperationsStatus isResident(Device *device, GraphicsAllocation &gfxAllocation) = 0;
 

shared/source/os_interface/aub_memory_operations_handler.cpp

@@ -69,6 +69,10 @@ MemoryOperationsStatus AubMemoryOperationsHandler::makeResident(Device *device,
     return MemoryOperationsStatus::success;
 }
 
+MemoryOperationsStatus AubMemoryOperationsHandler::lock(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) {
+    return makeResident(device, gfxAllocations);
+}
+
 MemoryOperationsStatus AubMemoryOperationsHandler::evict(Device *device, GraphicsAllocation &gfxAllocation) {
     auto lock = acquireLock(resourcesLock);
     auto itor = std::find(residentAllocations.begin(), residentAllocations.end(), &gfxAllocation);

shared/source/os_interface/aub_memory_operations_handler.h

@@ -23,6 +23,7 @@ class AubMemoryOperationsHandler : public MemoryOperationsHandler {
     ~AubMemoryOperationsHandler() override = default;
 
     MemoryOperationsStatus makeResident(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) override;
+    MemoryOperationsStatus lock(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) override;
     MemoryOperationsStatus evict(Device *device, GraphicsAllocation &gfxAllocation) override;
     MemoryOperationsStatus isResident(Device *device, GraphicsAllocation &gfxAllocation) override;
 

shared/source/os_interface/linux/drm_buffer_object.h

@@ -217,6 +217,9 @@ class BufferObject {
     std::vector<uint64_t> &getColourAddresses() {
         return this->bindAddresses;
     }
+    void requireExplicitLockedMemory(bool locked) { requiresLocked = locked; }
+    bool isExplicitLockedMemoryRequired() { return requiresLocked; }
+
     uint64_t peekPatIndex() const { return patIndex; }
     void setPatIndex(uint64_t newPatIndex) { this->patIndex = newPatIndex; }
     BOType peekBOType() const { return boType; }
@@ -266,6 +269,7 @@ class BufferObject {
     bool allowCapture = false;
     bool requiresImmediateBinding = false;
     bool requiresExplicitResidency = false;
+    bool requiresLocked = false;
     bool chunked = false;
     bool isReused = false;
     bool readOnlyGpuResource = false;

shared/source/os_interface/linux/drm_memory_operations_handler_bind.cpp

@@ -39,6 +39,13 @@ MemoryOperationsStatus DrmMemoryOperationsHandlerBind::makeResident(Device *devi
     return result;
 }
 
+MemoryOperationsStatus DrmMemoryOperationsHandlerBind::lock(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) {
+    for (auto gfxAllocation = gfxAllocations.begin(); gfxAllocation != gfxAllocations.end(); gfxAllocation++) {
+        (*gfxAllocation)->setLockedMemory(true);
+    }
+    return makeResident(device, gfxAllocations);
+}
+
 MemoryOperationsStatus DrmMemoryOperationsHandlerBind::makeResidentWithinOsContext(OsContext *osContext, ArrayRef<GraphicsAllocation *> gfxAllocations, bool evictable) {
     auto deviceBitfield = osContext->getDeviceBitfield();
 
@@ -59,12 +66,12 @@ MemoryOperationsStatus DrmMemoryOperationsHandlerBind::makeResidentWithinOsConte
             }
 
             if (!bo->getBindInfo()[bo->getOsContextId(osContext)][drmIterator]) {
+                bo->requireExplicitLockedMemory(drmAllocation->isLockedMemory());
                 int result = drmAllocation->makeBOsResident(osContext, drmIterator, nullptr, true);
                 if (result) {
                     return MemoryOperationsStatus::outOfMemory;
                 }
             }
-
             if (!evictable) {
                 drmAllocation->updateResidencyTaskCount(GraphicsAllocation::objectAlwaysResident, osContext->getContextId());
             }
@@ -77,6 +84,7 @@ MemoryOperationsStatus DrmMemoryOperationsHandlerBind::makeResidentWithinOsConte
 MemoryOperationsStatus DrmMemoryOperationsHandlerBind::evict(Device *device, GraphicsAllocation &gfxAllocation) {
     auto &engines = device->getAllEngines();
     auto retVal = MemoryOperationsStatus::success;
+    gfxAllocation.setLockedMemory(false);
     for (const auto &engine : engines) {
         retVal = this->evictWithinOsContext(engine.osContext, gfxAllocation);
         if (retVal != MemoryOperationsStatus::success) {
@@ -185,7 +193,10 @@ MemoryOperationsStatus DrmMemoryOperationsHandlerBind::evictUnusedAllocationsImp
                         evict = false;
                         break;
                     }
-
+                    if (allocation->isLockedMemory()) {
+                        evict = false;
+                        break;
+                    }
                     if (waitForCompletion) {
                         const auto waitStatus = engine.commandStreamReceiver->waitForCompletionWithTimeout(WaitParams{false, false, 0}, engine.commandStreamReceiver->peekLatestFlushedTaskCount());
                         if (waitStatus == WaitStatus::gpuHang) {

shared/source/os_interface/linux/drm_memory_operations_handler_bind.h

@@ -18,6 +18,7 @@ class DrmMemoryOperationsHandlerBind : public DrmMemoryOperationsHandler {
 
     MemoryOperationsStatus makeResidentWithinOsContext(OsContext *osContext, ArrayRef<GraphicsAllocation *> gfxAllocations, bool evictable) override;
     MemoryOperationsStatus makeResident(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) override;
+    MemoryOperationsStatus lock(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) override;
     MemoryOperationsStatus evict(Device *device, GraphicsAllocation &gfxAllocation) override;
     MemoryOperationsStatus evictWithinOsContext(OsContext *osContext, GraphicsAllocation &gfxAllocation) override;
     MemoryOperationsStatus isResident(Device *device, GraphicsAllocation &gfxAllocation) override;

shared/source/os_interface/linux/drm_memory_operations_handler_default.cpp

@@ -8,6 +8,8 @@
 #include "shared/source/os_interface/linux/drm_memory_operations_handler_default.h"
 
 #include "shared/source/debug_settings/debug_settings_manager.h"
+#include "shared/source/os_interface/linux/drm_allocation.h"
+#include "shared/source/os_interface/linux/drm_buffer_object.h"
 
 #include <algorithm>
 
@@ -29,6 +31,18 @@ MemoryOperationsStatus DrmMemoryOperationsHandlerDefault::makeResident(Device *d
     return this->makeResidentWithinOsContext(osContext, gfxAllocations, false);
 }
 
+MemoryOperationsStatus DrmMemoryOperationsHandlerDefault::lock(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) {
+    OsContext *osContext = nullptr;
+    for (auto gfxAllocation = gfxAllocations.begin(); gfxAllocation != gfxAllocations.end(); gfxAllocation++) {
+        auto drmAllocation = static_cast<DrmAllocation *>(*gfxAllocation);
+        drmAllocation->setLockedMemory(true);
+        for (auto bo : drmAllocation->getBOs()) {
+            bo->requireExplicitLockedMemory(true);
+        }
+    }
+    return this->makeResidentWithinOsContext(osContext, gfxAllocations, false);
+}
+
 MemoryOperationsStatus DrmMemoryOperationsHandlerDefault::evictWithinOsContext(OsContext *osContext, GraphicsAllocation &gfxAllocation) {
     std::lock_guard<std::mutex> lock(mutex);
     this->residency.erase(&gfxAllocation);
@@ -37,6 +51,16 @@ MemoryOperationsStatus DrmMemoryOperationsHandlerDefault::evictWithinOsContext(O
 
 MemoryOperationsStatus DrmMemoryOperationsHandlerDefault::evict(Device *device, GraphicsAllocation &gfxAllocation) {
     OsContext *osContext = nullptr;
+    auto drmAllocation = static_cast<DrmAllocation *>(&gfxAllocation);
+    drmAllocation->setLockedMemory(false);
+    if (drmAllocation->storageInfo.isChunked || drmAllocation->storageInfo.getNumBanks() == 1) {
+        auto bo = drmAllocation->getBO();
+        bo->requireExplicitLockedMemory(false);
+    } else {
+        for (auto bo : drmAllocation->getBOs()) {
+            bo->requireExplicitLockedMemory(false);
+        }
+    }
     return this->evictWithinOsContext(osContext, gfxAllocation);
 }
 

shared/source/os_interface/linux/drm_memory_operations_handler_default.h

@@ -21,6 +21,7 @@ class DrmMemoryOperationsHandlerDefault : public DrmMemoryOperationsHandler {
 
     MemoryOperationsStatus makeResidentWithinOsContext(OsContext *osContext, ArrayRef<GraphicsAllocation *> gfxAllocations, bool evictable) override;
     MemoryOperationsStatus makeResident(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) override;
+    MemoryOperationsStatus lock(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) override;
     MemoryOperationsStatus isResident(Device *device, GraphicsAllocation &gfxAllocation) override;
     MemoryOperationsStatus evictWithinOsContext(OsContext *osContext, GraphicsAllocation &gfxAllocation) override;
     MemoryOperationsStatus evict(Device *device, GraphicsAllocation &gfxAllocation) override;

shared/source/os_interface/linux/drm_memory_operations_handler_with_aub_dump.h

@@ -41,6 +41,11 @@ class DrmMemoryOperationsHandlerWithAubDump : public BaseOperationsHandler {
         return BaseOperationsHandler::makeResident(device, gfxAllocations);
     }
 
+    MemoryOperationsStatus lock(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) override {
+        aubMemoryOperationsHandler->makeResident(device, gfxAllocations);
+        return BaseOperationsHandler::lock(device, gfxAllocations);
+    }
+
     MemoryOperationsStatus evict(Device *device, GraphicsAllocation &gfxAllocation) override {
         aubMemoryOperationsHandler->evict(device, gfxAllocation);
         return BaseOperationsHandler::evict(device, gfxAllocation);

shared/source/os_interface/linux/drm_neo.cpp

@@ -1315,7 +1315,8 @@ int changeBufferObjectBinding(Drm *drm, OsContext *osContext, uint32_t vmHandleI
             bindMakeResident = bo->isExplicitResidencyRequired();
             bindImmediate = true;
         }
-        flags |= ioctlHelper->getFlagsForVmBind(bindCapture, bindImmediate, bindMakeResident, readOnlyResource);
+        bool bindLock = bo->isExplicitLockedMemoryRequired();
+        flags |= ioctlHelper->getFlagsForVmBind(bindCapture, bindImmediate, bindMakeResident, bindLock, readOnlyResource);
     }
 
     auto &bindAddresses = bo->getColourAddresses();

shared/source/os_interface/linux/ioctl_helper.h

@@ -116,7 +116,7 @@ class IoctlHelper {
     virtual bool getGemTiling(void *setTiling) = 0;
     virtual uint32_t getDirectSubmissionFlag() = 0;
     virtual std::unique_ptr<uint8_t[]> prepareVmBindExt(const StackVec<uint32_t, 2> &bindExtHandles) = 0;
-    virtual uint64_t getFlagsForVmBind(bool bindCapture, bool bindImmediate, bool bindMakeResident, bool readOnlyResource) = 0;
+    virtual uint64_t getFlagsForVmBind(bool bindCapture, bool bindImmediate, bool bindMakeResident, bool bindLockedMemory, bool readOnlyResource) = 0;
     virtual int queryDistances(std::vector<QueryItem> &queryItems, std::vector<DistanceInfo> &distanceInfos) = 0;
     virtual uint16_t getWaitUserFenceSoftFlag() = 0;
     virtual int execBuffer(ExecBuffer *execBuffer, uint64_t completionGpuAddress, TaskCountType counterValue) = 0;
@@ -259,7 +259,7 @@ class IoctlHelperUpstream : public IoctlHelperI915 {
     bool setVmPrefetch(uint64_t start, uint64_t length, uint32_t region, uint32_t vmId) override;
     uint32_t getDirectSubmissionFlag() override;
     std::unique_ptr<uint8_t[]> prepareVmBindExt(const StackVec<uint32_t, 2> &bindExtHandles) override;
-    uint64_t getFlagsForVmBind(bool bindCapture, bool bindImmediate, bool bindMakeResident, bool readOnlyResource) override;
+    uint64_t getFlagsForVmBind(bool bindCapture, bool bindImmediate, bool bindMakeResident, bool bindLockedMemory, bool readOnlyResource) override;
     int queryDistances(std::vector<QueryItem> &queryItems, std::vector<DistanceInfo> &distanceInfos) override;
     uint16_t getWaitUserFenceSoftFlag() override;
     int execBuffer(ExecBuffer *execBuffer, uint64_t completionGpuAddress, TaskCountType counterValue) override;
@@ -336,7 +336,7 @@ class IoctlHelperPrelim20 : public IoctlHelperI915 {
     bool setVmPrefetch(uint64_t start, uint64_t length, uint32_t region, uint32_t vmId) override;
     uint32_t getDirectSubmissionFlag() override;
     std::unique_ptr<uint8_t[]> prepareVmBindExt(const StackVec<uint32_t, 2> &bindExtHandles) override;
-    uint64_t getFlagsForVmBind(bool bindCapture, bool bindImmediate, bool bindMakeResident, bool readOnlyResource) override;
+    uint64_t getFlagsForVmBind(bool bindCapture, bool bindImmediate, bool bindMakeResident, bool bindLockedMemory, bool readOnlyResource) override;
     int queryDistances(std::vector<QueryItem> &queryItems, std::vector<DistanceInfo> &distanceInfos) override;
     uint16_t getWaitUserFenceSoftFlag() override;
     int execBuffer(ExecBuffer *execBuffer, uint64_t completionGpuAddress, TaskCountType counterValue) override;

shared/source/os_interface/linux/ioctl_helper_prelim.cpp

@@ -507,7 +507,7 @@ std::unique_ptr<uint8_t[]> IoctlHelperPrelim20::prepareVmBindExt(const StackVec<
     return extensionsBuffer;
 }
 
-uint64_t IoctlHelperPrelim20::getFlagsForVmBind(bool bindCapture, bool bindImmediate, bool bindMakeResident, bool readOnlyResource) {
+uint64_t IoctlHelperPrelim20::getFlagsForVmBind(bool bindCapture, bool bindImmediate, bool bindMakeResident, bool bindLockedMemory, bool readOnlyResource) {
     uint64_t flags = 0u;
     if (bindCapture) {
         flags |= PRELIM_I915_GEM_VM_BIND_CAPTURE;
@@ -515,7 +515,7 @@ uint64_t IoctlHelperPrelim20::getFlagsForVmBind(bool bindCapture, bool bindImmed
     if (bindImmediate) {
         flags |= PRELIM_I915_GEM_VM_BIND_IMMEDIATE;
     }
-    if (bindMakeResident) {
+    if (bindMakeResident || bindLockedMemory) { // lockedMemory is equal to residency in i915_prelim
         flags |= PRELIM_I915_GEM_VM_BIND_MAKE_RESIDENT;
     }
     if (readOnlyResource) {

shared/source/os_interface/linux/ioctl_helper_upstream.cpp

@@ -168,7 +168,7 @@ std::unique_ptr<uint8_t[]> IoctlHelperUpstream::prepareVmBindExt(const StackVec<
     return {};
 }
 
-uint64_t IoctlHelperUpstream::getFlagsForVmBind(bool bindCapture, bool bindImmediate, bool bindMakeResident, bool readOnlyResource) {
+uint64_t IoctlHelperUpstream::getFlagsForVmBind(bool bindCapture, bool bindImmediate, bool bindMakeResident, bool bindLockedMemory, bool readOnlyResource) {
     return 0u;
 }
 

shared/source/os_interface/linux/xe/CMakeLists.txt

@@ -12,6 +12,7 @@ set(NEO_CORE_OS_INTERFACE_LINUX_XE
     ${CMAKE_CURRENT_SOURCE_DIR}${BRANCH_DIR_SUFFIX}ioctl_helper_xe_vm_export.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}${BRANCH_DIR_SUFFIX}ioctl_helper_xe_context.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}${BRANCH_DIR_SUFFIX}ioctl_helper_xe_perf.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}${BRANCH_DIR_SUFFIX}ioctl_helper_xe_vm_bind_flag.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/ioctl_helper_xe.h
 )
 

shared/source/os_interface/linux/xe/ioctl_helper_xe.cpp

@@ -730,12 +730,13 @@ bool IoctlHelperXe::completionFenceExtensionSupported(const bool isVmBindAvailab
     return isVmBindAvailable;
 }
 
-uint64_t IoctlHelperXe::getFlagsForVmBind(bool bindCapture, bool bindImmediate, bool bindMakeResident, bool readOnlyResource) {
+uint64_t IoctlHelperXe::getFlagsForVmBind(bool bindCapture, bool bindImmediate, bool bindMakeResident, bool bindLock, bool readOnlyResource) {
     uint64_t ret = 0;
-    xeLog(" -> IoctlHelperXe::%s %d %d %d %d\n", __FUNCTION__, bindCapture, bindImmediate, bindMakeResident, readOnlyResource);
+    xeLog(" -> IoctlHelperXe::%s %d %d %d %d %d\n", __FUNCTION__, bindCapture, bindImmediate, bindMakeResident, bindLock, readOnlyResource);
     if (bindCapture) {
         ret |= DRM_XE_VM_BIND_FLAG_DUMPABLE;
     }
+    ret |= getExtraFlagsForVmBind(bindCapture, bindImmediate, bindMakeResident, bindLock, readOnlyResource);
     return ret;
 }
 

shared/source/os_interface/linux/xe/ioctl_helper_xe.h

@@ -58,7 +58,7 @@ class IoctlHelperXe : public IoctlHelper {
     bool getGemTiling(void *setTiling) override;
     uint32_t getDirectSubmissionFlag() override;
     std::unique_ptr<uint8_t[]> prepareVmBindExt(const StackVec<uint32_t, 2> &bindExtHandles) override;
-    uint64_t getFlagsForVmBind(bool bindCapture, bool bindImmediate, bool bindMakeResident, bool readOnlyResource) override;
+    uint64_t getFlagsForVmBind(bool bindCapture, bool bindImmediate, bool bindMakeResident, bool bindLock, bool readOnlyResource) override;
     int queryDistances(std::vector<QueryItem> &queryItems, std::vector<DistanceInfo> &distanceInfos) override;
     uint16_t getWaitUserFenceSoftFlag() override;
     int execBuffer(ExecBuffer *execBuffer, uint64_t completionGpuAddress, TaskCountType counterValue) override;
@@ -134,6 +134,7 @@ class IoctlHelperXe : public IoctlHelper {
     const char *xeGetBindOperationName(int bindOperation);
     const char *xeGetBindFlagsName(int bindFlags);
 
+    uint64_t getExtraFlagsForVmBind(bool bindCapture, bool bindImmediate, bool bindMakeResident, bool bindLock, bool readOnlyResource);
     const char *xeGetengineClassName(uint32_t engineClass);
     template <typename DataType>
     std::vector<DataType> queryData(uint32_t queryId);

shared/source/os_interface/linux/xe/ioctl_helper_xe_vm_bind_flag.cpp

@@ -0,0 +1,16 @@
+/*
+ * Copyright (C) 2024 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "shared/source/os_interface/linux/xe/ioctl_helper_xe.h"
+
+namespace NEO {
+
+uint64_t IoctlHelperXe::getExtraFlagsForVmBind(bool bindCapture, bool bindImmediate, bool bindMakeResident, bool bindLock, bool readOnlyResource) {
+    return 0;
+}
+
+} // namespace NEO

shared/source/os_interface/windows/wddm_memory_operations_handler.h

@@ -27,6 +27,9 @@ class WddmMemoryOperationsHandler : public MemoryOperationsHandler {
     MemoryOperationsStatus evict(Device *device, GraphicsAllocation &gfxAllocation) override;
     MemoryOperationsStatus isResident(Device *device, GraphicsAllocation &gfxAllocation) override;
 
+    MemoryOperationsStatus lock(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) override {
+        return MemoryOperationsStatus::unsupported;
+    }
     MemoryOperationsStatus makeResidentWithinOsContext(OsContext *osContext, ArrayRef<GraphicsAllocation *> gfxAllocations, bool evictable) override {
         return makeResident(nullptr, gfxAllocations);
     }

shared/test/common/mocks/mock_memory_operations_handler.h

@@ -21,6 +21,7 @@ class MockMemoryOperationsHandler : public MemoryOperationsHandler {
   public:
     MockMemoryOperationsHandler() {}
     MemoryOperationsStatus makeResident(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) override { return MemoryOperationsStatus::unsupported; }
+    MemoryOperationsStatus lock(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) override { return MemoryOperationsStatus::unsupported; }
     MemoryOperationsStatus evict(Device *device, GraphicsAllocation &gfxAllocation) override { return MemoryOperationsStatus::unsupported; }
     MemoryOperationsStatus isResident(Device *device, GraphicsAllocation &gfxAllocation) override { return MemoryOperationsStatus::unsupported; }
     MemoryOperationsStatus makeResidentWithinOsContext(OsContext *osContext, ArrayRef<GraphicsAllocation *> gfxAllocations, bool evictable) override { return MemoryOperationsStatus::unsupported; }
@@ -31,6 +32,7 @@ class MockMemoryOperationsHandlerTests : public MemoryOperationsHandler {
   public:
     MockMemoryOperationsHandlerTests() {}
     ADDMETHOD_NOBASE(makeResident, MemoryOperationsStatus, MemoryOperationsStatus::unsupported, (Device * device, ArrayRef<GraphicsAllocation *> gfxAllocations));
+    ADDMETHOD_NOBASE(lock, MemoryOperationsStatus, MemoryOperationsStatus::unsupported, (Device * device, ArrayRef<GraphicsAllocation *> gfxAllocations));
     ADDMETHOD_NOBASE(evict, MemoryOperationsStatus, MemoryOperationsStatus::unsupported, (Device * device, GraphicsAllocation &gfxAllocation));
     ADDMETHOD_NOBASE(isResident, MemoryOperationsStatus, MemoryOperationsStatus::unsupported, (Device * device, GraphicsAllocation &gfxAllocation));
     ADDMETHOD_NOBASE(makeResidentWithinOsContext, MemoryOperationsStatus, MemoryOperationsStatus::unsupported, (OsContext * osContext, ArrayRef<GraphicsAllocation *> gfxAllocations, bool evictable));
@@ -52,6 +54,12 @@ class MockMemoryOperations : public MemoryOperationsHandler {
         }
         return MemoryOperationsStatus::success;
     }
+
+    MemoryOperationsStatus lock(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) override {
+        lockCalledCount++;
+        return MemoryOperationsStatus::success;
+    }
+
     MemoryOperationsStatus evict(Device *device, GraphicsAllocation &gfxAllocation) override {
         evictCalledCount++;
         if (captureGfxAllocationsForMakeResident) {
@@ -98,6 +106,7 @@ class MockMemoryOperations : public MemoryOperationsHandler {
     int makeResidentCalledCount = 0;
     int evictCalledCount = 0;
     uint32_t isResidentCalledCount = 0;
+    uint32_t lockCalledCount = 0;
     uint32_t makeResidentContextId = std::numeric_limits<uint32_t>::max();
     bool captureGfxAllocationsForMakeResident = false;
 };

shared/test/unit_test/os_interface/aub_memory_operations_handler_tests.cpp

@@ -44,6 +44,29 @@ TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerWhenMakeResidentCalledThe
     EXPECT_EQ(2u, memoryOperationsInterface->residentAllocations.size());
 }
 
+TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerWhenCallingLockThenTrueReturnedAndWriteCalled) {
+    MockAubManager aubManager;
+    getMemoryOperationsHandler()->setAubManager(&aubManager);
+    auto memoryOperationsInterface = getMemoryOperationsHandler();
+    auto result = memoryOperationsInterface->lock(device.get(), ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
+    EXPECT_EQ(result, MemoryOperationsStatus::success);
+    EXPECT_TRUE(aubManager.writeMemory2Called);
+
+    auto itor = std::find(memoryOperationsInterface->residentAllocations.begin(), memoryOperationsInterface->residentAllocations.end(), allocPtr);
+    EXPECT_NE(memoryOperationsInterface->residentAllocations.end(), itor);
+    EXPECT_EQ(1u, memoryOperationsInterface->residentAllocations.size());
+
+    aubManager.writeMemory2Called = false;
+
+    result = memoryOperationsInterface->lock(device.get(), ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
+    EXPECT_EQ(result, MemoryOperationsStatus::success);
+    EXPECT_TRUE(aubManager.writeMemory2Called);
+
+    itor = std::find(memoryOperationsInterface->residentAllocations.begin(), memoryOperationsInterface->residentAllocations.end(), allocPtr);
+    EXPECT_NE(memoryOperationsInterface->residentAllocations.end(), itor);
+    EXPECT_EQ(2u, memoryOperationsInterface->residentAllocations.size());
+}
+
 TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerAndAllocationOfOneTimeAubWritableAllocationTypeWhenMakeResidentCalledTwoTimesThenWriteMemoryOnce) {
     ASSERT_TRUE(AubHelper::isOneTimeAubWritableAllocationType(AllocationType::buffer));
     allocPtr->setAllocationType(AllocationType::buffer);

shared/test/unit_test/os_interface/linux/CMakeLists.txt

@@ -22,12 +22,12 @@ set(NEO_CORE_OS_INTERFACE_TESTS_LINUX
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_mapper_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_manager_bindless_heap_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_manager_tests.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_operations_handler_default_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_operations_handler_with_aub_dump_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_mock_impl.h
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_os_memory_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_pci_speed_info_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_query_topology_upstream_tests.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/drm_residency_handler_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_special_heap_test.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_system_info_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_tests.cpp
@@ -54,9 +54,9 @@ if(NEO_ENABLE_i915_PRELIM_DETECTION)
        ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_info_prelim_tests.cpp
        ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_manager_debug_surface_prelim_tests.cpp
        ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_manager_localmem_prelim_tests.cpp
+       ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_operations_handler_bind_tests.cpp
        ${CMAKE_CURRENT_SOURCE_DIR}/drm_query_prelim_tests.cpp
        ${CMAKE_CURRENT_SOURCE_DIR}/drm_query_topology_prelim_tests.cpp
-       ${CMAKE_CURRENT_SOURCE_DIR}/drm_residency_handler_prelim_tests.cpp
        ${CMAKE_CURRENT_SOURCE_DIR}/drm_vm_bind_prelim_tests.cpp
        ${CMAKE_CURRENT_SOURCE_DIR}/drm_with_prelim_tests.cpp
        ${CMAKE_CURRENT_SOURCE_DIR}/ioctl_helper_tests_prelim.cpp

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

@@ -1324,6 +1324,91 @@ TEST_F(DrmMemoryOperationsHandlerBindTest, givenClosEnabledAndAllocationToBeCach
     memoryManager->freeGraphicsMemory(allocation);
 }
 
+TEST_F(DrmMemoryOperationsHandlerBindTest, whenIoctlFailDuringLockingThenOutOfMemoryIsThrown) {
+    auto allocation = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{device->getRootDeviceIndex(), MemoryConstants::pageSize});
+
+    mock->context.vmBindReturn = -1;
+    EXPECT_EQ(operationHandler->isResident(device, *allocation), MemoryOperationsStatus::memoryNotFound);
+    EXPECT_EQ(operationHandler->lock(device, ArrayRef<GraphicsAllocation *>(&allocation, 1)), MemoryOperationsStatus::outOfMemory);
+    memoryManager->freeGraphicsMemory(allocation);
+}
+
+TEST_F(DrmMemoryOperationsHandlerBindTest, whenLockingDrmAllocationThenBosRequireExplicitLockedMemory) {
+
+    BufferObjects bos;
+    MockBufferObject mockBo1(0, mock, 3, 0, 0, 1), mockBo2(0, mock, 3, 0, 0, 1);
+    mockBo1.setSize(1024);
+    mockBo2.setSize(1024);
+    bos.push_back(&mockBo1);
+    bos.push_back(&mockBo2);
+    GraphicsAllocation *mockDrmAllocation = new MockDrmAllocation(AllocationType::unknown, MemoryPool::localMemory, bos);
+    mockDrmAllocation->storageInfo.memoryBanks = 3;
+    EXPECT_EQ(2u, mockDrmAllocation->storageInfo.getNumBanks());
+
+    mock->context.vmBindReturn = 0;
+    EXPECT_EQ(operationHandler->isResident(device, *mockDrmAllocation), MemoryOperationsStatus::memoryNotFound);
+
+    EXPECT_EQ(operationHandler->lock(device, ArrayRef<GraphicsAllocation *>(&mockDrmAllocation, 1)), MemoryOperationsStatus::success);
+    EXPECT_EQ(operationHandler->isResident(device, *mockDrmAllocation), MemoryOperationsStatus::success);
+    EXPECT_TRUE(mockDrmAllocation->isLockedMemory());
+    EXPECT_TRUE(mockBo1.isExplicitLockedMemoryRequired());
+    EXPECT_TRUE(mockBo2.isExplicitLockedMemoryRequired());
+
+    delete mockDrmAllocation;
+}
+
+TEST_F(DrmMemoryOperationsHandlerBindTest, givenPreviouslyLockedMemoryWhenCallingResidentMemoryThenBosDoNotRequireExplicitLockedMemory) {
+
+    BufferObjects bos;
+    MockBufferObject mockBo(0, mock, 3, 0, 0, 1);
+    mockBo.setSize(1024);
+    bos.push_back(&mockBo);
+    GraphicsAllocation *mockDrmAllocation = new MockDrmAllocation(AllocationType::unknown, MemoryPool::localMemory, bos);
+
+    mock->context.vmBindReturn = 0;
+    mock->context.vmUnbindReturn = 0;
+    EXPECT_EQ(operationHandler->isResident(device, *mockDrmAllocation), MemoryOperationsStatus::memoryNotFound);
+
+    EXPECT_EQ(operationHandler->lock(device, ArrayRef<GraphicsAllocation *>(&mockDrmAllocation, 1)), MemoryOperationsStatus::success);
+    EXPECT_EQ(operationHandler->isResident(device, *mockDrmAllocation), MemoryOperationsStatus::success);
+
+    EXPECT_EQ(operationHandler->evict(device, *mockDrmAllocation), MemoryOperationsStatus::success);
+    EXPECT_EQ(operationHandler->isResident(device, *mockDrmAllocation), MemoryOperationsStatus::memoryNotFound);
+    EXPECT_FALSE(mockDrmAllocation->isLockedMemory());
+
+    EXPECT_EQ(operationHandler->makeResident(device, ArrayRef<GraphicsAllocation *>(&mockDrmAllocation, 1)), MemoryOperationsStatus::success);
+    EXPECT_EQ(operationHandler->isResident(device, *mockDrmAllocation), MemoryOperationsStatus::success);
+    EXPECT_FALSE(mockDrmAllocation->isLockedMemory());
+    EXPECT_FALSE(mockBo.isExplicitLockedMemoryRequired());
+
+    delete mockDrmAllocation;
+}
+
+TEST_F(DrmMemoryOperationsHandlerBindTest, givenLockedAndResidentAllocationsWhenCallingEvictUnusedMemoryThenBothAllocationsAreNotEvicted) {
+    auto allocation1 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{device->getRootDeviceIndex(), MemoryConstants::pageSize});
+    auto allocation2 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{device->getRootDeviceIndex(), MemoryConstants::pageSize});
+
+    EXPECT_EQ(operationHandler->isResident(device, *allocation1), MemoryOperationsStatus::memoryNotFound);
+    EXPECT_EQ(operationHandler->isResident(device, *allocation2), MemoryOperationsStatus::memoryNotFound);
+
+    EXPECT_EQ(operationHandler->lock(device, ArrayRef<GraphicsAllocation *>(&allocation1, 1)), MemoryOperationsStatus::success);
+    EXPECT_EQ(operationHandler->isResident(device, *allocation1), MemoryOperationsStatus::success);
+
+    EXPECT_EQ(operationHandler->makeResident(device, ArrayRef<GraphicsAllocation *>(&allocation2, 1)), MemoryOperationsStatus::success);
+    EXPECT_EQ(operationHandler->isResident(device, *allocation2), MemoryOperationsStatus::success);
+
+    operationHandler->useBaseEvictUnused = true;
+    EXPECT_EQ(operationHandler->evictUnusedCalled, 0u);
+    EXPECT_EQ(operationHandler->evictUnusedAllocations(false, true), MemoryOperationsStatus::success);
+    EXPECT_EQ(operationHandler->evictUnusedCalled, 1u);
+
+    EXPECT_EQ(operationHandler->isResident(device, *allocation1), MemoryOperationsStatus::success);
+    EXPECT_EQ(operationHandler->isResident(device, *allocation2), MemoryOperationsStatus::success);
+
+    memoryManager->freeGraphicsMemory(allocation2);
+    memoryManager->freeGraphicsMemory(allocation1);
+}
+
 using DrmResidencyHandlerTests = ::testing::Test;
 
 HWTEST2_F(DrmResidencyHandlerTests, givenClosIndexAndMemoryTypeWhenAskingForPatIndexThenReturnCorrectValue, IsWithinXeGfxFamily) {

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

@@ -0,0 +1,163 @@
+/*
+ * Copyright (C) 2019-2024 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "shared/source/os_interface/linux/drm_memory_operations_handler_default.h"
+#include "shared/test/common/libult/linux/drm_query_mock.h"
+#include "shared/test/common/mocks/linux/mock_drm_allocation.h"
+#include "shared/test/common/mocks/mock_execution_environment.h"
+#include "shared/test/common/mocks/mock_graphics_allocation.h"
+#include "shared/test/common/test_macros/test.h"
+
+#include <memory>
+
+using namespace NEO;
+
+struct MockDrmMemoryOperationsHandlerDefault : public DrmMemoryOperationsHandlerDefault {
+    using DrmMemoryOperationsHandlerDefault::DrmMemoryOperationsHandlerDefault;
+    using DrmMemoryOperationsHandlerDefault::residency;
+};
+struct DrmMemoryOperationsHandlerBaseTest : public ::testing::Test {
+    void SetUp() override {
+        executionEnvironment = new ExecutionEnvironment;
+        executionEnvironment->prepareRootDeviceEnvironments(1);
+        executionEnvironment->rootDeviceEnvironments[0]->setHwInfoAndInitHelpers(defaultHwInfo.get());
+        executionEnvironment->rootDeviceEnvironments[0]->initGmm();
+        executionEnvironment->calculateMaxOsContextCount();
+        mock = new DrmQueryMock(*executionEnvironment->rootDeviceEnvironments[0]);
+        mock->setBindAvailable();
+
+        drmMemoryOperationsHandler = std::make_unique<MockDrmMemoryOperationsHandlerDefault>(0);
+    }
+    void initializeAllocation(int numBos) {
+        if (!drmAllocation) {
+            for (auto i = 0; i < numBos; i++) {
+                mockBos.push_back(new MockBufferObject(0, mock, 3, 0, 0, 1));
+            }
+            drmAllocation = new MockDrmAllocation(AllocationType::unknown, MemoryPool::localMemory, mockBos);
+            for (auto i = 0; i < numBos; i++) {
+                drmAllocation->storageInfo.memoryBanks[i] = 1;
+            }
+            allocationPtr = drmAllocation;
+        }
+    }
+    void TearDown() override {
+        for (auto i = 0u; i < mockBos.size(); i++) {
+            delete mockBos[i];
+        }
+        mockBos.clear();
+        delete drmAllocation;
+        delete mock;
+        delete executionEnvironment;
+    }
+
+    ExecutionEnvironment *executionEnvironment;
+    DrmQueryMock *mock;
+    BufferObjects mockBos;
+    MockDrmAllocation *drmAllocation = nullptr;
+    GraphicsAllocation *allocationPtr = nullptr;
+    std::unique_ptr<MockDrmMemoryOperationsHandlerDefault> drmMemoryOperationsHandler;
+};
+
+TEST_F(DrmMemoryOperationsHandlerBaseTest, whenMakingAllocationResidentThenAllocationIsResident) {
+    initializeAllocation(1);
+    EXPECT_EQ(1u, drmAllocation->storageInfo.getNumBanks());
+
+    EXPECT_EQ(drmMemoryOperationsHandler->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocationPtr, 1)), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 1u);
+    EXPECT_TRUE(drmMemoryOperationsHandler->residency.find(allocationPtr) != drmMemoryOperationsHandler->residency.end());
+    EXPECT_EQ(drmMemoryOperationsHandler->isResident(nullptr, *allocationPtr), MemoryOperationsStatus::success);
+}
+
+TEST_F(DrmMemoryOperationsHandlerBaseTest, whenEvictingResidentAllocationThenAllocationIsNotResident) {
+    initializeAllocation(1);
+    EXPECT_EQ(1u, drmAllocation->storageInfo.getNumBanks());
+
+    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 0u);
+    EXPECT_EQ(drmMemoryOperationsHandler->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocationPtr, 1)), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandler->isResident(nullptr, *allocationPtr), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 1u);
+    EXPECT_TRUE(drmMemoryOperationsHandler->residency.find(allocationPtr) != drmMemoryOperationsHandler->residency.end());
+    EXPECT_EQ(drmMemoryOperationsHandler->evict(nullptr, *allocationPtr), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandler->isResident(nullptr, *allocationPtr), MemoryOperationsStatus::memoryNotFound);
+    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 0u);
+}
+
+TEST_F(DrmMemoryOperationsHandlerBaseTest, whenLockingAllocationThenAllocationIsResident) {
+    initializeAllocation(2);
+    EXPECT_EQ(2u, drmAllocation->storageInfo.getNumBanks());
+
+    EXPECT_EQ(drmMemoryOperationsHandler->lock(nullptr, ArrayRef<GraphicsAllocation *>(&allocationPtr, 1)), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 1u);
+    EXPECT_TRUE(drmMemoryOperationsHandler->residency.find(drmAllocation) != drmMemoryOperationsHandler->residency.end());
+    EXPECT_EQ(drmMemoryOperationsHandler->isResident(nullptr, *drmAllocation), MemoryOperationsStatus::success);
+    EXPECT_TRUE(drmAllocation->isLockedMemory());
+    EXPECT_TRUE(mockBos[0]->isExplicitLockedMemoryRequired());
+    EXPECT_TRUE(mockBos[1]->isExplicitLockedMemoryRequired());
+}
+
+TEST_F(DrmMemoryOperationsHandlerBaseTest, whenEvictingLockedAllocationThenAllocationIsNotResident) {
+    initializeAllocation(1);
+    EXPECT_EQ(1u, drmAllocation->storageInfo.getNumBanks());
+
+    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 0u);
+
+    EXPECT_EQ(drmMemoryOperationsHandler->lock(nullptr, ArrayRef<GraphicsAllocation *>(&allocationPtr, 1)), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandler->isResident(nullptr, *drmAllocation), MemoryOperationsStatus::success);
+    EXPECT_TRUE(drmMemoryOperationsHandler->residency.find(drmAllocation) != drmMemoryOperationsHandler->residency.end());
+    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 1u);
+    EXPECT_TRUE(drmAllocation->isLockedMemory());
+    EXPECT_TRUE(mockBos[0]->isExplicitLockedMemoryRequired());
+
+    EXPECT_EQ(drmMemoryOperationsHandler->evict(nullptr, *drmAllocation), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandler->isResident(nullptr, *drmAllocation), MemoryOperationsStatus::memoryNotFound);
+    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 0u);
+    EXPECT_FALSE(drmAllocation->isLockedMemory());
+    EXPECT_FALSE(mockBos[0]->isExplicitLockedMemoryRequired());
+}
+
+TEST_F(DrmMemoryOperationsHandlerBaseTest, whenEvictingLockedAllocationWithMultipleBOsThenAllocationIsNotResident) {
+    initializeAllocation(2);
+    EXPECT_EQ(2u, drmAllocation->storageInfo.getNumBanks());
+
+    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 0u);
+
+    EXPECT_EQ(drmMemoryOperationsHandler->lock(nullptr, ArrayRef<GraphicsAllocation *>(&allocationPtr, 1)), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandler->isResident(nullptr, *drmAllocation), MemoryOperationsStatus::success);
+    EXPECT_TRUE(drmMemoryOperationsHandler->residency.find(drmAllocation) != drmMemoryOperationsHandler->residency.end());
+    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 1u);
+    EXPECT_TRUE(drmAllocation->isLockedMemory());
+    EXPECT_TRUE(mockBos[0]->isExplicitLockedMemoryRequired());
+    EXPECT_TRUE(mockBos[1]->isExplicitLockedMemoryRequired());
+
+    EXPECT_EQ(drmMemoryOperationsHandler->evict(nullptr, *drmAllocation), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandler->isResident(nullptr, *drmAllocation), MemoryOperationsStatus::memoryNotFound);
+    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 0u);
+    EXPECT_FALSE(drmAllocation->isLockedMemory());
+    EXPECT_FALSE(mockBos[0]->isExplicitLockedMemoryRequired());
+    EXPECT_FALSE(mockBos[1]->isExplicitLockedMemoryRequired());
+}
+
+TEST_F(DrmMemoryOperationsHandlerBaseTest, whenEvictingLockedAllocationWithChunkedThenAllocationIsNotResident) {
+    initializeAllocation(1);
+    EXPECT_EQ(1u, drmAllocation->storageInfo.getNumBanks());
+    drmAllocation->storageInfo.isChunked = true;
+
+    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 0u);
+
+    EXPECT_EQ(drmMemoryOperationsHandler->lock(nullptr, ArrayRef<GraphicsAllocation *>(&allocationPtr, 1)), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandler->isResident(nullptr, *drmAllocation), MemoryOperationsStatus::success);
+    EXPECT_TRUE(drmMemoryOperationsHandler->residency.find(drmAllocation) != drmMemoryOperationsHandler->residency.end());
+    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 1u);
+    EXPECT_TRUE(drmAllocation->isLockedMemory());
+    EXPECT_TRUE(mockBos[0]->isExplicitLockedMemoryRequired());
+
+    EXPECT_EQ(drmMemoryOperationsHandler->evict(nullptr, *drmAllocation), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandler->isResident(nullptr, *drmAllocation), MemoryOperationsStatus::memoryNotFound);
+    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 0u);
+    EXPECT_FALSE(drmAllocation->isLockedMemory());
+    EXPECT_FALSE(mockBos[0]->isExplicitLockedMemoryRequired());
+}

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

@@ -9,6 +9,8 @@
 #include "shared/source/execution_environment/root_device_environment.h"
 #include "shared/source/os_interface/linux/drm_memory_operations_handler_default.h"
 #include "shared/source/os_interface/linux/drm_memory_operations_handler_with_aub_dump.h"
+#include "shared/test/common/libult/linux/drm_query_mock.h"
+#include "shared/test/common/mocks/linux/mock_drm_allocation.h"
 #include "shared/test/common/mocks/mock_aub_memory_operations_handler.h"
 #include "shared/test/common/mocks/mock_device.h"
 #include "shared/test/common/mocks/mock_graphics_allocation.h"
@@ -59,17 +61,55 @@ TEST_F(DrmMemoryOperationsHandlerWithAubDumpTest, whenMakingAllocationResidentTh
 }
 
 TEST_F(DrmMemoryOperationsHandlerWithAubDumpTest, whenEvictingResidentAllocationThenAllocationIsNotResident) {
+    auto mock = new DrmQueryMock(*device->executionEnvironment->rootDeviceEnvironments[0]);
+    mock->setBindAvailable();
+
+    BufferObjects bos;
+    MockBufferObject mockBo(0, mock, 3, 0, 0, 1);
+    mockBo.setSize(1024);
+    bos.push_back(&mockBo);
+    GraphicsAllocation *mockDrmAllocation = new MockDrmAllocation(AllocationType::unknown, MemoryPool::localMemory, bos);
+
     EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->residency.size(), 0u);
-    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocationPtr, 1)), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&mockDrmAllocation, 1)), MemoryOperationsStatus::success);
-    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, graphicsAllocation), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, *mockDrmAllocation), MemoryOperationsStatus::success);
     EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->residency.size(), 1u);
-    EXPECT_TRUE(drmMemoryOperationsHandlerWithAubDumpMock->residency.find(allocationPtr) != drmMemoryOperationsHandlerWithAubDumpMock->residency.end());
+    EXPECT_TRUE(drmMemoryOperationsHandlerWithAubDumpMock->residency.find(mockDrmAllocation) != drmMemoryOperationsHandlerWithAubDumpMock->residency.end());
-    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->evict(nullptr, graphicsAllocation), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->evict(nullptr, *mockDrmAllocation), MemoryOperationsStatus::success);
-    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, graphicsAllocation), MemoryOperationsStatus::memoryNotFound);
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, *mockDrmAllocation), MemoryOperationsStatus::memoryNotFound);
     EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->residency.size(), 0u);
     EXPECT_TRUE(mockAubMemoryOperationsHandler->makeResidentCalled);
     EXPECT_TRUE(mockAubMemoryOperationsHandler->isResidentCalled);
     EXPECT_TRUE(mockAubMemoryOperationsHandler->evictCalled);
+
+    delete mockDrmAllocation;
+    delete mock;
+}
+
+TEST_F(DrmMemoryOperationsHandlerWithAubDumpTest, whenEvictingLockedAllocationThenAllocationIsNotResident) {
+    auto mock = new DrmQueryMock(*device->executionEnvironment->rootDeviceEnvironments[0]);
+    mock->setBindAvailable();
+
+    BufferObjects bos;
+    MockBufferObject mockBo(0, mock, 3, 0, 0, 1);
+    mockBo.setSize(1024);
+    bos.push_back(&mockBo);
+    GraphicsAllocation *mockDrmAllocation = new MockDrmAllocation(AllocationType::unknown, MemoryPool::localMemory, bos);
+
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->residency.size(), 0u);
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->lock(nullptr, ArrayRef<GraphicsAllocation *>(&mockDrmAllocation, 1)), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, *mockDrmAllocation), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->residency.size(), 1u);
+    EXPECT_TRUE(drmMemoryOperationsHandlerWithAubDumpMock->residency.find(mockDrmAllocation) != drmMemoryOperationsHandlerWithAubDumpMock->residency.end());
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->evict(nullptr, *mockDrmAllocation), MemoryOperationsStatus::success);
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, *mockDrmAllocation), MemoryOperationsStatus::memoryNotFound);
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->residency.size(), 0u);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->makeResidentCalled);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->isResidentCalled);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->evictCalled);
+
+    delete mockDrmAllocation;
+    delete mock;
 }
 
 TEST_F(DrmMemoryOperationsHandlerWithAubDumpTest, whenConstructingDrmMemoryOperationsHandlerWithAubDumpWithoutAubCenterThenAubCenterIsInitialized) {

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

@@ -1,48 +0,0 @@
-/*
- * Copyright (C) 2019-2024 Intel Corporation
- *
- * SPDX-License-Identifier: MIT
- *
- */
-
-#include "shared/source/os_interface/linux/drm_memory_operations_handler_default.h"
-#include "shared/test/common/mocks/mock_graphics_allocation.h"
-#include "shared/test/common/test_macros/test.h"
-
-#include <memory>
-
-using namespace NEO;
-
-struct MockDrmMemoryOperationsHandlerDefault : public DrmMemoryOperationsHandlerDefault {
-    using DrmMemoryOperationsHandlerDefault::DrmMemoryOperationsHandlerDefault;
-    using DrmMemoryOperationsHandlerDefault::residency;
-};
-
-struct DrmMemoryOperationsHandlerBaseTest : public ::testing::Test {
-    void SetUp() override {
-        drmMemoryOperationsHandler = std::make_unique<MockDrmMemoryOperationsHandlerDefault>(0);
-        allocationPtr = &graphicsAllocation;
-    }
-
-    MockGraphicsAllocation graphicsAllocation;
-    GraphicsAllocation *allocationPtr;
-    std::unique_ptr<MockDrmMemoryOperationsHandlerDefault> drmMemoryOperationsHandler;
-};
-
-TEST_F(DrmMemoryOperationsHandlerBaseTest, whenMakingAllocationResidentThenAllocationIsResident) {
-    EXPECT_EQ(drmMemoryOperationsHandler->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocationPtr, 1)), MemoryOperationsStatus::success);
-    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 1u);
-    EXPECT_TRUE(drmMemoryOperationsHandler->residency.find(allocationPtr) != drmMemoryOperationsHandler->residency.end());
-    EXPECT_EQ(drmMemoryOperationsHandler->isResident(nullptr, graphicsAllocation), MemoryOperationsStatus::success);
-}
-
-TEST_F(DrmMemoryOperationsHandlerBaseTest, whenEvictingResidentAllocationThenAllocationIsNotResident) {
-    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 0u);
-    EXPECT_EQ(drmMemoryOperationsHandler->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocationPtr, 1)), MemoryOperationsStatus::success);
-    EXPECT_EQ(drmMemoryOperationsHandler->isResident(nullptr, graphicsAllocation), MemoryOperationsStatus::success);
-    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 1u);
-    EXPECT_TRUE(drmMemoryOperationsHandler->residency.find(allocationPtr) != drmMemoryOperationsHandler->residency.end());
-    EXPECT_EQ(drmMemoryOperationsHandler->evict(nullptr, graphicsAllocation), MemoryOperationsStatus::success);
-    EXPECT_EQ(drmMemoryOperationsHandler->isResident(nullptr, graphicsAllocation), MemoryOperationsStatus::memoryNotFound);
-    EXPECT_EQ(drmMemoryOperationsHandler->residency.size(), 0u);
-}

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

@@ -266,15 +266,16 @@ TEST_F(IoctlPrelimHelperTests, whenGettingFlagsForVmBindThenProperValuesAreRetur
     for (auto &bindCapture : ::testing::Bool()) {
         for (auto &bindImmediate : ::testing::Bool()) {
             for (auto &bindMakeResident : ::testing::Bool()) {
+                for (auto &bindLockedMemory : ::testing::Bool()) {
                     for (auto &readOnlyResource : ::testing::Bool()) {
-                    auto flags = ioctlHelper.getFlagsForVmBind(bindCapture, bindImmediate, bindMakeResident, readOnlyResource);
+                        auto flags = ioctlHelper.getFlagsForVmBind(bindCapture, bindImmediate, bindMakeResident, bindLockedMemory, readOnlyResource);
                         if (bindCapture) {
                             EXPECT_EQ(PRELIM_I915_GEM_VM_BIND_CAPTURE, (flags & PRELIM_I915_GEM_VM_BIND_CAPTURE));
                         }
                         if (bindImmediate) {
                             EXPECT_EQ(PRELIM_I915_GEM_VM_BIND_IMMEDIATE, (flags & PRELIM_I915_GEM_VM_BIND_IMMEDIATE));
                         }
-                    if (bindMakeResident) {
+                        if (bindMakeResident || bindLockedMemory) {
                             EXPECT_EQ(PRELIM_I915_GEM_VM_BIND_MAKE_RESIDENT, (flags & PRELIM_I915_GEM_VM_BIND_MAKE_RESIDENT));
                         }
                         if (readOnlyResource) {
@@ -284,6 +285,7 @@ TEST_F(IoctlPrelimHelperTests, whenGettingFlagsForVmBindThenProperValuesAreRetur
                             EXPECT_FALSE(bindCapture);
                             EXPECT_FALSE(bindImmediate);
                             EXPECT_FALSE(bindMakeResident);
+                            EXPECT_FALSE(bindLockedMemory);
                             EXPECT_FALSE(readOnlyResource);
                         }
                     }
@@ -291,6 +293,7 @@ TEST_F(IoctlPrelimHelperTests, whenGettingFlagsForVmBindThenProperValuesAreRetur
             }
         }
     }
+}
 
 TEST_F(IoctlPrelimHelperTests, givenIoctlHelperisVmBindPatIndexExtSupportedReturnsTrue) {
     ASSERT_EQ(true, ioctlHelper.isVmBindPatIndexExtSupported());

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

@@ -318,14 +318,16 @@ TEST(IoctlHelperUpstreamTest, whenGettingFlagsForVmBindThenZeroIsReturned) {
     for (auto &bindCapture : ::testing::Bool()) {
         for (auto &bindImmediate : ::testing::Bool()) {
             for (auto &bindMakeResident : ::testing::Bool()) {
+                for (auto &bindLock : ::testing::Bool()) {
                     for (auto &readOnlyResource : ::testing::Bool()) {
-                    auto flags = ioctlHelper.getFlagsForVmBind(bindCapture, bindImmediate, bindMakeResident, readOnlyResource);
+                        auto flags = ioctlHelper.getFlagsForVmBind(bindCapture, bindImmediate, bindMakeResident, bindLock, readOnlyResource);
                         EXPECT_EQ(0u, flags);
                     }
                 }
             }
         }
     }
+}
 
 TEST(IoctlHelperUpstreamTest, whenGettingVmBindExtFromHandlesThenNullptrIsReturned) {
     auto executionEnvironment = std::make_unique<MockExecutionEnvironment>();

shared/test/unit_test/os_interface/linux/xe/CMakeLists.txt

@@ -9,6 +9,7 @@ set(NEO_CORE_OS_INTERFACE_TESTS_LINUX_XE
     ${CMAKE_CURRENT_SOURCE_DIR}/ioctl_helper_xe_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/ioctl_helper_xe_vm_export_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/ioctl_helper_xe_perf_tests.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/${BRANCH_TYPE}/ioctl_helper_xe_vm_bind_flag_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/ioctl_helper_xe_tests.h
 )
 if(NEO_ENABLE_XE_EU_DEBUG_SUPPORT)

shared/test/unit_test/os_interface/linux/xe/ioctl_helper_xe_tests.cpp

@@ -261,7 +261,8 @@ TEST(IoctlHelperXeTest, givenIoctlHelperXeWhenCallingAnyMethodThenDummyValueIsRe
 
     EXPECT_EQ(0u, xeIoctlHelper->getDirectSubmissionFlag());
 
-    EXPECT_EQ(0u, xeIoctlHelper->getFlagsForVmBind(false, false, false, false));
+    EXPECT_EQ(0u, xeIoctlHelper->getFlagsForVmBind(false, false, false, false, false));
+
     std::vector<QueryItem> queryItems;
     std::vector<DistanceInfo> distanceInfos;
     EXPECT_EQ(0, xeIoctlHelper->queryDistances(queryItems, distanceInfos));
@@ -418,8 +419,8 @@ TEST(IoctlHelperXeTest, whenGettingFlagsForVmBindThenPropertValueIsReturned) {
     DrmMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
     auto xeIoctlHelper = std::make_unique<IoctlHelperXe>(drm);
 
-    EXPECT_EQ(static_cast<uint64_t>(DRM_XE_VM_BIND_FLAG_DUMPABLE), xeIoctlHelper->getFlagsForVmBind(true, true, true, true));
+    EXPECT_EQ(static_cast<uint64_t>(DRM_XE_VM_BIND_FLAG_DUMPABLE), xeIoctlHelper->getFlagsForVmBind(true, false, false, false, false));
-    EXPECT_EQ(static_cast<uint64_t>(0), xeIoctlHelper->getFlagsForVmBind(false, true, true, true));
+    EXPECT_EQ(static_cast<uint64_t>(0), xeIoctlHelper->getFlagsForVmBind(false, false, false, false, false));
 }
 
 TEST(IoctlHelperXeTest, whenGettingIoctlRequestValueThenPropertValueIsReturned) {

shared/test/unit_test/os_interface/linux/xe/ioctl_helper_xe_vm_bind_flag_tests.cpp

@@ -0,0 +1,35 @@
+/*
+ * Copyright (C) 2024 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "shared/test/unit_test/os_interface/linux/xe/ioctl_helper_xe_tests.h"
+
+using namespace NEO;
+
+TEST(IoctlHelperXeTest, givenIoctlHelperXeWhenCallingGetFlagsForVmBindThenExpectedValueIsReturned) {
+    auto executionEnvironment = std::make_unique<MockExecutionEnvironment>();
+    DrmMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
+    auto xeIoctlHelper = std::make_unique<IoctlHelperXe>(drm);
+    ASSERT_NE(nullptr, xeIoctlHelper);
+
+    for (auto &bindCapture : ::testing::Bool()) {
+        for (auto &bindImmediate : ::testing::Bool()) {
+            for (auto &bindMakeResident : ::testing::Bool()) {
+                for (auto &bindLockedMemory : ::testing::Bool()) {
+                    for (auto &readOnlyResource : ::testing::Bool()) {
+                        auto flags = xeIoctlHelper->getFlagsForVmBind(bindCapture, bindImmediate, bindMakeResident, bindLockedMemory, readOnlyResource);
+                        if (bindCapture) {
+                            EXPECT_EQ(static_cast<uint32_t>(DRM_XE_VM_BIND_FLAG_DUMPABLE), (flags & DRM_XE_VM_BIND_FLAG_DUMPABLE));
+                        }
+                        if (flags == 0) {
+                            EXPECT_FALSE(bindCapture);
+                        }
+                    }
+                }
+            }
+        }
+    }
+}

shared/test/unit_test/os_interface/windows/wddm_memory_manager_tests.cpp

@@ -844,6 +844,23 @@ TEST_F(WddmMemoryManagerSimpleTest, givenAllocationDataWithFlagsWhenAllocateGrap
     memoryManager->freeGraphicsMemory(allocation);
 }
 
+TEST_F(WddmMemoryManagerSimpleTest, givenAllocationWhenCallingSetLockedMemoryThenFlagIsSetCorrectly) {
+    class MockGraphicsAllocation : public GraphicsAllocation {
+      public:
+        using GraphicsAllocation::allocationInfo;
+    };
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
+    AllocationData allocationData;
+    auto allocation = static_cast<MockGraphicsAllocation *>(memoryManager->allocateGraphicsMemory64kb(allocationData));
+    EXPECT_NE(nullptr, allocation);
+    EXPECT_FALSE(allocation->allocationInfo.flags.lockedMemory);
+    allocation->setLockedMemory(true);
+    EXPECT_TRUE(allocation->allocationInfo.flags.lockedMemory);
+    EXPECT_TRUE(allocation->isLockedMemory());
+
+    memoryManager->freeGraphicsMemory(allocation);
+}
+
 TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWhenAllocateGraphicsMemoryWithPtrIsCalledThenMemoryPoolIsSystem4KBPages) {
     memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     if (memoryManager->isLimitedGPU(0)) {

shared/test/unit_test/os_interface/windows/wddm_memory_operations_handler_with_aub_dump_tests.cpp

@@ -122,3 +122,7 @@ TEST_F(WddmMemoryOperationsHandlerWithAubDumpTest, whenConstructingWddmMemoryOpe
     auto wddmMemoryOperationsHandlerWithAubDump = std::make_unique<WddmMemoryOperationsHandlerWithAubDump<WddmMemoryOperationsHandler>>(wddm, *rootDeviceEnvironment);
     EXPECT_NE(nullptr, rootDeviceEnvironment->aubCenter.get());
 }
+
+TEST_F(WddmMemoryOperationsHandlerWithAubDumpTest, givenRegularAllocationWhenLockingAllocationThenUnsupportIsReturned) {
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->lock(nullptr, ArrayRef<GraphicsAllocation *>(&allocationPtr, 1)), MemoryOperationsStatus::unsupported);
+}