Graphics Allocation: store task count per context id

Move definition of allocations list method to internal_allocation_storage.cpp

Change-Id: I4c6038df8fd1b9335e8a74edbab33b78f9293d8f
Signed-off-by: Mateusz Jablonski <mateusz.jablonski@intel.com>

runtime/command_stream/command_stream_receiver.cpp

@@ -61,7 +61,7 @@ void CommandStreamReceiver::makeResident(GraphicsAllocation &gfxAllocation) {
     auto submissionTaskCount = this->taskCount + 1;
     if (gfxAllocation.residencyTaskCount[deviceIndex] < (int)submissionTaskCount) {
         this->getResidencyAllocations().push_back(&gfxAllocation);
-        gfxAllocation.taskCount = submissionTaskCount;
+        gfxAllocation.updateTaskCount(submissionTaskCount, deviceIndex);
         if (gfxAllocation.residencyTaskCount[deviceIndex] == ObjectNotResident) {
             this->totalMemoryUsed += gfxAllocation.getUnderlyingBufferSize();
         }
@@ -103,13 +103,13 @@ void CommandStreamReceiver::makeResidentHostPtrAllocation(GraphicsAllocation *gf
     }
 }
 
-void CommandStreamReceiver::waitForTaskCountAndCleanAllocationList(uint32_t requiredTaskCount, uint32_t allocationType) {
+void CommandStreamReceiver::waitForTaskCountAndCleanAllocationList(uint32_t requiredTaskCount, uint32_t allocationUsage) {
     auto address = getTagAddress();
     if (address && requiredTaskCount != ObjectNotUsed) {
         while (*address < requiredTaskCount)
             ;
     }
-    internalAllocationStorage->cleanAllocationList(requiredTaskCount, allocationType);
+    internalAllocationStorage->cleanAllocationList(requiredTaskCount, allocationUsage);
 }
 
 MemoryManager *CommandStreamReceiver::getMemoryManager() const {
@@ -359,7 +359,7 @@ bool CommandStreamReceiver::createAllocationForHostSurface(HostPtrSurface &surfa
     if (allocation == nullptr) {
         return false;
     }
-    allocation->taskCount = Event::eventNotReady;
+    allocation->updateTaskCount(Event::eventNotReady, deviceIndex);
     surface.setAllocation(allocation);
     internalAllocationStorage->storeAllocation(std::unique_ptr<GraphicsAllocation>(allocation), TEMPORARY_ALLOCATION);
     return true;

runtime/command_stream/command_stream_receiver.h

@@ -82,7 +82,7 @@ class CommandStreamReceiver {
 
     virtual GmmPageTableMngr *createPageTableManager() { return nullptr; }
 
-    MOCKABLE_VIRTUAL void waitForTaskCountAndCleanAllocationList(uint32_t requiredTaskCount, uint32_t allocationType);
+    MOCKABLE_VIRTUAL void waitForTaskCountAndCleanAllocationList(uint32_t requiredTaskCount, uint32_t allocationUsage);
 
     LinearStream &getCS(size_t minRequiredSize = 1024u);
     OSInterface *getOSInterface() { return osInterface; };
@@ -149,6 +149,7 @@ 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.inl

@@ -248,7 +248,7 @@ CompletionStamp CommandStreamReceiverHw<GfxFamily>::flushTask(
 
     if (requiredScratchSize && (!scratchAllocation || scratchAllocation->getUnderlyingBufferSize() < requiredScratchSizeInBytes)) {
         if (scratchAllocation) {
-            scratchAllocation->taskCount = this->taskCount;
+            scratchAllocation->updateTaskCount(this->taskCount, this->deviceIndex);
             internalAllocationStorage->storeAllocation(std::unique_ptr<GraphicsAllocation>(scratchAllocation), TEMPORARY_ALLOCATION);
         }
         createScratchSpaceAllocation(requiredScratchSizeInBytes);

runtime/mem_obj/mem_obj.cpp

@@ -7,6 +7,7 @@
 
 #include "runtime/context/context.h"
 #include "runtime/command_queue/command_queue.h"
+#include "runtime/command_stream/command_stream_receiver.h"
 #include "runtime/device/device.h"
 #include "runtime/mem_obj/mem_obj.h"
 #include "runtime/memory_manager/deferred_deleter.h"
@@ -15,7 +16,6 @@
 #include "runtime/gmm_helper/gmm.h"
 #include "runtime/helpers/aligned_memory.h"
 #include "runtime/helpers/get_info.h"
-#include "runtime/command_stream/command_stream_receiver.h"
 #include <algorithm>
 
 namespace OCLRT {
@@ -64,7 +64,7 @@ MemObj::~MemObj() {
             if (!doAsyncDestrucions) {
                 needWait = true;
             }
-            if (needWait && graphicsAllocation->taskCount != ObjectNotUsed) {
+            if (needWait && graphicsAllocation->peekWasUsed()) {
                 waitForCsrCompletion();
             }
             destroyGraphicsAllocation(graphicsAllocation, doAsyncDestrucions);
@@ -288,22 +288,15 @@ void MemObj::releaseAllocatedMapPtr() {
 }
 
 void MemObj::waitForCsrCompletion() {
-    if (graphicsAllocation) {
+    memoryManager->getCommandStreamReceiver(0)->waitForCompletionWithTimeout(false, TimeoutControls::maxTimeout, graphicsAllocation->getTaskCount(0u));
-        memoryManager->getCommandStreamReceiver(0)->waitForCompletionWithTimeout(false, TimeoutControls::maxTimeout, graphicsAllocation->taskCount);
-    }
 }
 
 void MemObj::destroyGraphicsAllocation(GraphicsAllocation *allocation, bool asyncDestroy) {
-    if (asyncDestroy && allocation->taskCount != ObjectNotUsed) {
+    if (asyncDestroy) {
-        auto commandStreamReceiver = memoryManager->getCommandStreamReceiver(0);
+        memoryManager->checkGpuUsageAndDestroyGraphicsAllocations(allocation);
-        auto currentTag = *commandStreamReceiver->getTagAddress();
+    } else {
-        if (currentTag < allocation->taskCount) {
-            auto storageForAllocation = commandStreamReceiver->getInternalAllocationStorage();
-            storageForAllocation->storeAllocation(std::unique_ptr<GraphicsAllocation>(allocation), TEMPORARY_ALLOCATION);
-            return;
-        }
-    }
         memoryManager->freeGraphicsMemory(allocation);
+    }
 }
 
 bool MemObj::checkIfMemoryTransferIsRequired(size_t offsetInMemObjest, size_t offsetInHostPtr, const void *hostPtr, cl_command_type cmdType) {

runtime/memory_manager/allocations_list.h

@@ -11,10 +11,11 @@
 
 namespace OCLRT {
 class GraphicsAllocation;
+class CommandStreamReceiver;
 
 class AllocationsList : public IDList<GraphicsAllocation, true, true> {
   public:
-    std::unique_ptr<GraphicsAllocation> detachAllocation(size_t requiredMinimalSize, volatile uint32_t *csrTagAddress, bool internalAllocationRequired);
+    std::unique_ptr<GraphicsAllocation> detachAllocation(size_t requiredMinimalSize, CommandStreamReceiver &commandStreamReceiver, bool internalAllocationRequired);
 
   private:
     GraphicsAllocation *detachAllocationImpl(GraphicsAllocation *, void *);

runtime/memory_manager/graphics_allocation.cpp

@@ -5,13 +5,54 @@
  *
  */
 
-#include "graphics_allocation.h"
 #include "runtime/helpers/aligned_memory.h"
+#include "runtime/memory_manager/graphics_allocation.h"
 
-bool OCLRT::GraphicsAllocation::isL3Capable() {
+namespace OCLRT {
+bool GraphicsAllocation::isL3Capable() {
     auto ptr = ptrOffset(cpuPtr, static_cast<size_t>(this->allocationOffset));
     if (alignUp(ptr, MemoryConstants::cacheLineSize) == ptr && alignUp(this->size, MemoryConstants::cacheLineSize) == this->size) {
         return true;
     }
     return false;
 }
+GraphicsAllocation::GraphicsAllocation(void *cpuPtrIn, uint64_t gpuAddress, uint64_t baseAddress, size_t sizeIn) : gpuBaseAddress(baseAddress),
+                                                                                                                   size(sizeIn),
+                                                                                                                   cpuPtr(cpuPtrIn),
+                                                                                                                   gpuAddress(gpuAddress),
+                                                                                                                   taskCounts(maxOsContextCount) {
+    initTaskCounts();
+}
+
+GraphicsAllocation::GraphicsAllocation(void *cpuPtrIn, size_t sizeIn, osHandle sharedHandleIn) : size(sizeIn),
+                                                                                                 cpuPtr(cpuPtrIn),
+                                                                                                 gpuAddress(castToUint64(cpuPtrIn)),
+                                                                                                 sharedHandle(sharedHandleIn),
+                                                                                                 taskCounts(maxOsContextCount) {
+    initTaskCounts();
+}
+GraphicsAllocation::~GraphicsAllocation() = default;
+
+bool GraphicsAllocation::peekWasUsed() const { return registeredContextsNum > 0; }
+
+void GraphicsAllocation::updateTaskCount(uint32_t newTaskCount, uint32_t contextId) {
+    UNRECOVERABLE_IF(contextId >= taskCounts.size());
+    if (taskCounts[contextId] == ObjectNotUsed) {
+        registeredContextsNum++;
+    }
+    if (newTaskCount == ObjectNotUsed) {
+        registeredContextsNum--;
+    }
+    taskCounts[contextId] = newTaskCount;
+}
+
+uint32_t GraphicsAllocation::getTaskCount(uint32_t contextId) const {
+    UNRECOVERABLE_IF(contextId >= taskCounts.size());
+    return taskCounts[contextId];
+}
+void GraphicsAllocation::initTaskCounts() {
+    for (auto i = 0u; i < taskCounts.size(); i++) {
+        taskCounts[i] = ObjectNotUsed;
+    }
+}
+} // namespace OCLRT

runtime/memory_manager/graphics_allocation.h

@@ -17,6 +17,7 @@
 #include "runtime/memory_manager/memory_pool.h"
 #include "runtime/memory_manager/residency_container.h"
 #include "runtime/utilities/idlist.h"
+#include "runtime/utilities/stackvec.h"
 
 namespace OCLRT {
 
@@ -34,19 +35,7 @@ const uint32_t ObjectNotUsed = (uint32_t)-1;
 class Gmm;
 
 class GraphicsAllocation : public IDNode<GraphicsAllocation> {
-  protected:
-    size_t size = 0;
-    void *cpuPtr = nullptr;
-    uint64_t gpuAddress = 0;
-    bool coherent = false;
-    osHandle sharedHandle;
-    bool locked = false;
-    uint32_t reuseCount = 0; // GraphicsAllocation can be reused by shared resources
-    bool evictable = true;
-    MemoryPool::Type memoryPool = MemoryPool::MemoryNull;
-
   public:
-    uint32_t taskCount = ObjectNotUsed;
     OsHandleStorage fragmentsStorage;
     bool is32BitAllocation = false;
     uint64_t gpuBaseAddress = 0;
@@ -81,20 +70,13 @@ class GraphicsAllocation : public IDNode<GraphicsAllocation> {
         SHARED_RESOURCE,
     };
 
-    virtual ~GraphicsAllocation() = default;
+    virtual ~GraphicsAllocation();
     GraphicsAllocation &operator=(const GraphicsAllocation &) = delete;
     GraphicsAllocation(const GraphicsAllocation &) = delete;
 
-    GraphicsAllocation(void *cpuPtrIn, uint64_t gpuAddress, uint64_t baseAddress, size_t sizeIn) : size(sizeIn),
+    GraphicsAllocation(void *cpuPtrIn, uint64_t gpuAddress, uint64_t baseAddress, size_t sizeIn);
-                                                                                                   cpuPtr(cpuPtrIn),
-                                                                                                   gpuAddress(gpuAddress),
-                                                                                                   sharedHandle(Sharing::nonSharedResource),
-                                                                                                   gpuBaseAddress(baseAddress) {}
 
-    GraphicsAllocation(void *cpuPtrIn, size_t sizeIn, osHandle sharedHandleIn) : size(sizeIn),
+    GraphicsAllocation(void *cpuPtrIn, size_t sizeIn, osHandle sharedHandleIn);
-                                                                                 cpuPtr(cpuPtrIn),
-                                                                                 gpuAddress(castToUint64(cpuPtrIn)),
-                                                                                 sharedHandle(sharedHandleIn) {}
 
     void *getUnderlyingBuffer() const { return cpuPtr; }
     void setCpuPtrAndGpuAddress(void *cpuPtr, uint64_t gpuAddress) {
@@ -141,14 +123,30 @@ class GraphicsAllocation : public IDNode<GraphicsAllocation> {
     MemoryPool::Type getMemoryPool() {
         return memoryPool;
     }
+    bool peekWasUsed() const;
+    void updateTaskCount(uint32_t newTaskCount, uint32_t contextId);
+    uint32_t getTaskCount(uint32_t contextId) const;
 
   protected:
+    void initTaskCounts();
+
     //this variable can only be modified from SubmissionAggregator
     friend class SubmissionAggregator;
+    size_t size = 0;
+    void *cpuPtr = nullptr;
+    uint64_t gpuAddress = 0;
+    bool coherent = false;
+    osHandle sharedHandle = Sharing::nonSharedResource;
+    bool locked = false;
+    uint32_t reuseCount = 0; // GraphicsAllocation can be reused by shared resources
+    bool evictable = true;
+    MemoryPool::Type memoryPool = MemoryPool::MemoryNull;
     uint32_t inspectionId = 0;
     AllocationType allocationType = AllocationType::UNKNOWN;
     bool aubWritable = true;
     bool allocDumpable = false;
     bool memObjectsAllocationWithWritableFlags = false;
+    StackVec<uint32_t, maxOsContextCount> taskCounts;
+    std::atomic<uint32_t> registeredContextsNum{0};
 };
 } // namespace OCLRT

runtime/memory_manager/internal_allocation_storage.cpp

@@ -10,9 +10,9 @@
 #include "runtime/memory_manager/memory_manager.h"
 
 namespace OCLRT {
-InternalAllocationStorage::InternalAllocationStorage(CommandStreamReceiver &commandStreamReceiver) : commandStreamReceiver(commandStreamReceiver){};
+InternalAllocationStorage::InternalAllocationStorage(CommandStreamReceiver &commandStreamReceiver) : commandStreamReceiver(commandStreamReceiver), contextId(commandStreamReceiver.getDeviceIndex()){};
 void InternalAllocationStorage::storeAllocation(std::unique_ptr<GraphicsAllocation> gfxAllocation, uint32_t allocationUsage) {
-    uint32_t taskCount = gfxAllocation->taskCount;
+    uint32_t taskCount = gfxAllocation->getTaskCount(contextId);
 
     if (allocationUsage == REUSABLE_ALLOCATION) {
         taskCount = commandStreamReceiver.peekTaskCount();
@@ -28,7 +28,7 @@ void InternalAllocationStorage::storeAllocationWithTaskCount(std::unique_ptr<Gra
         }
     }
     auto &allocationsList = (allocationUsage == TEMPORARY_ALLOCATION) ? temporaryAllocations : allocationsForReuse;
-    gfxAllocation->taskCount = taskCount;
+    gfxAllocation->updateTaskCount(taskCount, contextId);
     allocationsList.pushTailOne(*gfxAllocation.release());
 }
 
@@ -43,7 +43,7 @@ void InternalAllocationStorage::freeAllocationsList(uint32_t waitTaskCount, Allo
     IDList<GraphicsAllocation, false, true> allocationsLeft;
     while (curr != nullptr) {
         auto *next = curr->next;
-        if (curr->taskCount <= waitTaskCount) {
+        if (curr->getTaskCount(contextId) <= waitTaskCount) {
             memoryManager->freeGraphicsMemory(curr);
         } else {
             allocationsLeft.pushTailOne(*curr);
@@ -57,8 +57,41 @@ void InternalAllocationStorage::freeAllocationsList(uint32_t waitTaskCount, Allo
 }
 
 std::unique_ptr<GraphicsAllocation> InternalAllocationStorage::obtainReusableAllocation(size_t requiredSize, bool internalAllocation) {
-    auto allocation = allocationsForReuse.detachAllocation(requiredSize, commandStreamReceiver.getTagAddress(), internalAllocation);
+    auto allocation = allocationsForReuse.detachAllocation(requiredSize, commandStreamReceiver, internalAllocation);
     return allocation;
 }
 
+struct ReusableAllocationRequirements {
+    size_t requiredMinimalSize;
+    volatile uint32_t *csrTagAddress;
+    bool internalAllocationRequired;
+    uint32_t contextId;
+};
+
+std::unique_ptr<GraphicsAllocation> AllocationsList::detachAllocation(size_t requiredMinimalSize, CommandStreamReceiver &commandStreamReceiver, bool internalAllocationRequired) {
+    ReusableAllocationRequirements req;
+    req.requiredMinimalSize = requiredMinimalSize;
+    req.csrTagAddress = commandStreamReceiver.getTagAddress();
+    req.internalAllocationRequired = internalAllocationRequired;
+    req.contextId = commandStreamReceiver.getDeviceIndex();
+    GraphicsAllocation *a = nullptr;
+    GraphicsAllocation *retAlloc = processLocked<AllocationsList, &AllocationsList::detachAllocationImpl>(a, static_cast<void *>(&req));
+    return std::unique_ptr<GraphicsAllocation>(retAlloc);
+}
+
+GraphicsAllocation *AllocationsList::detachAllocationImpl(GraphicsAllocation *, void *data) {
+    ReusableAllocationRequirements *req = static_cast<ReusableAllocationRequirements *>(data);
+    auto *curr = head;
+    while (curr != nullptr) {
+        auto currentTagValue = *req->csrTagAddress;
+        if ((req->internalAllocationRequired == curr->is32BitAllocation) &&
+            (curr->getUnderlyingBufferSize() >= req->requiredMinimalSize) &&
+            ((currentTagValue > curr->getTaskCount(req->contextId)) || (curr->getTaskCount(req->contextId) == 0))) {
+            return removeOneImpl(curr, nullptr);
+        }
+        curr = curr->next;
+    }
+    return nullptr;
+}
+
 } // namespace OCLRT

runtime/memory_manager/internal_allocation_storage.h

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

runtime/memory_manager/memory_manager.cpp

@@ -27,37 +27,6 @@
 namespace OCLRT {
 constexpr size_t TagCount = 512;
 
-struct ReusableAllocationRequirements {
-    size_t requiredMinimalSize;
-    volatile uint32_t *csrTagAddress;
-    bool internalAllocationRequired;
-};
-
-std::unique_ptr<GraphicsAllocation> AllocationsList::detachAllocation(size_t requiredMinimalSize, volatile uint32_t *csrTagAddress, bool internalAllocationRequired) {
-    ReusableAllocationRequirements req;
-    req.requiredMinimalSize = requiredMinimalSize;
-    req.csrTagAddress = csrTagAddress;
-    req.internalAllocationRequired = internalAllocationRequired;
-    GraphicsAllocation *a = nullptr;
-    GraphicsAllocation *retAlloc = processLocked<AllocationsList, &AllocationsList::detachAllocationImpl>(a, static_cast<void *>(&req));
-    return std::unique_ptr<GraphicsAllocation>(retAlloc);
-}
-
-GraphicsAllocation *AllocationsList::detachAllocationImpl(GraphicsAllocation *, void *data) {
-    ReusableAllocationRequirements *req = static_cast<ReusableAllocationRequirements *>(data);
-    auto *curr = head;
-    while (curr != nullptr) {
-        auto currentTagValue = req->csrTagAddress ? *req->csrTagAddress : -1;
-        if ((req->internalAllocationRequired == curr->is32BitAllocation) &&
-            (curr->getUnderlyingBufferSize() >= req->requiredMinimalSize) &&
-            ((currentTagValue > curr->taskCount) || (curr->taskCount == 0))) {
-            return removeOneImpl(curr, nullptr);
-        }
-        curr = curr->next;
-    }
-    return nullptr;
-}
-
 MemoryManager::MemoryManager(bool enable64kbpages, bool enableLocalMemory,
                              ExecutionEnvironment &executionEnvironment) : allocator32Bit(nullptr), enable64kbpages(enable64kbpages),
                                                                            localMemorySupported(enableLocalMemory),
@@ -198,7 +167,7 @@ void MemoryManager::freeGraphicsMemory(GraphicsAllocation *gfxAllocation) {
 //if not in use destroy in place
 //if in use pass to temporary allocation list that is cleaned on blocking calls
 void MemoryManager::checkGpuUsageAndDestroyGraphicsAllocations(GraphicsAllocation *gfxAllocation) {
-    if (gfxAllocation->taskCount == ObjectNotUsed || gfxAllocation->taskCount <= *getCommandStreamReceiver(0)->getTagAddress()) {
+    if (!gfxAllocation->peekWasUsed() || gfxAllocation->getTaskCount(0u) <= *getCommandStreamReceiver(0)->getTagAddress()) {
         freeGraphicsMemory(gfxAllocation);
     } else {
         getCommandStreamReceiver(0)->getInternalAllocationStorage()->storeAllocation(std::unique_ptr<GraphicsAllocation>(gfxAllocation), TEMPORARY_ALLOCATION);

runtime/os_interface/windows/wddm_memory_manager.cpp

@@ -294,9 +294,9 @@ void WddmMemoryManager::freeGraphicsMemoryImpl(GraphicsAllocation *gfxAllocation
     }
 
     UNRECOVERABLE_IF(DebugManager.flags.CreateMultipleDevices.get() == 0 &&
-                     gfxAllocation->taskCount != ObjectNotUsed && this->executionEnvironment.commandStreamReceivers.size() > 0 &&
+                     gfxAllocation->peekWasUsed() && this->executionEnvironment.commandStreamReceivers.size() > 0 &&
                      this->getCommandStreamReceiver(0) && this->getCommandStreamReceiver(0)->getTagAddress() &&
-                     gfxAllocation->taskCount > *this->getCommandStreamReceiver(0)->getTagAddress());
+                     gfxAllocation->getTaskCount(0u) > *this->getCommandStreamReceiver(0)->getTagAddress());
 
     if (input->gmm) {
         if (input->gmm->isRenderCompressed && wddm->getPageTableManager()) {

unit_tests/command_queue/enqueue_api_tests_mt_with_asyncGPU.cpp

@@ -1,23 +1,8 @@
 /*
- * Copyright (c) 2017, Intel Corporation
+ * Copyright (C) 2017-2018 Intel Corporation
  *
- * Permission is hereby granted, free of charge, to any person obtaining a
+ * SPDX-License-Identifier: MIT
- * copy of this software and associated documentation files (the "Software"),
- * to deal in the Software without restriction, including without limitation
- * the rights to use, copy, modify, merge, publish, distribute, sublicense,
- * and/or sell copies of the Software, and to permit persons to whom the
- * Software is furnished to do so, subject to the following conditions:
  *
- * The above copyright notice and this permission notice shall be included
- * in all copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
- * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
- * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
- * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
- * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
- * OTHER DEALINGS IN THE SOFTWARE.
  */
 
 #include "buffer_operations_withAsyncGPU_fixture.h"
@@ -73,7 +58,7 @@ HWTEST_F(AsyncGPUoperations, MapBufferAfterWriteBuffer) {
     }
     t.join();
 
-    srcBuffer->getGraphicsAllocation()->taskCount = ObjectNotUsed;
+    srcBuffer->getGraphicsAllocation()->updateTaskCount(ObjectNotUsed, 0);
 
     alignedFree(ptrMemory);
 }

unit_tests/command_queue/enqueue_fill_buffer_tests.cpp

@@ -376,7 +376,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->taskCount)) {
+            (pCmdQ->taskCount == allocation->getTaskCount(0))) {
             break;
         }
         allocation = allocation->next;
@@ -397,7 +397,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->taskCount)) {
+            (pCmdQ->taskCount == allocation->getTaskCount(0))) {
             break;
         }
         allocation = allocation->next;

unit_tests/command_queue/enqueue_kernel_tests.cpp

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

unit_tests/command_stream/aub_command_stream_receiver_tests.cpp

@@ -218,14 +218,14 @@ HWTEST_F(AubCommandStreamReceiverTests, givenGraphicsAllocationWhenMakeResidentC
     // First makeResident marks the allocation resident
     aubCsr->makeResident(*gfxAllocation);
     EXPECT_NE(ObjectNotResident, gfxAllocation->residencyTaskCount[0u]);
-    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->taskCount);
+    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->getTaskCount(0));
     EXPECT_EQ((int)aubCsr->peekTaskCount() + 1, gfxAllocation->residencyTaskCount[0u]);
     EXPECT_EQ(1u, aubCsr->getResidencyAllocations().size());
 
     // Second makeResident should have no impact
     aubCsr->makeResident(*gfxAllocation);
     EXPECT_NE(ObjectNotResident, gfxAllocation->residencyTaskCount[0u]);
-    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->taskCount);
+    EXPECT_EQ(aubCsr->peekTaskCount() + 1, gfxAllocation->getTaskCount(0));
     EXPECT_EQ((int)aubCsr->peekTaskCount() + 1, gfxAllocation->residencyTaskCount[0u]);
     EXPECT_EQ(1u, aubCsr->getResidencyAllocations().size());
 

unit_tests/command_stream/command_stream_receiver_flush_task_tests.cpp

@@ -1697,7 +1697,7 @@ HWTEST_F(CommandStreamReceiverFlushTaskTests, InForced32BitAllocationsModeStore3
         auto newScratchAllocation = commandStreamReceiver->getScratchAllocation();
         EXPECT_NE(scratchAllocation, newScratchAllocation); // Allocation changed
 
-        std::unique_ptr<GraphicsAllocation> allocationTemporary = commandStreamReceiver->getTemporaryAllocations().detachAllocation(0, nullptr, true);
+        std::unique_ptr<GraphicsAllocation> allocationTemporary = commandStreamReceiver->getTemporaryAllocations().detachAllocation(0, *commandStreamReceiver, true);
 
         EXPECT_EQ(scratchAllocation, allocationTemporary.get());
         pDevice->getMemoryManager()->freeGraphicsMemory(allocationTemporary.release());
@@ -3125,11 +3125,11 @@ 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->taskCount = 1;
+    temporaryToClean->updateTaskCount(1, 0u);
-    reusableToClean->taskCount = 1;
+    reusableToClean->updateTaskCount(1, 0u);
 
-    temporaryToHold->taskCount = 10;
+    temporaryToHold->updateTaskCount(10, 0u);
-    reusableToHold->taskCount = 10;
+    reusableToHold->updateTaskCount(10, 0u);
 
     commandStreamReceiver.latestFlushedTaskCount = 9;
     commandStreamReceiver.cleanupResources();

unit_tests/command_stream/command_stream_receiver_tests.cpp

@@ -177,13 +177,13 @@ HWTEST_F(CommandStreamReceiverTest, whenStoreAllocationThenStoredAllocationHasTa
     void *host_ptr = (void *)0x1234;
     auto allocation = memoryManager->allocateGraphicsMemory(1, host_ptr);
 
-    EXPECT_EQ(ObjectNotUsed, allocation->taskCount);
+    EXPECT_EQ(ObjectNotUsed, allocation->getTaskCount(0));
 
     csr.taskCount = 2u;
 
     csr.getInternalAllocationStorage()->storeAllocation(std::unique_ptr<GraphicsAllocation>(allocation), REUSABLE_ALLOCATION);
 
-    EXPECT_EQ(csr.peekTaskCount(), allocation->taskCount);
+    EXPECT_EQ(csr.peekTaskCount(), allocation->getTaskCount(0));
 }
 
 HWTEST_F(CommandStreamReceiverTest, givenCommandStreamReceiverWhenCheckedForInitialStatusOfStatelessMocsIndexThenUnknownMocsIsReturend) {
@@ -417,7 +417,7 @@ TEST_F(CreateAllocationForHostSurfaceTest, givenReadOnlyHostPointerWhenAllocatio
     EXPECT_NE(memory, allocation->getUnderlyingBuffer());
     EXPECT_THAT(allocation->getUnderlyingBuffer(), MemCompare(memory, size));
 
-    allocation->taskCount = commandStreamReceiver->peekLatestFlushedTaskCount();
+    allocation->updateTaskCount(commandStreamReceiver->peekLatestFlushedTaskCount(), 0u);
 }
 
 TEST_F(CreateAllocationForHostSurfaceTest, givenReadOnlyHostPointerWhenAllocationForHostSurfaceWithPtrCopyNotAllowedIsCreatedThenCopyAllocationIsNotCreated) {

unit_tests/event/event_tests.cpp

@@ -410,7 +410,7 @@ TEST_F(UpdateEventTest, givenEventContainingCommandQueueWhenItsStatusIsUpdatedTo
     void *ptr = (void *)0x1000;
     size_t size = 4096;
     auto temporary = memoryManager->allocateGraphicsMemory(size, ptr);
-    temporary->taskCount = 3;
+    temporary->updateTaskCount(3, 0);
     device->getCommandStreamReceiver().getInternalAllocationStorage()->storeAllocation(std::unique_ptr<GraphicsAllocation>(temporary), TEMPORARY_ALLOCATION);
     Event event(commandQueue.get(), CL_COMMAND_NDRANGE_KERNEL, 3, 3);
 

unit_tests/event/user_events_tests.cpp

@@ -904,7 +904,7 @@ TEST_F(EventTests, waitForEventsDestroysTemporaryAllocations) {
 
     EXPECT_EQ(temporaryAllocation, csr.getTemporaryAllocations().peekHead());
 
-    temporaryAllocation->taskCount = 10;
+    temporaryAllocation->updateTaskCount(10, 0u);
 
     Event event(pCmdQ, CL_COMMAND_NDRANGE_KERNEL, 3, 11);
 

unit_tests/kernel/kernel_tests.cpp

@@ -586,7 +586,7 @@ TEST_F(KernelPrivateSurfaceTest, givenKernelWithPrivateSurfaceThatIsInUseByGpuWh
     auto privateSurface = pKernel->getPrivateSurface();
     auto tagAddress = context.getDevice(0)->getTagAddress();
 
-    privateSurface->taskCount = *tagAddress + 1;
+    privateSurface->updateTaskCount(*tagAddress + 1, 0u);
 
     EXPECT_TRUE(csr.getTemporaryAllocations().peekIsEmpty());
     pKernel.reset(nullptr);

unit_tests/kernel/substitute_kernel_heap_tests.cpp

@@ -117,7 +117,9 @@ TEST_F(KernelSubstituteTest, givenKernelWithUsedKernelAllocationWhenSubstituteKe
     kernel.kernelInfo.createKernelAllocation(memoryManager);
     auto firstAllocation = kernel.kernelInfo.kernelAllocation;
 
-    firstAllocation->taskCount = ObjectNotUsed - 1;
+    uint32_t notReadyTaskCount = *commandStreamReceiver.getTagAddress() + 1u;
+
+    firstAllocation->updateTaskCount(notReadyTaskCount, 0u);
 
     const size_t newHeapSize = initialHeapSize + 1;
     char newHeap[newHeapSize];
@@ -130,5 +132,5 @@ TEST_F(KernelSubstituteTest, givenKernelWithUsedKernelAllocationWhenSubstituteKe
     EXPECT_FALSE(commandStreamReceiver.getTemporaryAllocations().peekIsEmpty());
     EXPECT_EQ(commandStreamReceiver.getTemporaryAllocations().peekHead(), firstAllocation);
     memoryManager->checkGpuUsageAndDestroyGraphicsAllocations(secondAllocation);
-    commandStreamReceiver.getInternalAllocationStorage()->cleanAllocationList(firstAllocation->taskCount, TEMPORARY_ALLOCATION);
+    commandStreamReceiver.getInternalAllocationStorage()->cleanAllocationList(notReadyTaskCount, TEMPORARY_ALLOCATION);
 }

unit_tests/mem_obj/mem_obj_destruction_tests.cpp

@@ -46,19 +46,20 @@ class MemObjDestructionTest : public ::testing::TestWithParam<bool> {
     }
 
     void makeMemObjUsed() {
-        memObj->getGraphicsAllocation()->taskCount = 3;
+        memObj->getGraphicsAllocation()->updateTaskCount(taskCountReady, 0u);
     }
 
     void makeMemObjNotReady() {
         makeMemObjUsed();
-        *context->getDevice(0)->getTagAddress() = memObj->getGraphicsAllocation()->taskCount - 1;
+        *context->getDevice(0)->getTagAddress() = taskCountReady - 1;
     }
 
     void makeMemObjReady() {
         makeMemObjUsed();
-        *context->getDevice(0)->getTagAddress() = memObj->getGraphicsAllocation()->taskCount;
+        *context->getDevice(0)->getTagAddress() = taskCountReady;
     }
 
+    constexpr static uint32_t taskCountReady = 3u;
     MockDevice *device;
     MockMemoryManager *memoryManager;
     std::unique_ptr<MockContext> context;
@@ -135,7 +136,7 @@ HWTEST_P(MemObjAsyncDestructionTest, givenUsedMemObjWithAsyncDestructionsEnabled
         .WillByDefault(::testing::Invoke(waitForCompletionWithTimeoutMock));
 
     if (hasCallbacks) {
-        EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, TimeoutControls::maxTimeout, allocation->taskCount))
+        EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, TimeoutControls::maxTimeout, allocation->getTaskCount(0)))
             .Times(1);
     } else {
         EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, ::testing::_, ::testing::_))
@@ -165,7 +166,7 @@ HWTEST_P(MemObjAsyncDestructionTest, givenUsedMemObjWithAsyncDestructionsEnabled
         .WillByDefault(::testing::Invoke(waitForCompletionWithTimeoutMock));
 
     if (hasAllocatedMappedPtr) {
-        EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, TimeoutControls::maxTimeout, allocation->taskCount))
+        EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, TimeoutControls::maxTimeout, allocation->getTaskCount(0)))
             .Times(1);
     } else {
         EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, ::testing::_, ::testing::_))
@@ -206,7 +207,7 @@ HWTEST_P(MemObjAsyncDestructionTest, givenUsedMemObjWithAsyncDestructionsEnabled
         .WillByDefault(::testing::Invoke(waitForCompletionWithTimeoutMock));
 
     if (hasAllocatedMappedPtr) {
-        EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, TimeoutControls::maxTimeout, allocation->taskCount))
+        EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, TimeoutControls::maxTimeout, allocation->getTaskCount(0)))
             .Times(1);
     } else {
         EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, ::testing::_, ::testing::_))
@@ -238,7 +239,7 @@ HWTEST_P(MemObjSyncDestructionTest, givenMemObjWithDestructableAllocationWhenAsy
     ON_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, ::testing::_, ::testing::_))
         .WillByDefault(::testing::Invoke(waitForCompletionWithTimeoutMock));
 
-    EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, TimeoutControls::maxTimeout, allocation->taskCount))
+    EXPECT_CALL(*mockCsr, waitForCompletionWithTimeout(::testing::_, TimeoutControls::maxTimeout, allocation->getTaskCount(0)))
         .Times(1);
 
     delete memObj;

unit_tests/mem_obj/mem_obj_tests.cpp

@@ -162,7 +162,7 @@ TEST(MemObj, givenNotReadyGraphicsAllocationWhenMemObjDestroysAllocationAsyncThe
     context.setMemoryManager(&memoryManager);
 
     auto allocation = memoryManager.allocateGraphicsMemory(MemoryConstants::pageSize);
-    allocation->taskCount = 2;
+    allocation->updateTaskCount(2, 0);
     *(memoryManager.getCommandStreamReceiver(0)->getTagAddress()) = 1;
     MemObj memObj(&context, CL_MEM_OBJECT_BUFFER, CL_MEM_COPY_HOST_PTR,
                   MemoryConstants::pageSize, nullptr, nullptr, nullptr, true, false, false);
@@ -180,7 +180,7 @@ TEST(MemObj, givenReadyGraphicsAllocationWhenMemObjDestroysAllocationAsyncThenAl
     context.setMemoryManager(&memoryManager);
 
     auto allocation = memoryManager.allocateGraphicsMemory(MemoryConstants::pageSize);
-    allocation->taskCount = 1;
+    allocation->updateTaskCount(1, 0);
     *context.getDevice(0)->getTagAddress() = 1;
     MemObj memObj(&context, CL_MEM_OBJECT_BUFFER, CL_MEM_COPY_HOST_PTR,
                   MemoryConstants::pageSize, nullptr, nullptr, nullptr, true, false, false);
@@ -199,7 +199,6 @@ TEST(MemObj, givenNotUsedGraphicsAllocationWhenMemObjDestroysAllocationAsyncThen
     context.setMemoryManager(&memoryManager);
 
     auto allocation = memoryManager.allocateGraphicsMemory(MemoryConstants::pageSize);
-    allocation->taskCount = ObjectNotUsed;
     MemObj memObj(&context, CL_MEM_OBJECT_BUFFER, CL_MEM_COPY_HOST_PTR,
                   MemoryConstants::pageSize, nullptr, nullptr, nullptr, true, false, false);
 
@@ -228,21 +227,6 @@ TEST(MemObj, givenMemoryManagerWithoutDeviceWhenMemObjDestroysAllocationAsyncThe
     EXPECT_TRUE(allocationList.peekIsEmpty());
 }
 
-TEST(MemObj, givenMemObjWhenItDoesntHaveGraphicsAllocationThenWaitForCsrCompletionDoesntCrash) {
-    MockContext context;
-    MockMemoryManager memoryManager(*context.getDevice(0)->getExecutionEnvironment());
-
-    context.setMemoryManager(&memoryManager);
-
-    MemObj memObj(&context, CL_MEM_OBJECT_BUFFER, CL_MEM_COPY_HOST_PTR,
-                  MemoryConstants::pageSize, nullptr, nullptr, nullptr, true, false, false);
-
-    EXPECT_EQ(nullptr, memObj.getGraphicsAllocation());
-    memObj.waitForCsrCompletion();
-
-    EXPECT_EQ(nullptr, memObj.getGraphicsAllocation());
-    memObj.waitForCsrCompletion();
-}
 TEST(MemObj, givenMemObjAndPointerToObjStorageWithProperCommandWhenCheckIfMemTransferRequiredThenReturnFalse) {
     MockMemoryManager memoryManager;
     MockContext context;

unit_tests/memory_manager/CMakeLists.txt

@@ -8,6 +8,7 @@ set(IGDRCL_SRCS_tests_memory_manager
   ${CMAKE_CURRENT_SOURCE_DIR}/CMakeLists.txt
   ${CMAKE_CURRENT_SOURCE_DIR}/address_mapper_tests.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/deferred_deleter_mt_tests.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/graphics_allocation_tests.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/host_ptr_manager_tests.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/internal_allocation_storage_tests.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/memory_manager_tests.cpp

unit_tests/memory_manager/graphics_allocation_tests.cpp

@@ -0,0 +1,69 @@
+/*
+ * Copyright (C) 2018 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "gtest/gtest.h"
+#include "runtime/memory_manager/graphics_allocation.h"
+
+using namespace OCLRT;
+
+TEST(GraphicsAllocationTest, givenGraphicsAllocationWhenIsCreatedThenTaskCountsAreInitializedProperly) {
+    GraphicsAllocation graphicsAllocation1(nullptr, 0u, 0u, 0u);
+    GraphicsAllocation graphicsAllocation2(nullptr, 0u, 0u);
+    for (auto i = 0u; i < maxOsContextCount; i++) {
+        EXPECT_EQ(ObjectNotUsed, graphicsAllocation1.getTaskCount(i));
+        EXPECT_EQ(ObjectNotUsed, graphicsAllocation2.getTaskCount(i));
+    }
+}
+TEST(GraphicsAllocationTest, givenGraphicsAllocationWhenAccessTaskCountForInvalidContextThenAbort) {
+    GraphicsAllocation graphicsAllocation(nullptr, 0u, 0u);
+    EXPECT_THROW(graphicsAllocation.getTaskCount(maxOsContextCount), std::exception);
+}
+TEST(GraphicsAllocationTest, givenGraphicsAllocationWhenUpdateTaskCountForInvalidContextThenAbort) {
+    GraphicsAllocation graphicsAllocation(nullptr, 0u, 0u);
+    EXPECT_THROW(graphicsAllocation.updateTaskCount(0u, maxOsContextCount), std::exception);
+}
+TEST(GraphicsAllocationTest, givenGraphicsAllocationWhenUpdatedTaskCountThenAllocationWasUsed) {
+    GraphicsAllocation graphicsAllocation(nullptr, 0u, 0u);
+    EXPECT_FALSE(graphicsAllocation.peekWasUsed());
+    graphicsAllocation.updateTaskCount(0u, 0u);
+    EXPECT_TRUE(graphicsAllocation.peekWasUsed());
+}
+TEST(GraphicsAllocationTest, givenGraphicsAllocationWhenUpdatedTaskCountThenOnlyOneTaskCountIsUpdated) {
+    GraphicsAllocation graphicsAllocation(nullptr, 0u, 0u);
+    graphicsAllocation.updateTaskCount(1u, 0u);
+    EXPECT_EQ(1u, graphicsAllocation.getTaskCount(0u));
+    for (auto i = 1u; i < maxOsContextCount; i++) {
+        EXPECT_EQ(ObjectNotUsed, graphicsAllocation.getTaskCount(i));
+    }
+    graphicsAllocation.updateTaskCount(2u, 1u);
+    EXPECT_EQ(1u, graphicsAllocation.getTaskCount(0u));
+    EXPECT_EQ(2u, graphicsAllocation.getTaskCount(1u));
+    for (auto i = 2u; i < maxOsContextCount; i++) {
+        EXPECT_EQ(ObjectNotUsed, graphicsAllocation.getTaskCount(i));
+    }
+}
+TEST(GraphicsAllocationTest, givenGraphicsAllocationWhenUpdatedTaskCountToObjectNotUsedValueThenUnregisterContext) {
+    GraphicsAllocation graphicsAllocation(nullptr, 0u, 0u);
+    EXPECT_FALSE(graphicsAllocation.peekWasUsed());
+    graphicsAllocation.updateTaskCount(0u, 0u);
+    EXPECT_TRUE(graphicsAllocation.peekWasUsed());
+    graphicsAllocation.updateTaskCount(ObjectNotUsed, 0u);
+    EXPECT_FALSE(graphicsAllocation.peekWasUsed());
+}
+TEST(GraphicsAllocationTest, whenTwoContextsUpdatedTaskCountAndOneOfThemUnregisteredThenOneContextUsageRemains) {
+    GraphicsAllocation graphicsAllocation(nullptr, 0u, 0u);
+    EXPECT_FALSE(graphicsAllocation.peekWasUsed());
+    graphicsAllocation.updateTaskCount(0u, 0u);
+    graphicsAllocation.updateTaskCount(0u, 1u);
+    EXPECT_TRUE(graphicsAllocation.peekWasUsed());
+    graphicsAllocation.updateTaskCount(ObjectNotUsed, 0u);
+    EXPECT_TRUE(graphicsAllocation.peekWasUsed());
+    graphicsAllocation.updateTaskCount(ObjectNotUsed, 0u);
+    EXPECT_TRUE(graphicsAllocation.peekWasUsed());
+    graphicsAllocation.updateTaskCount(ObjectNotUsed, 1u);
+    EXPECT_FALSE(graphicsAllocation.peekWasUsed());
+}

unit_tests/memory_manager/internal_allocation_storage_tests.cpp

@@ -41,9 +41,9 @@ TEST_F(InternalAllocationStorageTest, whenCleanAllocationListThenRemoveOnlyCompl
     auto allocation2 = memoryManager->allocateGraphicsMemory(1, host_ptr);
     auto allocation3 = memoryManager->allocateGraphicsMemory(1, host_ptr);
 
-    allocation->taskCount = 10;
+    allocation->updateTaskCount(10, 0);
-    allocation2->taskCount = 5;
+    allocation2->updateTaskCount(5, 0);
-    allocation3->taskCount = 15;
+    allocation3->updateTaskCount(15, 0);
 
     storage->storeAllocation(std::unique_ptr<GraphicsAllocation>(allocation), TEMPORARY_ALLOCATION);
     storage->storeAllocation(std::unique_ptr<GraphicsAllocation>(allocation2), TEMPORARY_ALLOCATION);

unit_tests/memory_manager/memory_manager_tests.cpp

@@ -200,7 +200,7 @@ TEST_F(MemoryAllocatorTest, allocateGraphics) {
     auto allocation = memoryManager->allocateGraphicsMemory(sizeof(char));
     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->taskCount);
+    EXPECT_EQ((uint32_t)-1, allocation->getTaskCount(0));
     // 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());
@@ -1207,7 +1207,7 @@ TEST_F(MemoryManagerWithCsrTest, givenAllocationThatWasUsedAndIsCompletedWhenche
     auto tagAddress = csr->getTagAddress();
     ASSERT_NE(0u, *tagAddress);
 
-    usedAllocationButGpuCompleted->taskCount = *tagAddress - 1;
+    usedAllocationButGpuCompleted->updateTaskCount(*tagAddress - 1, 0);
 
     memoryManager->checkGpuUsageAndDestroyGraphicsAllocations(usedAllocationButGpuCompleted);
     EXPECT_TRUE(csr->getTemporaryAllocations().peekIsEmpty());
@@ -1218,14 +1218,14 @@ TEST_F(MemoryManagerWithCsrTest, givenAllocationThatWasUsedAndIsNotCompletedWhen
 
     auto tagAddress = csr->getTagAddress();
 
-    usedAllocationAndNotGpuCompleted->taskCount = *tagAddress + 1;
+    usedAllocationAndNotGpuCompleted->updateTaskCount(*tagAddress + 1, 0);
 
     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->taskCount = ObjectNotUsed;
+    usedAllocationAndNotGpuCompleted->updateTaskCount(ObjectNotUsed, 0);
 }
 
 class MockAlignMallocMemoryManager : public MockMemoryManager {

unit_tests/mocks/mock_csr.h

@@ -76,9 +76,9 @@ class MockCsrBase : public UltCommandStreamReceiver<GfxFamily> {
         processEvictionCalled = true;
     }
 
-    void waitForTaskCountAndCleanAllocationList(uint32_t requiredTaskCount, uint32_t allocationType) override {
+    void waitForTaskCountAndCleanAllocationList(uint32_t requiredTaskCount, uint32_t allocationUsage) override {
         waitForTaskCountRequiredTaskCount = requiredTaskCount;
-        BaseUltCsrClass::waitForTaskCountAndCleanAllocationList(requiredTaskCount, allocationType);
+        BaseUltCsrClass::waitForTaskCountAndCleanAllocationList(requiredTaskCount, allocationUsage);
     }
 
     ResidencyContainer madeResidentGfxAllocations;

unit_tests/os_interface/windows/device_command_stream_tests.cpp

@@ -638,8 +638,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->taskCount = 1;
+    graphicsAllocation->updateTaskCount(1, 0u);
-    graphicsAllocation2->taskCount = 100;
+    graphicsAllocation2->updateTaskCount(100, 0u);
 
     csr->waitForTaskCountAndCleanAllocationList(1, TEMPORARY_ALLOCATION);
     // graphicsAllocation2 still lives

unit_tests/os_interface/windows/wddm_memory_manager_tests.cpp

@@ -1602,7 +1602,7 @@ TEST(WddmMemoryManagerCleanupTest, givenUsedTagAllocationInWddmMemoryManagerWhen
     EXPECT_EQ(executionEnvironment.commandStreamReceivers[0].get(), executionEnvironment.memoryManager->getCommandStreamReceiver(0));
     auto tagAllocator = executionEnvironment.memoryManager->getEventPerfCountAllocator();
     auto allocation = tagAllocator->getTag()->getGraphicsAllocation();
-    allocation->taskCount = 1;
+    allocation->updateTaskCount(1, 0);
     executionEnvironment.commandStreamReceivers.clear();
     EXPECT_THROW(executionEnvironment.memoryManager->getCommandStreamReceiver(0), std::exception);
     EXPECT_NO_THROW(executionEnvironment.memoryManager.reset());

unit_tests/program/program_data_tests.cpp

@@ -179,7 +179,7 @@ TEST_F(ProgramDataTest, givenConstantAllocationThatIsInUseByGpuWhenProgramIsBein
     auto &csr = pPlatform->getDevice(0)->getCommandStreamReceiver();
     auto tagAddress = csr.getTagAddress();
     auto constantSurface = pProgram->getConstantSurface();
-    constantSurface->taskCount = *tagAddress + 1;
+    constantSurface->updateTaskCount(*tagAddress + 1, 0);
 
     EXPECT_TRUE(csr.getTemporaryAllocations().peekIsEmpty());
     delete pProgram;
@@ -196,7 +196,7 @@ TEST_F(ProgramDataTest, givenGlobalAllocationThatIsInUseByGpuWhenProgramIsBeingD
     auto &csr = pPlatform->getDevice(0)->getCommandStreamReceiver();
     auto tagAddress = csr.getTagAddress();
     auto globalSurface = pProgram->getGlobalSurface();
-    globalSurface->taskCount = *tagAddress + 1;
+    globalSurface->updateTaskCount(*tagAddress + 1, 0);
 
     EXPECT_TRUE(csr.getTemporaryAllocations().peekIsEmpty());
     delete pProgram;

unit_tests/program/program_tests.cpp

@@ -664,7 +664,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->taskCount = 100;
+    kernelAllocation->updateTaskCount(100, 0);
     *pDevice->getTagAddress() = 0;
     pProgram->cleanCurrentKernelInfo();
     EXPECT_FALSE(csr.getTemporaryAllocations().peekIsEmpty());