No Description
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

RewindableAllocator.cs 15KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358
  1. using AOT;
  2. using System;
  3. using System.Threading;
  4. using Unity.Burst;
  5. using Unity.Collections.LowLevel.Unsafe;
  6. using Unity.Jobs.LowLevel.Unsafe;
  7. using Unity.Mathematics;
  8. namespace Unity.Collections
  9. {
  10. struct Spinner
  11. {
  12. int m_value;
  13. public void Lock()
  14. {
  15. while (0 != Interlocked.CompareExchange(ref m_value, 1, 0))
  16. {
  17. }
  18. Interlocked.MemoryBarrier();
  19. }
  20. public void Unlock()
  21. {
  22. Interlocked.MemoryBarrier();
  23. while (1 != Interlocked.CompareExchange(ref m_value, 0, 1))
  24. {
  25. }
  26. }
  27. }
  28. internal struct UnmanagedArray<T> : IDisposable where T : unmanaged
  29. {
  30. IntPtr m_pointer;
  31. int m_length;
  32. AllocatorManager.AllocatorHandle m_allocator;
  33. public UnmanagedArray(int length, AllocatorManager.AllocatorHandle allocator)
  34. {
  35. unsafe
  36. {
  37. m_pointer = (IntPtr)Memory.Unmanaged.Array.Allocate<T>(length, allocator);
  38. }
  39. m_length = length;
  40. m_allocator = allocator;
  41. }
  42. public void Dispose()
  43. {
  44. unsafe
  45. {
  46. Memory.Unmanaged.Free((T*)m_pointer, Allocator.Persistent);
  47. }
  48. }
  49. public unsafe T* GetUnsafePointer()
  50. {
  51. return (T*)m_pointer;
  52. }
  53. public ref T this[int index]
  54. {
  55. get { unsafe { return ref ((T*)m_pointer)[index]; } }
  56. }
  57. }
  58. /// <summary>
  59. /// An allocator that is fast like a linear allocator, is threadsafe, and automatically invalidates
  60. /// all allocations made from it, when "rewound" by the user.
  61. /// </summary>
  62. [BurstCompile]
  63. public struct RewindableAllocator : AllocatorManager.IAllocator
  64. {
  65. [BurstCompatible]
  66. internal unsafe struct MemoryBlock : IDisposable
  67. {
  68. public const int kMaximumAlignment = 16384; // can't align any coarser than this many bytes
  69. public byte* m_pointer; // pointer to contiguous memory
  70. public long m_bytes; // how many bytes of contiguous memory it points to
  71. public long m_current; // next byte to give out, when people "allocate" from this block
  72. public long m_allocations; // how many allocations have been made from this block, so far?
  73. public MemoryBlock(long bytes)
  74. {
  75. m_pointer = (byte*)Memory.Unmanaged.Allocate(bytes, kMaximumAlignment, Allocator.Persistent);
  76. m_bytes = bytes;
  77. m_current = 0;
  78. m_allocations = 0;
  79. }
  80. public void Rewind()
  81. {
  82. m_current = 0;
  83. m_allocations = 0;
  84. }
  85. public void Dispose()
  86. {
  87. Memory.Unmanaged.Free(m_pointer, Allocator.Persistent);
  88. m_pointer = null;
  89. m_bytes = 0;
  90. m_current = 0;
  91. m_allocations = 0;
  92. }
  93. public int TryAllocate(ref AllocatorManager.Block block)
  94. {
  95. // Make the alignment multiple of cacheline size
  96. var alignment = math.max(JobsUtility.CacheLineSize, block.Alignment);
  97. var extra = alignment != JobsUtility.CacheLineSize ? 1 : 0;
  98. var cachelineMask = JobsUtility.CacheLineSize - 1;
  99. if (extra == 1)
  100. {
  101. alignment = (alignment + cachelineMask) & ~cachelineMask;
  102. }
  103. // Adjust the size to be multiple of alignment, add extra alignment
  104. // to size if alignment is more than cacheline size
  105. var mask = alignment - 1L;
  106. var size = (block.Bytes + extra * alignment + mask) & ~mask;
  107. var begin = Interlocked.Add(ref m_current, size) - size;
  108. begin = (begin + mask) & ~mask; // align the offset here
  109. if (begin + block.Bytes > m_bytes)
  110. return AllocatorManager.kErrorBufferOverflow;
  111. block.Range.Pointer = (IntPtr)(m_pointer + begin);
  112. block.AllocatedItems = block.Range.Items;
  113. Interlocked.Increment(ref m_allocations);
  114. return AllocatorManager.kErrorNone;
  115. }
  116. public bool Contains(IntPtr ptr)
  117. {
  118. unsafe
  119. {
  120. void* pointer = (void*)ptr;
  121. return (pointer >= m_pointer) && (pointer < m_pointer + m_current);
  122. }
  123. }
  124. };
  125. Spinner m_spinner;
  126. AllocatorManager.AllocatorHandle m_handle;
  127. UnmanagedArray<MemoryBlock> m_block;
  128. int m_best; // block we expect is best to allocate from next
  129. int m_last; // highest-index block that has memory to allocate from
  130. int m_used; // highest-index block that we actually allocated from, since last rewind
  131. bool m_enableBlockFree; // flag indicating if allocator enables individual block free
  132. /// <summary>
  133. /// Initializes the allocator. Must be called before first use.
  134. /// </summary>
  135. /// <param name="initialSizeInBytes">The initial capacity of the allocator, in bytes</param>
  136. public void Initialize(int initialSizeInBytes, bool enableBlockFree = false)
  137. {
  138. m_spinner = default;
  139. m_block = new UnmanagedArray<MemoryBlock>(64, Allocator.Persistent);
  140. m_block[0] = new MemoryBlock(initialSizeInBytes);
  141. m_last = m_used = m_best = 0;
  142. m_enableBlockFree = enableBlockFree;
  143. }
  144. /// <summary>
  145. /// Property to get and set enable block free flag, a flag indicating whether allocator enables individual block free.
  146. /// </summary>
  147. public bool EnableBlockFree
  148. {
  149. get => m_enableBlockFree;
  150. set => m_enableBlockFree = value;
  151. }
  152. /// <summary>
  153. /// Retrieves the number of memory blocks that the allocator has requested from the system.
  154. /// </summary>
  155. public int BlocksAllocated => (int)(m_last + 1);
  156. /// <summary>
  157. /// Retrieves the size of the initial memory block, as requested in the Initialize function.
  158. /// </summary>
  159. public int InitialSizeInBytes => (int)(m_block[0].m_bytes);
  160. /// <summary>
  161. /// Rewind the allocator; invalidate all allocations made from it, and potentially also free memory blocks
  162. /// it has allocated from the system.
  163. /// </summary>
  164. public void Rewind()
  165. {
  166. if (JobsUtility.IsExecutingJob)
  167. throw new InvalidOperationException("You cannot Rewind a RewindableAllocator from a Job.");
  168. m_handle.Rewind(); // bump the allocator handle version, invalidate all dependents
  169. while (m_last > m_used) // *delete* all blocks we didn't even allocate from this time around.
  170. m_block[m_last--].Dispose();
  171. while (m_used > 0) // simply *rewind* all blocks we used in this update, to avoid allocating again, every update.
  172. m_block[m_used--].Rewind();
  173. m_block[0].Rewind();
  174. }
  175. /// <summary>
  176. /// Dispose the allocator. This must be called to free the memory blocks that were allocated from the system.
  177. /// </summary>
  178. public void Dispose()
  179. {
  180. if (JobsUtility.IsExecutingJob)
  181. throw new InvalidOperationException("You cannot Dispose a RewindableAllocator from a Job.");
  182. m_used = 0; // so that we delete all blocks in Rewind() on the next line
  183. Rewind();
  184. m_block[0].Dispose();
  185. m_block.Dispose();
  186. m_last = m_used = m_best = 0;
  187. }
  188. /// <summary>
  189. /// All allocators must implement this property, in order to be installed in the custom allocator table.
  190. /// </summary>
  191. [NotBurstCompatible]
  192. public AllocatorManager.TryFunction Function => Try;
  193. /// <summary>
  194. /// Try to allocate, free, or reallocate a block of memory. This is an internal function, and
  195. /// is not generally called by the user.
  196. /// </summary>
  197. /// <param name="block">The memory block to allocate, free, or reallocate</param>
  198. public int Try(ref AllocatorManager.Block block)
  199. {
  200. if (block.Range.Pointer == IntPtr.Zero)
  201. {
  202. // first, try to allocate from the block that succeeded last time, which we expect is likely to succeed again.
  203. var error = m_block[m_best].TryAllocate(ref block);
  204. if (error == AllocatorManager.kErrorNone)
  205. return error;
  206. // if that fails, check all the blocks to see if any of them have enough memory
  207. m_spinner.Lock();
  208. int best;
  209. for (best = 0; best <= m_last; ++best)
  210. {
  211. error = m_block[best].TryAllocate(ref block);
  212. if (error == AllocatorManager.kErrorNone)
  213. {
  214. m_used = best > m_used ? best : m_used;
  215. m_best = best;
  216. m_spinner.Unlock();
  217. return error;
  218. }
  219. }
  220. // if that fails, allocate another block that's guaranteed big enough, and allocate from it.
  221. var bytes = math.max(m_block[0].m_bytes << best, math.ceilpow2(block.Bytes)); // if user suddenly asks for 1GB, skip smaller sizes
  222. m_block[best] = new MemoryBlock(bytes);
  223. error = m_block[best].TryAllocate(ref block);
  224. m_best = best;
  225. m_used = best;
  226. m_last = best;
  227. m_spinner.Unlock();
  228. return error;
  229. }
  230. // To free memory, no-op unless allocator enables individual block to be freed
  231. if (block.Range.Items == 0)
  232. {
  233. if (m_enableBlockFree)
  234. {
  235. m_spinner.Lock();
  236. if (m_block[m_best].Contains(block.Range.Pointer))
  237. if (0 == Interlocked.Decrement(ref m_block[m_best].m_allocations))
  238. m_block[m_best].Rewind();
  239. m_spinner.Unlock();
  240. }
  241. return 0; // we could check to see if the pointer belongs to us, if we want to be strict about it.
  242. }
  243. return -1;
  244. }
  245. [BurstCompile]
  246. [MonoPInvokeCallback(typeof(AllocatorManager.TryFunction))]
  247. internal static int Try(IntPtr state, ref AllocatorManager.Block block)
  248. {
  249. unsafe { return ((RewindableAllocator*)state)->Try(ref block); }
  250. }
  251. /// <summary>
  252. /// Retrieve the AllocatorHandle associated with this allocator. The handle is used as an index into a
  253. /// global table, for times when a reference to the allocator object isn't available.
  254. /// </summary>
  255. /// <value>The AllocatorHandle retrieved.</value>
  256. public AllocatorManager.AllocatorHandle Handle { get { return m_handle; } set { m_handle = value; } }
  257. /// <summary>
  258. /// Retrieve the Allocator associated with this allocator.
  259. /// </summary>
  260. /// <value>The Allocator retrieved.</value>
  261. public Allocator ToAllocator { get { return m_handle.ToAllocator; } }
  262. /// <summary>
  263. /// Check whether this AllocatorHandle is a custom allocator.
  264. /// </summary>
  265. /// <value>True if this AllocatorHandle is a custom allocator.</value>
  266. public bool IsCustomAllocator { get { return m_handle.IsCustomAllocator; } }
  267. /// <summary>
  268. /// Allocate a NativeArray of type T from memory that is guaranteed to remain valid until the end of the
  269. /// next Update of this World. There is no need to Dispose the NativeArray so allocated. It is not possible
  270. /// to free the memory by Disposing it - it is automatically freed after the end of the next Update for this
  271. /// World.
  272. /// </summary>
  273. /// <typeparam name="T">The element type of the NativeArray to allocate.</typeparam>
  274. /// <param name="length">The length of the NativeArray to allocate, measured in elements.</param>
  275. /// <returns>The NativeArray allocated by this function.</returns>
  276. [BurstCompatible(GenericTypeArguments = new[] { typeof(int) })]
  277. public NativeArray<T> AllocateNativeArray<T>(int length) where T : struct
  278. {
  279. var container = new NativeArray<T>();
  280. unsafe
  281. {
  282. container.m_Buffer = this.AllocateStruct(default(T), length);
  283. }
  284. container.m_Length = length;
  285. container.m_AllocatorLabel = Allocator.None;
  286. #if ENABLE_UNITY_COLLECTIONS_CHECKS
  287. container.m_MinIndex = 0;
  288. container.m_MaxIndex = length - 1;
  289. container.m_Safety = CollectionHelper.CreateSafetyHandle(ToAllocator);
  290. #if REMOVE_DISPOSE_SENTINEL
  291. #else
  292. container.m_DisposeSentinel = null;
  293. #endif
  294. CollectionHelper.SetStaticSafetyId<NativeArray<T>>(ref container.m_Safety, ref NativeArrayExtensions.NativeArrayStaticId<T>.s_staticSafetyId.Data);
  295. Handle.AddSafetyHandle(container.m_Safety);
  296. #endif
  297. return container;
  298. }
  299. /// <summary>
  300. /// Allocate a NativeList of type T from memory that is guaranteed to remain valid until the end of the
  301. /// next Update of this World. There is no need to Dispose the NativeList so allocated. It is not possible
  302. /// to free the memory by Disposing it - it is automatically freed after the end of the next Update for this
  303. /// World. The NativeList must be initialized with its maximum capacity; if it were to dynamically resize,
  304. /// up to 1/2 of the total final capacity would be wasted, because the memory can't be dynamically freed.
  305. /// </summary>
  306. /// <typeparam name="T">The element type of the NativeList to allocate.</typeparam>
  307. /// <param name="capacity">The capacity of the NativeList to allocate, measured in elements.</param>
  308. /// <returns>The NativeList allocated by this function.</returns>
  309. [BurstCompatible(GenericTypeArguments = new[] { typeof(int) })]
  310. public NativeList<T> AllocateNativeList<T>(int capacity) where T : unmanaged
  311. {
  312. var container = new NativeList<T>();
  313. unsafe
  314. {
  315. container.m_ListData = this.Allocate(default(UnsafeList<T>), 1);
  316. container.m_ListData->Ptr = this.Allocate(default(T), capacity);
  317. container.m_ListData->m_capacity = capacity;
  318. container.m_ListData->m_length = 0;
  319. container.m_ListData->Allocator = Allocator.None;
  320. }
  321. container.m_DeprecatedAllocator = Allocator.None;
  322. #if ENABLE_UNITY_COLLECTIONS_CHECKS
  323. container.m_Safety = CollectionHelper.CreateSafetyHandle(ToAllocator);
  324. #if REMOVE_DISPOSE_SENTINEL
  325. #else
  326. container.m_DisposeSentinel = null;
  327. #endif
  328. CollectionHelper.SetStaticSafetyId<NativeList<T>>(ref container.m_Safety, ref NativeList<T>.s_staticSafetyId.Data);
  329. AtomicSafetyHandle.SetBumpSecondaryVersionOnScheduleWrite(container.m_Safety, true);
  330. Handle.AddSafetyHandle(container.m_Safety);
  331. #endif
  332. return container;
  333. }
  334. }
  335. }