Use different DRM Context for each OsContext on Linux

Change-Id: I543df4accdeba6c69b7dcf86d4238d12dafe92fe Signed-off-by: Dunajski, Bartosz <bartosz.dunajski@intel.com>
2025-12-24 12:23:05 +08:00 · 2018-12-11 08:21:56 +01:00
parent 9fcc06f0c5
commit cfafe943eb
18 changed files with 311 additions and 178 deletions
--- a/2
+++ b/2
@@ -1,5 +1,5 @@
 #!groovy
 dependenciesRevision='0069cbdd1559b47c9de98ca934b334b12fa59703-1174'
 strategy='EQUAL'
-allowedCD=274
+allowedCD=275
 allowedF=4
--- a/runtime/dll/linux/drm_neo_create.cpp
+++ b/runtime/dll/linux/drm_neo_create.cpp
@@ -28,9 +28,6 @@ const DeviceDescriptor deviceDescriptorTable[] = {
 static std::array<Drm *, 1> drms = {{nullptr}};

 Drm::~Drm() {
-    if (lowPriorityContextId != 0) {
-        destroyLowPriorityContext();
-    }
    close(fd);
    fd = -1;
 }
--- a/runtime/os_interface/linux/drm_buffer_object.cpp
+++ b/runtime/os_interface/linux/drm_buffer_object.cpp
@@ -106,7 +106,7 @@ bool BufferObject::setTiling(uint32_t mode, uint32_t stride) {
    return set_tiling.tiling_mode == mode;
 }

-void BufferObject::fillExecObject(drm_i915_gem_exec_object2 &execObject) {
+void BufferObject::fillExecObject(drm_i915_gem_exec_object2 &execObject, uint32_t drmContextId) {
    execObject.handle = this->handle;
    execObject.relocation_count = 0; //No relocations, we are SoftPinning
    execObject.relocs_ptr = 0ul;
@@ -116,23 +116,23 @@ void BufferObject::fillExecObject(drm_i915_gem_exec_object2 &execObject) {
 #ifdef __x86_64__
    execObject.flags |= reinterpret_cast<uint64_t>(this->address) & MemoryConstants::zoneHigh ? EXEC_OBJECT_SUPPORTS_48B_ADDRESS : 0;
 #endif
-    execObject.rsvd1 = this->drm->lowPriorityContextId;
+    execObject.rsvd1 = drmContextId;
    execObject.rsvd2 = 0;
 }

-void BufferObject::processRelocs(int &idx) {
+void BufferObject::processRelocs(int &idx, uint32_t drmContextId) {
    for (size_t i = 0; i < this->residency.size(); i++) {
-        residency[i]->fillExecObject(execObjectsStorage[idx]);
+        residency[i]->fillExecObject(execObjectsStorage[idx], drmContextId);
        idx++;
    }
 }

-int BufferObject::exec(uint32_t used, size_t startOffset, unsigned int flags, bool requiresCoherency, bool lowPriority) {
+int BufferObject::exec(uint32_t used, size_t startOffset, unsigned int flags, bool requiresCoherency, uint32_t drmContextId) {
    drm_i915_gem_execbuffer2 execbuf = {};

    int idx = 0;
-    processRelocs(idx);
-    this->fillExecObject(execObjectsStorage[idx]);
+    processRelocs(idx, drmContextId);
+    this->fillExecObject(execObjectsStorage[idx], drmContextId);
    idx++;

    execbuf.buffers_ptr = reinterpret_cast<uintptr_t>(execObjectsStorage);
@@ -140,10 +140,7 @@ int BufferObject::exec(uint32_t used, size_t startOffset, unsigned int flags, bo
    execbuf.batch_start_offset = static_cast<uint32_t>(startOffset);
    execbuf.batch_len = alignUp(used, 8);
    execbuf.flags = flags;
-
-    if (lowPriority) {
-        execbuf.rsvd1 = this->drm->lowPriorityContextId & I915_EXEC_CONTEXT_ID_MASK;
-    }
+    execbuf.rsvd1 = drmContextId;

    int ret = this->drm->ioctl(DRM_IOCTL_I915_GEM_EXECBUFFER2, &execbuf);
    if (ret != 0) {
@@ -155,7 +152,7 @@ int BufferObject::exec(uint32_t used, size_t startOffset, unsigned int flags, bo
    return ret;
 }

-int BufferObject::pin(BufferObject *boToPin[], size_t numberOfBos) {
+int BufferObject::pin(BufferObject *boToPin[], size_t numberOfBos, uint32_t drmContextId) {
    drm_i915_gem_execbuffer2 execbuf = {};
    StackVec<drm_i915_gem_exec_object2, maxFragmentsCount + 1> execObject;

@@ -166,10 +163,10 @@ int BufferObject::pin(BufferObject *boToPin[], size_t numberOfBos) {

    uint32_t boIndex = 0;
    for (boIndex = 0; boIndex < (uint32_t)numberOfBos; boIndex++) {
-        boToPin[boIndex]->fillExecObject(execObject[boIndex]);
+        boToPin[boIndex]->fillExecObject(execObject[boIndex], drmContextId);
    }

-    this->fillExecObject(execObject[boIndex]);
+    this->fillExecObject(execObject[boIndex], drmContextId);

    execbuf.buffers_ptr = reinterpret_cast<uintptr_t>(&execObject[0]);
    execbuf.buffer_count = boIndex + 1;
--- a/runtime/os_interface/linux/drm_buffer_object.h
+++ b/runtime/os_interface/linux/drm_buffer_object.h
@@ -44,9 +44,9 @@ class BufferObject {

    bool setTiling(uint32_t mode, uint32_t stride);

-    MOCKABLE_VIRTUAL int pin(BufferObject *boToPin[], size_t numberOfBos);
+    MOCKABLE_VIRTUAL int pin(BufferObject *boToPin[], size_t numberOfBos, uint32_t drmContextId);

-    int exec(uint32_t used, size_t startOffset, unsigned int flags, bool requiresCoherency = false, bool lowPriority = false);
+    int exec(uint32_t used, size_t startOffset, unsigned int flags, bool requiresCoherency, uint32_t drmContextId);

    int wait(int64_t timeoutNs);
    bool close();
@@ -94,8 +94,8 @@ class BufferObject {
    uint32_t tiling_mode;
    uint32_t stride;

-    MOCKABLE_VIRTUAL void fillExecObject(drm_i915_gem_exec_object2 &execObject);
-    void processRelocs(int &idx);
+    MOCKABLE_VIRTUAL void fillExecObject(drm_i915_gem_exec_object2 &execObject, uint32_t drmContextId);
+    void processRelocs(int &idx, uint32_t drmContextId);

    uint64_t offset64; // last-seen GPU offset
    size_t size;
--- a/runtime/os_interface/linux/drm_command_stream.inl
+++ b/runtime/os_interface/linux/drm_command_stream.inl
@@ -68,7 +68,7 @@ FlushStamp DrmCommandStreamReceiver<GfxFamily>::flush(BatchBuffer &batchBuffer,
        bb->exec(static_cast<uint32_t>(alignUp(batchBuffer.usedSize - batchBuffer.startOffset, 8)),
                 alignedStart, engineFlag | I915_EXEC_NO_RELOC,
                 batchBuffer.requiresCoherency,
-                 batchBuffer.low_priority);
+                 osContext->get()->getDrmContextId());

        bb->getResidency()->clear();

--- a/runtime/os_interface/linux/drm_memory_manager.cpp
+++ b/runtime/os_interface/linux/drm_memory_manager.cpp
@@ -5,6 +5,7 @@
 *
 */

+#include "runtime/command_stream/command_stream_receiver.h"
 #include "runtime/device/device.h"
 #include "runtime/helpers/ptr_math.h"
 #include "runtime/helpers/options.h"
@@ -13,6 +14,7 @@
 #include "runtime/os_interface/linux/drm_allocation.h"
 #include "runtime/os_interface/linux/drm_buffer_object.h"
 #include "runtime/os_interface/linux/drm_memory_manager.h"
+#include "runtime/os_interface/linux/os_context_linux.h"
 #include "runtime/helpers/surface_formats.h"
 #include <cstring>
 #include <iostream>
@@ -224,7 +226,7 @@ DrmAllocation *DrmMemoryManager::allocateGraphicsMemoryWithAlignment(const Alloc

    bo->isAllocated = true;
    if (forcePinEnabled && pinBB != nullptr && allocationData.flags.forcePin && allocationData.size >= this->pinThreshold) {
-        pinBB->pin(&bo, 1);
+        pinBB->pin(&bo, 1, getDefaultCommandStreamReceiver(0)->getOsContext().get()->getDrmContextId());
    }
    return new DrmAllocation(bo, res, cSize, MemoryPool::System4KBPages, getOsContextCount(), false);
 }
@@ -235,7 +237,7 @@ DrmAllocation *DrmMemoryManager::allocateGraphicsMemoryWithHostPtr(const Allocat
    bool forcePinAllowed = res != nullptr && pinBB != nullptr && forcePinEnabled && allocationData.flags.forcePin && allocationData.size >= this->pinThreshold;
    if (!validateHostPtrMemory && forcePinAllowed) {
        BufferObject *boArray[] = {res->getBO()};
-        pinBB->pin(boArray, 1);
+        pinBB->pin(boArray, 1, getDefaultCommandStreamReceiver(0)->getOsContext().get()->getDrmContextId());
    }
    return res;
 }
@@ -578,7 +580,7 @@ MemoryManager::AllocationStatus DrmMemoryManager::populateOsHandles(OsHandleStor
    }

    if (validateHostPtrMemory) {
-        int result = pinBB->pin(allocatedBos, numberOfBosAllocated);
+        int result = pinBB->pin(allocatedBos, numberOfBosAllocated, getDefaultCommandStreamReceiver(0)->getOsContext().get()->getDrmContextId());

        if (result == EFAULT) {
            for (uint32_t i = 0; i < numberOfBosAllocated; i++) {
--- a/runtime/os_interface/linux/drm_neo.cpp
+++ b/runtime/os_interface/linux/drm_neo.cpp
@@ -137,24 +137,27 @@ void Drm::checkPreemptionSupport() {
 #endif
 }

-void Drm::createLowPriorityContext() {
-    drm_i915_gem_context_create gcc = {};
-    auto retVal = ioctl(DRM_IOCTL_I915_GEM_CONTEXT_CREATE, &gcc);
-    UNRECOVERABLE_IF(retVal != 0);
-    lowPriorityContextId = gcc.ctx_id;
-
+void Drm::setLowPriorityContextParam(uint32_t drmContextId) {
    drm_i915_gem_context_param gcp = {};
-    gcp.ctx_id = lowPriorityContextId;
+    gcp.ctx_id = drmContextId;
    gcp.param = I915_CONTEXT_PARAM_PRIORITY;
    gcp.value = -1023;

-    retVal = ioctl(DRM_IOCTL_I915_GEM_CONTEXT_SETPARAM, &gcp);
+    auto retVal = ioctl(DRM_IOCTL_I915_GEM_CONTEXT_SETPARAM, &gcp);
    UNRECOVERABLE_IF(retVal != 0);
 }

-void Drm::destroyLowPriorityContext() {
+uint32_t Drm::createDrmContext() {
+    drm_i915_gem_context_create gcc = {};
+    auto retVal = ioctl(DRM_IOCTL_I915_GEM_CONTEXT_CREATE, &gcc);
+    UNRECOVERABLE_IF(retVal != 0);
+
+    return gcc.ctx_id;
+}
+
+void Drm::destroyDrmContext(uint32_t drmContextId) {
    drm_i915_gem_context_destroy destroy = {};
-    destroy.ctx_id = lowPriorityContextId;
+    destroy.ctx_id = drmContextId;
    auto retVal = ioctl(DRM_IOCTL_I915_GEM_CONTEXT_DESTROY, &destroy);
    UNRECOVERABLE_IF(retVal != 0);
 }
--- a/runtime/os_interface/linux/drm_neo.h
+++ b/runtime/os_interface/linux/drm_neo.h
@@ -37,7 +37,6 @@ class Drm {
    friend DeviceFactory;

  public:
-    uint32_t lowPriorityContextId;
    static Drm *get(int32_t deviceOrdinal);

    virtual int ioctl(unsigned long request, void *arg);
@@ -57,8 +56,9 @@ class Drm {
    bool isPreemptionSupported() const { return preemptionSupported; }
    MOCKABLE_VIRTUAL void checkPreemptionSupport();
    int getFileDescriptor() const { return fd; }
-    void createLowPriorityContext();
-    void destroyLowPriorityContext();
+    uint32_t createDrmContext();
+    void destroyDrmContext(uint32_t drmContextId);
+    void setLowPriorityContextParam(uint32_t drmContextId);

    void setGtType(GTTYPE eGtType) { this->eGtType = eGtType; }
    GTTYPE getGtType() const { return this->eGtType; }
@@ -73,7 +73,7 @@ class Drm {
    int deviceId;
    int revisionId;
    GTTYPE eGtType;
-    Drm(int fd) : lowPriorityContextId(0), fd(fd), deviceId(0), revisionId(0), eGtType(GTTYPE_UNDEFINED) {}
+    Drm(int fd) : fd(fd), deviceId(0), revisionId(0), eGtType(GTTYPE_UNDEFINED) {}
    virtual ~Drm();

    static bool isi915Version(int fd);
--- a/runtime/os_interface/linux/hw_info_config.cpp
+++ b/runtime/os_interface/linux/hw_info_config.cpp
@@ -162,9 +162,6 @@ int HwInfoConfig::configureHwInfo(const HardwareInfo *inHwInfo, HardwareInfo *ou

    drm->checkPreemptionSupport();
    bool preemption = drm->isPreemptionSupported();
-    if (preemption) {
-        drm->createLowPriorityContext();
-    }
    preemption = hwHelper.setupPreemptionRegisters(outHwInfo, preemption);
    PreemptionHelper::adjustDefaultPreemptionMode(outHwInfo->capabilityTable,
                                                  static_cast<bool>(outHwInfo->pSkuTable->ftrGpGpuMidThreadLevelPreempt) && preemption,
--- a/runtime/os_interface/linux/os_context_linux.cpp
+++ b/runtime/os_interface/linux/os_context_linux.cpp
@@ -24,5 +24,15 @@ OsContext::~OsContext() = default;

 OsContextLinux::OsContextImpl(Drm &drm, EngineInstanceT engineType) : drm(drm) {
    engineFlag = DrmEngineMapper::engineNodeMap(engineType.type);
+    this->drmContextId = drm.createDrmContext();
+    if (drm.isPreemptionSupported() &&
+        engineType.type == gpgpuEngineInstances[EngineInstanceConstants::lowPriorityGpgpuEngineIndex].type &&
+        engineType.id == gpgpuEngineInstances[EngineInstanceConstants::lowPriorityGpgpuEngineIndex].id) {
+        drm.setLowPriorityContextParam(this->drmContextId);
+    }
+}
+
+OsContextLinux::~OsContextImpl() {
+    drm.destroyDrmContext(drmContextId);
 }
 } // namespace OCLRT
--- a/runtime/os_interface/linux/os_context_linux.h
+++ b/runtime/os_interface/linux/os_context_linux.h
@@ -13,11 +13,14 @@ using OsContextLinux = OsContext::OsContextImpl;

 class OsContext::OsContextImpl {
  public:
+    ~OsContextImpl();
    OsContextImpl(Drm &drm, EngineInstanceT engineType);
    unsigned int getEngineFlag() const { return engineFlag; }
+    uint32_t getDrmContextId() const { return drmContextId; }

  protected:
    unsigned int engineFlag = 0;
+    uint32_t drmContextId = 0;
    Drm &drm;
 };
 } // namespace OCLRT
--- a/unit_tests/linux/main_linux_dll.cpp
+++ b/unit_tests/linux/main_linux_dll.cpp
@@ -282,7 +282,7 @@ TEST_F(DrmTests, failOnContextCreate) {
    auto drm = DrmWrap::createDrm(0);
    EXPECT_NE(drm, nullptr);
    failOnContextCreate = -1;
-    EXPECT_THROW(drm->createLowPriorityContext(), std::exception);
+    EXPECT_THROW(drm->createDrmContext(), std::exception);
    EXPECT_FALSE(drm->isPreemptionSupported());
    failOnContextCreate = 0;
    DrmWrap::closeDevice(0);
@@ -295,7 +295,8 @@ TEST_F(DrmTests, failOnSetPriority) {
    auto drm = DrmWrap::createDrm(0);
    EXPECT_NE(drm, nullptr);
    failOnSetPriority = -1;
-    EXPECT_THROW(drm->createLowPriorityContext(), std::exception);
+    auto drmContext = drm->createDrmContext();
+    EXPECT_THROW(drm->setLowPriorityContextParam(drmContext), std::exception);
    EXPECT_FALSE(drm->isPreemptionSupported());
    failOnSetPriority = 0;
    DrmWrap::closeDevice(0);
--- a/unit_tests/os_interface/linux/device_command_stream_fixture.h
+++ b/unit_tests/os_interface/linux/device_command_stream_fixture.h
@@ -10,6 +10,7 @@
 #include "drm/i915_drm.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include "engine_node.h"
 #include "runtime/helpers/aligned_memory.h"
 #include "runtime/os_interface/linux/drm_memory_manager.h"
 #include "runtime/os_interface/linux/drm_neo.h"
@@ -81,6 +82,8 @@ class DrmMockCustom : public Drm {
            gemClose = 0;
            regRead = 0;
            contextGetParam = 0;
+            contextCreate = 0;
+            contextDestroy = 0;
        }

        std::atomic<int32_t> total;
@@ -95,6 +98,8 @@ class DrmMockCustom : public Drm {
        std::atomic<int32_t> gemClose;
        std::atomic<int32_t> regRead;
        std::atomic<int32_t> contextGetParam;
+        std::atomic<int32_t> contextCreate;
+        std::atomic<int32_t> contextDestroy;
    };

    std::atomic<int> ioctl_res;
@@ -121,6 +126,8 @@ class DrmMockCustom : public Drm {
        NEO_IOCTL_EXPECT_EQ(gemClose);
        NEO_IOCTL_EXPECT_EQ(regRead);
        NEO_IOCTL_EXPECT_EQ(contextGetParam);
+        NEO_IOCTL_EXPECT_EQ(contextCreate);
+        NEO_IOCTL_EXPECT_EQ(contextDestroy);
 #undef NEO_IOCTL_EXPECT_EQ
    }

@@ -231,6 +238,13 @@ class DrmMockCustom : public Drm {
            getContextParam->value = getContextParamRetValue;
        } break;

+        case DRM_IOCTL_I915_GEM_CONTEXT_CREATE: {
+            ioctl_cnt.contextCreate++;
+        } break;
+        case DRM_IOCTL_I915_GEM_CONTEXT_DESTROY: {
+            ioctl_cnt.contextDestroy++;
+        } break;
+
        default:
            std::cout << std::hex << DRM_IOCTL_I915_GEM_WAIT << std::endl;
            std::cout << "unexpected IOCTL: " << std::hex << request << std::endl;
@@ -255,6 +269,8 @@ class DrmMockCustom : public Drm {

    DrmMockCustom() : Drm(mockFd) {
        reset();
+        ioctl_expected.contextCreate = OCLRT::EngineInstanceConstants::numGpgpuEngineInstances;
+        ioctl_expected.contextDestroy = ioctl_expected.contextCreate.load();
    }
    int getErrno() override {
        return errnoValue;
--- a/unit_tests/os_interface/linux/drm_buffer_object_tests.cpp
+++ b/unit_tests/os_interface/linux/drm_buffer_object_tests.cpp
@@ -24,8 +24,8 @@ class TestedBufferObject : public BufferObject {
        this->tiling_mode = mode;
    }

-    void fillExecObject(drm_i915_gem_exec_object2 &execObject) override {
-        BufferObject::fillExecObject(execObject);
+    void fillExecObject(drm_i915_gem_exec_object2 &execObject, uint32_t drmContextId) override {
+        BufferObject::fillExecObject(execObject, drmContextId);
        execObjectPointerFilled = &execObject;
    }

@@ -63,7 +63,7 @@ TEST_F(DrmBufferObjectTest, exec) {
    mock->ioctl_expected.total = 1;
    mock->ioctl_res = 0;

-    auto ret = bo->exec(0, 0, 0);
+    auto ret = bo->exec(0, 0, 0, false, 1);
    EXPECT_EQ(mock->ioctl_res, ret);
    EXPECT_EQ(0u, mock->execBuffer.flags);
 }
@@ -71,7 +71,7 @@ TEST_F(DrmBufferObjectTest, exec) {
 TEST_F(DrmBufferObjectTest, exec_ioctlFailed) {
    mock->ioctl_expected.total = 1;
    mock->ioctl_res = -1;
-    EXPECT_THROW(bo->exec(0, 0, 0), std::exception);
+    EXPECT_THROW(bo->exec(0, 0, 0, false, 1), std::exception);
 }

 TEST_F(DrmBufferObjectTest, setTiling_success) {
@@ -100,13 +100,13 @@ TEST_F(DrmBufferObjectTest, testExecObjectFlags) {
 #ifdef __x86_64__
    memset(&execObject, 0, sizeof(execObject));
    bo->setAddress((void *)((uint64_t)1u << 34)); //anything above 4GB
-    bo->fillExecObject(execObject);
+    bo->fillExecObject(execObject, 1);
    EXPECT_TRUE(execObject.flags & EXEC_OBJECT_SUPPORTS_48B_ADDRESS);
 #endif

    memset(&execObject, 0, sizeof(execObject));
    bo->setAddress((void *)((uint64_t)1u << 31)); //anything below 4GB
-    bo->fillExecObject(execObject);
+    bo->fillExecObject(execObject, 1);
    EXPECT_FALSE(execObject.flags & EXEC_OBJECT_SUPPORTS_48B_ADDRESS);
 }

@@ -122,7 +122,7 @@ TEST_F(DrmBufferObjectTest, onPinIoctlFailed) {

    bo->setAddress(buff.get());
    BufferObject *boArray[1] = {boToPin.get()};
-    auto ret = bo->pin(boArray, 1);
+    auto ret = bo->pin(boArray, 1, 1);
    EXPECT_EQ(EINVAL, ret);
 }

@@ -145,7 +145,7 @@ TEST(DrmBufferObjectSimpleTest, givenInvalidBoWhenPinIsCalledThenErrorIsReturned
    mock->errnoValue = EFAULT;

    BufferObject *boArray[1] = {boToPin.get()};
-    auto ret = bo->pin(boArray, 1);
+    auto ret = bo->pin(boArray, 1, 1);
    EXPECT_EQ(EFAULT, ret);
 }

@@ -170,7 +170,7 @@ TEST(DrmBufferObjectSimpleTest, givenArrayOfBosWhenPinnedThenAllBosArePinned) {
    BufferObject *array[3] = {boToPin.get(), boToPin2.get(), boToPin3.get()};

    bo->setAddress(buff.get());
-    auto ret = bo->pin(array, 3);
+    auto ret = bo->pin(array, 3, 1);
    EXPECT_EQ(mock->ioctl_res, ret);
    uint32_t bb_end = 0x05000000;
    EXPECT_EQ(buff[0], bb_end);
--- a/unit_tests/os_interface/linux/drm_command_stream_tests.cpp
+++ b/unit_tests/os_interface/linux/drm_command_stream_tests.cpp
@@ -36,7 +36,7 @@ class DrmCommandStreamFixture {
        //make sure this is disabled, we don't want test this now
        DebugManager.flags.EnableForcePin.set(false);

-        mock = std::make_unique<DrmMockImpl>(mockFd);
+        mock = std::make_unique<::testing::NiceMock<DrmMockImpl>>(mockFd);

        executionEnvironment.osInterface = std::make_unique<OSInterface>();
        executionEnvironment.osInterface->get()->setDrm(mock.get());
@@ -73,7 +73,7 @@ class DrmCommandStreamFixture {

    DeviceCommandStreamReceiver<DEFAULT_TEST_FAMILY_NAME> *csr = nullptr;
    DrmMemoryManager *memoryManager = nullptr;
-    std::unique_ptr<DrmMockImpl> mock;
+    std::unique_ptr<::testing::NiceMock<DrmMockImpl>> mock;
    const int mockFd = 33;
    static const uint64_t alignment = MemoryConstants::allocationAlignment;
    DebugManagerStateRestore dbgState;
@@ -170,6 +170,21 @@ MATCHER_P2(BoExecFlushEq, batch_start_offset, batch_len, "") {
    return (exec2->batch_start_offset == batch_start_offset) && (exec2->batch_len == batch_len);
 }

+// matcher to check DrmContextId
+MATCHER_P2(BoExecFlushContextEq, drmContextId, numExecs, "") {
+    auto execBuff2 = reinterpret_cast<drm_i915_gem_execbuffer2 *>(arg);
+
+    bool allExecsWithTheSameId = (execBuff2->buffer_count == numExecs);
+    allExecsWithTheSameId &= (execBuff2->rsvd1 == drmContextId);
+
+    auto execObjects = reinterpret_cast<drm_i915_gem_exec_object2 *>(execBuff2->buffers_ptr);
+    for (uint32_t i = 0; i < execBuff2->buffer_count - 1; i++) {
+        allExecsWithTheSameId &= (execObjects[i].rsvd1 == drmContextId);
+    }
+
+    return allExecsWithTheSameId;
+}
+
 TEST_F(DrmCommandStreamTest, Flush) {
    auto expectedSize = alignUp(8u, MemoryConstants::cacheLineSize); // bbEnd
    int boHandle = 123;
@@ -211,6 +226,48 @@ TEST_F(DrmCommandStreamTest, Flush) {
    EXPECT_EQ(availableSpacePriorToFlush, cs.getAvailableSpace());
 }

+TEST_F(DrmCommandStreamTest, givenDrmContextIdWhenFlushingThenSetIdToAllExecBuffersAndObjects) {
+    uint32_t expectedDrmContextId = 321;
+    uint32_t numAllocations = 3;
+
+    auto createdContextId = [&expectedDrmContextId](unsigned long request, void *arg) {
+        auto contextCreate = static_cast<drm_i915_gem_context_create *>(arg);
+        contextCreate->ctx_id = expectedDrmContextId;
+        return 0;
+    };
+
+    auto allocation1 = memoryManager->allocate32BitGraphicsMemory(1, reinterpret_cast<void *>(1), AllocationOrigin::INTERNAL_ALLOCATION);
+    auto allocation2 = memoryManager->allocate32BitGraphicsMemory(1, reinterpret_cast<void *>(2), AllocationOrigin::INTERNAL_ALLOCATION);
+    csr->makeResident(*allocation1);
+    csr->makeResident(*allocation2);
+
+    EXPECT_CALL(*mock, ioctl(::testing::_, ::testing::_)).WillRepeatedly(::testing::Return(0)).RetiresOnSaturation();
+
+    EXPECT_CALL(*mock, ioctl(DRM_IOCTL_I915_GEM_CONTEXT_CREATE, ::testing::_))
+        .Times(1)
+        .WillRepeatedly(::testing::Invoke(createdContextId))
+        .RetiresOnSaturation();
+
+    EXPECT_CALL(*mock, ioctl(DRM_IOCTL_I915_GEM_EXECBUFFER2, BoExecFlushContextEq(expectedDrmContextId, numAllocations)))
+        .Times(1)
+        .WillRepeatedly(::testing::Return(0))
+        .RetiresOnSaturation();
+
+    osContext = std::make_unique<OsContext>(executionEnvironment.osInterface.get(), 1,
+                                            gpgpuEngineInstances[0], PreemptionHelper::getDefaultPreemptionMode(*platformDevices[0]));
+    csr->setOsContext(*osContext);
+
+    auto &cs = csr->getCS();
+
+    csr->addBatchBufferEnd(cs, nullptr);
+    csr->alignToCacheLine(cs);
+    BatchBuffer batchBuffer{cs.getGraphicsAllocation(), 0, 0, nullptr, false, false, QueueThrottle::MEDIUM, cs.getUsed(), &cs};
+    csr->flush(batchBuffer, csr->getResidencyAllocations());
+
+    memoryManager->freeGraphicsMemory(allocation1);
+    memoryManager->freeGraphicsMemory(allocation2);
+}
+
 TEST_F(DrmCommandStreamTest, FlushWithLowPriorityContext) {
    auto expectedSize = alignUp(8u, MemoryConstants::cacheLineSize); // bbEnd

@@ -655,43 +712,39 @@ class DrmCommandStreamEnhancedFixture

 {
  public:
-    DrmMockCustom *mock;
-    DeviceCommandStreamReceiver<DEFAULT_TEST_FAMILY_NAME> *csr = nullptr;
-    DrmMemoryManager *mm = nullptr;
-    MockDevice *device = nullptr;
-    DebugManagerStateRestore *dbgState;
+    std::unique_ptr<DebugManagerStateRestore> dbgState;
    ExecutionEnvironment *executionEnvironment;
-    std::unique_ptr<OsContext> osContext;
+    std::unique_ptr<DrmMockCustom> mock;
+    DeviceCommandStreamReceiver<DEFAULT_TEST_FAMILY_NAME> *csr = nullptr;

-    void SetUp() {
+    DrmMemoryManager *mm = nullptr;
+    std::unique_ptr<MockDevice> device;
+
+    virtual void SetUp() {
        executionEnvironment = new ExecutionEnvironment;
+        executionEnvironment->incRefInternal();
        executionEnvironment->initGmm(*platformDevices);
-        this->dbgState = new DebugManagerStateRestore();
+        this->dbgState = std::make_unique<DebugManagerStateRestore>();
        //make sure this is disabled, we don't want test this now
        DebugManager.flags.EnableForcePin.set(false);

-        mock = new DrmMockCustom();
+        mock = std::make_unique<DrmMockCustom>();
        executionEnvironment->osInterface = std::make_unique<OSInterface>();
-        executionEnvironment->osInterface->get()->setDrm(mock);
-        osContext = std::make_unique<OsContext>(executionEnvironment->osInterface.get(), 0u, gpgpuEngineInstances[0], PreemptionHelper::getDefaultPreemptionMode(*platformDevices[0]));
+        executionEnvironment->osInterface->get()->setDrm(mock.get());

        tCsr = new TestedDrmCommandStreamReceiver<DEFAULT_TEST_FAMILY_NAME>(*executionEnvironment);
        csr = tCsr;
-        csr->setOsContext(*osContext);
        ASSERT_NE(nullptr, csr);
        mm = reinterpret_cast<DrmMemoryManager *>(csr->createMemoryManager(false, false));
        ASSERT_NE(nullptr, mm);
        executionEnvironment->memoryManager.reset(mm);
-        device = Device::create<MockDevice>(platformDevices[0], executionEnvironment, 0u);
+        device.reset(MockDevice::create<MockDevice>(platformDevices[0], executionEnvironment, 0u));
+        device->resetCommandStreamReceiver(tCsr);
        ASSERT_NE(nullptr, device);
    }

-    void TearDown() {
-        //And close at destruction
-        delete csr;
-        delete device;
-        delete mock;
-        delete dbgState;
+    virtual void TearDown() {
+        executionEnvironment->decRefInternal();
    }

    bool isResident(BufferObject *bo) {
@@ -716,7 +769,7 @@ class DrmCommandStreamEnhancedFixture
    };

    MockBufferObject *createBO(size_t size) {
-        return new MockBufferObject(this->mock, size);
+        return new MockBufferObject(this->mock.get(), size);
    }
 };

@@ -859,7 +912,7 @@ TEST_F(DrmCommandStreamGemWorkerTests, givenCommandStreamWithDuplicatesWhenItIsF

 TEST_F(DrmCommandStreamGemWorkerTests, givenDrmCsrCreatedWithInactiveGemCloseWorkerPolicyThenThreadIsNotCreated) {
    this->executionEnvironment->osInterface = std::make_unique<OSInterface>();
-    this->executionEnvironment->osInterface->get()->setDrm(mock);
+    this->executionEnvironment->osInterface->get()->setDrm(mock.get());
    TestedDrmCommandStreamReceiver<DEFAULT_TEST_FAMILY_NAME> testedCsr(gemCloseWorkerMode::gemCloseWorkerInactive,
                                                                       *this->executionEnvironment);
    EXPECT_EQ(gemCloseWorkerMode::gemCloseWorkerInactive, testedCsr.peekGemCloseWorkerOperationMode());
@@ -889,6 +942,7 @@ class DrmCommandStreamBatchingTests : public Test<DrmCommandStreamEnhancedFixtur
 };

 TEST_F(DrmCommandStreamBatchingTests, givenCSRWhenFlushIsCalledThenProperFlagsArePassed) {
+    mock->reset();
    auto commandBuffer = mm->allocateGraphicsMemory(1024);
    auto dummyAllocation = mm->allocateGraphicsMemory(1024);
    ASSERT_NE(nullptr, commandBuffer);
@@ -902,13 +956,14 @@ TEST_F(DrmCommandStreamBatchingTests, givenCSRWhenFlushIsCalledThenProperFlagsAr
    BatchBuffer batchBuffer{cs.getGraphicsAllocation(), 0, 0, nullptr, false, false, QueueThrottle::MEDIUM, cs.getUsed(), &cs};
    csr->flush(batchBuffer, csr->getResidencyAllocations());

-    //preemption allocation + Sip Kernel
-    int ioctlPreemptionCnt = (PreemptionHelper::getDefaultPreemptionMode(*platformDevices[0]) == PreemptionMode::MidThread) ? 2 : 0;
-    int ioctlTagAllocCnt = gpgpuEngineInstances.size();
+    int ioctlExecCnt = 1;
+    int ioctlUserPtrCnt = 2;

    auto engineFlag = csr->getOsContext().get()->getEngineFlag();

-    EXPECT_EQ(5 + ioctlPreemptionCnt + ioctlTagAllocCnt, this->mock->ioctl_cnt.total);
+    EXPECT_EQ(ioctlExecCnt + ioctlUserPtrCnt, this->mock->ioctl_cnt.total);
+    EXPECT_EQ(ioctlExecCnt, this->mock->ioctl_cnt.execbuffer2);
+    EXPECT_EQ(ioctlUserPtrCnt, this->mock->ioctl_cnt.gemUserptr);
    uint64_t flags = engineFlag | I915_EXEC_NO_RELOC;
    EXPECT_EQ(flags, this->mock->execBuffer.flags);

@@ -917,6 +972,7 @@ TEST_F(DrmCommandStreamBatchingTests, givenCSRWhenFlushIsCalledThenProperFlagsAr
 }

 TEST_F(DrmCommandStreamBatchingTests, givenCsrWhenDispatchPolicyIsSetToBatchingThenCommandBufferIsNotSubmitted) {
+    mock->reset();
    tCsr->overrideDispatchPolicy(DispatchMode::BatchedDispatch);

    auto mockedSubmissionsAggregator = new mockSubmissionsAggregator();
@@ -930,8 +986,6 @@ TEST_F(DrmCommandStreamBatchingTests, givenCsrWhenDispatchPolicyIsSetToBatchingT

    tCsr->makeResident(*dummyAllocation);

-    std::unique_ptr<Device> device(MockDevice::createWithNewExecutionEnvironment<MockDevice>(nullptr));
-
    tCsr->setTagAllocation(tagAllocation);
    tCsr->setPreemptionCsrAllocation(preemptionAllocation);
    DispatchFlags dispatchFlags;
@@ -946,10 +1000,6 @@ TEST_F(DrmCommandStreamBatchingTests, givenCsrWhenDispatchPolicyIsSetToBatchingT
    //preemption allocation
    size_t csrSurfaceCount = (device->getPreemptionMode() == PreemptionMode::MidThread) ? 2 : 0;

-    //preemption allocation + sipKernel
-    int ioctlPreemptionCnt = (PreemptionHelper::getDefaultPreemptionMode(*platformDevices[0]) == PreemptionMode::MidThread) ? 2 : 0;
-    int ioctlTagAllocCnt = gpgpuEngineInstances.size();
-
    auto recordedCmdBuffer = cmdBuffers.peekHead();
    EXPECT_EQ(3u + csrSurfaceCount, recordedCmdBuffer->surfaces.size());

@@ -965,26 +1015,28 @@ TEST_F(DrmCommandStreamBatchingTests, givenCsrWhenDispatchPolicyIsSetToBatchingT

    EXPECT_EQ(tCsr->commandStream.getGraphicsAllocation(), recordedCmdBuffer->batchBuffer.commandBufferAllocation);

-    EXPECT_EQ(5 + ioctlPreemptionCnt + ioctlTagAllocCnt, this->mock->ioctl_cnt.total);
+    int ioctlUserPtrCnt = 3;
+
+    EXPECT_EQ(ioctlUserPtrCnt, this->mock->ioctl_cnt.total);
+    EXPECT_EQ(ioctlUserPtrCnt, this->mock->ioctl_cnt.gemUserptr);

    EXPECT_EQ(0u, this->mock->execBuffer.flags);

+    tCsr->flushBatchedSubmissions();
+
    mm->freeGraphicsMemory(dummyAllocation);
    mm->freeGraphicsMemory(commandBuffer);
-    tCsr->setTagAllocation(nullptr);
-    tCsr->setPreemptionCsrAllocation(nullptr);
 }

 TEST_F(DrmCommandStreamBatchingTests, givenRecordedCommandBufferWhenItIsSubmittedThenFlushTaskIsProperlyCalled) {
+    mock->reset();
    tCsr->overrideDispatchPolicy(DispatchMode::BatchedDispatch);

    auto mockedSubmissionsAggregator = new mockSubmissionsAggregator();
    tCsr->overrideSubmissionAggregator(mockedSubmissionsAggregator);

    auto commandBuffer = mm->allocateGraphicsMemory(1024);
-    auto dummyAllocation = mm->allocateGraphicsMemory(1024);
    IndirectHeap cs(commandBuffer);
-    std::unique_ptr<Device> device(MockDevice::createWithNewExecutionEnvironment<MockDevice>(nullptr));

    tCsr->setTagAllocation(tagAllocation);
    tCsr->setPreemptionCsrAllocation(preemptionAllocation);
@@ -1013,10 +1065,6 @@ TEST_F(DrmCommandStreamBatchingTests, givenRecordedCommandBufferWhenItIsSubmitte
    //preemption allocation
    size_t csrSurfaceCount = (device->getPreemptionMode() == PreemptionMode::MidThread) ? 2 : 0;

-    //preemption allocation +sip Kernel
-    int ioctlPreemptionCnt = (PreemptionHelper::getDefaultPreemptionMode(*platformDevices[0]) == PreemptionMode::MidThread) ? 2 : 0;
-    int ioctlTagAllocCnt = gpgpuEngineInstances.size();
-
    //validate that submited command buffer has what we want
    EXPECT_EQ(3u + csrSurfaceCount, this->mock->execBuffer.buffer_count);
    EXPECT_EQ(4u, this->mock->execBuffer.batch_start_offset);
@@ -1037,11 +1085,13 @@ TEST_F(DrmCommandStreamBatchingTests, givenRecordedCommandBufferWhenItIsSubmitte
        EXPECT_TRUE(handleFound);
    }

-    EXPECT_EQ(6 + ioctlPreemptionCnt + ioctlTagAllocCnt, this->mock->ioctl_cnt.total);
+    int ioctlExecCnt = 1;
+    int ioctlUserPtrCnt = 2;
+    EXPECT_EQ(ioctlExecCnt, this->mock->ioctl_cnt.execbuffer2);
+    EXPECT_EQ(ioctlUserPtrCnt, this->mock->ioctl_cnt.gemUserptr);
+    EXPECT_EQ(ioctlExecCnt + ioctlUserPtrCnt, this->mock->ioctl_cnt.total);

-    mm->freeGraphicsMemory(dummyAllocation);
    mm->freeGraphicsMemory(commandBuffer);
-    tCsr->setTagAllocation(nullptr);
    tCsr->setPreemptionCsrAllocation(nullptr);
 }

@@ -1180,10 +1230,12 @@ TEST_F(DrmCommandStreamLeaksTest, givenFragmentedAllocationsWithResuedFragmentsW

    tCsr->processResidency(csr->getResidencyAllocations());

-    EXPECT_TRUE(graphicsAllocation->fragmentsStorage.fragmentStorageData[0].residency->resident[osContext->getContextId()]);
-    EXPECT_TRUE(graphicsAllocation->fragmentsStorage.fragmentStorageData[1].residency->resident[osContext->getContextId()]);
-    EXPECT_TRUE(graphicsAllocation->fragmentsStorage.fragmentStorageData[2].residency->resident[osContext->getContextId()]);
-    EXPECT_TRUE(graphicsAllocation2->fragmentsStorage.fragmentStorageData[0].residency->resident[osContext->getContextId()]);
+    auto &osContext = tCsr->getOsContext();
+
+    EXPECT_TRUE(graphicsAllocation->fragmentsStorage.fragmentStorageData[0].residency->resident[osContext.getContextId()]);
+    EXPECT_TRUE(graphicsAllocation->fragmentsStorage.fragmentStorageData[1].residency->resident[osContext.getContextId()]);
+    EXPECT_TRUE(graphicsAllocation->fragmentsStorage.fragmentStorageData[2].residency->resident[osContext.getContextId()]);
+    EXPECT_TRUE(graphicsAllocation2->fragmentsStorage.fragmentStorageData[0].residency->resident[osContext.getContextId()]);

    auto residency = tCsr->getResidencyVector();

@@ -1192,10 +1244,10 @@ TEST_F(DrmCommandStreamLeaksTest, givenFragmentedAllocationsWithResuedFragmentsW
    tCsr->makeSurfacePackNonResident(tCsr->getResidencyAllocations());

    //check that each packet is not resident
-    EXPECT_FALSE(graphicsAllocation->fragmentsStorage.fragmentStorageData[0].residency->resident[osContext->getContextId()]);
-    EXPECT_FALSE(graphicsAllocation->fragmentsStorage.fragmentStorageData[1].residency->resident[osContext->getContextId()]);
-    EXPECT_FALSE(graphicsAllocation->fragmentsStorage.fragmentStorageData[2].residency->resident[osContext->getContextId()]);
-    EXPECT_FALSE(graphicsAllocation2->fragmentsStorage.fragmentStorageData[0].residency->resident[osContext->getContextId()]);
+    EXPECT_FALSE(graphicsAllocation->fragmentsStorage.fragmentStorageData[0].residency->resident[osContext.getContextId()]);
+    EXPECT_FALSE(graphicsAllocation->fragmentsStorage.fragmentStorageData[1].residency->resident[osContext.getContextId()]);
+    EXPECT_FALSE(graphicsAllocation->fragmentsStorage.fragmentStorageData[2].residency->resident[osContext.getContextId()]);
+    EXPECT_FALSE(graphicsAllocation2->fragmentsStorage.fragmentStorageData[0].residency->resident[osContext.getContextId()]);

    EXPECT_EQ(0u, residency->size());

@@ -1204,10 +1256,10 @@ TEST_F(DrmCommandStreamLeaksTest, givenFragmentedAllocationsWithResuedFragmentsW

    tCsr->processResidency(csr->getResidencyAllocations());

-    EXPECT_TRUE(graphicsAllocation->fragmentsStorage.fragmentStorageData[0].residency->resident[osContext->getContextId()]);
-    EXPECT_TRUE(graphicsAllocation->fragmentsStorage.fragmentStorageData[1].residency->resident[osContext->getContextId()]);
-    EXPECT_TRUE(graphicsAllocation->fragmentsStorage.fragmentStorageData[2].residency->resident[osContext->getContextId()]);
-    EXPECT_TRUE(graphicsAllocation2->fragmentsStorage.fragmentStorageData[0].residency->resident[osContext->getContextId()]);
+    EXPECT_TRUE(graphicsAllocation->fragmentsStorage.fragmentStorageData[0].residency->resident[osContext.getContextId()]);
+    EXPECT_TRUE(graphicsAllocation->fragmentsStorage.fragmentStorageData[1].residency->resident[osContext.getContextId()]);
+    EXPECT_TRUE(graphicsAllocation->fragmentsStorage.fragmentStorageData[2].residency->resident[osContext.getContextId()]);
+    EXPECT_TRUE(graphicsAllocation2->fragmentsStorage.fragmentStorageData[0].residency->resident[osContext.getContextId()]);

    EXPECT_EQ(3u, residency->size());

@@ -1215,10 +1267,10 @@ TEST_F(DrmCommandStreamLeaksTest, givenFragmentedAllocationsWithResuedFragmentsW

    EXPECT_EQ(0u, residency->size());

-    EXPECT_FALSE(graphicsAllocation->fragmentsStorage.fragmentStorageData[0].residency->resident[osContext->getContextId()]);
-    EXPECT_FALSE(graphicsAllocation->fragmentsStorage.fragmentStorageData[1].residency->resident[osContext->getContextId()]);
-    EXPECT_FALSE(graphicsAllocation->fragmentsStorage.fragmentStorageData[2].residency->resident[osContext->getContextId()]);
-    EXPECT_FALSE(graphicsAllocation2->fragmentsStorage.fragmentStorageData[0].residency->resident[osContext->getContextId()]);
+    EXPECT_FALSE(graphicsAllocation->fragmentsStorage.fragmentStorageData[0].residency->resident[osContext.getContextId()]);
+    EXPECT_FALSE(graphicsAllocation->fragmentsStorage.fragmentStorageData[1].residency->resident[osContext.getContextId()]);
+    EXPECT_FALSE(graphicsAllocation->fragmentsStorage.fragmentStorageData[2].residency->resident[osContext.getContextId()]);
+    EXPECT_FALSE(graphicsAllocation2->fragmentsStorage.fragmentStorageData[0].residency->resident[osContext.getContextId()]);

    mm->freeGraphicsMemory(graphicsAllocation);
    mm->freeGraphicsMemory(graphicsAllocation2);
--- a/unit_tests/os_interface/linux/drm_memory_manager_tests.cpp
+++ b/unit_tests/os_interface/linux/drm_memory_manager_tests.cpp
@@ -60,21 +60,27 @@ class DrmMemoryManagerFixture : public MemoryManagementFixture {
  public:
    TestedDrmMemoryManager *memoryManager = nullptr;
    DrmMockCustom *mock;
+    MockDevice *device = nullptr;

    void SetUp() override {
        MemoryManagementFixture::SetUp();
        this->mock = new DrmMockCustom;
        executionEnvironment = new ExecutionEnvironment;
        executionEnvironment->incRefInternal();
+        executionEnvironment->osInterface = std::make_unique<OSInterface>();
+        executionEnvironment->osInterface->get()->setDrm(mock);
+
        memoryManager = new (std::nothrow) TestedDrmMemoryManager(this->mock, *executionEnvironment);
        //assert we have memory manager
        ASSERT_NE(nullptr, memoryManager);
        if (memoryManager->getgemCloseWorker()) {
            memoryManager->getgemCloseWorker()->close(true);
        }
+        device = MockDevice::createWithExecutionEnvironment<MockDevice>(*platformDevices, executionEnvironment, 0);
    }

    void TearDown() override {
+        delete device;
        delete memoryManager;
        executionEnvironment->decRefInternal();

@@ -98,23 +104,30 @@ class DrmMemoryManagerFixtureWithoutQuietIoctlExpectation : public MemoryManagem

    void SetUp() override {
        MemoryManagementFixture::SetUp();
+        executionEnvironment = new ExecutionEnvironment;
        this->mock = new DrmMockCustom;
-        memoryManager = new (std::nothrow) TestedDrmMemoryManager(this->mock, executionEnvironment);
+        memoryManager = new (std::nothrow) TestedDrmMemoryManager(this->mock, *executionEnvironment);
        ASSERT_NE(nullptr, memoryManager);
        if (memoryManager->getgemCloseWorker()) {
            memoryManager->getgemCloseWorker()->close(true);
        }
+        executionEnvironment->osInterface = std::make_unique<OSInterface>();
+        executionEnvironment->osInterface->get()->setDrm(mock);
+        device = MockDevice::createWithExecutionEnvironment<MockDevice>(*platformDevices, executionEnvironment, 0);
    }

    void TearDown() override {
+        delete device;
        delete memoryManager;
        delete this->mock;
        this->mock = nullptr;
+
        MemoryManagementFixture::TearDown();
    }

  protected:
-    ExecutionEnvironment executionEnvironment;
+    ExecutionEnvironment *executionEnvironment = nullptr;
+    MockDevice *device = nullptr;
    DrmMockCustom::IoctlResExt ioctlResExt = {0, 0};
 };

@@ -164,7 +177,7 @@ TEST_F(DrmMemoryManagerTest, GivenGraphicsAllocationWhenAddAndRemoveAllocationTo
 }

 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenforcePinAllowedWhenMemoryManagerIsCreatedThenPinBbIsCreated) {
-    auto memoryManager = std::make_unique<TestedDrmMemoryManager>(this->mock, true, false, executionEnvironment);
+    auto memoryManager = std::make_unique<TestedDrmMemoryManager>(this->mock, true, false, *executionEnvironment);
    EXPECT_NE(nullptr, memoryManager->getPinBB());
 }

@@ -186,6 +199,7 @@ TEST_F(DrmMemoryManagerTest, pinBBnotCreatedWhenIoctlFailed) {
    mock->ioctl_res = -1;
    auto memoryManager = new (std::nothrow) TestedDrmMemoryManager(this->mock, true, false, *executionEnvironment);
    EXPECT_EQ(nullptr, memoryManager->getPinBB());
+    mock->ioctl_res = 0;
    delete memoryManager;
 }

@@ -343,12 +357,12 @@ TEST_F(DrmMemoryManagerTest, UnreferenceNullPtr) {
 }

 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDrmMemoryManagerCreatedWithGemCloseWorkerModeInactiveThenGemCloseWorkerIsNotCreated) {
-    DrmMemoryManager drmMemoryManger(this->mock, gemCloseWorkerMode::gemCloseWorkerInactive, false, false, executionEnvironment);
+    DrmMemoryManager drmMemoryManger(this->mock, gemCloseWorkerMode::gemCloseWorkerInactive, false, false, *executionEnvironment);
    EXPECT_EQ(nullptr, drmMemoryManger.peekGemCloseWorker());
 }

 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDrmMemoryManagerCreatedWithGemCloseWorkerActiveThenGemCloseWorkerIsCreated) {
-    DrmMemoryManager drmMemoryManger(this->mock, gemCloseWorkerMode::gemCloseWorkerActive, false, false, executionEnvironment);
+    DrmMemoryManager drmMemoryManger(this->mock, gemCloseWorkerMode::gemCloseWorkerActive, false, false, *executionEnvironment);
    EXPECT_NE(nullptr, drmMemoryManger.peekGemCloseWorker());
 }

@@ -411,6 +425,7 @@ TEST_F(DrmMemoryManagerTest, AllocateUserptrFail) {

    auto ptr = memoryManager->allocateGraphicsMemory(3);
    EXPECT_EQ(nullptr, ptr);
+    mock->ioctl_res = 0;
 }

 TEST_F(DrmMemoryManagerTest, FreeNullPtr) {
@@ -494,6 +509,7 @@ TEST_F(DrmMemoryManagerTest, Allocate_HostPtr_UserptrFail) {
    EXPECT_EQ(nullptr, alloc);

    ::alignedFree(ptrT);
+    mock->ioctl_res = 0;
 }

 TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerWhengetSystemSharedMemoryIsCalledThenContextGetParamIsCalled) {
@@ -523,6 +539,7 @@ TEST_F(DrmMemoryManagerTest, getMinimumSystemSharedMemory) {
    mock->ioctl_expected.contextGetParam = 1;

    EXPECT_EQ(gpuMemorySize, systemSharedMemorySize);
+    mock->ioctl_expected.contextDestroy = 0;
    mock->testIoctls();

    // gpuMemSize > hostMemSize
@@ -532,6 +549,7 @@ TEST_F(DrmMemoryManagerTest, getMinimumSystemSharedMemory) {
    mock->ioctl_expected.contextGetParam = 2;
    EXPECT_EQ(hostMemorySize, systemSharedMemorySize);
    mock->testIoctls();
+    mock->ioctl_expected.contextDestroy = OCLRT::EngineInstanceConstants::numGpgpuEngineInstances;
 }

 TEST_F(DrmMemoryManagerTest, BoWaitFailure) {
@@ -546,6 +564,7 @@ TEST_F(DrmMemoryManagerTest, BoWaitFailure) {
    mock->ioctl_res = 1;

    memoryManager->unreference(bo);
+    mock->ioctl_res = 0;
 }

 TEST_F(DrmMemoryManagerTest, NullOsHandleStorageAskedForPopulationReturnsFilledPointer) {
@@ -565,7 +584,7 @@ TEST_F(DrmMemoryManagerTest, NullOsHandleStorageAskedForPopulationReturnsFilledP
 }

 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenEnabledHostMemoryValidationWhenReadOnlyPointerCausesPinningFailWithEfaultThenPopulateOsHandlesReturnsInvalidHostPointerError) {
-    std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new (std::nothrow) TestedDrmMemoryManager(this->mock, false, true, executionEnvironment));
+    std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new (std::nothrow) TestedDrmMemoryManager(this->mock, false, true, *executionEnvironment));

    OsHandleStorage storage;
    storage.fragmentStorageData[0].cpuPtr = reinterpret_cast<void *>(0x1000);
@@ -594,7 +613,7 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenEnabledHostMemoryValid
 }

 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenEnabledHostMemoryValidationWhenPinningFailWithErrorDifferentThanEfaultThenPopulateOsHandlesReturnsError) {
-    std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new (std::nothrow) TestedDrmMemoryManager(this->mock, false, true, executionEnvironment));
+    std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new (std::nothrow) TestedDrmMemoryManager(this->mock, false, true, *executionEnvironment));

    OsHandleStorage storage;
    storage.fragmentStorageData[0].cpuPtr = reinterpret_cast<void *>(0x1000);
@@ -721,8 +740,6 @@ TEST_F(DrmMemoryManagerTest, givenMemoryManagerWhenAskedFor32BitAllocationThen32
 }

 TEST_F(DrmMemoryManagerTest, givenMemoryManagerWhensetForce32BitAllocationsIsCalledWithTrueMutlipleTimesThenAllocatorIsReused) {
-    mock->ioctl_expected.reset();
-
    EXPECT_EQ(nullptr, memoryManager->allocator32Bit.get());
    memoryManager->setForce32BitAllocations(true);
    EXPECT_NE(nullptr, memoryManager->allocator32Bit.get());
@@ -732,8 +749,6 @@ TEST_F(DrmMemoryManagerTest, givenMemoryManagerWhensetForce32BitAllocationsIsCal
 }

 TEST_F(DrmMemoryManagerTest, givenMemoryManagerWhensetForce32BitAllocationsIsCalledWithFalseThenAllocatorIsNotDeleted) {
-    mock->ioctl_expected.reset();
-
    memoryManager->setForce32BitAllocations(true);
    EXPECT_NE(nullptr, memoryManager->allocator32Bit.get());
    memoryManager->setForce32BitAllocations(false);
@@ -920,7 +935,7 @@ TEST_F(DrmMemoryManagerTest, Given32bitAllocatorWhenAskedForBufferCreatedFrom64B
 TEST_F(DrmMemoryManagerTest, givenMemoryManagerWhenAskedFor32BitAllocationWithHostPtrAndAllocUserptrFailsThenFails) {
    mock->ioctl_expected.gemUserptr = 1;

-    this->ioctlResExt = {0, -1};
+    this->ioctlResExt = {mock->ioctl_cnt.total, -1};
    mock->ioctl_res_ext = &ioctlResExt;

    auto size = 10u;
@@ -935,7 +950,7 @@ TEST_F(DrmMemoryManagerTest, givenMemoryManagerWhenAskedFor32BitAllocationWithHo
 TEST_F(DrmMemoryManagerTest, givenMemoryManagerWhenAskedFor32BitAllocationAndAllocUserptrFailsThenFails) {
    mock->ioctl_expected.gemUserptr = 1;

-    this->ioctlResExt = {0, -1};
+    this->ioctlResExt = {mock->ioctl_cnt.total, -1};
    mock->ioctl_res_ext = &ioctlResExt;

    auto size = 10u;
@@ -1230,6 +1245,7 @@ TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerWhenTiledImageIsBeingCreatedAn

    injectFailures(method);
    mock->reset();
+    mock->ioctl_expected.contextDestroy = OCLRT::EngineInstanceConstants::numGpgpuEngineInstances;
 }

 TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerWhenTiledImageIsBeingCreatedFromHostPtrThenallocateGraphicsMemoryForImageIsUsed) {
@@ -1238,8 +1254,7 @@ TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerWhenTiledImageIsBeingCreatedFr
    mock->ioctl_expected.gemWait = 1;
    mock->ioctl_expected.gemClose = 1;

-    std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(*platformDevices, executionEnvironment, 0u));
-    MockContext context(device.get());
+    MockContext context(device);
    context.setMemoryManager(memoryManager);

    cl_image_format imageFormat;
@@ -1459,6 +1474,7 @@ TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerAndOsHandleWhenAllocationFails

    injectFailures(method);
    mock->reset();
+    mock->ioctl_expected.contextDestroy = OCLRT::EngineInstanceConstants::numGpgpuEngineInstances;
 }

 TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerAndThreeOsHandlesWhenReuseCreatesAreCalledThenGraphicsAllocationsAreReturned) {
@@ -1603,7 +1619,7 @@ TEST_F(DrmMemoryManagerTest, givenNon32BitAddressingWhenBufferFromSharedHandleIs

 TEST_F(DrmMemoryManagerTest, givenSharedHandleWhenAllocationIsCreatedAndIoctlPrimeFdToHandleFailsThenNullPtrIsReturned) {
    mock->ioctl_expected.primeFdToHandle = 1;
-    this->ioctlResExt = {0, -1};
+    this->ioctlResExt = {mock->ioctl_cnt.total, -1};
    mock->ioctl_res_ext = &this->ioctlResExt;

    osHandle handle = 1u;
@@ -1625,9 +1641,7 @@ TEST_F(DrmMemoryManagerTest, givenTwoGraphicsAllocationsThatShareTheSameBufferOb
    executionEnvironment->osInterface = std::make_unique<OSInterface>();
    executionEnvironment->osInterface->get()->setDrm(mock);
    auto testedCsr = new TestedDrmCommandStreamReceiver<DEFAULT_TEST_FAMILY_NAME>(*executionEnvironment);
-    executionEnvironment->commandStreamReceivers.resize(1);
-    executionEnvironment->commandStreamReceivers[0][0].reset(testedCsr);
-    testedCsr->setOsContext(*memoryManager->createAndRegisterOsContext(gpgpuEngineInstances[0], PreemptionHelper::getDefaultPreemptionMode(*platformDevices[0])));
+    device->resetCommandStreamReceiver(testedCsr);

    testedCsr->makeResident(*graphicsAllocation);
    testedCsr->makeResident(*graphicsAllocation2);
@@ -1654,9 +1668,7 @@ TEST_F(DrmMemoryManagerTest, givenTwoGraphicsAllocationsThatDoesnShareTheSameBuf
    executionEnvironment->osInterface = std::make_unique<OSInterface>();
    executionEnvironment->osInterface->get()->setDrm(mock);
    auto testedCsr = new TestedDrmCommandStreamReceiver<DEFAULT_TEST_FAMILY_NAME>(*executionEnvironment);
-    executionEnvironment->commandStreamReceivers.resize(1);
-    executionEnvironment->commandStreamReceivers[0][0].reset(testedCsr);
-    testedCsr->setOsContext(*memoryManager->createAndRegisterOsContext(gpgpuEngineInstances[0], PreemptionHelper::getDefaultPreemptionMode(*platformDevices[0])));
+    device->resetCommandStreamReceiver(testedCsr);

    testedCsr->makeResident(*graphicsAllocation);
    testedCsr->makeResident(*graphicsAllocation2);
@@ -1784,7 +1796,7 @@ TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerWhenLockUnlockIsCalledOnAlloca

 TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerWhenLockUnlockIsCalledButFailsOnIoctlMmapThenReturnNullPtr) {
    mock->ioctl_expected.gemMmap = 1;
-    this->ioctlResExt = {0, -1};
+    this->ioctlResExt = {mock->ioctl_cnt.total, -1};
    mock->ioctl_res_ext = &ioctlResExt;

    DrmMockCustom drmMock;
@@ -1812,7 +1824,7 @@ TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerWhenSetDomainCpuIsCalledOnAllo

 TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerWhenSetDomainCpuIsCalledButFailsOnIoctlSetDomainThenReturnFalse) {
    mock->ioctl_expected.gemSetDomain = 1;
-    this->ioctlResExt = {0, -1};
+    this->ioctlResExt = {mock->ioctl_cnt.total, -1};
    mock->ioctl_res_ext = &ioctlResExt;

    DrmMockCustom drmMock;
@@ -2037,7 +2049,7 @@ TEST_F(DrmMemoryManagerTest, givenMemoryManagerSupportingVirutalPaddingWhenAlloc
    mock->ioctl_expected.gemUserptr = 3;
    mock->ioctl_expected.gemWait = 2;
    mock->ioctl_expected.gemClose = 2;
-    this->ioctlResExt = {2, -1};
+    this->ioctlResExt = {mock->ioctl_cnt.total + 2, -1};
    mock->ioctl_res_ext = &ioctlResExt;

    //first let's create normal buffer
@@ -2415,9 +2427,9 @@ TEST(DrmMemoryManager, givenMemoryManagerWhenAllocateGraphicsMemoryIsCalledThenM
    memoryManager->freeGraphicsMemory(allocation);
 }

-TEST(DrmMemoryManager, givenMemoryManagerWhenAllocateGraphicsMemoryWithPtrIsCalledThenMemoryPoolIsSystem4KBPages) {
-    ExecutionEnvironment executionEnvironment;
-    std::unique_ptr<TestedDrmMemoryManager> memoryManager(new (std::nothrow) TestedDrmMemoryManager(Drm::get(0), false, true, executionEnvironment));
+TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenMemoryManagerWhenAllocateGraphicsMemoryWithPtrIsCalledThenMemoryPoolIsSystem4KBPages) {
+    std::unique_ptr<TestedDrmMemoryManager> memoryManager(new (std::nothrow) TestedDrmMemoryManager(Drm::get(0), false, true, *executionEnvironment));
+
    void *ptr = reinterpret_cast<void *>(0x1001);
    auto size = 4096u;
    auto allocation = memoryManager->allocateGraphicsMemory(size, ptr);
@@ -2429,6 +2441,7 @@ TEST(DrmMemoryManager, givenMemoryManagerWhenAllocateGraphicsMemoryWithPtrIsCall
 TEST(DrmMemoryManager, givenMemoryManagerWhenAllocate32BitGraphicsMemoryWithPtrIsCalledThenMemoryPoolIsSystem4KBPagesWith32BitGpuAddressing) {
    ExecutionEnvironment executionEnvironment;
    std::unique_ptr<TestedDrmMemoryManager> memoryManager(new (std::nothrow) TestedDrmMemoryManager(Drm::get(0), false, true, executionEnvironment));
+
    memoryManager->setForce32BitAllocations(true);

    void *ptr = reinterpret_cast<void *>(0x1001);
@@ -2490,17 +2503,18 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDisabledForcePinAndEna
    this->mock->ioctl_expected.gemUserptr = 1;
    EXPECT_THROW(
        {
-            std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new TestedDrmMemoryManager(this->mock, false, true, executionEnvironment));
+            std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new TestedDrmMemoryManager(this->mock, false, true, *executionEnvironment));
            EXPECT_NE(nullptr, testedMemoryManager.get());
        },
        std::exception);
    this->mock->ioctl_res = 0;
+    this->mock->ioctl_expected.contextDestroy = 0;
    this->mock->testIoctls();
 }

 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDisabledForcePinAndEnabledValidateHostMemoryWhenPopulateOsHandlesIsCalledThenHostMemoryIsValidated) {

-    std::unique_ptr<DrmMemoryManager> testedMemoryManager(new TestedDrmMemoryManager(this->mock, false, true, executionEnvironment));
+    std::unique_ptr<DrmMemoryManager> testedMemoryManager(new TestedDrmMemoryManager(this->mock, false, true, *executionEnvironment));
    ASSERT_NE(nullptr, testedMemoryManager.get());
    ASSERT_NE(nullptr, testedMemoryManager->getPinBB());

@@ -2529,7 +2543,7 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDisabledForcePinAndEna
        PinBufferObject(Drm *drm) : BufferObject(drm, 1, true) {
        }

-        int pin(BufferObject *boToPin[], size_t numberOfBos) override {
+        int pin(BufferObject *boToPin[], size_t numberOfBos, uint32_t drmContextId) override {
            for (size_t i = 0; i < numberOfBos; i++) {
                pinnedBoArray[i] = boToPin[i];
            }
@@ -2540,7 +2554,7 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDisabledForcePinAndEna
        size_t numberOfBosPinned;
    };

-    std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new TestedDrmMemoryManager(this->mock, false, true, executionEnvironment));
+    std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new TestedDrmMemoryManager(this->mock, false, true, *executionEnvironment));
    ASSERT_NE(nullptr, testedMemoryManager.get());
    ASSERT_NE(nullptr, testedMemoryManager->getPinBB());

@@ -2589,7 +2603,7 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenValidateHostPtrMemoryE
    mock->ioctl_expected.gemWait = 1;
    mock->ioctl_expected.gemClose = 2;

-    std::unique_ptr<TestedDrmMemoryManager> memoryManager(new (std::nothrow) TestedDrmMemoryManager(this->mock, true, true, executionEnvironment));
+    std::unique_ptr<TestedDrmMemoryManager> memoryManager(new (std::nothrow) TestedDrmMemoryManager(this->mock, true, true, *executionEnvironment));
    ASSERT_NE(nullptr, memoryManager->getPinBB());

    size_t size = 10 * 1024 * 1024;
@@ -2603,7 +2617,7 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenValidateHostPtrMemoryE
 }

 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenEnabledHostMemoryValidationWhenValidHostPointerIsPassedToPopulateThenSuccessIsReturned) {
-    std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new (std::nothrow) TestedDrmMemoryManager(this->mock, false, true, executionEnvironment));
+    std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new (std::nothrow) TestedDrmMemoryManager(this->mock, false, true, *executionEnvironment));

    OsHandleStorage storage;
    storage.fragmentStorageData[0].cpuPtr = reinterpret_cast<void *>(0x1000);
@@ -2693,7 +2707,7 @@ TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerWhenAllocateGraphicsMemoryForN
    memoryManager->forceLimitedRangeAllocator(0xFFFFFFFFF);

    mock->ioctl_expected.gemUserptr = 1;
-    this->ioctlResExt = {0, -1};
+    this->ioctlResExt = {mock->ioctl_cnt.total, -1};
    mock->ioctl_res_ext = &ioctlResExt;

    auto allocation = memoryManager->allocateGraphicsMemoryForNonSvmHostPtr(size, hostPtr);
@@ -2703,7 +2717,7 @@ TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerWhenAllocateGraphicsMemoryForN
 }

 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenForcePinNotAllowedAndHostMemoryValidationEnabledWhenAllocationIsCreatedThenBufferObjectIsPinnedOnlyOnce) {
-    std::unique_ptr<TestedDrmMemoryManager> memoryManager(new TestedDrmMemoryManager(this->mock, false, true, executionEnvironment));
+    std::unique_ptr<TestedDrmMemoryManager> memoryManager(new TestedDrmMemoryManager(this->mock, false, true, *executionEnvironment));
    mock->reset();
    mock->ioctl_expected.gemUserptr = 1;
    mock->ioctl_expected.execbuffer2 = 1;
@@ -2725,7 +2739,7 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenForcePinNotAllowedAndH
 }

 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenForcePinNotAllowedAndHostMemoryValidationDisabledWhenAllocationIsCreatedThenBufferObjectIsNotPinned) {
-    std::unique_ptr<TestedDrmMemoryManager> memoryManager(new TestedDrmMemoryManager(this->mock, false, false, executionEnvironment));
+    std::unique_ptr<TestedDrmMemoryManager> memoryManager(new TestedDrmMemoryManager(this->mock, false, false, *executionEnvironment));
    mock->reset();
    mock->ioctl_expected.gemUserptr = 1;
    mock->ioctl_expected.gemClose = 1;
@@ -2746,7 +2760,7 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenForcePinNotAllowedAndH
 }

 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenEnabledValidateHostMemoryWhenReadOnlyPointerCausesPinningFailWithEfaultThenPopulateOsHandlesMarksFragmentsToFree) {
-    std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new TestedDrmMemoryManager(this->mock, false, true, executionEnvironment));
+    std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new TestedDrmMemoryManager(this->mock, false, true, *executionEnvironment));
    ASSERT_NE(nullptr, testedMemoryManager.get());
    ASSERT_NE(nullptr, testedMemoryManager->getPinBB());

@@ -2792,7 +2806,7 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenEnabledValidateHostMem
 }

 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenEnabledValidateHostMemoryWhenReadOnlyPointerCausesPinningFailWithEfaultThenPopulateOsHandlesDoesNotStoreTheFragments) {
-    std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new TestedDrmMemoryManager(this->mock, false, true, executionEnvironment));
+    std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new TestedDrmMemoryManager(this->mock, false, true, *executionEnvironment));
    ASSERT_NE(nullptr, testedMemoryManager->getPinBB());

    mock->reset();
@@ -2836,7 +2850,7 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenEnabledValidateHostMem
 }

 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenEnabledValidateHostMemoryWhenPopulateOsHandlesSucceedsThenFragmentIsStoredInHostPtrManager) {
-    std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new TestedDrmMemoryManager(this->mock, false, true, executionEnvironment));
+    std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new TestedDrmMemoryManager(this->mock, false, true, *executionEnvironment));
    ASSERT_NE(nullptr, testedMemoryManager->getPinBB());

    mock->reset();
@@ -2862,7 +2876,7 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenEnabledValidateHostMem
 }

 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDrmMemoryManagerWhenCleanOsHandlesDeletesHandleDataThenOsHandleStorageAndResidencyIsSetToNullptr) {
-    std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new TestedDrmMemoryManager(this->mock, false, true, executionEnvironment));
+    std::unique_ptr<TestedDrmMemoryManager> testedMemoryManager(new TestedDrmMemoryManager(this->mock, false, true, *executionEnvironment));
    ASSERT_NE(nullptr, testedMemoryManager->getPinBB());

    OsHandleStorage handleStorage;
--- a/unit_tests/os_interface/linux/drm_mock.h
+++ b/unit_tests/os_interface/linux/drm_mock.h
@@ -27,6 +27,8 @@ static const int mockFd = 33;
 // Mock DRM class that responds to DRM_IOCTL_I915_GETPARAMs
 class DrmMock : public Drm {
  public:
+    using Drm::preemptionSupported;
+
    DrmMock() : Drm(mockFd) {
        sysFsDefaultGpuPathToRestore = nullptr;
    }
@@ -117,24 +119,30 @@ class DrmMock : public Drm {
                return this->StoredRetVal;
            }
        }
-#if defined(I915_PARAM_HAS_PREEMPTION)
+
        if ((request == DRM_IOCTL_I915_GEM_CONTEXT_CREATE) && (arg != nullptr)) {
            drm_i915_gem_context_create *create = (drm_i915_gem_context_create *)arg;
            create->ctx_id = this->StoredCtxId;
            return this->StoredRetVal;
        }

+        if ((request == DRM_IOCTL_I915_GEM_CONTEXT_DESTROY) && (arg != nullptr)) {
+            drm_i915_gem_context_destroy *destroy = (drm_i915_gem_context_destroy *)arg;
+            this->receivedDestroyContextId = destroy->ctx_id;
+            return this->StoredRetVal;
+        }
+
        if ((request == DRM_IOCTL_I915_GEM_CONTEXT_SETPARAM) && (arg != nullptr)) {
-            drm_i915_gem_context_param *gp = (drm_i915_gem_context_param *)arg;
-            if ((gp->param == I915_CONTEXT_PARAM_PRIORITY) && (gp->ctx_id == this->StoredCtxId)) {
-                EXPECT_EQ(0u, gp->size);
+            receivedContextParamRequestCount++;
+            receivedContextParamRequest = *reinterpret_cast<drm_i915_gem_context_param *>(arg);
+            if (receivedContextParamRequest.param == I915_CONTEXT_PARAM_PRIORITY) {
                return this->StoredRetVal;
            }
-            if ((gp->param == I915_CONTEXT_PRIVATE_PARAM_BOOST) && (gp->value == 1)) {
+            if ((receivedContextParamRequest.param == I915_CONTEXT_PRIVATE_PARAM_BOOST) && (receivedContextParamRequest.value == 1)) {
                return this->StoredRetVal;
            }
        }
-#endif
+
        if (request == DRM_IOCTL_I915_GEM_EXECBUFFER2) {
            drm_i915_gem_execbuffer2 *execbuf = (drm_i915_gem_execbuffer2 *)arg;
            this->execBuffer = *execbuf;
@@ -219,6 +227,10 @@ class DrmMock : public Drm {
    int StoredPreemptionSupport = 0;
    int StoredExecSoftPin = 0;
    uint32_t StoredCtxId = 1;
+    uint32_t receivedDestroyContextId = 0;
+
+    uint32_t receivedContextParamRequestCount = 0;
+    drm_i915_gem_context_param receivedContextParamRequest = {};

    //DRM_IOCTL_I915_GEM_EXECBUFFER2
    drm_i915_gem_execbuffer2 execBuffer = {0};
--- a/unit_tests/os_interface/linux/drm_tests.cpp
+++ b/unit_tests/os_interface/linux/drm_tests.cpp
@@ -7,6 +7,7 @@

 #include "runtime/helpers/options.h"
 #include "runtime/os_interface/device_factory.h"
+#include "runtime/os_interface/linux/os_context_linux.h"
 #include "runtime/helpers/file_io.h"
 #include "unit_tests/os_interface/linux/drm_mock.h"
 #include "unit_tests/fixtures/memory_management_fixture.h"
@@ -137,30 +138,58 @@ TEST(DrmTest, GivenDrmWhenAskedForPreemptionCorrectValueReturned) {
    delete pDrm;
 }

-TEST(DrmTest, GivenDrmWhenAskedForContextThatPassedThenValidContextIdsReturned) {
-    DrmMock *pDrm = new DrmMock;
-    EXPECT_EQ(0u, pDrm->lowPriorityContextId);
-    pDrm->StoredRetVal = 0;
-    pDrm->StoredCtxId = 2;
-    pDrm->createLowPriorityContext();
-    EXPECT_EQ(2u, pDrm->lowPriorityContextId);
-    pDrm->StoredRetVal = 0;
-    pDrm->StoredCtxId = 1;
-    delete pDrm;
-}
-
 TEST(DrmTest, GivenDrmWhenAskedForContextThatFailsThenFalseIsReturned) {
    DrmMock *pDrm = new DrmMock;
    pDrm->StoredRetVal = -1;
-    EXPECT_THROW(pDrm->createLowPriorityContext(), std::exception);
+    EXPECT_THROW(pDrm->createDrmContext(), std::exception);
    pDrm->StoredRetVal = 0;
    delete pDrm;
 }

-TEST(DrmTest, GivenDrmWhenContextDestroyIsCalledThenThereAreNoLeaksOrCrashes) {
+TEST(DrmTest, givenDrmWhenOsContextIsCreatedThenCreateAndDestroyNewDrmOsContext) {
    DrmMock drmMock;
-    drmMock.createLowPriorityContext();
-    drmMock.destroyLowPriorityContext();
+    uint32_t drmContextId1 = 123;
+    uint32_t drmContextId2 = 456;
+
+    {
+        drmMock.StoredCtxId = drmContextId1;
+        OsContextLinux osContext1(drmMock, gpgpuEngineInstances[0]);
+        EXPECT_EQ(drmContextId1, osContext1.getDrmContextId());
+        EXPECT_EQ(0u, drmMock.receivedDestroyContextId);
+
+        {
+            drmMock.StoredCtxId = drmContextId2;
+            OsContextLinux osContext2(drmMock, gpgpuEngineInstances[1]);
+            EXPECT_EQ(drmContextId2, osContext2.getDrmContextId());
+            EXPECT_EQ(0u, drmMock.receivedDestroyContextId);
+        }
+        EXPECT_EQ(drmContextId2, drmMock.receivedDestroyContextId);
+    }
+
+    EXPECT_EQ(drmContextId1, drmMock.receivedDestroyContextId);
+    EXPECT_EQ(0u, drmMock.receivedContextParamRequestCount);
+}
+
+TEST(DrmTest, givenDrmPreemptionEnabledAndLowPriorityEngineWhenCreatingOsContextThenCallSetContextPriorityIoctl) {
+    DrmMock drmMock;
+    drmMock.StoredCtxId = 123;
+    drmMock.preemptionSupported = false;
+
+    OsContextLinux osContext1(drmMock, gpgpuEngineInstances[0]);
+    OsContextLinux osContext2(drmMock, gpgpuEngineInstances[EngineInstanceConstants::lowPriorityGpgpuEngineIndex]);
+    EXPECT_EQ(0u, drmMock.receivedContextParamRequestCount);
+
+    drmMock.preemptionSupported = true;
+
+    OsContextLinux osContext3(drmMock, gpgpuEngineInstances[0]);
+    EXPECT_EQ(0u, drmMock.receivedContextParamRequestCount);
+
+    OsContextLinux osContext4(drmMock, gpgpuEngineInstances[EngineInstanceConstants::lowPriorityGpgpuEngineIndex]);
+    EXPECT_EQ(1u, drmMock.receivedContextParamRequestCount);
+    EXPECT_EQ(drmMock.StoredCtxId, drmMock.receivedContextParamRequest.ctx_id);
+    EXPECT_EQ(static_cast<uint64_t>(I915_CONTEXT_PARAM_PRIORITY), drmMock.receivedContextParamRequest.param);
+    EXPECT_EQ(static_cast<uint64_t>(-1023), drmMock.receivedContextParamRequest.value);
+    EXPECT_EQ(0u, drmMock.receivedContextParamRequest.size);
 }

 TEST(DrmTest, getExecSoftPin) {