1 /*
2  * Copyright (C) 2019-2021 Intel Corporation
3  *
4  * SPDX-License-Identifier: MIT
5  *
6  */
7 
8 #pragma once
9 #include "shared/source/helpers/timestamp_packet.h"
10 #include "shared/test/common/mocks/mock_graphics_allocation.h"
11 
12 namespace NEO {
13 
14 template <typename TagType = TimestampPackets<uint32_t>>
15 class MockTagAllocator : public TagAllocator<TagType> {
16   public:
17     using BaseClass = TagAllocator<TagType>;
18     using BaseClass::freeTags;
19     using BaseClass::usedTags;
20     using NodeType = typename BaseClass::NodeType;
21 
MockTagAllocator(uint32_t rootDeviceIndex,MemoryManager * memoryManager,size_t tagCount,size_t tagAlignment,size_t tagSize,bool doNotReleaseNodes,DeviceBitfield deviceBitfield)22     MockTagAllocator(uint32_t rootDeviceIndex, MemoryManager *memoryManager, size_t tagCount,
23                      size_t tagAlignment, size_t tagSize, bool doNotReleaseNodes, DeviceBitfield deviceBitfield)
24         : BaseClass(std::vector<uint32_t>{rootDeviceIndex}, memoryManager, tagCount, tagAlignment, tagSize, doNotReleaseNodes, deviceBitfield) {
25     }
26 
27     MockTagAllocator(uint32_t rootDeviceIndex, MemoryManager *memoryManager, size_t tagCount = 10)
MockTagAllocator(rootDeviceIndex,memoryManager,tagCount,MemoryConstants::cacheLineSize,sizeof (TagType),false,mockDeviceBitfield)28         : MockTagAllocator(rootDeviceIndex, memoryManager, tagCount, MemoryConstants::cacheLineSize, sizeof(TagType), false, mockDeviceBitfield) {
29     }
30 
31     MockTagAllocator(const std::vector<uint32_t> &rootDeviceIndices, MemoryManager *memoryManager, size_t tagCount = 10)
BaseClass(rootDeviceIndices,memoryManager,tagCount,MemoryConstants::cacheLineSize,sizeof (TagType),false,mockDeviceBitfield)32         : BaseClass(rootDeviceIndices, memoryManager, tagCount, MemoryConstants::cacheLineSize, sizeof(TagType), false, mockDeviceBitfield) {}
33 
returnTag(TagNodeBase * node)34     void returnTag(TagNodeBase *node) override {
35         releaseReferenceNodes.push_back(static_cast<NodeType *>(node));
36         BaseClass::returnTag(node);
37     }
38 
returnTagToFreePool(TagNodeBase * node)39     void returnTagToFreePool(TagNodeBase *node) override {
40         returnedToFreePoolNodes.push_back(static_cast<NodeType *>(node));
41         BaseClass::returnTagToFreePool(node);
42     }
43 
44     std::vector<NodeType *> releaseReferenceNodes;
45     std::vector<NodeType *> returnedToFreePoolNodes;
46 };
47 
48 class MockTimestampPacketContainer : public TimestampPacketContainer {
49   public:
50     using TimestampPacketContainer::timestampPacketNodes;
51 
MockTimestampPacketContainer(TagAllocatorBase & tagAllocator,size_t numberOfPreallocatedTags)52     MockTimestampPacketContainer(TagAllocatorBase &tagAllocator, size_t numberOfPreallocatedTags) {
53         for (size_t i = 0; i < numberOfPreallocatedTags; i++) {
54             add(tagAllocator.getTag());
55         }
56     }
57 
getNode(size_t position)58     TagNodeBase *getNode(size_t position) {
59         return timestampPacketNodes.at(position);
60     }
61 };
62 } // namespace NEO
63