Allow Device creating multiple CSRs [8/n]

Use OsContextId instead of DeviceIndex for residency

Change-Id: Ib2367b32b5b3e320252d8254f1042f1c3d497068

runtime/command_stream/aub_command_stream_receiver_hw.h

@@ -25,6 +25,7 @@ class AubSubCaptureManager;
 
 template <typename GfxFamily>
 class AUBCommandStreamReceiverHw : public CommandStreamReceiverSimulatedHw<GfxFamily> {
+  protected:
     typedef CommandStreamReceiverSimulatedHw<GfxFamily> BaseClass;
     typedef typename AUBFamilyMapper<GfxFamily>::AUB AUB;
     typedef typename AUB::MiContextDescriptorReg MiContextDescriptorReg;

runtime/command_stream/aub_command_stream_receiver_hw.inl

@@ -368,7 +368,7 @@ FlushStamp AUBCommandStreamReceiverHw<GfxFamily>::flush(BatchBuffer &batchBuffer
             }
         } else {
             allocationsForResidency.push_back(batchBuffer.commandBufferAllocation);
-            batchBuffer.commandBufferAllocation->updateResidencyTaskCount(this->taskCount, this->deviceIndex);
+            batchBuffer.commandBufferAllocation->updateResidencyTaskCount(this->taskCount, this->osContext->getContextId());
         }
     }
     processResidency(allocationsForResidency);
@@ -760,7 +760,7 @@ void AUBCommandStreamReceiverHw<GfxFamily>::processResidency(ResidencyContainer
         if (!writeMemory(*gfxAllocation)) {
             DEBUG_BREAK_IF(!((gfxAllocation->getUnderlyingBufferSize() == 0) || !gfxAllocation->isAubWritable()));
         }
-        gfxAllocation->updateResidencyTaskCount(this->taskCount + 1, this->deviceIndex);
+        gfxAllocation->updateResidencyTaskCount(this->taskCount + 1, this->osContext->getContextId());
     }
 
     dumpAubNonWritable = false;
@@ -768,9 +768,9 @@ void AUBCommandStreamReceiverHw<GfxFamily>::processResidency(ResidencyContainer
 
 template <typename GfxFamily>
 void AUBCommandStreamReceiverHw<GfxFamily>::makeNonResident(GraphicsAllocation &gfxAllocation) {
-    if (gfxAllocation.isResident(this->deviceIndex)) {
+    if (gfxAllocation.isResident(this->osContext->getContextId())) {
         this->getEvictionAllocations().push_back(&gfxAllocation);
-        gfxAllocation.resetResidencyTaskCount(this->deviceIndex);
+        gfxAllocation.resetResidencyTaskCount(this->osContext->getContextId());
     }
 }
 

runtime/command_stream/command_stream_receiver.cpp

@@ -21,6 +21,7 @@
 #include "runtime/memory_manager/internal_allocation_storage.h"
 #include "runtime/memory_manager/memory_manager.h"
 #include "runtime/memory_manager/surface.h"
+#include "runtime/os_interface/os_context.h"
 #include "runtime/os_interface/os_interface.h"
 #include "runtime/utilities/tag_allocator.h"
 
@@ -62,15 +63,15 @@ CommandStreamReceiver::~CommandStreamReceiver() {
 
 void CommandStreamReceiver::makeResident(GraphicsAllocation &gfxAllocation) {
     auto submissionTaskCount = this->taskCount + 1;
-    bool isNotResident = !gfxAllocation.isResident(deviceIndex);
+    bool isNotResident = !gfxAllocation.isResident(osContext->getContextId());
-    if (isNotResident || gfxAllocation.getResidencyTaskCount(deviceIndex) < submissionTaskCount) {
+    if (isNotResident || gfxAllocation.getResidencyTaskCount(osContext->getContextId()) < submissionTaskCount) {
         this->getResidencyAllocations().push_back(&gfxAllocation);
-        gfxAllocation.updateTaskCount(submissionTaskCount, deviceIndex);
+        gfxAllocation.updateTaskCount(submissionTaskCount, osContext->getContextId());
         if (isNotResident) {
             this->totalMemoryUsed += gfxAllocation.getUnderlyingBufferSize();
         }
     }
-    gfxAllocation.updateResidencyTaskCount(submissionTaskCount, deviceIndex);
+    gfxAllocation.updateResidencyTaskCount(submissionTaskCount, osContext->getContextId());
 }
 
 void CommandStreamReceiver::processEviction() {
@@ -78,7 +79,7 @@ void CommandStreamReceiver::processEviction() {
 }
 
 void CommandStreamReceiver::makeNonResident(GraphicsAllocation &gfxAllocation) {
-    if (gfxAllocation.isResident(deviceIndex)) {
+    if (gfxAllocation.isResident(osContext->getContextId())) {
         makeCoherent(gfxAllocation);
         if (gfxAllocation.peekEvictable()) {
             this->getEvictionAllocations().push_back(&gfxAllocation);
@@ -87,7 +88,7 @@ void CommandStreamReceiver::makeNonResident(GraphicsAllocation &gfxAllocation) {
         }
     }
 
-    gfxAllocation.resetResidencyTaskCount(this->deviceIndex);
+    gfxAllocation.resetResidencyTaskCount(this->osContext->getContextId());
 }
 
 void CommandStreamReceiver::makeSurfacePackNonResident(ResidencyContainer &allocationsForResidency) {
@@ -363,7 +364,7 @@ bool CommandStreamReceiver::createAllocationForHostSurface(HostPtrSurface &surfa
     if (allocation == nullptr) {
         return false;
     }
-    allocation->updateTaskCount(Event::eventNotReady, deviceIndex);
+    allocation->updateTaskCount(Event::eventNotReady, osContext->getContextId());
     surface.setAllocation(allocation);
     internalAllocationStorage->storeAllocation(std::unique_ptr<GraphicsAllocation>(allocation), TEMPORARY_ALLOCATION);
     return true;

runtime/command_stream/command_stream_receiver.h

@@ -156,7 +156,6 @@ class CommandStreamReceiver {
     size_t defaultSshSize;
 
     void setDeviceIndex(uint32_t deviceIndex) { this->deviceIndex = deviceIndex; }
-    uint32_t getDeviceIndex() const { return this->deviceIndex; }
     AllocationsList &getTemporaryAllocations();
     AllocationsList &getAllocationsForReuse();
     InternalAllocationStorage *getInternalAllocationStorage() const { return internalAllocationStorage.get(); }

runtime/command_stream/command_stream_receiver_hw.h

@@ -72,6 +72,8 @@ class CommandStreamReceiverHw : public CommandStreamReceiver {
     }
 
   protected:
+    using CommandStreamReceiver::osContext;
+
     void programPreemption(LinearStream &csr, Device &device, DispatchFlags &dispatchFlags);
     void programL3(LinearStream &csr, DispatchFlags &dispatchFlags, uint32_t &newL3Config);
     void programPreamble(LinearStream &csr, Device &device, DispatchFlags &dispatchFlags, uint32_t &newL3Config);

runtime/command_stream/command_stream_receiver_hw.inl

@@ -247,7 +247,7 @@ CompletionStamp CommandStreamReceiverHw<GfxFamily>::flushTask(
         scratchSpaceController->setRequiredScratchSpace(ssh.getCpuBase(),
                                                         requiredScratchSize,
                                                         this->taskCount,
-                                                        this->deviceIndex,
+                                                        this->osContext->getContextId(),
                                                         stateBaseAddressDirty,
                                                         checkVfeStateDirty);
         if (checkVfeStateDirty) {

runtime/command_stream/command_stream_receiver_simulated_common_hw.h

@@ -19,7 +19,9 @@ namespace OCLRT {
 class GraphicsAllocation;
 template <typename GfxFamily>
 class CommandStreamReceiverSimulatedCommonHw : public CommandStreamReceiverHw<GfxFamily> {
+  protected:
     using CommandStreamReceiverHw<GfxFamily>::CommandStreamReceiverHw;
+    using CommandStreamReceiverHw<GfxFamily>::osContext;
 
   public:
     uint64_t getGTTBits() const {

runtime/command_stream/command_stream_receiver_with_aub_dump.h

@@ -12,6 +12,9 @@ namespace OCLRT {
 
 template <typename BaseCSR>
 class CommandStreamReceiverWithAUBDump : public BaseCSR {
+  protected:
+    using BaseCSR::osContext;
+
   public:
     using BaseCSR::createMemoryManager;
     CommandStreamReceiverWithAUBDump(const HardwareInfo &hwInfoIn, ExecutionEnvironment &executionEnvironment);

runtime/command_stream/command_stream_receiver_with_aub_dump.inl

@@ -34,9 +34,9 @@ FlushStamp CommandStreamReceiverWithAUBDump<BaseCSR>::flush(BatchBuffer &batchBu
 
 template <typename BaseCSR>
 void CommandStreamReceiverWithAUBDump<BaseCSR>::makeNonResident(GraphicsAllocation &gfxAllocation) {
-    uint32_t residencyTaskCount = gfxAllocation.getResidencyTaskCount(this->deviceIndex);
+    uint32_t residencyTaskCount = gfxAllocation.getResidencyTaskCount(this->osContext->getContextId());
     BaseCSR::makeNonResident(gfxAllocation);
-    gfxAllocation.updateResidencyTaskCount(residencyTaskCount, this->deviceIndex);
+    gfxAllocation.updateResidencyTaskCount(residencyTaskCount, this->osContext->getContextId());
     if (aubCSR) {
         aubCSR->makeNonResident(gfxAllocation);
     }

runtime/command_stream/definitions/command_stream_receiver_simulated_hw.h

@@ -13,7 +13,9 @@ namespace OCLRT {
 class GraphicsAllocation;
 template <typename GfxFamily>
 class CommandStreamReceiverSimulatedHw : public CommandStreamReceiverSimulatedCommonHw<GfxFamily> {
+  protected:
     using CommandStreamReceiverSimulatedCommonHw<GfxFamily>::CommandStreamReceiverSimulatedCommonHw;
+    using CommandStreamReceiverSimulatedCommonHw<GfxFamily>::osContext;
 
   public:
     uint32_t getMemoryBank(GraphicsAllocation *allocation) const {

runtime/command_stream/scratch_space_controller_base.cpp

@@ -22,13 +22,13 @@ ScratchSpaceControllerBase::ScratchSpaceControllerBase(const HardwareInfo &info,
 void ScratchSpaceControllerBase::setRequiredScratchSpace(void *sshBaseAddress,
                                                          uint32_t requiredPerThreadScratchSize,
                                                          uint32_t currentTaskCount,
-                                                         uint32_t deviceIdx,
+                                                         uint32_t contextId,
                                                          bool &stateBaseAddressDirty,
                                                          bool &vfeStateDirty) {
     size_t requiredScratchSizeInBytes = requiredPerThreadScratchSize * computeUnitsUsedForScratch;
     if (requiredScratchSizeInBytes && (!scratchAllocation || scratchSizeBytes < requiredScratchSizeInBytes)) {
         if (scratchAllocation) {
-            scratchAllocation->updateTaskCount(currentTaskCount, deviceIdx);
+            scratchAllocation->updateTaskCount(currentTaskCount, contextId);
             csrAllocationStorage.storeAllocation(std::unique_ptr<GraphicsAllocation>(scratchAllocation), TEMPORARY_ALLOCATION);
         }
         scratchSizeBytes = requiredScratchSizeInBytes;

runtime/command_stream/scratch_space_controller_base.h

@@ -17,7 +17,7 @@ class ScratchSpaceControllerBase : public ScratchSpaceController {
     void setRequiredScratchSpace(void *sshBaseAddress,
                                  uint32_t requiredPerThreadScratchSize,
                                  uint32_t currentTaskCount,
-                                 uint32_t deviceIdx,
+                                 uint32_t contextId,
                                  bool &stateBaseAddressDirty,
                                  bool &vfeStateDirty) override;
     uint64_t calculateNewGSH() override;

runtime/command_stream/tbx_command_stream_receiver_hw.h

@@ -27,6 +27,7 @@ class TbxMemoryManager : public OsAgnosticMemoryManager {
 
 template <typename GfxFamily>
 class TbxCommandStreamReceiverHw : public CommandStreamReceiverSimulatedHw<GfxFamily> {
+  protected:
     typedef CommandStreamReceiverSimulatedHw<GfxFamily> BaseClass;
     typedef typename OCLRT::AUBFamilyMapper<GfxFamily>::AUB AUB;
     typedef typename AUB::MiContextDescriptorReg MiContextDescriptorReg;

runtime/command_stream/tbx_command_stream_receiver_hw.inl

@@ -371,7 +371,7 @@ void TbxCommandStreamReceiverHw<GfxFamily>::processResidency(ResidencyContainer
         if (!writeMemory(*gfxAllocation)) {
             DEBUG_BREAK_IF(!(gfxAllocation->getUnderlyingBufferSize() == 0));
         }
-        gfxAllocation->updateResidencyTaskCount(this->taskCount + 1, this->deviceIndex);
+        gfxAllocation->updateResidencyTaskCount(this->taskCount + 1, this->osContext->getContextId());
     }
 }
 

runtime/mem_obj/mem_obj.cpp

@@ -16,6 +16,7 @@
 #include "runtime/gmm_helper/gmm.h"
 #include "runtime/helpers/aligned_memory.h"
 #include "runtime/helpers/get_info.h"
+#include "runtime/os_interface/os_context.h"
 #include <algorithm>
 
 namespace OCLRT {
@@ -288,7 +289,8 @@ void MemObj::releaseAllocatedMapPtr() {
 }
 
 void MemObj::waitForCsrCompletion() {
-    memoryManager->getDefaultCommandStreamReceiver(0)->waitForCompletionWithTimeout(false, TimeoutControls::maxTimeout, graphicsAllocation->getTaskCount(0u));
+    auto osContextId = context->getDevice(0)->getDefaultEngine().osContext->getContextId();
+    memoryManager->getDefaultCommandStreamReceiver(0)->waitForCompletionWithTimeout(false, TimeoutControls::maxTimeout, graphicsAllocation->getTaskCount(osContextId));
 }
 
 void MemObj::destroyGraphicsAllocation(GraphicsAllocation *allocation, bool asyncDestroy) {

runtime/memory_manager/internal_allocation_storage.cpp

@@ -8,11 +8,12 @@
 #include "runtime/command_stream/command_stream_receiver.h"
 #include "runtime/memory_manager/internal_allocation_storage.h"
 #include "runtime/memory_manager/memory_manager.h"
+#include "runtime/os_interface/os_context.h"
 
 namespace OCLRT {
-InternalAllocationStorage::InternalAllocationStorage(CommandStreamReceiver &commandStreamReceiver) : commandStreamReceiver(commandStreamReceiver), contextId(commandStreamReceiver.getDeviceIndex()){};
+InternalAllocationStorage::InternalAllocationStorage(CommandStreamReceiver &commandStreamReceiver) : commandStreamReceiver(commandStreamReceiver){};
 void InternalAllocationStorage::storeAllocation(std::unique_ptr<GraphicsAllocation> gfxAllocation, uint32_t allocationUsage) {
-    uint32_t taskCount = gfxAllocation->getTaskCount(contextId);
+    uint32_t taskCount = gfxAllocation->getTaskCount(commandStreamReceiver.getOsContext().getContextId());
 
     if (allocationUsage == REUSABLE_ALLOCATION) {
         taskCount = commandStreamReceiver.peekTaskCount();
@@ -28,7 +29,7 @@ void InternalAllocationStorage::storeAllocationWithTaskCount(std::unique_ptr<Gra
         }
     }
     auto &allocationsList = (allocationUsage == TEMPORARY_ALLOCATION) ? temporaryAllocations : allocationsForReuse;
-    gfxAllocation->updateTaskCount(taskCount, contextId);
+    gfxAllocation->updateTaskCount(taskCount, commandStreamReceiver.getOsContext().getContextId());
     allocationsList.pushTailOne(*gfxAllocation.release());
 }
 
@@ -43,7 +44,7 @@ void InternalAllocationStorage::freeAllocationsList(uint32_t waitTaskCount, Allo
     IDList<GraphicsAllocation, false, true> allocationsLeft;
     while (curr != nullptr) {
         auto *next = curr->next;
-        if (curr->getTaskCount(contextId) <= waitTaskCount) {
+        if (curr->getTaskCount(commandStreamReceiver.getOsContext().getContextId()) <= waitTaskCount) {
             memoryManager->freeGraphicsMemory(curr);
         } else {
             allocationsLeft.pushTailOne(*curr);
@@ -73,7 +74,7 @@ std::unique_ptr<GraphicsAllocation> AllocationsList::detachAllocation(size_t req
     req.requiredMinimalSize = requiredMinimalSize;
     req.csrTagAddress = commandStreamReceiver.getTagAddress();
     req.internalAllocationRequired = internalAllocationRequired;
-    req.contextId = commandStreamReceiver.getDeviceIndex();
+    req.contextId = commandStreamReceiver.getOsContext().getContextId();
     GraphicsAllocation *a = nullptr;
     GraphicsAllocation *retAlloc = processLocked<AllocationsList, &AllocationsList::detachAllocationImpl>(a, static_cast<void *>(&req));
     return std::unique_ptr<GraphicsAllocation>(retAlloc);

runtime/memory_manager/internal_allocation_storage.h

@@ -27,7 +27,6 @@ class InternalAllocationStorage {
   protected:
     void freeAllocationsList(uint32_t waitTaskCount, AllocationsList &allocationsList);
     CommandStreamReceiver &commandStreamReceiver;
-    const uint32_t contextId;
 
     AllocationsList temporaryAllocations;
     AllocationsList allocationsForReuse;

runtime/os_interface/linux/drm_command_stream.inl

@@ -129,7 +129,7 @@ void DrmCommandStreamReceiver<GfxFamily>::makeNonResident(GraphicsAllocation &gf
     // Vector is moved to command buffer inside flush.
     // If flush wasn't called we need to make all objects non-resident.
     // If makeNonResident is called before flush, vector will be cleared.
-    if (gfxAllocation.isResident(this->deviceIndex)) {
+    if (gfxAllocation.isResident(this->osContext->getContextId())) {
         if (this->residency.size() != 0) {
             this->residency.clear();
         }
@@ -139,7 +139,7 @@ void DrmCommandStreamReceiver<GfxFamily>::makeNonResident(GraphicsAllocation &gf
             }
         }
     }
-    gfxAllocation.resetResidencyTaskCount(this->deviceIndex);
+    gfxAllocation.resetResidencyTaskCount(this->osContext->getContextId());
 }
 
 template <typename GfxFamily>

runtime/os_interface/windows/wddm_device_command_stream.inl

@@ -71,7 +71,7 @@ FlushStamp WddmCommandStreamReceiver<GfxFamily>::flush(BatchBuffer &batchBuffer,
         makeResident(*batchBuffer.commandBufferAllocation);
     } else {
         allocationsForResidency.push_back(batchBuffer.commandBufferAllocation);
-        batchBuffer.commandBufferAllocation->updateResidencyTaskCount(this->taskCount, this->deviceIndex);
+        batchBuffer.commandBufferAllocation->updateResidencyTaskCount(this->taskCount, this->osContext->getContextId());
     }
 
     this->processResidency(allocationsForResidency);

runtime/os_interface/windows/wddm_memory_manager.cpp

@@ -289,7 +289,7 @@ void WddmMemoryManager::freeGraphicsMemoryImpl(GraphicsAllocation *gfxAllocation
     UNRECOVERABLE_IF(DebugManager.flags.CreateMultipleDevices.get() == 0 &&
                      gfxAllocation->isUsed() && this->executionEnvironment.commandStreamReceivers.size() > 0 &&
                      this->getDefaultCommandStreamReceiver(0) && this->getDefaultCommandStreamReceiver(0)->getTagAddress() &&
-                     gfxAllocation->getTaskCount(0u) > *this->getDefaultCommandStreamReceiver(0)->getTagAddress());
+                     gfxAllocation->getTaskCount(defaultEngineIndex) > *this->getDefaultCommandStreamReceiver(0)->getTagAddress());
 
     if (input->gmm) {
         if (input->gmm->isRenderCompressed && wddm->getPageTableManager()) {

unit_tests/command_queue/enqueue_api_tests_mt_with_asyncGPU.cpp

@@ -8,6 +8,7 @@
 #include "buffer_operations_withAsyncGPU_fixture.h"
 #include "runtime/memory_manager/graphics_allocation.h"
 #include "runtime/helpers/aligned_memory.h"
+#include "runtime/os_interface/os_context.h"
 #include "test.h"
 #include <thread>
 
@@ -58,7 +59,7 @@ HWTEST_F(AsyncGPUoperations, MapBufferAfterWriteBuffer) {
     }
     t.join();
 
-    srcBuffer->getGraphicsAllocation()->updateTaskCount(0u, 0u);
+    srcBuffer->getGraphicsAllocation()->updateTaskCount(0u, pCmdQ->getCommandStreamReceiver().getOsContext().getContextId());
 
     alignedFree(ptrMemory);
 }

unit_tests/command_queue/enqueue_debug_kernel_tests.cpp

@@ -9,6 +9,7 @@
 #include "runtime/command_queue/command_queue.h"
 #include "runtime/program/program.h"
 #include "runtime/source_level_debugger/source_level_debugger.h"
+#include "runtime/os_interface/os_context.h"
 #include "unit_tests/fixtures/enqueue_handler_fixture.h"
 #include "unit_tests/helpers/kernel_binary_helper.h"
 #include "unit_tests/helpers/kernel_filename_helper.h"
@@ -104,7 +105,7 @@ HWTEST_F(EnqueueDebugKernelTest, givenDebugKernelWhenEnqueuedThenSSHAndBtiAreCor
 
         auto &commandStreamReceiver = mockCmdQ->getCommandStreamReceiver();
         auto debugSurface = commandStreamReceiver.getDebugSurfaceAllocation();
-        EXPECT_EQ(1u, debugSurface->getTaskCount(0u));
+        EXPECT_EQ(1u, debugSurface->getTaskCount(commandStreamReceiver.getOsContext().getContextId()));
 
         EXPECT_EQ(debugSurface->getGpuAddress(), debugSurfaceState->getSurfaceBaseAddress());
     }

unit_tests/command_queue/enqueue_fill_buffer_tests.cpp

@@ -9,6 +9,7 @@
 #include "runtime/built_ins/builtins_dispatch_builder.h"
 #include "runtime/command_queue/command_queue.h"
 #include "runtime/command_stream/command_stream_receiver.h"
+#include "runtime/os_interface/os_context.h"
 #include "reg_configs_common.h"
 #include "runtime/helpers/ptr_math.h"
 #include "runtime/helpers/aligned_memory.h"
@@ -376,7 +377,7 @@ HWTEST_F(EnqueueFillBufferCmdTests, patternShouldBeCopied) {
         if ((allocation->getUnderlyingBufferSize() >= sizeof(float)) &&
             (allocation->getUnderlyingBuffer() != nullptr) &&
             (*(static_cast<float *>(allocation->getUnderlyingBuffer())) == EnqueueFillBufferHelper<>::Traits::pattern[0]) &&
-            (pCmdQ->taskCount == allocation->getTaskCount(0))) {
+            (pCmdQ->taskCount == allocation->getTaskCount(csr.getOsContext().getContextId()))) {
             break;
         }
         allocation = allocation->next;
@@ -397,7 +398,7 @@ HWTEST_F(EnqueueFillBufferCmdTests, patternShouldBeAligned) {
         if ((allocation->getUnderlyingBufferSize() >= sizeof(float)) &&
             (allocation->getUnderlyingBuffer() != nullptr) &&
             (*(static_cast<float *>(allocation->getUnderlyingBuffer())) == EnqueueFillBufferHelper<>::Traits::pattern[0]) &&
-            (pCmdQ->taskCount == allocation->getTaskCount(0))) {
+            (pCmdQ->taskCount == allocation->getTaskCount(csr.getOsContext().getContextId()))) {
             break;
         }
         allocation = allocation->next;

unit_tests/command_queue/enqueue_kernel_1_tests.cpp

@@ -917,7 +917,7 @@ HWTEST_F(EnqueueKernelTest, givenKernelWhenItIsEnqueuedThenAllResourceGraphicsAl
     auto csrTaskCount = mockCsr->peekTaskCount();
     auto &passedAllocationPack = mockCsr->copyOfAllocations;
     for (auto &allocation : passedAllocationPack) {
-        EXPECT_EQ(csrTaskCount, allocation->getTaskCount(0u));
+        EXPECT_EQ(csrTaskCount, allocation->getTaskCount(mockCsr->getOsContext().getContextId()));
     }
 }
 

unit_tests/command_stream/aub_command_stream_receiver_1_tests.cpp

@@ -235,30 +235,33 @@ HWTEST_F(AubCommandStreamReceiverTests, givenGraphicsAllocationWhenMakeResidentC
     std::unique_ptr<MemoryManager> memoryManager(nullptr);
     std::unique_ptr<AUBCommandStreamReceiverHw<FamilyType>> aubCsr(new AUBCommandStreamReceiverHw<FamilyType>(*platformDevices[0], "", true, *pDevice->executionEnvironment));
     memoryManager.reset(aubCsr->createMemoryManager(false, false));
+    aubCsr->setOsContext(*pDevice->getDefaultEngine().osContext);
     auto gfxAllocation = memoryManager->allocateGraphicsMemory(sizeof(uint32_t), sizeof(uint32_t), false, false);
 
+    auto osContextId = aubCsr->getOsContext().getContextId();
+
     // First makeResident marks the allocation resident
     aubCsr->makeResident(*gfxAllocation);
-    EXPECT_TRUE(gfxAllocation->isResident(0u));
+    EXPECT_TRUE(gfxAllocation->isResident(osContextId));
-    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->getTaskCount(0));
+    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->getTaskCount(osContextId));
-    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->getResidencyTaskCount(0u));
+    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->getResidencyTaskCount(osContextId));
     EXPECT_EQ(1u, aubCsr->getResidencyAllocations().size());
 
     // Second makeResident should have no impact
     aubCsr->makeResident(*gfxAllocation);
-    EXPECT_TRUE(gfxAllocation->isResident(0u));
+    EXPECT_TRUE(gfxAllocation->isResident(osContextId));
-    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->getTaskCount(0));
+    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->getTaskCount(osContextId));
-    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->getResidencyTaskCount(0u));
+    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->getResidencyTaskCount(osContextId));
     EXPECT_EQ(1u, aubCsr->getResidencyAllocations().size());
 
     // First makeNonResident marks the allocation as nonresident
     aubCsr->makeNonResident(*gfxAllocation);
-    EXPECT_FALSE(gfxAllocation->isResident(0u));
+    EXPECT_FALSE(gfxAllocation->isResident(osContextId));
     EXPECT_EQ(1u, aubCsr->getEvictionAllocations().size());
 
     // Second makeNonResident should have no impact
     aubCsr->makeNonResident(*gfxAllocation);
-    EXPECT_FALSE(gfxAllocation->isResident(0u));
+    EXPECT_FALSE(gfxAllocation->isResident(osContextId));
     EXPECT_EQ(1u, aubCsr->getEvictionAllocations().size());
 
     memoryManager->freeGraphicsMemoryImpl(gfxAllocation);
@@ -502,17 +505,17 @@ HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverInStandalon
 
     BatchBuffer batchBuffer{cs.getGraphicsAllocation(), 0, 0, nullptr, false, false, QueueThrottle::MEDIUM, cs.getUsed(), &cs};
 
-    EXPECT_FALSE(commandBuffer->isResident(0u));
+    EXPECT_FALSE(commandBuffer->isResident(aubCsr->getOsContext().getContextId()));
 
     aubCsr->overrideDispatchPolicy(DispatchMode::ImmediateDispatch);
     aubCsr->flush(batchBuffer, allocationsForResidency);
 
-    EXPECT_TRUE(commandBuffer->isResident(0u));
+    EXPECT_TRUE(commandBuffer->isResident(aubCsr->getOsContext().getContextId()));
-    EXPECT_EQ(aubCsr->peekTaskCount() + 1, commandBuffer->getResidencyTaskCount(0u));
+    EXPECT_EQ(aubCsr->peekTaskCount() + 1, commandBuffer->getResidencyTaskCount(aubCsr->getOsContext().getContextId()));
 
     aubCsr->makeSurfacePackNonResident(aubCsr->getResidencyAllocations());
 
-    EXPECT_FALSE(commandBuffer->isResident(0u));
+    EXPECT_FALSE(commandBuffer->isResident(aubCsr->getOsContext().getContextId()));
 }
 
 HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverInNonStandaloneModeWhenFlushIsCalledThenItShouldNotCallMakeResidentOnCommandBufferAllocation) {
@@ -523,11 +526,11 @@ HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverInNonStanda
 
     BatchBuffer batchBuffer{cs.getGraphicsAllocation(), 0, 0, nullptr, false, false, QueueThrottle::MEDIUM, cs.getUsed(), &cs};
 
-    EXPECT_FALSE(aubExecutionEnvironment->commandBuffer->isResident(0u));
+    EXPECT_FALSE(aubExecutionEnvironment->commandBuffer->isResident(aubCsr->getOsContext().getContextId()));
 
     aubCsr->flush(batchBuffer, allocationsForResidency);
 
-    EXPECT_FALSE(aubExecutionEnvironment->commandBuffer->isResident(0u));
+    EXPECT_FALSE(aubExecutionEnvironment->commandBuffer->isResident(aubCsr->getOsContext().getContextId()));
 }
 
 HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverInStandaloneModeWhenFlushIsCalledThenItShouldCallMakeResidentOnResidencyAllocations) {
@@ -544,22 +547,22 @@ HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverInStandalon
 
     ResidencyContainer allocationsForResidency = {gfxAllocation};
 
-    EXPECT_FALSE(gfxAllocation->isResident(0u));
+    EXPECT_FALSE(gfxAllocation->isResident(aubCsr->getOsContext().getContextId()));
-    EXPECT_FALSE(commandBuffer->isResident(0u));
+    EXPECT_FALSE(commandBuffer->isResident(aubCsr->getOsContext().getContextId()));
 
     aubCsr->overrideDispatchPolicy(DispatchMode::BatchedDispatch);
     aubCsr->flush(batchBuffer, allocationsForResidency);
 
-    EXPECT_TRUE(gfxAllocation->isResident(0u));
+    EXPECT_TRUE(gfxAllocation->isResident(aubCsr->getOsContext().getContextId()));
-    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->getResidencyTaskCount(0u));
+    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->getResidencyTaskCount(aubCsr->getOsContext().getContextId()));
 
-    EXPECT_TRUE(commandBuffer->isResident(0u));
+    EXPECT_TRUE(commandBuffer->isResident(aubCsr->getOsContext().getContextId()));
-    EXPECT_EQ(aubCsr->peekTaskCount() + 1, commandBuffer->getResidencyTaskCount(0u));
+    EXPECT_EQ(aubCsr->peekTaskCount() + 1, commandBuffer->getResidencyTaskCount(aubCsr->getOsContext().getContextId()));
 
     aubCsr->makeSurfacePackNonResident(allocationsForResidency);
 
-    EXPECT_FALSE(gfxAllocation->isResident(0u));
+    EXPECT_FALSE(gfxAllocation->isResident(aubCsr->getOsContext().getContextId()));
-    EXPECT_FALSE(commandBuffer->isResident(0u));
+    EXPECT_FALSE(commandBuffer->isResident(aubCsr->getOsContext().getContextId()));
 
     memoryManager->freeGraphicsMemory(gfxAllocation);
 }
@@ -577,15 +580,15 @@ HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverInNonStanda
 
     ResidencyContainer allocationsForResidency = {gfxAllocation};
 
-    EXPECT_FALSE(gfxAllocation->isResident(0u));
+    EXPECT_FALSE(gfxAllocation->isResident(aubCsr->getOsContext().getContextId()));
-    EXPECT_FALSE(commandBuffer->isResident(0u));
+    EXPECT_FALSE(commandBuffer->isResident(aubCsr->getOsContext().getContextId()));
 
     aubCsr->flush(batchBuffer, allocationsForResidency);
 
-    EXPECT_TRUE(gfxAllocation->isResident(0u));
+    EXPECT_TRUE(gfxAllocation->isResident(aubCsr->getOsContext().getContextId()));
-    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->getResidencyTaskCount(0u));
+    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->getResidencyTaskCount(aubCsr->getOsContext().getContextId()));
 
-    EXPECT_FALSE(commandBuffer->isResident(0u));
+    EXPECT_FALSE(commandBuffer->isResident(aubCsr->getOsContext().getContextId()));
 
     memoryManager->freeGraphicsMemoryImpl(gfxAllocation);
 }
@@ -615,22 +618,22 @@ HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverInStandalon
 
     ResidencyContainer allocationsForResidency = {gfxAllocation};
 
-    EXPECT_FALSE(gfxAllocation->isResident(0u));
+    EXPECT_FALSE(gfxAllocation->isResident(aubCsr->getOsContext().getContextId()));
-    EXPECT_FALSE(commandBuffer->isResident(0u));
+    EXPECT_FALSE(commandBuffer->isResident(aubCsr->getOsContext().getContextId()));
 
     aubCsr->overrideDispatchPolicy(DispatchMode::BatchedDispatch);
     aubCsr->flush(batchBuffer, allocationsForResidency);
 
-    EXPECT_TRUE(gfxAllocation->isResident(0u));
+    EXPECT_TRUE(gfxAllocation->isResident(aubCsr->getOsContext().getContextId()));
-    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->getResidencyTaskCount(0u));
+    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->getResidencyTaskCount(aubCsr->getOsContext().getContextId()));
 
-    EXPECT_TRUE(commandBuffer->isResident(0u));
+    EXPECT_TRUE(commandBuffer->isResident(aubCsr->getOsContext().getContextId()));
-    EXPECT_EQ(aubCsr->peekTaskCount() + 1, commandBuffer->getResidencyTaskCount(0u));
+    EXPECT_EQ(aubCsr->peekTaskCount() + 1, commandBuffer->getResidencyTaskCount(aubCsr->getOsContext().getContextId()));
 
     aubCsr->makeSurfacePackNonResident(allocationsForResidency);
 
-    EXPECT_FALSE(gfxAllocation->isResident(0u));
+    EXPECT_FALSE(gfxAllocation->isResident(aubCsr->getOsContext().getContextId()));
-    EXPECT_FALSE(commandBuffer->isResident(0u));
+    EXPECT_FALSE(commandBuffer->isResident(aubCsr->getOsContext().getContextId()));
 
     memoryManager->freeGraphicsMemory(gfxAllocation);
 }
@@ -689,6 +692,7 @@ HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverWhenProcess
     std::unique_ptr<MemoryManager> memoryManager(nullptr);
     std::unique_ptr<AUBCommandStreamReceiverHw<FamilyType>> aubCsr(new AUBCommandStreamReceiverHw<FamilyType>(*platformDevices[0], "", true, *pDevice->executionEnvironment));
     memoryManager.reset(aubCsr->createMemoryManager(false, false));
+    aubCsr->setOsContext(*pDevice->getDefaultEngine().osContext);
 
     auto gfxBufferAllocation = memoryManager->allocateGraphicsMemory(sizeof(uint32_t), sizeof(uint32_t), false, false);
     gfxBufferAllocation->setAllocationType(GraphicsAllocation::AllocationType::BUFFER);
@@ -711,6 +715,7 @@ HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverInSubCaptur
     std::unique_ptr<MemoryManager> memoryManager(nullptr);
     std::unique_ptr<MockAubCsrToTestDumpAubNonWritable<FamilyType>> aubCsr(new MockAubCsrToTestDumpAubNonWritable<FamilyType>(*platformDevices[0], "", true, *pDevice->executionEnvironment));
     memoryManager.reset(aubCsr->createMemoryManager(false, false));
+    aubCsr->setOsContext(*pDevice->getDefaultEngine().osContext);
 
     auto gfxBufferAllocation = memoryManager->allocateGraphicsMemory(sizeof(uint32_t), sizeof(uint32_t), false, false);
     gfxBufferAllocation->setAllocationType(GraphicsAllocation::AllocationType::BUFFER);
@@ -737,6 +742,7 @@ HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverWhenProcess
     std::unique_ptr<MemoryManager> memoryManager(nullptr);
     std::unique_ptr<MockAubCsrToTestDumpAubNonWritable<FamilyType>> aubCsr(new MockAubCsrToTestDumpAubNonWritable<FamilyType>(*platformDevices[0], "", true, *pDevice->executionEnvironment));
     memoryManager.reset(aubCsr->createMemoryManager(false, false));
+    aubCsr->setOsContext(*pDevice->getDefaultEngine().osContext);
 
     auto gfxBufferAllocation = memoryManager->allocateGraphicsMemory(sizeof(uint32_t), sizeof(uint32_t), false, false);
     gfxBufferAllocation->setAllocationType(GraphicsAllocation::AllocationType::BUFFER);

unit_tests/command_stream/command_stream_receiver_flush_task_1_tests.cpp

@@ -751,6 +751,8 @@ HWTEST_F(CommandStreamReceiverFlushTaskTests, flushTaskWithBothCSCallsFlushOnce)
     commandStreamReceiver.initializeTagAllocation();
     commandStream.getSpace(sizeof(typename FamilyType::MI_NOOP));
 
+    commandStreamReceiver.setOsContext(*pDevice->getDefaultEngine().osContext);
+
     flushTask(commandStreamReceiver);
     EXPECT_EQ(1, commandStreamReceiver.flushCount);
 }

unit_tests/command_stream/command_stream_receiver_flush_task_3_tests.cpp

@@ -378,8 +378,8 @@ HWTEST_F(CommandStreamReceiverFlushTaskTests, givenCsrInBatchingModeWhenRecorded
     std::vector<GraphicsAllocation *> residentSurfaces = cmdBuffer->surfaces;
 
     for (auto &graphicsAllocation : residentSurfaces) {
-        EXPECT_TRUE(graphicsAllocation->isResident(0u));
+        EXPECT_TRUE(graphicsAllocation->isResident(mockCsr->getOsContext().getContextId()));
-        EXPECT_EQ(1u, graphicsAllocation->getResidencyTaskCount(0u));
+        EXPECT_EQ(1u, graphicsAllocation->getResidencyTaskCount(mockCsr->getOsContext().getContextId()));
     }
 
     mockCsr->flushBatchedSubmissions();
@@ -395,7 +395,7 @@ HWTEST_F(CommandStreamReceiverFlushTaskTests, givenCsrInBatchingModeWhenRecorded
     EXPECT_EQ(0u, surfacesForResidency.size());
 
     for (auto &graphicsAllocation : residentSurfaces) {
-        EXPECT_FALSE(graphicsAllocation->isResident(0u));
+        EXPECT_FALSE(graphicsAllocation->isResident(mockCsr->getOsContext().getContextId()));
     }
 }
 
@@ -665,6 +665,7 @@ HWTEST_F(CommandStreamReceiverFlushTaskTests, givenCsrInBatchingModeWhenTotalRes
     mockCsr->initializeTagAllocation();
     mockCsr->setPreemptionCsrAllocation(pDevice->getPreemptionAllocation());
     mockCsr->overrideDispatchPolicy(DispatchMode::BatchedDispatch);
+    mockCsr->setOsContext(*pDevice->getDefaultEngine().osContext);
 
     auto mockedSubmissionsAggregator = new mockSubmissionsAggregator();
     mockCsr->overrideSubmissionAggregator(mockedSubmissionsAggregator);
@@ -1193,11 +1194,13 @@ HWTEST_F(CommandStreamReceiverFlushTaskTests, givenCsrWhenTemporaryAndReusableAl
     commandStreamReceiver.getInternalAllocationStorage()->storeAllocation(std::unique_ptr<GraphicsAllocation>(reusableToClean), REUSABLE_ALLOCATION);
     commandStreamReceiver.getInternalAllocationStorage()->storeAllocation(std::unique_ptr<GraphicsAllocation>(reusableToHold), REUSABLE_ALLOCATION);
 
-    temporaryToClean->updateTaskCount(1, 0u);
+    auto osContextId = commandStreamReceiver.getOsContext().getContextId();
-    reusableToClean->updateTaskCount(1, 0u);
 
-    temporaryToHold->updateTaskCount(10, 0u);
+    temporaryToClean->updateTaskCount(1, osContextId);
-    reusableToHold->updateTaskCount(10, 0u);
+    reusableToClean->updateTaskCount(1, osContextId);
+
+    temporaryToHold->updateTaskCount(10, osContextId);
+    reusableToHold->updateTaskCount(10, osContextId);
 
     commandStreamReceiver.latestFlushedTaskCount = 9;
     commandStreamReceiver.cleanupResources();

unit_tests/command_stream/command_stream_receiver_tests.cpp

@@ -82,11 +82,11 @@ TEST_F(CommandStreamReceiverTest, makeResident_setsBufferResidencyFlag) {
         srcMemory,
         retVal);
     ASSERT_NE(nullptr, buffer);
-    EXPECT_FALSE(buffer->getGraphicsAllocation()->isResident(0u));
+    EXPECT_FALSE(buffer->getGraphicsAllocation()->isResident(commandStreamReceiver->getOsContext().getContextId()));
 
     commandStreamReceiver->makeResident(*buffer->getGraphicsAllocation());
 
-    EXPECT_TRUE(buffer->getGraphicsAllocation()->isResident(0u));
+    EXPECT_TRUE(buffer->getGraphicsAllocation()->isResident(commandStreamReceiver->getOsContext().getContextId()));
 
     delete buffer;
 }
@@ -185,7 +185,7 @@ HWTEST_F(CommandStreamReceiverTest, whenStoreAllocationThenStoredAllocationHasTa
 
     csr.getInternalAllocationStorage()->storeAllocation(std::unique_ptr<GraphicsAllocation>(allocation), REUSABLE_ALLOCATION);
 
-    EXPECT_EQ(csr.peekTaskCount(), allocation->getTaskCount(0));
+    EXPECT_EQ(csr.peekTaskCount(), allocation->getTaskCount(csr.getOsContext().getContextId()));
 }
 
 HWTEST_F(CommandStreamReceiverTest, givenCommandStreamReceiverWhenCheckedForInitialStatusOfStatelessMocsIndexThenUnknownMocsIsReturend) {
@@ -392,6 +392,9 @@ TEST(CommandStreamReceiverMultiContextTests, givenMultipleCsrsWhenSameResourcesA
     auto &commandStreamReceiver0 = device0->getCommandStreamReceiver();
     auto &commandStreamReceiver1 = device1->getCommandStreamReceiver();
 
+    auto csr0ContextId = commandStreamReceiver0.getOsContext().getContextId();
+    auto csr1ContextId = commandStreamReceiver1.getOsContext().getContextId();
+
     MockGraphicsAllocation graphicsAllocation;
 
     commandStreamReceiver0.makeResident(graphicsAllocation);
@@ -402,16 +405,16 @@ TEST(CommandStreamReceiverMultiContextTests, givenMultipleCsrsWhenSameResourcesA
     EXPECT_EQ(1u, commandStreamReceiver0.getResidencyAllocations().size());
     EXPECT_EQ(1u, commandStreamReceiver1.getResidencyAllocations().size());
 
-    EXPECT_EQ(1u, graphicsAllocation.getResidencyTaskCount(0u));
+    EXPECT_EQ(1u, graphicsAllocation.getResidencyTaskCount(csr0ContextId));
-    EXPECT_EQ(1u, graphicsAllocation.getResidencyTaskCount(1u));
+    EXPECT_EQ(1u, graphicsAllocation.getResidencyTaskCount(csr1ContextId));
 
     commandStreamReceiver0.makeNonResident(graphicsAllocation);
-    EXPECT_FALSE(graphicsAllocation.isResident(0u));
+    EXPECT_FALSE(graphicsAllocation.isResident(csr0ContextId));
-    EXPECT_TRUE(graphicsAllocation.isResident(1u));
+    EXPECT_TRUE(graphicsAllocation.isResident(csr1ContextId));
 
     commandStreamReceiver1.makeNonResident(graphicsAllocation);
-    EXPECT_FALSE(graphicsAllocation.isResident(0u));
+    EXPECT_FALSE(graphicsAllocation.isResident(csr0ContextId));
-    EXPECT_FALSE(graphicsAllocation.isResident(1u));
+    EXPECT_FALSE(graphicsAllocation.isResident(csr1ContextId));
 
     EXPECT_EQ(1u, commandStreamReceiver0.getEvictionAllocations().size());
     EXPECT_EQ(1u, commandStreamReceiver1.getEvictionAllocations().size());
@@ -456,7 +459,7 @@ TEST_F(CreateAllocationForHostSurfaceTest, givenReadOnlyHostPointerWhenAllocatio
     EXPECT_NE(memory, allocation->getUnderlyingBuffer());
     EXPECT_THAT(allocation->getUnderlyingBuffer(), MemCompare(memory, size));
 
-    allocation->updateTaskCount(commandStreamReceiver->peekLatestFlushedTaskCount(), 0u);
+    allocation->updateTaskCount(commandStreamReceiver->peekLatestFlushedTaskCount(), commandStreamReceiver->getOsContext().getContextId());
 }
 
 TEST_F(CreateAllocationForHostSurfaceTest, givenReadOnlyHostPointerWhenAllocationForHostSurfaceWithPtrCopyNotAllowedIsCreatedThenCopyAllocationIsNotCreated) {

unit_tests/command_stream/command_stream_receiver_with_aub_dump_tests.cpp

@@ -33,7 +33,7 @@ struct MyMockCsr : UltCommandStreamReceiver<DEFAULT_TEST_FAMILY_NAME> {
     void makeResident(GraphicsAllocation &gfxAllocation) override {
         makeResidentParameterization.wasCalled = true;
         makeResidentParameterization.receivedGfxAllocation = &gfxAllocation;
-        gfxAllocation.updateResidencyTaskCount(1, deviceIndex);
+        gfxAllocation.updateResidencyTaskCount(1, osContext->getContextId());
     }
 
     void processResidency(ResidencyContainer &allocationsForResidency) override {
@@ -42,10 +42,10 @@ struct MyMockCsr : UltCommandStreamReceiver<DEFAULT_TEST_FAMILY_NAME> {
     }
 
     void makeNonResident(GraphicsAllocation &gfxAllocation) override {
-        if (gfxAllocation.isResident(this->deviceIndex)) {
+        if (gfxAllocation.isResident(this->osContext->getContextId())) {
             makeNonResidentParameterization.wasCalled = true;
             makeNonResidentParameterization.receivedGfxAllocation = &gfxAllocation;
-            gfxAllocation.resetResidencyTaskCount(this->deviceIndex);
+            gfxAllocation.resetResidencyTaskCount(this->osContext->getContextId());
         }
     }
 

unit_tests/command_stream/tbx_command_stream_tests.cpp

@@ -12,6 +12,7 @@
 #include "runtime/helpers/ptr_math.h"
 #include "runtime/memory_manager/memory_banks.h"
 #include "runtime/os_interface/debug_settings_manager.h"
+#include "runtime/os_interface/os_context.h"
 #include "unit_tests/command_queue/command_queue_fixture.h"
 #include "unit_tests/fixtures/device_fixture.h"
 #include "unit_tests/gen_common/gen_cmd_parse.h"
@@ -227,13 +228,13 @@ HWTEST_F(TbxCommandStreamTests, givenTbxCommandStreamReceiverWhenProcessResidenc
     auto graphicsAllocation = memoryManager->allocateGraphicsMemory(4096);
     ASSERT_NE(nullptr, graphicsAllocation);
 
-    EXPECT_FALSE(graphicsAllocation->isResident(0u));
+    EXPECT_FALSE(graphicsAllocation->isResident(tbxCsr->getOsContext().getContextId()));
 
     ResidencyContainer allocationsForResidency = {graphicsAllocation};
     tbxCsr->processResidency(allocationsForResidency);
 
-    EXPECT_TRUE(graphicsAllocation->isResident(0u));
+    EXPECT_TRUE(graphicsAllocation->isResident(tbxCsr->getOsContext().getContextId()));
-    EXPECT_EQ(tbxCsr->peekTaskCount() + 1, graphicsAllocation->getResidencyTaskCount(0u));
+    EXPECT_EQ(tbxCsr->peekTaskCount() + 1, graphicsAllocation->getResidencyTaskCount(tbxCsr->getOsContext().getContextId()));
 
     memoryManager->freeGraphicsMemory(graphicsAllocation);
 }
@@ -246,13 +247,13 @@ HWTEST_F(TbxCommandStreamTests, givenTbxCommandStreamReceiverWhenProcessResidenc
     auto graphicsAllocation = memoryManager->allocateGraphicsMemory(4096);
     ASSERT_NE(nullptr, graphicsAllocation);
 
-    EXPECT_FALSE(graphicsAllocation->isResident(0u));
+    EXPECT_FALSE(graphicsAllocation->isResident(tbxCsr->getOsContext().getContextId()));
 
     ResidencyContainer allocationsForResidency = {graphicsAllocation};
     tbxCsr->processResidency(allocationsForResidency);
 
-    EXPECT_TRUE(graphicsAllocation->isResident(0u));
+    EXPECT_TRUE(graphicsAllocation->isResident(tbxCsr->getOsContext().getContextId()));
-    EXPECT_EQ(tbxCsr->peekTaskCount() + 1, graphicsAllocation->getResidencyTaskCount(0u));
+    EXPECT_EQ(tbxCsr->peekTaskCount() + 1, graphicsAllocation->getResidencyTaskCount(tbxCsr->getOsContext().getContextId()));
 
     memoryManager->freeGraphicsMemory(graphicsAllocation);
 }
@@ -273,12 +274,12 @@ HWTEST_F(TbxCommandStreamTests, givenTbxCommandStreamReceiverWhenFlushIsCalledTh
 
     ResidencyContainer allocationsForResidency = {graphicsAllocation};
 
-    EXPECT_FALSE(graphicsAllocation->isResident(0u));
+    EXPECT_FALSE(graphicsAllocation->isResident(tbxCsr->getOsContext().getContextId()));
 
     tbxCsr->flush(batchBuffer, allocationsForResidency);
 
-    EXPECT_TRUE(graphicsAllocation->isResident(0u));
+    EXPECT_TRUE(graphicsAllocation->isResident(tbxCsr->getOsContext().getContextId()));
-    EXPECT_EQ(tbxCsr->peekTaskCount() + 1, graphicsAllocation->getResidencyTaskCount(0u));
+    EXPECT_EQ(tbxCsr->peekTaskCount() + 1, graphicsAllocation->getResidencyTaskCount(tbxCsr->getOsContext().getContextId()));
 
     memoryManager->freeGraphicsMemory(commandBuffer);
     memoryManager->freeGraphicsMemory(graphicsAllocation);

unit_tests/event/event_tests.cpp

@@ -411,7 +411,7 @@ TEST_F(UpdateEventTest, givenEventContainingCommandQueueWhenItsStatusIsUpdatedTo
     void *ptr = (void *)0x1000;
     size_t size = 4096;
     auto temporary = memoryManager->allocateGraphicsMemory(size, ptr);
-    temporary->updateTaskCount(3, 0);
+    temporary->updateTaskCount(3, commandQueue->getCommandStreamReceiver().getOsContext().getContextId());
     commandQueue->getCommandStreamReceiver().getInternalAllocationStorage()->storeAllocation(std::unique_ptr<GraphicsAllocation>(temporary), TEMPORARY_ALLOCATION);
     Event event(commandQueue.get(), CL_COMMAND_NDRANGE_KERNEL, 3, 3);
 
@@ -488,7 +488,7 @@ TEST_F(InternalsEventTest, processBlockedCommandsKernelOperation) {
     EXPECT_EQ(taskLevelBefore + 1, taskLevelAfter);
 
     EXPECT_EQ(surface->resident, 1u);
-    EXPECT_FALSE(surface->graphicsAllocation->isResident(0u));
+    EXPECT_FALSE(surface->graphicsAllocation->isResident(csr.getOsContext().getContextId()));
     delete surface->graphicsAllocation;
 }
 
@@ -579,7 +579,7 @@ TEST_F(InternalsEventTest, givenBlockedKernelWithPrintfWhenSubmittedThenPrintOut
 
     std::string output = testing::internal::GetCapturedStdout();
     EXPECT_STREQ("test", output.c_str());
-    EXPECT_FALSE(surface->isResident(0u));
+    EXPECT_FALSE(surface->isResident(pDevice->getDefaultEngine().osContext->getContextId()));
 
     delete pPrintfSurface;
 }
@@ -906,6 +906,7 @@ HWTEST_F(InternalsEventTest, GivenBufferWithoutZeroCopyOnCommandMapOrUnmapFlushe
     MockNonZeroCopyBuff buffer(executionStamp);
     MockCsr<FamilyType> csr(executionStamp, *pDevice->executionEnvironment);
     csr.setTagAllocation(pDevice->getDefaultEngine().commandStreamReceiver->getTagAllocation());
+    csr.setOsContext(*pDevice->getDefaultEngine().osContext);
 
     MemObjSizeArray size = {{4, 1, 1}};
     MemObjOffsetArray offset = {{0, 0, 0}};

unit_tests/event/user_events_tests.cpp

@@ -12,6 +12,7 @@
 #include "unit_tests/mocks/mock_event.h"
 #include "runtime/memory_manager/internal_allocation_storage.h"
 #include "runtime/memory_manager/memory_manager.h"
+#include "runtime/os_interface/os_context.h"
 
 TEST(UserEvent, testInitialStatusOfUserEventCmdQueue) {
     UserEvent uEvent;
@@ -900,7 +901,7 @@ TEST_F(EventTests, waitForEventsDestroysTemporaryAllocations) {
 
     EXPECT_EQ(temporaryAllocation, csr.getTemporaryAllocations().peekHead());
 
-    temporaryAllocation->updateTaskCount(10, 0u);
+    temporaryAllocation->updateTaskCount(10, csr.getOsContext().getContextId());
 
     Event event(pCmdQ, CL_COMMAND_NDRANGE_KERNEL, 3, 11);
 

unit_tests/execution_model/enqueue_execution_model_kernel_tests.cpp

@@ -115,7 +115,7 @@ HWTEST_P(ParentKernelEnqueueTest, GivenParentKernelWithPrivateSurfaceWhenEnqueue
         pKernel->getProgram()->getBlockKernelManager()->pushPrivateSurface(privateSurface, 0);
         pCmdQ->enqueueKernel(pKernel, 1, offset, gws, gws, 0, nullptr, nullptr);
 
-        EXPECT_TRUE(privateSurface->isResident(0u));
+        EXPECT_TRUE(privateSurface->isResident(mockCSR->getOsContext().getContextId()));
     }
 }
 

unit_tests/fixtures/memory_allocator_fixture.h

@@ -23,6 +23,7 @@ class MemoryAllocatorFixture : public MemoryManagementFixture {
         memoryManager = new OsAgnosticMemoryManager(false, false, *executionEnvironment);
         executionEnvironment->memoryManager.reset(memoryManager);
         csr = memoryManager->getDefaultCommandStreamReceiver(0);
+        csr->setOsContext(*memoryManager->createAndRegisterOsContext(gpgpuEngineInstances[0]));
     }
 
     void TearDown() override {

unit_tests/fixtures/memory_manager_fixture.cpp

@@ -5,6 +5,7 @@
  *
  */
 
+#include "runtime/os_interface/os_context.h"
 #include "unit_tests/fixtures/memory_manager_fixture.h"
 #include "unit_tests/mocks/mock_csr.h"
 #include "unit_tests/mocks/mock_memory_manager.h"
@@ -18,6 +19,8 @@ void MemoryManagerWithCsrFixture::SetUp() {
     csr->tagAddress = &currentGpuTag;
     executionEnvironment.commandStreamReceivers.resize(1);
     executionEnvironment.commandStreamReceivers[0][0].reset(csr);
+
+    csr->setOsContext(*memoryManager->createAndRegisterOsContext(gpgpuEngineInstances[0]));
 }
 
 void MemoryManagerWithCsrFixture::TearDown() {

unit_tests/gtpin/gtpin_tests.cpp

@@ -19,6 +19,7 @@
 #include "runtime/kernel/kernel.h"
 #include "runtime/mem_obj/buffer.h"
 #include "runtime/memory_manager/surface.h"
+#include "runtime/os_interface/os_context.h"
 #include "unit_tests/fixtures/context_fixture.h"
 #include "unit_tests/fixtures/memory_management_fixture.h"
 #include "unit_tests/fixtures/platform_fixture.h"
@@ -1668,11 +1669,11 @@ TEST_F(GTPinTests, givenInitializedGTPinInterfaceWhenKernelIsCreatedThenAllKerne
     auto pBuffer1 = castToObject<Buffer>(gtpinBuffer1);
     GraphicsAllocation *pGfxAlloc1 = pBuffer1->getGraphicsAllocation();
     CommandStreamReceiver &csr = pCmdQueue->getCommandStreamReceiver();
-    EXPECT_FALSE(pGfxAlloc0->isResident(0u));
+    EXPECT_FALSE(pGfxAlloc0->isResident(csr.getOsContext().getContextId()));
-    EXPECT_FALSE(pGfxAlloc1->isResident(0u));
+    EXPECT_FALSE(pGfxAlloc1->isResident(csr.getOsContext().getContextId()));
     gtpinNotifyMakeResident(pKernel, &csr);
-    EXPECT_TRUE(pGfxAlloc0->isResident(0u));
+    EXPECT_TRUE(pGfxAlloc0->isResident(csr.getOsContext().getContextId()));
-    EXPECT_FALSE(pGfxAlloc1->isResident(0u));
+    EXPECT_FALSE(pGfxAlloc1->isResident(csr.getOsContext().getContextId()));
 
     // Cancel information about second submitted kernel
     kernelExecQueue.pop_back();
@@ -1839,24 +1840,24 @@ TEST_F(GTPinTests, givenInitializedGTPinInterfaceWhenOneKernelIsSubmittedSeveral
     GraphicsAllocation *pGfxAlloc1 = pBuffer1->getGraphicsAllocation();
     CommandStreamReceiver &csr = pCmdQueue->getCommandStreamReceiver();
     // Make resident resource of first submitted kernel
-    EXPECT_FALSE(pGfxAlloc0->isResident(0u));
+    EXPECT_FALSE(pGfxAlloc0->isResident(csr.getOsContext().getContextId()));
-    EXPECT_FALSE(pGfxAlloc1->isResident(0u));
+    EXPECT_FALSE(pGfxAlloc1->isResident(csr.getOsContext().getContextId()));
     gtpinNotifyMakeResident(pKernel, &csr);
-    EXPECT_TRUE(pGfxAlloc0->isResident(0u));
+    EXPECT_TRUE(pGfxAlloc0->isResident(csr.getOsContext().getContextId()));
-    EXPECT_FALSE(pGfxAlloc1->isResident(0u));
+    EXPECT_FALSE(pGfxAlloc1->isResident(csr.getOsContext().getContextId()));
     // Make resident resource of second submitted kernel
     gtpinNotifyMakeResident(pKernel, &csr);
-    EXPECT_TRUE(pGfxAlloc0->isResident(0u));
+    EXPECT_TRUE(pGfxAlloc0->isResident(csr.getOsContext().getContextId()));
-    EXPECT_TRUE(pGfxAlloc1->isResident(0u));
+    EXPECT_TRUE(pGfxAlloc1->isResident(csr.getOsContext().getContextId()));
 
     // Verify that correct GT-Pin resource is added to residency list.
     // This simulates enqueuing blocked kernels
     kernelExecQueue[0].isResourceResident = false;
     kernelExecQueue[1].isResourceResident = false;
-    pGfxAlloc0->resetResidencyTaskCount(0u);
+    pGfxAlloc0->resetResidencyTaskCount(csr.getOsContext().getContextId());
-    pGfxAlloc1->resetResidencyTaskCount(0u);
+    pGfxAlloc1->resetResidencyTaskCount(csr.getOsContext().getContextId());
-    EXPECT_FALSE(pGfxAlloc0->isResident(0u));
+    EXPECT_FALSE(pGfxAlloc0->isResident(csr.getOsContext().getContextId()));
-    EXPECT_FALSE(pGfxAlloc1->isResident(0u));
+    EXPECT_FALSE(pGfxAlloc1->isResident(csr.getOsContext().getContextId()));
     std::vector<Surface *> residencyVector;
     EXPECT_EQ(0u, residencyVector.size());
     // Add to residency list resource of first submitted kernel
@@ -1865,16 +1866,16 @@ TEST_F(GTPinTests, givenInitializedGTPinInterfaceWhenOneKernelIsSubmittedSeveral
     // Make resident first resource on residency list
     GeneralSurface *pSurf1 = (GeneralSurface *)residencyVector[0];
     pSurf1->makeResident(csr);
-    EXPECT_TRUE(pGfxAlloc0->isResident(0u));
+    EXPECT_TRUE(pGfxAlloc0->isResident(csr.getOsContext().getContextId()));
-    EXPECT_FALSE(pGfxAlloc1->isResident(0u));
+    EXPECT_FALSE(pGfxAlloc1->isResident(csr.getOsContext().getContextId()));
     // Add to residency list resource of second submitted kernel
     gtpinNotifyUpdateResidencyList(pKernel, &residencyVector);
     EXPECT_EQ(2u, residencyVector.size());
     // Make resident second resource on residency list
     GeneralSurface *pSurf2 = (GeneralSurface *)residencyVector[1];
     pSurf2->makeResident(csr);
-    EXPECT_TRUE(pGfxAlloc0->isResident(0u));
+    EXPECT_TRUE(pGfxAlloc0->isResident(csr.getOsContext().getContextId()));
-    EXPECT_TRUE(pGfxAlloc1->isResident(0u));
+    EXPECT_TRUE(pGfxAlloc1->isResident(csr.getOsContext().getContextId()));
 
     // Cleanup
     delete pSurf1;

unit_tests/helpers/timestamp_packet_tests.cpp

@@ -346,10 +346,8 @@ HWCMDTEST_F(IGFX_GEN8_CORE, TimestampPacketTests, givenTimestampPacketDisabledWh
 HWTEST_F(TimestampPacketTests, givenTimestampPacketWriteEnabledWhenEnqueueingThenObtainNewStampAndPassToEvent) {
     auto &csr = device->getUltCommandStreamReceiver<FamilyType>();
     csr.timestampPacketWriteEnabled = true;
-    auto mockMemoryManager = new MockMemoryManager(*device->getExecutionEnvironment());
+
-    device->injectMemoryManager(mockMemoryManager);
+    auto mockTagAllocator = new MockTagAllocator<>(executionEnvironment.memoryManager.get());
-    context->setMemoryManager(mockMemoryManager);
-    auto mockTagAllocator = new MockTagAllocator<>(mockMemoryManager);
     csr.timestampPacketAllocator.reset(mockTagAllocator);
     auto cmdQ = std::make_unique<MockCommandQueueHw<FamilyType>>(context.get(), device.get(), nullptr);
 
@@ -609,10 +607,7 @@ HWTEST_F(TimestampPacketTests, givenTimestampPacketWriteEnabledWhenDispatchingTh
 }
 
 HWTEST_F(TimestampPacketTests, givenAlreadyAssignedNodeWhenEnqueueingNonBlockedThenMakeItResident) {
-    auto mockMemoryManager = new MockMemoryManager(*device->getExecutionEnvironment());
+    auto mockTagAllocator = new MockTagAllocator<>(executionEnvironment.memoryManager.get(), 1);
-    device->injectMemoryManager(mockMemoryManager);
-    context->setMemoryManager(mockMemoryManager);
-    auto mockTagAllocator = new MockTagAllocator<>(mockMemoryManager, 1);
 
     auto &csr = device->getUltCommandStreamReceiver<FamilyType>();
     csr.timestampPacketAllocator.reset(mockTagAllocator);
@@ -634,10 +629,7 @@ HWTEST_F(TimestampPacketTests, givenAlreadyAssignedNodeWhenEnqueueingNonBlockedT
 }
 
 HWTEST_F(TimestampPacketTests, givenAlreadyAssignedNodeWhenEnqueueingBlockedThenMakeItResident) {
-    auto mockMemoryManager = new MockMemoryManager(*device->getExecutionEnvironment());
+    auto mockTagAllocator = new MockTagAllocator<>(executionEnvironment.memoryManager.get(), 1);
-    device->injectMemoryManager(mockMemoryManager);
-    context->setMemoryManager(mockMemoryManager);
-    auto mockTagAllocator = new MockTagAllocator<>(mockMemoryManager, 1);
 
     auto &csr = device->getUltCommandStreamReceiver<FamilyType>();
     csr.timestampPacketAllocator.reset(mockTagAllocator);

unit_tests/kernel/kernel_image_arg_tests.cpp

@@ -190,6 +190,7 @@ HWTEST_F(KernelImageArgTest, givenImgWithMcsAllocWhenMakeResidentThenMakeMcsAllo
     pKernel->setArg(0, sizeof(memObj), &memObj);
 
     std::unique_ptr<MockCsr<FamilyType>> csr(new MockCsr<FamilyType>(execStamp, *pDevice->executionEnvironment));
+    csr->setOsContext(*pDevice->getDefaultEngine().osContext);
 
     pKernel->makeResident(*csr.get());
     EXPECT_TRUE(csr->isMadeResident(mcsAlloc));

unit_tests/kernel/kernel_tests.cpp

@@ -15,6 +15,7 @@
 #include "runtime/memory_manager/allocations_list.h"
 #include "runtime/memory_manager/os_agnostic_memory_manager.h"
 #include "runtime/os_interface/debug_settings_manager.h"
+#include "runtime/os_interface/os_context.h"
 #include "unit_tests/fixtures/device_fixture.h"
 #include "unit_tests/fixtures/execution_model_fixture.h"
 #include "unit_tests/fixtures/memory_management_fixture.h"
@@ -580,8 +581,8 @@ TEST_F(KernelPrivateSurfaceTest, testPrivateSurface) {
 
     // Test it
     auto executionEnvironment = pDevice->getExecutionEnvironment();
-    executionEnvironment->memoryManager = std::make_unique<OsAgnosticMemoryManager>(false, false, *executionEnvironment);
     std::unique_ptr<CommandStreamReceiverMock> csr(new CommandStreamReceiverMock(*executionEnvironment));
+    csr->setOsContext(*pDevice->getDefaultEngine().osContext);
     csr->residency.clear();
     EXPECT_EQ(0u, csr->residency.size());
 
@@ -621,7 +622,7 @@ TEST_F(KernelPrivateSurfaceTest, givenKernelWithPrivateSurfaceThatIsInUseByGpuWh
     auto privateSurface = pKernel->getPrivateSurface();
     auto tagAddress = csr.getTagAddress();
 
-    privateSurface->updateTaskCount(*tagAddress + 1, 0u);
+    privateSurface->updateTaskCount(*tagAddress + 1, csr.getOsContext().getContextId());
 
     EXPECT_TRUE(csr.getTemporaryAllocations().peekIsEmpty());
     pKernel.reset(nullptr);

unit_tests/kernel/substitute_kernel_heap_tests.cpp

@@ -6,6 +6,7 @@
  */
 
 #include "runtime/memory_manager/internal_allocation_storage.h"
+#include "runtime/os_interface/os_context.h"
 #include "unit_tests/fixtures/device_fixture.h"
 #include "unit_tests/mocks/mock_kernel.h"
 #include "test.h"
@@ -119,7 +120,7 @@ TEST_F(KernelSubstituteTest, givenKernelWithUsedKernelAllocationWhenSubstituteKe
 
     uint32_t notReadyTaskCount = *commandStreamReceiver.getTagAddress() + 1u;
 
-    firstAllocation->updateTaskCount(notReadyTaskCount, 0u);
+    firstAllocation->updateTaskCount(notReadyTaskCount, commandStreamReceiver.getOsContext().getContextId());
 
     const size_t newHeapSize = initialHeapSize + 1;
     char newHeap[newHeapSize];

unit_tests/mem_obj/buffer_tests.cpp

@@ -14,6 +14,7 @@
 #include "runtime/mem_obj/buffer.h"
 #include "mem_obj_types.h"
 #include "runtime/memory_manager/svm_memory_manager.h"
+#include "runtime/os_interface/os_context.h"
 #include "test.h"
 #include "unit_tests/fixtures/device_fixture.h"
 #include "unit_tests/fixtures/memory_management_fixture.h"
@@ -352,10 +353,13 @@ TEST(Buffer, givenZeroFlagsNoSharedContextAndRenderCompressedBuffersDisabledWhen
 }
 
 TEST(Buffer, givenClMemCopyHostPointerPassedToBufferCreateWhenAllocationIsNotInSystemMemoryPoolThenAllocationIsWrittenByEnqueueWriteBuffer) {
-    std::unique_ptr<MockDevice> device(MockDevice::createWithNewExecutionEnvironment<MockDevice>(nullptr));
+    ExecutionEnvironment executionEnvironment;
-    ::testing::NiceMock<GMockMemoryManagerFailFirstAllocation> *memoryManager = new ::testing::NiceMock<GMockMemoryManagerFailFirstAllocation>(*device->getExecutionEnvironment());
+    executionEnvironment.incRefInternal();
+
+    auto *memoryManager = new ::testing::NiceMock<GMockMemoryManagerFailFirstAllocation>(executionEnvironment);
+    executionEnvironment.memoryManager.reset(memoryManager);
+    std::unique_ptr<MockDevice> device(MockDevice::create<MockDevice>(*platformDevices, &executionEnvironment, 0));
 
-    device->injectMemoryManager(memoryManager);
     MockContext ctx(device.get());
 
     auto allocateNonSystemGraphicsAllocation = [memoryManager](AllocationFlags flags, DevicesBitfield devicesBitfield, const void *hostPtr, size_t size, GraphicsAllocation::AllocationType type) -> GraphicsAllocation * {
@@ -693,7 +697,7 @@ TEST_P(ValidHostPtr, isResident_defaultsToFalseAfterCreate) {
     buffer = createBuffer();
     ASSERT_NE(nullptr, buffer);
 
-    EXPECT_FALSE(buffer->getGraphicsAllocation()->isResident(0u));
+    EXPECT_FALSE(buffer->getGraphicsAllocation()->isResident(pDevice->getDefaultEngine().osContext->getContextId()));
 }
 
 TEST_P(ValidHostPtr, getAddress) {

unit_tests/mem_obj/image_tests.cpp

@@ -10,6 +10,7 @@
 #include "runtime/helpers/aligned_memory.h"
 #include "runtime/helpers/mipmap.h"
 #include "runtime/mem_obj/image.h"
+#include "runtime/os_interface/os_context.h"
 #include "unit_tests/command_queue/command_queue_fixture.h"
 #include "unit_tests/fixtures/device_fixture.h"
 #include "unit_tests/fixtures/image_fixture.h"
@@ -45,9 +46,6 @@ class CreateImageTest : public DeviceFixture,
   protected:
     void SetUp() override {
         DeviceFixture::SetUp();
-        memoryManager = new MockMemoryManager(*pDevice->getExecutionEnvironment());
-
-        pDevice->injectMemoryManager(memoryManager);
 
         CommandQueueFixture::SetUp(pDevice, 0);
         flags = GetParam();
@@ -74,7 +72,6 @@ class CreateImageTest : public DeviceFixture,
         DeviceFixture::TearDown();
     }
 
-    MockMemoryManager *memoryManager;
     cl_image_format imageFormat;
     cl_image_desc imageDesc;
     cl_int retVal = CL_SUCCESS;
@@ -499,7 +496,7 @@ TEST_P(CreateImageHostPtr, isResidentDefaultsToFalseAfterCreate) {
     image = createImage(retVal);
     ASSERT_NE(nullptr, image);
 
-    EXPECT_FALSE(image->getGraphicsAllocation()->isResident(0u));
+    EXPECT_FALSE(image->getGraphicsAllocation()->isResident(pDevice->getDefaultEngine().osContext->getContextId()));
 }
 
 TEST_P(CreateImageHostPtr, getAddress) {
@@ -610,7 +607,10 @@ TEST_P(CreateImageHostPtr, failedAllocationInjection) {
 }
 
 TEST_P(CreateImageHostPtr, checkWritingOutsideAllocatedMemoryWhileCreatingImage) {
-    memoryManager->redundancyRatio = 2;
+    auto mockMemoryManager = new MockMemoryManager(executionEnvironment);
+    pDevice->injectMemoryManager(mockMemoryManager);
+    context->setMemoryManager(mockMemoryManager);
+    mockMemoryManager->redundancyRatio = 2;
     memset(pHostPtr, 1, testImageDimensions * testImageDimensions * elementSize * 4);
     imageDesc.image_type = CL_MEM_OBJECT_IMAGE1D_ARRAY;
     imageDesc.image_height = 1;
@@ -630,7 +630,7 @@ TEST_P(CreateImageHostPtr, checkWritingOutsideAllocatedMemoryWhileCreatingImage)
         EXPECT_EQ(0, memory[memorySize + i]);
     }
 
-    memoryManager->redundancyRatio = 1;
+    mockMemoryManager->redundancyRatio = 1;
 }
 
 struct ModifyableImage {

unit_tests/mem_obj/mem_obj_destruction_tests.cpp

@@ -137,12 +137,13 @@ HWTEST_P(MemObjAsyncDestructionTest, givenUsedMemObjWithAsyncDestructionsEnabled
     bool desired = true;
 
     auto waitForCompletionWithTimeoutMock = [=](bool enableTimeout, int64_t timeoutMs, uint32_t taskCountToWait) -> bool { return desired; };
+    auto osContextId = mockCsr->getOsContext().getContextId();
 
     ON_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, ::testing::_, ::testing::_))
         .WillByDefault(::testing::Invoke(waitForCompletionWithTimeoutMock));
 
     if (hasCallbacks) {
-        EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, TimeoutControls::maxTimeout, allocation->getTaskCount(0)))
+        EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, TimeoutControls::maxTimeout, allocation->getTaskCount(osContextId)))
             .Times(1);
     } else {
         EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, ::testing::_, ::testing::_))
@@ -163,6 +164,7 @@ HWTEST_P(MemObjAsyncDestructionTest, givenUsedMemObjWithAsyncDestructionsEnabled
 
     auto mockCsr = new ::testing::NiceMock<MyCsr<FamilyType>>(device->getHardwareInfo(), *device->executionEnvironment);
     device->resetCommandStreamReceiver(mockCsr);
+    auto osContextId = mockCsr->getOsContext().getContextId();
 
     bool desired = true;
 
@@ -172,7 +174,7 @@ HWTEST_P(MemObjAsyncDestructionTest, givenUsedMemObjWithAsyncDestructionsEnabled
         .WillByDefault(::testing::Invoke(waitForCompletionWithTimeoutMock));
 
     if (hasAllocatedMappedPtr) {
-        EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, TimeoutControls::maxTimeout, allocation->getTaskCount(0)))
+        EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, TimeoutControls::maxTimeout, allocation->getTaskCount(osContextId)))
             .Times(1);
     } else {
         EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, ::testing::_, ::testing::_))
@@ -208,12 +210,13 @@ HWTEST_P(MemObjAsyncDestructionTest, givenUsedMemObjWithAsyncDestructionsEnabled
     bool desired = true;
 
     auto waitForCompletionWithTimeoutMock = [=](bool enableTimeout, int64_t timeoutMs, uint32_t taskCountToWait) -> bool { return desired; };
+    auto osContextId = mockCsr->getOsContext().getContextId();
 
     ON_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, ::testing::_, ::testing::_))
         .WillByDefault(::testing::Invoke(waitForCompletionWithTimeoutMock));
 
     if (hasAllocatedMappedPtr) {
-        EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, TimeoutControls::maxTimeout, allocation->getTaskCount(0)))
+        EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, TimeoutControls::maxTimeout, allocation->getTaskCount(osContextId)))
             .Times(1);
     } else {
         EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, ::testing::_, ::testing::_))
@@ -241,11 +244,12 @@ HWTEST_P(MemObjSyncDestructionTest, givenMemObjWithDestructableAllocationWhenAsy
     bool desired = true;
 
     auto waitForCompletionWithTimeoutMock = [=](bool enableTimeout, int64_t timeoutMs, uint32_t taskCountToWait) -> bool { return desired; };
+    auto osContextId = mockCsr->getOsContext().getContextId();
 
     ON_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, ::testing::_, ::testing::_))
         .WillByDefault(::testing::Invoke(waitForCompletionWithTimeoutMock));
 
-    EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, TimeoutControls::maxTimeout, allocation->getTaskCount(0)))
+    EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, TimeoutControls::maxTimeout, allocation->getTaskCount(osContextId)))
         .Times(1);
 
     delete memObj;

unit_tests/mem_obj/mem_obj_tests.cpp

@@ -11,8 +11,10 @@
 #include "runtime/gmm_helper/gmm.h"
 #include "runtime/helpers/properties_helper.h"
 #include "runtime/memory_manager/allocations_list.h"
+#include "runtime/os_interface/os_context.h"
 #include "unit_tests/mocks/mock_context.h"
 #include "unit_tests/mocks/mock_deferred_deleter.h"
+#include "unit_tests/mocks/mock_device.h"
 #include "unit_tests/mocks/mock_graphics_allocation.h"
 #include "unit_tests/mocks/mock_memory_manager.h"
 #include "gtest/gtest.h"
@@ -147,7 +149,7 @@ TEST(MemObj, givenNotReadyGraphicsAllocationWhenMemObjDestroysAllocationAsyncThe
     context.setMemoryManager(&memoryManager);
 
     auto allocation = memoryManager.allocateGraphicsMemory(MemoryConstants::pageSize);
-    allocation->updateTaskCount(2, 0);
+    allocation->updateTaskCount(2, context.getDevice(0)->getDefaultEngine().osContext->getContextId());
     *(memoryManager.getDefaultCommandStreamReceiver(0)->getTagAddress()) = 1;
     MemObj memObj(&context, CL_MEM_OBJECT_BUFFER, CL_MEM_COPY_HOST_PTR,
                   MemoryConstants::pageSize, nullptr, nullptr, nullptr, true, false, false);
@@ -159,18 +161,19 @@ TEST(MemObj, givenNotReadyGraphicsAllocationWhenMemObjDestroysAllocationAsyncThe
 }
 
 TEST(MemObj, givenReadyGraphicsAllocationWhenMemObjDestroysAllocationAsyncThenAllocationIsNotAddedToMemoryManagerAllocationList) {
-    MockContext context;
+    ExecutionEnvironment executionEnvironment;
-    MockMemoryManager memoryManager(*context.getDevice(0)->getExecutionEnvironment());
+    executionEnvironment.incRefInternal();
+    auto device = std::unique_ptr<MockDevice>(MockDevice::create<MockDevice>(*platformDevices, &executionEnvironment, 0));
+    MockContext context(device.get());
+    auto memoryManager = executionEnvironment.memoryManager.get();
 
-    context.setMemoryManager(&memoryManager);
+    auto allocation = memoryManager->allocateGraphicsMemory(MemoryConstants::pageSize);
-
+    allocation->updateTaskCount(1, device->getDefaultEngine().osContext->getContextId());
-    auto allocation = memoryManager.allocateGraphicsMemory(MemoryConstants::pageSize);
+    *device->getDefaultEngine().commandStreamReceiver->getTagAddress() = 1;
-    allocation->updateTaskCount(1, 0);
-    *memoryManager.getDefaultCommandStreamReceiver(0)->getTagAddress() = 1;
     MemObj memObj(&context, CL_MEM_OBJECT_BUFFER, CL_MEM_COPY_HOST_PTR,
                   MemoryConstants::pageSize, nullptr, nullptr, nullptr, true, false, false);
 
-    auto &allocationList = memoryManager.getDefaultCommandStreamReceiver(0)->getTemporaryAllocations();
+    auto &allocationList = device->getDefaultEngine().commandStreamReceiver->getTemporaryAllocations();
     EXPECT_TRUE(allocationList.peekIsEmpty());
     memObj.destroyGraphicsAllocation(allocation, true);
 

unit_tests/memory_manager/internal_allocation_storage_tests.cpp

@@ -6,6 +6,7 @@
  */
 
 #include "runtime/memory_manager/internal_allocation_storage.h"
+#include "runtime/os_interface/os_context.h"
 #include "unit_tests/fixtures/memory_allocator_fixture.h"
 #include "unit_tests/helpers/debug_manager_state_restore.h"
 #include "unit_tests/utilities/containers_tests_helpers.h"
@@ -41,9 +42,9 @@ TEST_F(InternalAllocationStorageTest, whenCleanAllocationListThenRemoveOnlyCompl
     auto allocation2 = memoryManager->allocateGraphicsMemory(1, host_ptr);
     auto allocation3 = memoryManager->allocateGraphicsMemory(1, host_ptr);
 
-    allocation->updateTaskCount(10, 0);
+    allocation->updateTaskCount(10, csr->getOsContext().getContextId());
-    allocation2->updateTaskCount(5, 0);
+    allocation2->updateTaskCount(5, csr->getOsContext().getContextId());
-    allocation3->updateTaskCount(15, 0);
+    allocation3->updateTaskCount(15, csr->getOsContext().getContextId());
 
     storage->storeAllocation(std::unique_ptr<GraphicsAllocation>(allocation), TEMPORARY_ALLOCATION);
     storage->storeAllocation(std::unique_ptr<GraphicsAllocation>(allocation2), TEMPORARY_ALLOCATION);
@@ -119,7 +120,7 @@ TEST_F(InternalAllocationStorageTest, whenNotUsedAllocationIsStoredAsReusableAnd
     *csr->getTagAddress() = 0; // initial hw tag for dll
 
     storage->storeAllocation(std::unique_ptr<GraphicsAllocation>(allocation), REUSABLE_ALLOCATION);
-    EXPECT_EQ(0u, allocation->getTaskCount(0u));
+    EXPECT_EQ(0u, allocation->getTaskCount(csr->getOsContext().getContextId()));
     EXPECT_FALSE(csr->getAllocationsForReuse().peekIsEmpty());
 
     auto reusedAllocation = storage->obtainReusableAllocation(1, false).release();
@@ -179,8 +180,8 @@ TEST_F(InternalAllocationStorageTest, whenObtainAllocationFromMidlleOfReusableLi
     EXPECT_TRUE(reusableAllocations.peekContains(*allocation3));
 
     memoryManager->freeGraphicsMemory(allocation2);
-    allocation->updateTaskCount(0u, 0u);
+    allocation->updateTaskCount(0u, csr->getOsContext().getContextId());
-    allocation3->updateTaskCount(0u, 0u);
+    allocation3->updateTaskCount(0u, csr->getOsContext().getContextId());
 }
 
 TEST_F(InternalAllocationStorageTest, givenNonInternalAllocationWhenItIsPutOnReusableListWhenInternalAllocationIsRequestedThenNullIsReturned) {

unit_tests/memory_manager/memory_manager_tests.cpp

@@ -200,7 +200,7 @@ TEST_F(MemoryAllocatorTest, allocateGraphics) {
 
     ASSERT_NE(nullptr, allocation);
     // initial taskCount must be -1. if not, we may kill allocation before it will be used
-    EXPECT_EQ((uint32_t)-1, allocation->getTaskCount(0));
+    EXPECT_EQ((uint32_t)-1, allocation->getTaskCount(csr->getOsContext().getContextId()));
     // We know we want graphics memory to be page aligned
     EXPECT_EQ(0u, reinterpret_cast<uintptr_t>(allocation->getUnderlyingBuffer()) & (alignment - 1));
     EXPECT_EQ(Sharing::nonSharedResource, allocation->peekSharedHandle());
@@ -1179,7 +1179,7 @@ TEST_F(MemoryManagerWithCsrTest, givenAllocationThatWasUsedAndIsCompletedWhenche
     auto tagAddress = csr->getTagAddress();
     ASSERT_NE(0u, *tagAddress);
 
-    usedAllocationButGpuCompleted->updateTaskCount(*tagAddress - 1, 0);
+    usedAllocationButGpuCompleted->updateTaskCount(*tagAddress - 1, csr->getOsContext().getContextId());
 
     memoryManager->checkGpuUsageAndDestroyGraphicsAllocations(usedAllocationButGpuCompleted);
     EXPECT_TRUE(csr->getTemporaryAllocations().peekIsEmpty());
@@ -1191,14 +1191,14 @@ TEST_F(MemoryManagerWithCsrTest, givenAllocationThatWasUsedAndIsNotCompletedWhen
 
     auto tagAddress = csr->getTagAddress();
 
-    usedAllocationAndNotGpuCompleted->updateTaskCount(*tagAddress + 1, 0);
+    usedAllocationAndNotGpuCompleted->updateTaskCount(*tagAddress + 1, csr->getOsContext().getContextId());
 
     memoryManager->checkGpuUsageAndDestroyGraphicsAllocations(usedAllocationAndNotGpuCompleted);
     EXPECT_FALSE(csr->getTemporaryAllocations().peekIsEmpty());
     EXPECT_EQ(csr->getTemporaryAllocations().peekHead(), usedAllocationAndNotGpuCompleted);
 
     //change task count so cleanup will not clear alloc in use
-    usedAllocationAndNotGpuCompleted->updateTaskCount(csr->peekLatestFlushedTaskCount(), 0);
+    usedAllocationAndNotGpuCompleted->updateTaskCount(csr->peekLatestFlushedTaskCount(), csr->getOsContext().getContextId());
 }
 
 class MockAlignMallocMemoryManager : public MockMemoryManager {

unit_tests/memory_manager/surface_tests.cpp

@@ -61,6 +61,7 @@ HWTEST_TYPED_TEST(SurfaceTest, GivenSurfaceWhenInterfaceIsUsedThenSurfaceBehaves
     MockCsr<FamilyType> *csr = new MockCsr<FamilyType>(execStamp, executionEnvironment);
     executionEnvironment.commandStreamReceivers[0][0].reset(csr);
     executionEnvironment.memoryManager.reset(csr->createMemoryManager(false, false));
+    csr->setOsContext(*executionEnvironment.memoryManager->createAndRegisterOsContext(gpgpuEngineInstances[0]));
 
     Surface *surface = createSurface::Create<TypeParam>(this->data,
                                                         &this->buffer,

unit_tests/mocks/mock_csr.h

@@ -14,6 +14,7 @@
 #include "runtime/helpers/options.h"
 #include "runtime/helpers/flush_stamp.h"
 #include "runtime/helpers/string.h"
+#include "runtime/os_interface/os_context.h"
 #include "unit_tests/libult/ult_command_stream_receiver.h"
 #include "gmock/gmock.h"
 #include <vector>
@@ -44,7 +45,7 @@ class MockCsrBase : public UltCommandStreamReceiver<GfxFamily> {
         if (this->getMemoryManager()) {
             this->getResidencyAllocations().push_back(&gfxAllocation);
         }
-        gfxAllocation.updateResidencyTaskCount(this->taskCount, this->deviceIndex);
+        gfxAllocation.updateResidencyTaskCount(this->taskCount, this->osContext->getContextId());
     }
     void makeNonResident(GraphicsAllocation &gfxAllocation) override {
         madeNonResidentGfxAllocations.push_back(&gfxAllocation);

unit_tests/os_interface/linux/drm_command_stream_tests.cpp

@@ -805,21 +805,23 @@ TEST_F(DrmCommandStreamGemWorkerTests, GivenTwoAllocationsWhenBackingStorageIsDi
     csr->makeResident(*allocation);
     csr->makeResident(*allocation2);
 
-    EXPECT_TRUE(allocation->isResident(0u));
+    auto osContextId = csr->getOsContext().getContextId();
-    EXPECT_TRUE(allocation2->isResident(0u));
+
+    EXPECT_TRUE(allocation->isResident(osContextId));
+    EXPECT_TRUE(allocation2->isResident(osContextId));
 
     csr->processResidency(csr->getResidencyAllocations());
 
-    EXPECT_TRUE(allocation->isResident(0u));
+    EXPECT_TRUE(allocation->isResident(osContextId));
-    EXPECT_TRUE(allocation2->isResident(0u));
+    EXPECT_TRUE(allocation2->isResident(osContextId));
 
     EXPECT_EQ(tCsr->getResidencyVector()->size(), 2u);
 
     csr->makeNonResident(*allocation);
     csr->makeNonResident(*allocation2);
 
-    EXPECT_FALSE(allocation->isResident(0u));
+    EXPECT_FALSE(allocation->isResident(osContextId));
-    EXPECT_FALSE(allocation2->isResident(0u));
+    EXPECT_FALSE(allocation2->isResident(osContextId));
 
     EXPECT_EQ(tCsr->getResidencyVector()->size(), 0u);
     mm->freeGraphicsMemory(allocation);
@@ -1006,7 +1008,7 @@ TEST_F(DrmCommandStreamBatchingTests, givenRecordedCommandBufferWhenItIsSubmitte
     EXPECT_TRUE(cmdBuffers.peekIsEmpty());
 
     auto commandBufferGraphicsAllocation = submittedCommandBuffer.getGraphicsAllocation();
-    EXPECT_FALSE(commandBufferGraphicsAllocation->isResident(0u));
+    EXPECT_FALSE(commandBufferGraphicsAllocation->isResident(tCsr->getOsContext().getContextId()));
 
     //preemption allocation
     size_t csrSurfaceCount = (device->getPreemptionMode() == PreemptionMode::MidThread) ? 2 : 0;
@@ -1565,7 +1567,7 @@ TEST_F(DrmCommandStreamLeaksTest, givenMultipleMakeResidentWhenMakeNonResidentIs
     csr->makeSurfacePackNonResident(csr->getResidencyAllocations());
 
     EXPECT_EQ(0u, csr->getResidencyAllocations().size());
-    EXPECT_FALSE(allocation1->isResident(0u));
+    EXPECT_FALSE(allocation1->isResident(csr->getOsContext().getContextId()));
 
     mm->freeGraphicsMemory(allocation1);
 }
@@ -1616,22 +1618,24 @@ class DrmMockBuffer : public Buffer {
 TEST_F(DrmCommandStreamLeaksTest, BufferResidency) {
     std::unique_ptr<Buffer> buffer(DrmMockBuffer::create());
 
-    ASSERT_FALSE(buffer->getGraphicsAllocation()->isResident(0u));
+    auto osContextId = csr->getOsContext().getContextId();
+
+    ASSERT_FALSE(buffer->getGraphicsAllocation()->isResident(osContextId));
     ASSERT_GT(buffer->getSize(), 0u);
 
     //make it resident 8 times
     for (int c = 0; c < 8; c++) {
         csr->makeResident(*buffer->getGraphicsAllocation());
         csr->processResidency(csr->getResidencyAllocations());
-        EXPECT_TRUE(buffer->getGraphicsAllocation()->isResident(0u));
+        EXPECT_TRUE(buffer->getGraphicsAllocation()->isResident(osContextId));
-        EXPECT_EQ(buffer->getGraphicsAllocation()->getResidencyTaskCount(0u), csr->peekTaskCount() + 1);
+        EXPECT_EQ(buffer->getGraphicsAllocation()->getResidencyTaskCount(osContextId), csr->peekTaskCount() + 1);
     }
 
     csr->makeNonResident(*buffer->getGraphicsAllocation());
-    EXPECT_FALSE(buffer->getGraphicsAllocation()->isResident(0u));
+    EXPECT_FALSE(buffer->getGraphicsAllocation()->isResident(osContextId));
 
     csr->makeNonResident(*buffer->getGraphicsAllocation());
-    EXPECT_FALSE(buffer->getGraphicsAllocation()->isResident(0u));
+    EXPECT_FALSE(buffer->getGraphicsAllocation()->isResident(osContextId));
 }
 
 typedef Test<DrmCommandStreamEnhancedFixture> DrmCommandStreamMemoryManagerTest;

unit_tests/os_interface/linux/drm_memory_manager_tests.cpp

@@ -1247,6 +1247,7 @@ TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerWhenTiledImageIsBeingCreatedFr
     EXPECT_EQ(CL_SUCCESS, retVal);
     ASSERT_NE(nullptr, dstImage);
     auto imageGraphicsAllocation = dstImage->getGraphicsAllocation();
+    imageGraphicsAllocation->updateTaskCount(0, 0);
     ASSERT_NE(nullptr, imageGraphicsAllocation);
     EXPECT_TRUE(imageGraphicsAllocation->gmm->resourceParams.Usage ==
                 GMM_RESOURCE_USAGE_TYPE::GMM_RESOURCE_USAGE_OCL_IMAGE);
@@ -1611,6 +1612,7 @@ TEST_F(DrmMemoryManagerTest, givenTwoGraphicsAllocationsThatShareTheSameBufferOb
     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]));
 
     testedCsr->makeResident(*graphicsAllocation);
     testedCsr->makeResident(*graphicsAllocation2);
@@ -1639,6 +1641,7 @@ TEST_F(DrmMemoryManagerTest, givenTwoGraphicsAllocationsThatDoesnShareTheSameBuf
     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]));
 
     testedCsr->makeResident(*graphicsAllocation);
     testedCsr->makeResident(*graphicsAllocation2);

unit_tests/os_interface/windows/device_command_stream_tests.cpp

@@ -49,7 +49,7 @@ using namespace ::testing;
 
 class WddmCommandStreamFixture {
   public:
-    std::unique_ptr<Device> device;
+    std::unique_ptr<MockDevice> device;
     DeviceCommandStreamReceiver<DEFAULT_TEST_FAMILY_NAME> *csr;
     MockWddmMemoryManager *memoryManager = nullptr;
     WddmMock *wddm = nullptr;
@@ -62,13 +62,12 @@ class WddmCommandStreamFixture {
         wddm = static_cast<WddmMock *>(executionEnvironment->osInterface->get()->getWddm());
 
         csr = new WddmCommandStreamReceiver<DEFAULT_TEST_FAMILY_NAME>(*platformDevices[0], *executionEnvironment);
-        executionEnvironment->commandStreamReceivers.resize(1);
-        executionEnvironment->commandStreamReceivers[0][0].reset(csr);
 
         memoryManager = new MockWddmMemoryManager(wddm, *executionEnvironment);
         executionEnvironment->memoryManager.reset(memoryManager);
 
-        device.reset(Device::create<Device>(platformDevices[0], executionEnvironment, 0u));
+        device.reset(MockDevice::create<MockDevice>(platformDevices[0], executionEnvironment, 0u));
+        device->resetCommandStreamReceiver(csr);
         ASSERT_NE(nullptr, device);
     }
 
@@ -622,8 +621,8 @@ TEST_F(WddmCommandStreamTest, givenTwoTemporaryAllocationsWhenCleanTemporaryAllo
     csr->getInternalAllocationStorage()->storeAllocation(std::unique_ptr<GraphicsAllocation>(graphicsAllocation), TEMPORARY_ALLOCATION);
     csr->getInternalAllocationStorage()->storeAllocation(std::unique_ptr<GraphicsAllocation>(graphicsAllocation2), TEMPORARY_ALLOCATION);
 
-    graphicsAllocation->updateTaskCount(1, 0u);
+    graphicsAllocation->updateTaskCount(1, csr->getOsContext().getContextId());
-    graphicsAllocation2->updateTaskCount(100, 0u);
+    graphicsAllocation2->updateTaskCount(100, csr->getOsContext().getContextId());
 
     csr->waitForTaskCountAndCleanAllocationList(1, TEMPORARY_ALLOCATION);
     // graphicsAllocation2 still lives

unit_tests/os_interface/windows/wddm_memory_manager_tests.cpp

@@ -1431,17 +1431,23 @@ TEST_F(WddmMemoryManagerTest2, givenReadOnlyMemoryPassedToPopulateOsHandlesWhenC
 }
 
 TEST(WddmMemoryManagerCleanupTest, givenUsedTagAllocationInWddmMemoryManagerWhenCleanupMemoryManagerThenDontAccessCsr) {
-    auto wddm = std::make_unique<WddmMock>();
-    EXPECT_TRUE(wddm->init());
     ExecutionEnvironment executionEnvironment;
+    auto csr = createCommandStream(*platformDevices, executionEnvironment);
+    auto wddm = new WddmMock();
+    EXPECT_TRUE(wddm->init());
+
+    executionEnvironment.osInterface = std::make_unique<OSInterface>();
+    executionEnvironment.osInterface->get()->setWddm(wddm);
     executionEnvironment.commandStreamReceivers.resize(1);
-    executionEnvironment.commandStreamReceivers[0][0] =
+    executionEnvironment.commandStreamReceivers[0][0] = std::unique_ptr<CommandStreamReceiver>(csr);
-        std::unique_ptr<CommandStreamReceiver>(createCommandStream(*platformDevices, executionEnvironment));
+
-    executionEnvironment.memoryManager = std::make_unique<WddmMemoryManager>(false, false, wddm.get(), executionEnvironment);
+    executionEnvironment.memoryManager = std::make_unique<WddmMemoryManager>(false, false, wddm, executionEnvironment);
-    EXPECT_EQ(executionEnvironment.commandStreamReceivers[0][0].get(), executionEnvironment.memoryManager->getDefaultCommandStreamReceiver(0));
+    csr->setOsContext(*executionEnvironment.memoryManager->createAndRegisterOsContext(gpgpuEngineInstances[0]));
-    auto tagAllocator = executionEnvironment.commandStreamReceivers[0][0]->getEventPerfCountAllocator();
+    EXPECT_EQ(csr, executionEnvironment.memoryManager->getDefaultCommandStreamReceiver(0));
+
+    auto tagAllocator = csr->getEventPerfCountAllocator();
     auto allocation = tagAllocator->getTag()->getGraphicsAllocation();
-    allocation->updateTaskCount(1, 0);
+    allocation->updateTaskCount(1, csr->getOsContext().getContextId());
     executionEnvironment.commandStreamReceivers.clear();
     EXPECT_NO_THROW(executionEnvironment.memoryManager.reset());
 }

unit_tests/program/program_data_tests.cpp

@@ -178,7 +178,7 @@ TEST_F(ProgramDataTest, givenConstantAllocationThatIsInUseByGpuWhenProgramIsBein
     auto &csr = *pPlatform->getDevice(0)->getDefaultEngine().commandStreamReceiver;
     auto tagAddress = csr.getTagAddress();
     auto constantSurface = pProgram->getConstantSurface();
-    constantSurface->updateTaskCount(*tagAddress + 1, 0);
+    constantSurface->updateTaskCount(*tagAddress + 1, csr.getOsContext().getContextId());
 
     EXPECT_TRUE(csr.getTemporaryAllocations().peekIsEmpty());
     delete pProgram;
@@ -195,7 +195,7 @@ TEST_F(ProgramDataTest, givenGlobalAllocationThatIsInUseByGpuWhenProgramIsBeingD
     auto &csr = *pPlatform->getDevice(0)->getDefaultEngine().commandStreamReceiver;
     auto tagAddress = csr.getTagAddress();
     auto globalSurface = pProgram->getGlobalSurface();
-    globalSurface->updateTaskCount(*tagAddress + 1, 0);
+    globalSurface->updateTaskCount(*tagAddress + 1, csr.getOsContext().getContextId());
 
     EXPECT_TRUE(csr.getTemporaryAllocations().peekIsEmpty());
     delete pProgram;

unit_tests/program/program_tests.cpp

@@ -22,6 +22,7 @@
 #include "runtime/memory_manager/graphics_allocation.h"
 #include "runtime/memory_manager/surface.h"
 #include "runtime/program/create.inl"
+#include "runtime/os_interface/os_context.h"
 #include "unit_tests/fixtures/device_fixture.h"
 #include "unit_tests/fixtures/program_fixture.inl"
 #include "unit_tests/global_environment.h"
@@ -664,7 +665,7 @@ TEST_P(ProgramFromBinaryTest, givenProgramWhenCleanCurrentKernelInfoIsCalledButG
     EXPECT_TRUE(csr.getTemporaryAllocations().peekIsEmpty());
     pProgram->build(1, &device, nullptr, nullptr, nullptr, true);
     auto kernelAllocation = pProgram->getKernelInfo(size_t(0))->getGraphicsAllocation();
-    kernelAllocation->updateTaskCount(100, 0);
+    kernelAllocation->updateTaskCount(100, csr.getOsContext().getContextId());
     *csr.getTagAddress() = 0;
     pProgram->cleanCurrentKernelInfo();
     EXPECT_FALSE(csr.getTemporaryAllocations().peekIsEmpty());

unit_tests/sharings/gl/gl_texture_tests.cpp

@@ -42,31 +42,37 @@ class GlSharingTextureTests : public ::testing::Test {
     };
 
     void SetUp() override {
+        executionEnvironment.incRefInternal();
         imgDesc = {};
         imgDesc.image_type = CL_MEM_OBJECT_IMAGE1D;
         imgDesc.image_width = 10;
         auto imgInfo = MockGmm::initImgInfo(imgDesc, 0, nullptr);
         glSharing = new MockGlSharing;
-        clContext.setSharingFunctions(&glSharing->m_sharingFunctions);
+
+        tempMM = new TempMM(executionEnvironment);
+        executionEnvironment.memoryManager.reset(tempMM);
+        device.reset(MockDevice::create<MockDevice>(*platformDevices, &executionEnvironment, 0));
+        clContext = std::make_unique<MockContext>(device.get());
+
+        clContext->setSharingFunctions(&glSharing->m_sharingFunctions);
         ASSERT_FALSE(overrideCommandStreamReceiverCreation);
-        tempMM = new TempMM(*clContext.getDevice(0)->getExecutionEnvironment());
+
         tempMM->forceGmm = MockGmm::queryImgParams(imgInfo);
         tempMM->forceAllocationSize = textureSize;
-        MockDevice *device = static_cast<MockDevice *>(clContext.getDevice(0));
-        device->injectMemoryManager(tempMM);
-        clContext.setMemoryManager(tempMM);
         textureSize = imgInfo.size;
         textureId = 1;
     }
 
     void TearDown() override {
-        clContext.releaseSharingFunctions(SharingType::CLGL_SHARING);
+        clContext->releaseSharingFunctions(SharingType::CLGL_SHARING);
         delete glSharing;
     }
 
+    ExecutionEnvironment executionEnvironment;
     cl_image_desc imgDesc;
     TempMM *tempMM;
-    MockContext clContext;
+    std::unique_ptr<MockDevice> device;
+    std::unique_ptr<MockContext> clContext;
     MockGlSharing *glSharing;
     size_t textureSize;
     unsigned int textureId;
@@ -80,7 +86,7 @@ TEST_F(GlSharingTextureTests, givenMockGlWhen1dGlTextureIsCreatedThenMemObjectHa
 
     glSharing->uploadDataToTextureInfo(textureId);
 
-    auto glTexture = GlTexture::createSharedGlTexture(&clContext, (cl_mem_flags)0, GL_TEXTURE_1D, 0, textureId, &retVal);
+    auto glTexture = GlTexture::createSharedGlTexture(clContext.get(), (cl_mem_flags)0, GL_TEXTURE_1D, 0, textureId, &retVal);
     ASSERT_NE(nullptr, glTexture);
     EXPECT_NE(nullptr, glTexture->getGraphicsAllocation());
     EXPECT_EQ(textureSize, glTexture->getGraphicsAllocation()->getUnderlyingBufferSize());
@@ -106,18 +112,18 @@ class FailingMemoryManager : public MockMemoryManager {
 };
 
 TEST_F(GlSharingTextureTests, givenMockGlWhenGlTextureIsCreatedFromWrongHandleThenErrorAndNoTextureIsReturned) {
-    auto tempMemoryManager = clContext.getMemoryManager();
+    auto tempMemoryManager = clContext->getMemoryManager();
     tempMM->useForcedGmm = false;
     auto memoryManager = std::unique_ptr<FailingMemoryManager>(new FailingMemoryManager());
-    clContext.setMemoryManager(memoryManager.get());
+    clContext->setMemoryManager(memoryManager.get());
 
     auto retVal = CL_SUCCESS;
-    auto glTexture = GlTexture::createSharedGlTexture(&clContext, (cl_mem_flags)0, GL_TEXTURE_1D, 0, textureId, &retVal);
+    auto glTexture = GlTexture::createSharedGlTexture(clContext.get(), (cl_mem_flags)0, GL_TEXTURE_1D, 0, textureId, &retVal);
 
     EXPECT_EQ(nullptr, glTexture);
     EXPECT_EQ(CL_INVALID_GL_OBJECT, retVal);
 
-    clContext.setMemoryManager(tempMemoryManager);
+    clContext->setMemoryManager(tempMemoryManager);
 }
 
 GLboolean OSAPI mockGLAcquireSharedTexture(GLDisplay, GLContext, GLContext, GLvoid *pResourceInfo) {
@@ -141,7 +147,7 @@ TEST_F(GlSharingTextureTests, givenMockGlWhenGlTextureIsCreatedFromIncorrectForm
     glSharing->m_sharingFunctions.setGLAcquireSharedTextureMock(mockGLAcquireSharedTexture);
 
     auto retVal = CL_SUCCESS;
-    auto glTexture = GlTexture::createSharedGlTexture(&clContext, (cl_mem_flags)0, GL_TEXTURE_1D, 0, textureId, &retVal);
+    auto glTexture = GlTexture::createSharedGlTexture(clContext.get(), (cl_mem_flags)0, GL_TEXTURE_1D, 0, textureId, &retVal);
 
     EXPECT_EQ(nullptr, glTexture);
     EXPECT_EQ(CL_INVALID_GL_OBJECT, retVal);
@@ -155,7 +161,7 @@ TEST_F(GlSharingTextureTests, givenMockGlWhenRenderBufferTextureIsCreatedThenMem
 
     glSharing->uploadDataToTextureInfo(textureId);
 
-    auto glTexture = GlTexture::createSharedGlTexture(&clContext, (cl_mem_flags)0, GL_RENDERBUFFER_EXT, 0, textureId, &retVal);
+    auto glTexture = GlTexture::createSharedGlTexture(clContext.get(), (cl_mem_flags)0, GL_RENDERBUFFER_EXT, 0, textureId, &retVal);
     ASSERT_NE(nullptr, glTexture);
     EXPECT_NE(nullptr, glTexture->getGraphicsAllocation());
     EXPECT_EQ(textureSize, glTexture->getGraphicsAllocation()->getUnderlyingBufferSize());
@@ -183,7 +189,7 @@ TEST_F(GlSharingTextureTests, givenGmmResourceAsInputeWhenTextureIsCreatedItHasG
 
     glSharing->uploadDataToTextureInfo();
 
-    auto glTexture = GlTexture::createSharedGlTexture(&clContext, (cl_mem_flags)0, GL_TEXTURE_1D, 0, textureId, &retVal);
+    auto glTexture = GlTexture::createSharedGlTexture(clContext.get(), (cl_mem_flags)0, GL_TEXTURE_1D, 0, textureId, &retVal);
     ASSERT_NE(nullptr, glTexture);
     EXPECT_NE(nullptr, glTexture->getGraphicsAllocation());
 
@@ -208,7 +214,7 @@ TEST_F(GlSharingTextureTests, givenDifferentHwFormatWhenSurfaceFormatInfoIsSetTh
 
     glSharing->uploadDataToTextureInfo();
 
-    auto glTexture = GlTexture::createSharedGlTexture(&clContext, CL_MEM_READ_ONLY, GL_TEXTURE_2D, 0, textureId, &retVal);
+    auto glTexture = GlTexture::createSharedGlTexture(clContext.get(), CL_MEM_READ_ONLY, GL_TEXTURE_2D, 0, textureId, &retVal);
     ASSERT_NE(nullptr, glTexture);
     EXPECT_TRUE(newHwFormat == glTexture->getSurfaceFormatInfo().GenxSurfaceFormat);
 
@@ -221,14 +227,14 @@ TEST_F(GlSharingTextureTests, givenGLRGB10FormatWhenSharedGlTextureIsCreatedThen
     glSharing->m_textureInfoOutput.glHWFormat = GFX3DSTATE_SURFACEFORMAT_R16G16B16X16_UNORM;
     glSharing->uploadDataToTextureInfo();
 
-    std::unique_ptr<Image> glTexture(GlTexture::createSharedGlTexture(&clContext, CL_MEM_READ_ONLY, GL_TEXTURE_2D, 0, textureId, &retVal));
+    std::unique_ptr<Image> glTexture(GlTexture::createSharedGlTexture(clContext.get(), CL_MEM_READ_ONLY, GL_TEXTURE_2D, 0, textureId, &retVal));
     ASSERT_NE(nullptr, glTexture);
     EXPECT_EQ(glTexture->getSurfaceFormatInfo().GenxSurfaceFormat, GFX3DSTATE_SURFACEFORMAT_R16G16B16A16_UNORM);
 }
 
 TEST_F(GlSharingTextureTests, givenContextAnd1dTextureWhenClCreateFromGlTextureIsCalledThenImageIsReturned) {
     cl_int retVal = CL_INVALID_GL_OBJECT;
-    auto glImage = clCreateFromGLTexture(&clContext, 0, GL_TEXTURE_1D, 0, textureId, &retVal);
+    auto glImage = clCreateFromGLTexture(clContext.get(), 0, GL_TEXTURE_1D, 0, textureId, &retVal);
     ASSERT_EQ(CL_SUCCESS, retVal);
     ASSERT_NE(nullptr, glImage);
 
@@ -238,16 +244,16 @@ TEST_F(GlSharingTextureTests, givenContextAnd1dTextureWhenClCreateFromGlTextureI
 
 TEST_F(GlSharingTextureTests, givenContextWithoutSharingAnd1dTextureWhenClCreateFromGlTextureIsCalledThenErrorIsReturned) {
     tempMM->useForcedGmm = false;
-    clContext.releaseSharingFunctions(CLGL_SHARING);
+    clContext->releaseSharingFunctions(CLGL_SHARING);
     cl_int retVal = CL_INVALID_GL_OBJECT;
-    auto glImage = clCreateFromGLTexture(&clContext, 0, GL_TEXTURE_1D, 0, textureId, &retVal);
+    auto glImage = clCreateFromGLTexture(clContext.get(), 0, GL_TEXTURE_1D, 0, textureId, &retVal);
     ASSERT_EQ(CL_INVALID_CONTEXT, retVal);
     ASSERT_EQ(nullptr, glImage);
 }
 
 TEST_F(GlSharingTextureTests, givenContextAndRenderBufferTextureWhenClCreateFromGlTextureIsCalledThenImageIsReturned) {
     cl_int retVal = CL_INVALID_GL_OBJECT;
-    auto glImage = clCreateFromGLRenderbuffer(&clContext, 0, textureId, &retVal);
+    auto glImage = clCreateFromGLRenderbuffer(clContext.get(), 0, textureId, &retVal);
     ASSERT_EQ(CL_SUCCESS, retVal);
     ASSERT_NE(nullptr, glImage);
 
@@ -257,16 +263,16 @@ TEST_F(GlSharingTextureTests, givenContextAndRenderBufferTextureWhenClCreateFrom
 
 TEST_F(GlSharingTextureTests, givenContextWithoutSharingAndRenderBufferTextureWhenClCreateFromGlTextureIsCalledThenErrorIsReturned) {
     tempMM->useForcedGmm = false;
-    clContext.releaseSharingFunctions(CLGL_SHARING);
+    clContext->releaseSharingFunctions(CLGL_SHARING);
     cl_int retVal = CL_INVALID_GL_OBJECT;
-    auto glImage = clCreateFromGLRenderbuffer(&clContext, 0, textureId, &retVal);
+    auto glImage = clCreateFromGLRenderbuffer(clContext.get(), 0, textureId, &retVal);
     ASSERT_EQ(CL_INVALID_CONTEXT, retVal);
     ASSERT_EQ(nullptr, glImage);
 }
 
 TEST_F(GlSharingTextureTests, givenGlCl1dTextureWhenAskedForCLGLGetInfoThenIdAndTypeIsReturned) {
     auto retVal = CL_SUCCESS;
-    auto glImage = clCreateFromGLTexture(&clContext, 0, GL_TEXTURE_1D, 0, textureId, &retVal);
+    auto glImage = clCreateFromGLTexture(clContext.get(), 0, GL_TEXTURE_1D, 0, textureId, &retVal);
     ASSERT_EQ(CL_SUCCESS, retVal);
     ASSERT_NE(nullptr, glImage);
 
@@ -284,7 +290,7 @@ TEST_F(GlSharingTextureTests, givenGlCl1dTextureWhenAskedForCLGLGetInfoThenIdAnd
 
 TEST_F(GlSharingTextureTests, givenHwCommandQueueAndGlTextureWhenItIsCreatedWithClCreateFromGlTexture2dThenImageObjectIsReturned) {
     auto retVal = CL_SUCCESS;
-    auto glImage = clCreateFromGLTexture2D(&clContext, 0, GL_TEXTURE_2D, 0, textureId, &retVal);
+    auto glImage = clCreateFromGLTexture2D(clContext.get(), 0, GL_TEXTURE_2D, 0, textureId, &retVal);
     ASSERT_EQ(CL_SUCCESS, retVal);
     ASSERT_NE(nullptr, glImage);
 
@@ -302,16 +308,16 @@ TEST_F(GlSharingTextureTests, givenHwCommandQueueAndGlTextureWhenItIsCreatedWith
 
 TEST_F(GlSharingTextureTests, givenContextWithoutSharingAndGlTextureWhenItIsCreatedWithClCreateFromGlTexture2dThenErrorIsReturned) {
     tempMM->useForcedGmm = false;
-    clContext.releaseSharingFunctions(CLGL_SHARING);
+    clContext->releaseSharingFunctions(CLGL_SHARING);
     auto retVal = CL_SUCCESS;
-    auto glImage = clCreateFromGLTexture2D(&clContext, 0, GL_TEXTURE_2D, 0, textureId, &retVal);
+    auto glImage = clCreateFromGLTexture2D(clContext.get(), 0, GL_TEXTURE_2D, 0, textureId, &retVal);
     ASSERT_EQ(CL_INVALID_CONTEXT, retVal);
     ASSERT_EQ(nullptr, glImage);
 }
 
 TEST_F(GlSharingTextureTests, givenHwCommandQueueAndGlTextureWhenItIsCreatedWithClCreateFromGlTexture3dThenImageObjectIsReturned) {
     auto retVal = CL_SUCCESS;
-    auto glImage = clCreateFromGLTexture3D(&clContext, 0, GL_TEXTURE_3D, 0, textureId, &retVal);
+    auto glImage = clCreateFromGLTexture3D(clContext.get(), 0, GL_TEXTURE_3D, 0, textureId, &retVal);
     ASSERT_EQ(CL_SUCCESS, retVal);
     ASSERT_NE(nullptr, glImage);
 
@@ -329,9 +335,9 @@ TEST_F(GlSharingTextureTests, givenHwCommandQueueAndGlTextureWhenItIsCreatedWith
 
 TEST_F(GlSharingTextureTests, givenContextWithoutSharingAndGlTextureWhenItIsCreatedWithClCreateFromGlTexture3dThenErrorIsReturned) {
     tempMM->useForcedGmm = false;
-    clContext.releaseSharingFunctions(CLGL_SHARING);
+    clContext->releaseSharingFunctions(CLGL_SHARING);
     auto retVal = CL_SUCCESS;
-    auto glImage = clCreateFromGLTexture3D(&clContext, 0, GL_TEXTURE_3D, 0, textureId, &retVal);
+    auto glImage = clCreateFromGLTexture3D(clContext.get(), 0, GL_TEXTURE_3D, 0, textureId, &retVal);
     ASSERT_EQ(CL_INVALID_CONTEXT, retVal);
     ASSERT_EQ(nullptr, glImage);
 }
@@ -340,10 +346,10 @@ TEST_F(GlSharingTextureTests, givenHwCommandQueueAndGlTextureWhenAcquireIsCalled
     glSharing->uploadDataToTextureInfo(textureId);
 
     auto retVal = CL_SUCCESS;
-    auto commandQueue = clCreateCommandQueue(&clContext, clContext.getDevice(0), 0, &retVal);
+    auto commandQueue = clCreateCommandQueue(clContext.get(), clContext->getDevice(0), 0, &retVal);
     ASSERT_EQ(CL_SUCCESS, retVal);
 
-    auto glImage = clCreateFromGLTexture(&clContext, 0, GL_TEXTURE_1D, 0, textureId, &retVal);
+    auto glImage = clCreateFromGLTexture(clContext.get(), 0, GL_TEXTURE_1D, 0, textureId, &retVal);
 
     EXPECT_EQ(1, GLAcquireSharedTextureCalled);
     retVal = clEnqueueAcquireGLObjects(commandQueue, 1, &glImage, 0, nullptr, nullptr);
@@ -367,10 +373,10 @@ TEST_F(GlSharingTextureTests, verifyGlTextureBufferOffset) {
     glSharing->uploadDataToTextureInfo(textureId);
 
     auto retVal = CL_SUCCESS;
-    auto commandQueue = clCreateCommandQueue(&clContext, clContext.getDevice(0), 0, &retVal);
+    auto commandQueue = clCreateCommandQueue(clContext.get(), clContext->getDevice(0), 0, &retVal);
     ASSERT_EQ(CL_SUCCESS, retVal);
 
-    auto glImage = clCreateFromGLTexture(&clContext, 0, GL_TEXTURE_1D, 0, textureId, &retVal);
+    auto glImage = clCreateFromGLTexture(clContext.get(), 0, GL_TEXTURE_1D, 0, textureId, &retVal);
     EXPECT_NE(glImage, nullptr);
 
     retVal = clEnqueueAcquireGLObjects(commandQueue, 1, &glImage, 0, nullptr, nullptr);
@@ -403,10 +409,10 @@ TEST_F(GlSharingTextureTests, givenHwCommandQueueAndGlRenderBufferWhenAcquireIsC
     glSharing->uploadDataToTextureInfo(textureId);
 
     auto retVal = CL_SUCCESS;
-    auto commandQueue = clCreateCommandQueue(&clContext, clContext.getDevice(0), 0, &retVal);
+    auto commandQueue = clCreateCommandQueue(clContext.get(), clContext->getDevice(0), 0, &retVal);
     ASSERT_EQ(CL_SUCCESS, retVal);
 
-    auto glImage = clCreateFromGLRenderbuffer(&clContext, 0, textureId, &retVal);
+    auto glImage = clCreateFromGLRenderbuffer(clContext.get(), 0, textureId, &retVal);
 
     EXPECT_EQ(1, GLAcquireSharedRenderBufferCalled);
     retVal = clEnqueueAcquireGLObjects(commandQueue, 1, &glImage, 0, nullptr, nullptr);
@@ -427,7 +433,7 @@ TEST_F(GlSharingTextureTests, givenHwCommandQueueAndGlRenderBufferWhenAcquireIsC
 }
 
 TEST_F(GlSharingTextureTests, givenSharedGlTextureWhenItIsAcquireCountIsDecrementedToZeroThenCallReleaseFunction) {
-    std::unique_ptr<Image> image(GlTexture::createSharedGlTexture(&clContext, CL_MEM_READ_ONLY, GL_TEXTURE_2D, 0, textureId, nullptr));
+    std::unique_ptr<Image> image(GlTexture::createSharedGlTexture(clContext.get(), CL_MEM_READ_ONLY, GL_TEXTURE_2D, 0, textureId, nullptr));
     auto sharingHandler = image->peekSharingHandler();
 
     sharingHandler->acquire(image.get());
@@ -443,7 +449,7 @@ TEST_F(GlSharingTextureTests, givenSharedGlTextureWhenItIsAcquireCountIsDecremen
 }
 
 TEST_F(GlSharingTextureTests, givenSharedRenderBufferWhenItIsAcquireCountIsDecrementedToZeroThenCallReleaseFunction) {
-    std::unique_ptr<Image> image(GlTexture::createSharedGlTexture(&clContext, CL_MEM_READ_WRITE, GL_RENDERBUFFER_EXT, 0, textureId, nullptr));
+    std::unique_ptr<Image> image(GlTexture::createSharedGlTexture(clContext.get(), CL_MEM_READ_WRITE, GL_RENDERBUFFER_EXT, 0, textureId, nullptr));
     auto sharingHandler = image->peekSharingHandler();
 
     sharingHandler->acquire(image.get());
@@ -462,7 +468,7 @@ TEST_F(GlSharingTextureTests, givenMultisampleTextureWithMoreThanOneSampleWhenAs
     GLsizei expectedNumSamples = 2;
     glSharing->m_textureInfoOutput.numberOfSamples = expectedNumSamples;
     glSharing->uploadDataToTextureInfo();
-    std::unique_ptr<Image> image(GlTexture::createSharedGlTexture(&clContext, CL_MEM_READ_WRITE, GL_TEXTURE_2D_MULTISAMPLE, 0, textureId, nullptr));
+    std::unique_ptr<Image> image(GlTexture::createSharedGlTexture(clContext.get(), CL_MEM_READ_WRITE, GL_TEXTURE_2D_MULTISAMPLE, 0, textureId, nullptr));
 
     GLsizei numSamples = 0;
     size_t retSize = 0;
@@ -476,7 +482,7 @@ TEST_F(GlSharingTextureTests, givenMultisampleTextureWithMoreThanOneSampleWhenAs
 TEST_F(GlSharingTextureTests, givenTextureWithOneSampleWhenAskedForNumSamplesThenReturnZero) {
     glSharing->m_textureInfoOutput.numberOfSamples = 1;
     glSharing->uploadDataToTextureInfo();
-    std::unique_ptr<Image> image(GlTexture::createSharedGlTexture(&clContext, CL_MEM_READ_WRITE, GL_TEXTURE_2D_MULTISAMPLE, 0, textureId, nullptr));
+    std::unique_ptr<Image> image(GlTexture::createSharedGlTexture(clContext.get(), CL_MEM_READ_WRITE, GL_TEXTURE_2D_MULTISAMPLE, 0, textureId, nullptr));
 
     GLenum numSamples = 0;
     size_t retSize = 0;
@@ -490,7 +496,7 @@ TEST_F(GlSharingTextureTests, givenTextureWithOneSampleWhenAskedForNumSamplesThe
 TEST_F(GlSharingTextureTests, givenTextureWithZeroSamplesWhenAskedForNumSamplesThenReturnZero) {
     glSharing->m_textureInfoOutput.numberOfSamples = 0;
     glSharing->uploadDataToTextureInfo();
-    std::unique_ptr<Image> image(GlTexture::createSharedGlTexture(&clContext, CL_MEM_READ_WRITE, GL_TEXTURE_2D_MULTISAMPLE, 0, textureId, nullptr));
+    std::unique_ptr<Image> image(GlTexture::createSharedGlTexture(clContext.get(), CL_MEM_READ_WRITE, GL_TEXTURE_2D_MULTISAMPLE, 0, textureId, nullptr));
 
     GLenum numSamples = 0;
     size_t retSize = 0;
@@ -504,7 +510,7 @@ TEST_F(GlSharingTextureTests, givenTextureWithZeroSamplesWhenAskedForNumSamplesT
 TEST_F(GlSharingTextureTests, givenMockGlWhenGlTextureIsCreatedFromFormatNotIncludedInSurfaceFormatsThenErrorAndNoTextureIsReturned) {
     cl_int retVal = CL_SUCCESS;
     textureInfoOutput.glInternalFormat = GL_SRGB8_ALPHA8;
-    auto glTexture = GlTexture::createSharedGlTexture(&clContext, CL_MEM_WRITE_ONLY, GL_SRGB8_ALPHA8, 0, textureId, &retVal);
+    auto glTexture = GlTexture::createSharedGlTexture(clContext.get(), CL_MEM_WRITE_ONLY, GL_SRGB8_ALPHA8, 0, textureId, &retVal);
 
     EXPECT_EQ(nullptr, glTexture);
     EXPECT_EQ(CL_INVALID_GL_OBJECT, retVal);
@@ -523,7 +529,7 @@ TEST_P(GetGlTextureInfoTests, givenGlTextureWhenAskedForCLGetGLTextureInfoThenRe
     auto retVal = CL_SUCCESS;
     GLenum expectedTarget = GetParam();
     GLint mipLevel = 1u;
-    auto glImage = clCreateFromGLTexture(&clContext, 0, expectedTarget, mipLevel, textureId, &retVal);
+    auto glImage = clCreateFromGLTexture(clContext.get(), 0, expectedTarget, mipLevel, textureId, &retVal);
     ASSERT_EQ(CL_SUCCESS, retVal);
     ASSERT_NE(nullptr, glImage);