Refactor poolbuffer check to fix refcount issue when using clang comp…

opencl/source/context/context.cpp

@@ -574,6 +574,7 @@ Context::BufferPool::BufferPool(Context *context) : BaseType(context->memoryMana
                                            bufferCreateArgs,
                                            errcodeRet));
     if (this->mainStorage) {
+        this->mainStorage->setAsPoolBuffer(true);
         this->chunkAllocator.reset(new HeapAllocator(params.startingOffset,
                                                      context->getBufferPoolAllocator().getParams().aggregatedSmallBuffersPoolSize,
                                                      context->getBufferPoolAllocator().getParams().chunkAlignment));
@@ -602,6 +603,7 @@ Buffer *Context::BufferPool::allocate(const MemoryProperties &memoryProperties,
     auto bufferFromPool = this->mainStorage->createSubBuffer(flags, flagsIntel, &bufferRegion, errcodeRet);
     bufferFromPool->createFunction = this->mainStorage->createFunction;
     bufferFromPool->setSizeInPoolAllocator(actualSize);
+    bufferFromPool->setAsPoolBuffer(true);
     return bufferFromPool;
 }
 

opencl/source/mem_obj/buffer.h

@@ -66,6 +66,7 @@ class Buffer : public MemObj {
     constexpr static cl_ulong maskMagic = 0xFFFFFFFFFFFFFFFFLL;
     constexpr static cl_ulong objectMagic = MemObj::objectMagic | 0x02;
     bool forceDisallowCPUCopy = false;
+    bool poolBuffer = false;
 
     ~Buffer() override;
 
@@ -158,6 +159,8 @@ class Buffer : public MemObj {
     BufferCreateFunc createFunction = nullptr;
     bool isSubBuffer();
     bool isValidSubBufferOffset(size_t offset);
+    void setAsPoolBuffer(bool value) { this->poolBuffer = value; }
+    bool isPoolBuffer() const { return poolBuffer; }
     uint64_t setArgStateless(void *memory, uint32_t patchSize, uint32_t rootDeviceIndex, bool set32BitAddressing);
     virtual void setArgStateful(void *memory, bool forceNonAuxMode, bool disableL3, bool alignSizeForAuxTranslation,
                                 bool isReadOnly, const Device &device, bool areMultipleSubDevicesInContext) = 0;

shared/source/memory_manager/graphics_allocation.h

@@ -313,6 +313,9 @@ class GraphicsAllocation : public IDNode<GraphicsAllocation>, NEO::NonCopyableAn
     bool isAllocatedInLocalMemoryPool() const { return (this->memoryPool == MemoryPool::localMemory); }
     bool isAllocationLockable() const;
 
+    void setAsPoolBuffer(bool value) { this->poolBuffer = value; }
+    bool isPoolBuffer() const { return poolBuffer; }
+
     const AubInfo &getAubInfo() const { return aubInfo; }
 
     bool isCompressionEnabled() const;
@@ -438,6 +441,7 @@ class GraphicsAllocation : public IDNode<GraphicsAllocation>, NEO::NonCopyableAn
     std::atomic<uint32_t> registeredContextsNum{0};
     bool shareableHostMemory = false;
     bool cantBeReadOnly = false;
+    bool poolBuffer = false;
     bool explicitlyMadeResident = false;
     bool isImported = false;
 };

shared/source/utilities/buffer_pool_allocator.inl

@@ -48,7 +48,8 @@ template <typename PoolT, typename BufferType, typename BufferParentType>
 bool AbstractBuffersPool<PoolT, BufferType, BufferParentType>::isPoolBuffer(const BufferParentType *buffer) const {
     static_assert(std::is_base_of_v<BufferParentType, BufferType>);
 
-    return (buffer && this->mainStorage.get() == buffer);
+    const auto *bufferObj = static_cast<const BufferType *>(buffer);
+    return bufferObj && bufferObj->isPoolBuffer();
 }
 
 template <typename PoolT, typename BufferType, typename BufferParentType>

shared/source/utilities/generic_pool_allocator.inl

@@ -28,6 +28,9 @@ GenericPool<Traits>::GenericPool(Device *device, size_t poolSize)
                                                       MemoryConstants::pageSize,
                                                       0u));
     this->mainStorage.reset(graphicsAllocation);
+    if (this->mainStorage) {
+        this->mainStorage->setAsPoolBuffer(true);
+    }
     this->mtx = std::make_unique<std::mutex>();
     stackVec.push_back(graphicsAllocation);
 }

shared/source/utilities/isa_pool_allocator.cpp

@@ -28,6 +28,9 @@ ISAPool::ISAPool(Device *device, bool isBuiltin, size_t storageSize)
                                                       MemoryConstants::pageSize,
                                                       0u));
     this->mainStorage.reset(graphicsAllocation);
+    if (this->mainStorage) {
+        this->mainStorage->setAsPoolBuffer(true);
+    }
     this->mtx = std::make_unique<std::mutex>();
     this->stackVec.push_back(graphicsAllocation);
 }