From d306f183ded45eb724a62998961a3cf6ced9bfa4 Mon Sep 17 00:00:00 2001 From: Misaki Date: Fri, 11 Apr 2025 16:41:28 +0900 Subject: [PATCH] Refactor memory management in AllocationManager Changed the `AllocationManager` to use `ThreadLocal` for thread-specific memory allocation. Changed the `Initialize` method to directly validate `initialSize` and initialize the thread-local arena. Changed the `Allocate` and `Realloc` methods to utilize the thread-local arena for memory operations. Changed the `Free` method to remove unnecessary locking due to thread-local management. Added a new private method `EnsureInitialized` to verify the initialization of the thread-local arena. Added a public static method `ResetAll` to reset all thread-local arenas. Changed the `ResetCurrent` method to reset only the current thread's arena. Updated the `Dispose` method to clean up all thread-local arenas and clear the allocated dictionary. Uncommented the `UNSAFE_COLLECTION_CHECK` directive for enhanced debugging checks in `AllocationManager`. Changed `CollectionBenchmark` to call `AllocationManager.ResetCurrent()` after populating the `UnsafeArray`. --- .../CollectionBenchmark.cs | 1 + .../Buffer/AllocationManager.cs | 84 +++++++++++-------- 2 files changed, 49 insertions(+), 36 deletions(-) diff --git a/Misaki.HighPerformance.Test/CollectionBenchmark.cs b/Misaki.HighPerformance.Test/CollectionBenchmark.cs index b3361fc..442510a 100644 --- a/Misaki.HighPerformance.Test/CollectionBenchmark.cs +++ b/Misaki.HighPerformance.Test/CollectionBenchmark.cs @@ -34,6 +34,7 @@ public unsafe class CollectionBenchmark { array[i] = i; } + AllocationManager.ResetCurrent(); } [Benchmark] diff --git a/Misaki.HighPerformance.Unsafe/Buffer/AllocationManager.cs b/Misaki.HighPerformance.Unsafe/Buffer/AllocationManager.cs index 4256551..111afa1 100644 --- a/Misaki.HighPerformance.Unsafe/Buffer/AllocationManager.cs +++ b/Misaki.HighPerformance.Unsafe/Buffer/AllocationManager.cs @@ -1,6 +1,7 @@ -//#define UNSAFE_COLLECTION_CHECK +#define UNSAFE_COLLECTION_CHECK using Misaki.HighPerformance.Unsafe.Collections; +using System.Runtime.CompilerServices; #if UNSAFE_COLLECTION_CHECK #if DEBUG using System.Diagnostics; @@ -12,34 +13,37 @@ namespace Misaki.HighPerformance.Unsafe.Buffer; // TODO: Custom allocator public static unsafe class AllocationManager { - private const uint _DEFAULT_ARENA_SIZE = 512 * 1024; // 512 KB - - private static DynamicArena _arena; - private static bool _initialized; + private const uint _DEFAULT_ARENA_SIZE = 512 * 1024; + private static ThreadLocal? _threadLocalArena; #if UNSAFE_COLLECTION_CHECK private static Dictionary _allocated = null!; #endif - //private static readonly Lock _lock = new(); - /// /// Initializes the AllocationManager with a specified initial size for the memory arena. /// /// The initial size in bytes for the memory arena. public static void Initialize(uint initialSize = _DEFAULT_ARENA_SIZE) { - if (_initialized || initialSize <= 0) + if (initialSize <= 0) { return; } - _arena = new DynamicArena(initialSize); + _threadLocalArena = new ThreadLocal(() => new DynamicArena(initialSize), true); #if UNSAFE_COLLECTION_CHECK _allocated = new Dictionary(32); #endif + } - _initialized = true; + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static void EnsureInitialized() + { + if (_threadLocalArena == null) + { + throw new InvalidOperationException("The AllocationManager has not been initialized."); + } } public static T* Allocate(uint size, uint alignSize, Allocator allocator, AllocationOption allocationOption) @@ -50,18 +54,13 @@ public static unsafe class AllocationManager return (T*)AlignedAlloc(size, alignSize); } - if (!_initialized) - { - Initialize(); - } + EnsureInitialized(); - //lock (_lock) - //{ T* buffer; switch (allocator) { case Allocator.Temp: - buffer = (T*)_arena.Allocate(size * (uint)sizeof(T), alignSize, allocationOption); + buffer = (T*)_threadLocalArena!.Value.Allocate(size * (uint)sizeof(T), alignSize, allocationOption); break; case Allocator.Persistent: @@ -90,24 +89,18 @@ public static unsafe class AllocationManager } return buffer; - //} } public static T* Realloc(T* buffer, uint size, uint alignSize, Allocator allocator) where T : unmanaged { - if (!_initialized) - { - throw new InvalidOperationException("The AllocationManager has not been initialized."); - } + EnsureInitialized(); - //lock (_lock) - //{ T* newBuffer; switch (allocator) { case Allocator.Temp: - newBuffer = (T*)_arena.Allocate(size * (uint)sizeof(T), alignSize, AllocationOption.None); + newBuffer = (T*)_threadLocalArena!.Value.Allocate(size * (uint)sizeof(T), alignSize, AllocationOption.None); break; case Allocator.Persistent: @@ -134,13 +127,10 @@ public static unsafe class AllocationManager } return newBuffer; - //} } public static void Free(void* ptr, Allocator allocator) { - //lock (_lock) - //{ if (allocator == Allocator.Persistent) { AlignedFree(ptr); @@ -148,15 +138,34 @@ public static unsafe class AllocationManager _allocated.Remove((IntPtr)ptr); #endif } - //} } /// - /// Resets the memory arena, optionally clearing the allocated memory. + /// Resets the memory arenas on all of the threads. /// - public static void Reset() + public static void ResetAll() { - _arena.Reset(); + if (_threadLocalArena == null) + { + return; + } + + foreach (var arena in _threadLocalArena.Values) + { + arena.Reset(); + } + } + + /// + /// Resets the current thread-local arena. + /// + public static void ResetCurrent() + { + if (_threadLocalArena == null) + { + return; + } + _threadLocalArena.Value.Reset(); } /// @@ -167,12 +176,17 @@ public static unsafe class AllocationManager #endif public static void Dispose() { - if (!_initialized) + if (_threadLocalArena == null) { return; } - _arena.Dispose(); + foreach (var arena in _threadLocalArena.Values) + { + arena.Dispose(); + } + + _threadLocalArena.Dispose(); #if UNSAFE_COLLECTION_CHECK nuint unfreeBytes = 0u; @@ -189,7 +203,5 @@ public static unsafe class AllocationManager _allocated.Clear(); #endif - - _initialized = false; } } \ No newline at end of file