WIP: Removed internal header tracking for WorkspaceHeap. Modified Handles to track the alignment that the underlying memory was aligned to. Re-worked WorkspaceHeap to use lists rather than Stack

This commit is contained in:
Jim
2026-03-23 22:26:01 +00:00
parent c2150acb2c
commit f73c09add5
6 changed files with 215 additions and 138 deletions

View File

@@ -3,7 +3,6 @@
using System;
using System.Collections.Generic;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using UnmanagedMMU.Allocators;
using UnmanagedMMU.Handles;
using UnmanagedMMU.Handles.Internal;
@@ -33,6 +32,16 @@
/// </remarks>
public unsafe sealed class WorkspaceHeap : IDisposable, IUnmanagedMemoryOwner
{
/// <summary>
/// Struct defining metadata to be stored with each allocation
/// </summary>
private struct AllocationMetadata
{
public IntPtr Block;
public nuint Alignment;
}
/// <summary>
/// The maximum size, in bytes, for "small" allocations. Uses fixed-size buckets
/// </summary>
@@ -61,12 +70,12 @@
/// Predefined bucket sizes used for small allocation reuse.
/// </summary>
private static readonly nuint[] _sizeClasses =
{
[
32, 64, 96, 128, 160, 192, 224, 256,
288, 320, 352, 384, 416, 448, 480, 512,
544, 576, 608, 640, 672, 704, 736, 768,
800, 832, 864, 896, 928, 960, 992, 1024
};
];
/// <summary>
/// Allocator interface used for all underlying unmanaged memory operations.
@@ -84,17 +93,17 @@
/// <remarks>
/// For a given bucket, the corresponding stack contains previously allocated blocks that are available to be used
/// </remarks>
private readonly Dictionary<nuint, Stack<IntPtr>> _smallFree = new();
private readonly Dictionary<nuint, List<AllocationMetadata>> _smallFree = [];
/// <summary>
/// Free lists for medium allocations keyed by the exactly allocated size, in bytes, and sorted for best-fit.
/// </summary>
private readonly SortedDictionary<nuint, Stack<IntPtr>> _mediumFree = new();
private readonly SortedDictionary<nuint, List<AllocationMetadata>> _mediumFree = [];
/// <summary>
/// Free lists for large allocations keyed by the exactly allocated size, in bytes, sorted for tolerance-based reuse.
/// </summary>
private readonly SortedDictionary<nuint, Stack<IntPtr>> _largeFree = new();
private readonly SortedDictionary<nuint, List<AllocationMetadata>> _largeFree = [];
/// <summary>
/// Tracks the total bytes of memory reserved from the provided <see cref="IUnmanagedAllocator"/>.
@@ -116,33 +125,6 @@
/// </summary>
private volatile bool _disposed;
/// <summary>
/// Internal header prepended to each allocation to track its size.
/// </summary>
[StructLayout(LayoutKind.Sequential)]
private struct BlockHeader
{
/// <summary>
/// Size of the allocation in bytes.
/// </summary>
public nuint Size;
/// <summary>
/// Padding to ensure <see cref="BlockHeader"></see> is 32-byte aligned
/// </summary>
private readonly nuint _pad1;
/// <summary>
/// Padding to ensure <see cref="BlockHeader"></see> is 32-byte aligned
/// </summary>
private readonly nuint _pad2;
/// <summary>
/// Padding to ensure <see cref="BlockHeader"></see> is 32-byte aligned
/// </summary>
private readonly nuint _pad3;
}
/// <summary>
/// Creates a new <see cref="WorkspaceHeap"/>.
/// </summary>
@@ -160,8 +142,10 @@
_allocator = allocator;
// Initialize small-size buckets
foreach (var size in _sizeClasses)
_smallFree[size] = new Stack<IntPtr>();
foreach (nuint size in _sizeClasses)
{
_smallFree[size] = [];
}
}
/// <summary>
@@ -243,16 +227,15 @@
/// Allocates a new block from the underlying allocator including a header.
/// </summary>
/// <param name="payloadSize">Requested payload size in bytes.</param>
/// <param name="alignment">FOO</param>
/// <returns>
/// Pointer to the allocated block (header included).
/// </returns>
private IntPtr AllocateNew(nuint payloadSize)
private IntPtr AllocateNewAligned(nuint payloadSize, nuint alignment)
{
nuint total = payloadSize + (nuint)sizeof(BlockHeader);
void* raw = _allocator.Alloc(total);
_totalReserved += total;
void* raw = _allocator.AllocAligned(payloadSize, alignment);
_totalReserved += payloadSize;
_totalAllocations++;
return (IntPtr)raw;
}
@@ -264,100 +247,154 @@
/// <returns> An <see cref="IMemoryHandle{T}"/> to the allocated memory.</returns>
/// <exception cref="ObjectDisposedException">Thrown if heap is disposed.</exception>
/// <exception cref="ArgumentOutOfRangeException">Thrown if size is zero.</exception>
/// <remarks>The memory returned by this method is always 16-byte aligned, for other alignemnts use </remarks>
public IMemoryHandle<T> Allocate<T>(int count, bool zero = false) where T : unmanaged
{
return AllocateAligned<T>(count, SegmentAlignment.Aligned16, zero);
}
/// <summary>
///
/// </summary>
/// <typeparam name="T"></typeparam>
/// <param name="count"></param>
/// <param name="alignment"></param>
/// <param name="zero"></param>
/// <returns></returns>
/// <exception cref="OverflowException"></exception>
public IMemoryHandle<T> AllocateAligned<T>(int count, SegmentAlignment alignment, bool zero = false) where T: unmanaged
{
ArgumentOutOfRangeException.ThrowIfNegative(count);
ArgumentOutOfRangeException.ThrowIfZero(count);
if ((nuint)count > nuint.MaxValue / (nuint)(sizeof(T)))
{
throw new OverflowException($"Requested allocation of {count} elements of type {typeof(T)} exceeds allowable maximum memory size.");
}
nuint size = (nuint)count * (nuint)sizeof(T);
nuint requestedAlignment = (nuint)alignment;
nuint effectiveAlignment = requestedAlignment < (nuint)sizeof(T) ? (nuint)sizeof(T) : requestedAlignment;
lock (_lock)
{
ThrowIfDisposed();
void* ptr = null;
if (size <= _smallThreshold)
{
ptr = AllocateSmall(size, zero);
}
else if (size <= _mediumThreshold)
{
ptr = AllocateMedium(size, zero);
}
else
{
ptr = AllocateLarge(size, zero);
}
void* ptr;
return new PersistentMemoryHandle<T>((T*)ptr, size, this);
if (size <= _smallThreshold)
ptr = AllocateSmallAligned(size, effectiveAlignment, zero);
else if (size <= _mediumThreshold)
ptr = AllocateMediumAligned(size, effectiveAlignment, zero);
else
ptr = AllocateLargeAligned(size, effectiveAlignment, zero);
return new PersistentMemoryHandle<T>((T*)ptr, size, effectiveAlignment, this);
}
}
/// <summary>Allocates a small-size block using bucketed free lists.</summary>
private void* AllocateSmall(nuint size, bool zero)
/// <summary>
/// TODO: Fill in
/// </summary>
/// <param name="alignment"></param>
/// <param name="list"></param>
/// <returns></returns>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
private static int FindIndexAlignmentMatch(nuint alignment, List<AllocationMetadata> list)
{
// Attempt to find an index match the requested size
int mIdx = -1;
for (int i = 0; i < list.Count; i++)
{
if (list[i].Alignment >= alignment)
{
mIdx = i;
break;
}
}
return mIdx;
}
/// <summary>
/// Allocates a small-size block using bucketed free lists.
/// </summary>
private void* AllocateSmallAligned(nuint size, nuint alignment, bool zero)
{
nuint bucket = GetSizeClass(size);
Stack<IntPtr> stack = _smallFree[bucket];
List<AllocationMetadata> free = _smallFree[bucket];
IntPtr block = stack.Count > 0
? stack.Pop()
: AllocateNew(bucket);
// Attempt to find an index match the requested size
int mIdx = FindIndexAlignmentMatch(alignment, free);
BlockHeader* header = (BlockHeader*)block;
header->Size = bucket;
if (mIdx >= 0)
{
AllocationMetadata meta = free[mIdx];
free.RemoveAt(mIdx);
_totalInUse += bucket;
void* user = (void*)meta.Block;
if (zero)
{
Unsafe.InitBlockUnaligned(user, 0, (uint)bucket);
}
return user;
}
// allocate a new small block
IntPtr block = AllocateNewAligned(bucket, alignment);
_totalInUse += bucket;
void* user = header + 1;
void* nUser = (void*)block;
if (zero)
Unsafe.InitBlockUnaligned(user, 0, (uint)bucket);
return user;
{
Unsafe.InitBlockUnaligned(nUser, 0, (uint)bucket);
}
return nUser;
}
/// <summary>Allocates a medium-size block using best-fit reuse.</summary>
private void* AllocateMedium(nuint size, bool zero)
private void* AllocateMediumAligned(nuint size, nuint alignment, bool zero)
{
foreach (var kv in _mediumFree)
foreach (KeyValuePair<nuint, List<AllocationMetadata>> kv in _mediumFree)
{
if (kv.Key >= size && kv.Value.Count > 0)
{
var block = kv.Value.Pop();
var header = (BlockHeader*)block;
_totalInUse += header->Size;
List<AllocationMetadata> free = kv.Value;
void* user = header + 1;
if (zero)
Unsafe.InitBlockUnaligned(user, 0, (uint)header->Size);
// Attempt to find an index match the requested size
int mIdx = FindIndexAlignmentMatch(alignment, free);
return user;
if (mIdx >= 0)
{
AllocationMetadata meta = free[mIdx];
free.RemoveAt(mIdx);
_totalInUse += kv.Key;
void* user = (void*)meta.Block;
if (zero)
{
Unsafe.InitBlockUnaligned(user, 0, (uint)kv.Key);
}
return user;
}
}
}
var newBlock = AllocateNew(size);
var newHeader = (BlockHeader*)newBlock;
newHeader->Size = size;
// no match
IntPtr newBlock = AllocateNewAligned(size, alignment);
_totalInUse += size;
void* newUser = newHeader + 1;
void* nUser = (void*)newBlock;
if (zero)
{
Unsafe.InitBlockUnaligned(newUser, 0, (uint)size);
}
Unsafe.InitBlockUnaligned(nUser, 0, (uint)size);
}
return newUser;
return nUser;
}
/// <summary>Allocates a large block using tolerance-based reuse (smallest ≥ requested within waste bounds).</summary>
private void* AllocateLarge(nuint size, bool zero)
/// <summary>
/// Allocates a large block using tolerance-based reuse (smallest ≥ requested within waste bounds).
/// </summary>
private void* AllocateLargeAligned(nuint size, nuint alignment, bool zero)
{
foreach (var kv in _largeFree)
foreach (KeyValuePair<nuint, List<AllocationMetadata>> kv in _largeFree)
{
nuint blockSize = kv.Key;
if (blockSize < size || kv.Value.Count == 0)
@@ -365,41 +402,46 @@
continue;
}
// check waste tolerance before attempting to find an alignment match
nuint waste = blockSize - size;
bool acceptable =
waste <= _largeMaxWasteBytes ||
((double)blockSize / size) <= _largeWasteRatioLimit;
bool acceptable = waste <= _largeMaxWasteBytes || ((double)blockSize / size) <= _largeWasteRatioLimit;
if (!acceptable)
{
continue;
}
var block = kv.Value.Pop();
var header = (BlockHeader*)block;
_totalInUse += header->Size;
List<AllocationMetadata> free = kv.Value;
void* user = header + 1;
if (zero)
// Attempt to find an index match the requested size
int mIdx = FindIndexAlignmentMatch(alignment, free);
if (mIdx >= 0)
{
Unsafe.InitBlockUnaligned(user, 0, (uint)header->Size);
}
AllocationMetadata meta = free[mIdx];
free.RemoveAt(mIdx);
return user;
_totalInUse += blockSize;
void* user = (void*)meta.Block;
if (zero)
{
Unsafe.InitBlockUnaligned(user, 0, (uint)blockSize);
}
return user;
}
}
var newBlock = AllocateNew(size);
var newHeader = (BlockHeader*)newBlock;
newHeader->Size = size;
_totalInUse += size;
// no match
void* newUser = newHeader + 1;
IntPtr block = AllocateNewAligned(size, alignment);
_totalInUse += size;
void* nUser = (void*)block;
if (zero)
{
Unsafe.InitBlockUnaligned(newUser, 0, (uint)size);
Unsafe.InitBlockUnaligned(nUser, 0, (uint)size);
}
return newUser;
return nUser;
}
/// <summary>
@@ -421,23 +463,41 @@
lock (_lock)
{
var header = ((BlockHeader*)handle.Pointer) - 1;
nuint size = header->Size;
nuint size = handle.ByteCount;
nuint alignment = handle.Alignment;
_totalInUse -= size;
if (size <= _smallThreshold)
_smallFree[size].Push((IntPtr)header);
{
_smallFree[size].Add(new AllocationMetadata
{
Block = (IntPtr)handle.Pointer,
Alignment = alignment
});
}
else if (size <= _mediumThreshold)
{
if (!_mediumFree.TryGetValue(size, out var stack))
_mediumFree[size] = stack = new Stack<IntPtr>();
stack.Push((IntPtr)header);
if (!_mediumFree.TryGetValue(size, out List<AllocationMetadata>? free))
{
_mediumFree[size] = free = [];
}
free.Add(new AllocationMetadata
{
Block = (IntPtr)handle.Pointer,
Alignment = alignment
});
}
else
{
if (!_largeFree.TryGetValue(size, out var stack))
_largeFree[size] = stack = new Stack<IntPtr>();
stack.Push((IntPtr)header);
if (!_largeFree.TryGetValue(size, out List<AllocationMetadata>? free))
{
_largeFree[size] = free = [];
}
free.Add(new AllocationMetadata
{
Block = (IntPtr)handle.Pointer,
Alignment = alignment
});
}
}
}
@@ -457,17 +517,17 @@
}
/// <summary>Helper to free all blocks in a dictionary of free stacks.</summary>
private void PruneDictionary(IDictionary<nuint, Stack<IntPtr>> dict)
private void PruneDictionary(IDictionary<nuint, List<AllocationMetadata>> dict)
{
foreach (var kv in dict)
foreach (KeyValuePair<nuint, List<AllocationMetadata>> kv in dict)
{
var stack = kv.Value;
while (stack.Count > 0)
List<AllocationMetadata> list = kv.Value;
while (list.Count > 0)
{
var block = stack.Pop();
var header = (BlockHeader*)block;
_allocator.Free((void*)block);
_totalReserved -= (header->Size + (nuint)sizeof(BlockHeader));
var item = list[^1];
list.RemoveAt(list.Count - 1);
_allocator.FreeAligned((void*)item.Block, item.Alignment);
_totalReserved -= kv.Key;
}
}
}