ArmNN
 24.02
SingleAxisPriorityList.cpp
Go to the documentation of this file.
1 //
2 // Copyright © 2021 Arm Ltd and Contributors. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5 
7 
8 #include <algorithm>
9 #include <cstdlib>
10 
11 #include <iostream>
12 
13 namespace armnn
14 {
15 
16 // This strategy uses a vector of size_ts/words to represent occupancy of a memBlock in a memBin.
17 // Where each bit represents occupancy of a single time-step in that lifetime.
18 // We can then use bit masks to check for overlaps of memBlocks along the lifetime
19 
20 // For more information on the algorithm itself see: https://arxiv.org/pdf/2001.03288.pdf
21 // This strategy is an implementation of 4.3 Greedy by Size
22 constexpr size_t wordSize = sizeof(size_t) * 8;
23 
24 std::string SingleAxisPriorityList::GetName() const {
25  return m_Name;
26 }
27 
29  return m_MemBlockStrategyType;
30 }
31 
32 struct SingleAxisPriorityList::BinTracker
33 {
34  // maxLifeTime is the number of operators in the model
35  // We then divide that by wordSize to get the number of words we need to store all the lifetimes
36  BinTracker(unsigned int maxLifeTime)
37  : m_OccupiedSpaces(1 + maxLifeTime/wordSize, 0)
38  {}
39 
40  // Add a block of a single word size to the bin
41  void AddBlock(MemBlock* block, const size_t word, const size_t index)
42  {
43  m_OccupiedSpaces[index] = m_OccupiedSpaces[index] | word;
44 
45  m_PlacedBlocks.push_back(block);
46  m_MemSize = std::max(m_MemSize, block->m_MemSize);
47  }
48 
49  // Add a block with a word size of two or more to the bin
50  void AddBlock(MemBlock* block,
51  const size_t startIndex,
52  const size_t endIndex,
53  const size_t firstWord,
54  const size_t lastWord)
55  {
56  m_OccupiedSpaces[startIndex] = m_OccupiedSpaces[startIndex] | firstWord;
57  m_OccupiedSpaces[endIndex] = m_OccupiedSpaces[endIndex] | lastWord;
58 
59  for (size_t i = startIndex +1; i <= endIndex -1; ++i)
60  {
61  m_OccupiedSpaces[i] = std::numeric_limits<size_t>::max();
62  }
63 
64  m_PlacedBlocks.push_back(block);
65  m_MemSize = std::max(m_MemSize, block->m_MemSize);
66  }
67 
68  size_t m_MemSize = 0;
69  std::vector<size_t> m_OccupiedSpaces;
70  std::vector<MemBlock*> m_PlacedBlocks;
71 };
72 
73 void SingleAxisPriorityList::PlaceBlocks(const std::list<MemBlock*>& priorityList,
74  std::vector<BinTracker>& placedBlocks,
75  const unsigned int maxLifetime)
76 {
77  // This function is used when the given block start and end lifetimes fit within a single word.
78  auto singleWordLoop = [&](MemBlock* curBlock, const size_t firstWord, const size_t index)
79  {
80  bool placed = false;
81  // loop through all existing bins
82  for (auto& blockList : placedBlocks)
83  {
84  // Check if the lifetimes at the given index overlap with the lifetimes of the block
85  if ((blockList.m_OccupiedSpaces[index] & firstWord) == 0)
86  {
87  // If the binary AND is 0 there is no overlap between the words and the block will fit
88  blockList.AddBlock(curBlock, firstWord, index);
89  placed = true;
90  break;
91  }
92  }
93  // No suitable bin was found, create a new bin/BinTracker and add it to the placedBlocks vector
94  if (!placed)
95  {
96  placedBlocks.emplace_back(BinTracker{maxLifetime});
97  placedBlocks.back().AddBlock(curBlock, firstWord, index);
98  }
99  };
100 
101  // This function is used when the given block start and end lifetimes fit within two words.
102  auto doubleWordLoop =[&](MemBlock* curBlock,
103  const size_t firstWord,
104  const size_t firstIndex,
105  const size_t lastWord,
106  const size_t lastIndex)
107  {
108  bool placed = false;
109  for (auto& blockList : placedBlocks)
110  {
111  // Check if the lifetimes at the given indexes overlap with the lifetimes of the block
112  if ((blockList.m_OccupiedSpaces[firstIndex] & firstWord) == 0 &&
113  (blockList.m_OccupiedSpaces[lastIndex] & lastWord) == 0)
114  {
115  blockList.AddBlock(curBlock, firstIndex, lastIndex, firstWord, lastWord);
116  placed = true;
117  break;
118  }
119  }
120  // No suitable bin was found, create a new bin/BinTracker and add it to the placedBlocks vector
121  if (!placed)
122  {
123  placedBlocks.emplace_back(BinTracker{maxLifetime});
124  placedBlocks.back().AddBlock(curBlock, firstIndex, lastIndex, firstWord, lastWord);
125  }
126  };
127 
128  // Loop through the blocks to place
129  for(const auto curBlock : priorityList)
130  {
131  // The lifetimes of both the block and bin are represented by single bits on a word/s
132  // Each bin has maxLifetime/wordSize words
133  // The number of words for a block depends on the size of the blocks lifetime
134  // and the alignment of the block's lifetimes
135  // This allows for checking sizeof(size_t) * 8 lifetimes at once with a binary AND
136  const size_t firstWordIndex = curBlock->m_StartOfLife/wordSize;
137  const size_t lastWordIndex = curBlock->m_EndOfLife/wordSize;
138 
139  // Align and right shift the first word
140  // This sets the bits before curBlock->m_StartOfLife to 0
141  size_t remainder = (curBlock->m_StartOfLife - firstWordIndex * wordSize);
142  size_t firstWord = std::numeric_limits<size_t>::max() >> remainder;
143 
144  // If the indexes match the block can fit into a single word
145  if(firstWordIndex == lastWordIndex)
146  {
147  // We then need to zero the bits to the right of curBlock->m_EndOfLife
148  remainder += curBlock->m_EndOfLife + 1 - curBlock->m_StartOfLife;
149  firstWord = firstWord >> (wordSize - remainder);
150  firstWord = firstWord << (wordSize - remainder);
151 
152  singleWordLoop(curBlock, firstWord, firstWordIndex);
153  continue;
154  }
155 
156  // The indexes don't match we need at least two words
157  // Zero the bits to the right of curBlock->m_EndOfLife
158  remainder = (curBlock->m_EndOfLife + 1 - lastWordIndex * wordSize);
159  size_t lastWord = std::numeric_limits<size_t>::max() << (wordSize - remainder);
160 
161  if(firstWordIndex + 1 == lastWordIndex)
162  {
163  doubleWordLoop(curBlock, firstWord, firstWordIndex, lastWord, lastWordIndex);
164  continue;
165  }
166 
167  // The block cannot fit into two words
168  // We don't need to create any more words to represent this,
169  // as any word between the first and last block would always equal the maximum value of size_t,
170  // all the lifetimes would be occupied
171 
172  // Instead, we just check that the corresponding word in the bin is completely empty
173 
174  bool placed = false;
175  for (auto& blockList : placedBlocks)
176  {
177  // Check the first and last word
178  if ((blockList.m_OccupiedSpaces[firstWordIndex] & firstWord) != 0 ||
179  (blockList.m_OccupiedSpaces[lastWordIndex] & lastWord) != 0)
180  {
181  continue;
182  }
183 
184  bool fits = true;
185  // Check that all spaces in between are clear
186  for (size_t i = firstWordIndex +1; i <= lastWordIndex -1; ++i)
187  {
188  if (blockList.m_OccupiedSpaces[i] != 0)
189  {
190  fits = false;
191  break;
192  }
193  }
194 
195  if (!fits)
196  {
197  continue;
198  }
199 
200  blockList.AddBlock(curBlock, firstWordIndex, lastWordIndex, firstWord, lastWord);
201  placed = true;
202  break;
203  }
204 
205  // No suitable bin was found, create a new bin/BinTracker and add it to the placedBlocks vector
206  if (!placed)
207  {
208  placedBlocks.emplace_back(BinTracker{maxLifetime});
209  placedBlocks.back().AddBlock(curBlock, firstWordIndex, lastWordIndex, firstWord, lastWord);
210  }
211  }
212 }
213 
214 std::vector<MemBin> SingleAxisPriorityList::Optimize(std::vector<MemBlock>& blocks)
215 {
216  unsigned int maxLifetime = 0;
217  std::list<MemBlock*> priorityList;
218  for (auto& block: blocks)
219  {
220  maxLifetime = std::max(maxLifetime, block.m_EndOfLife);
221  priorityList.emplace_back(&block);
222  }
223  maxLifetime++;
224 
225  // From testing ordering by m_MemSize in non-descending order gives the best results overall
226  priorityList.sort([](const MemBlock* lhs, const MemBlock* rhs)
227  {
228  return lhs->m_MemSize > rhs->m_MemSize ;
229  });
230 
231 
232  std::vector<BinTracker> placedBlocks;
233  placedBlocks.reserve(maxLifetime);
234  PlaceBlocks(priorityList, placedBlocks, maxLifetime);
235 
236  std::vector<MemBin> bins;
237  bins.reserve(placedBlocks.size());
238  for (auto blockList: placedBlocks)
239  {
240  MemBin bin;
241  bin.m_MemBlocks.reserve(blockList.m_PlacedBlocks.size());
242  bin.m_MemSize = blockList.m_MemSize;
243 
244  for (auto block : blockList.m_PlacedBlocks)
245  {
246  bin.m_MemBlocks.emplace_back(MemBlock{block->m_StartOfLife,
247  block->m_EndOfLife,
248  block->m_MemSize,
249  0,
250  block->m_Index,});
251  }
252  bins.push_back(std::move(bin));
253  }
254 
255  return bins;
256 }
257 
258 } // namespace armnn
259 
armnn::SingleAxisPriorityList::GetMemBlockStrategyType
MemBlockStrategyType GetMemBlockStrategyType() const override
Definition: SingleAxisPriorityList.cpp:28
armnn::MemBlock
Definition: IMemoryOptimizerStrategy.hpp:13
armnn::MemBlock::m_StartOfLife
const unsigned int m_StartOfLife
Definition: IMemoryOptimizerStrategy.hpp:22
armnn::wordSize
constexpr size_t wordSize
Definition: SingleAxisPriorityList.cpp:22
armnn::SingleAxisPriorityList::Optimize
std::vector< MemBin > Optimize(std::vector< MemBlock > &memBlocks) override
Definition: SingleAxisPriorityList.cpp:214
armnn::MemBlockStrategyType
MemBlockStrategyType
Definition: Types.hpp:253
SingleAxisPriorityList.hpp
armnn::MemBin::m_MemSize
size_t m_MemSize
Definition: IMemoryOptimizerStrategy.hpp:35
armnn::MemBlock::m_MemSize
const size_t m_MemSize
Definition: IMemoryOptimizerStrategy.hpp:25
armnn::MemBin::m_MemBlocks
std::vector< MemBlock > m_MemBlocks
Definition: IMemoryOptimizerStrategy.hpp:34
armnn
Copyright (c) 2021 ARM Limited and Contributors.
Definition: 01_00_quick_start.dox:6
armnn::MemBin
Definition: IMemoryOptimizerStrategy.hpp:32
armnn::SingleAxisPriorityList::GetName
std::string GetName() const override
Definition: SingleAxisPriorityList.cpp:24