[Memory] Add free block tracker to BaseHeap for O(log n) allocation

Replace linear page_table_ scans in AllocRange with a std::map-based
free block index that tracks contiguous free regions.

Insertions now coalesce with adjacent blocks on release and AllocFixed
uses the targeted tracker for pure reserves and falls back to a full rebuild
for mixed-state commits.

Also fixing PhysicalHeap leaking parent memory on child allocation failure,
Reset() not restoring unreserved_page_count_ and some incorrect method names
in PhysicalHeap error messages
This commit is contained in:
Herman S.
2026-04-08 12:07:50 +09:00
parent 80c6751b82
commit ea02e8d317
3 changed files with 599 additions and 79 deletions

View File

@@ -803,6 +803,10 @@ void BaseHeap::Initialize(Memory* memory, uint8_t* membase, HeapType heap_type,
host_address_offset_ = host_address_offset;
page_table_.resize(heap_size / page_size);
unreserved_page_count_ = uint32_t(page_table_.size());
// Initialize free block tracker with a single block covering the entire heap.
free_blocks_.clear();
free_blocks_[0] = uint32_t(page_table_.size());
}
void BaseHeap::Dispose() {
@@ -816,6 +820,7 @@ void BaseHeap::Dispose() {
page_number += page_entry.region_page_count;
}
}
free_blocks_.clear();
}
void BaseHeap::DumpMap() {
@@ -938,14 +943,98 @@ bool BaseHeap::Restore(ByteStream* stream) {
}
}
RebuildFreeBlocks();
return true;
}
void BaseHeap::RebuildFreeBlocks() {
free_blocks_.clear();
uint32_t run_start = UINT32_MAX;
for (uint32_t i = 0; i < uint32_t(page_table_.size()); ++i) {
if (page_table_[i].state == 0) {
if (run_start == UINT32_MAX) {
run_start = i;
}
} else {
if (run_start != UINT32_MAX) {
free_blocks_[run_start] = i - run_start;
run_start = UINT32_MAX;
}
}
}
if (run_start != UINT32_MAX) {
free_blocks_[run_start] = uint32_t(page_table_.size()) - run_start;
}
}
void BaseHeap::RemoveFreeBlock(uint32_t start_page, uint32_t page_count) {
if (free_blocks_.empty()) {
return;
}
// Find the free block that contains the allocated range.
auto it = free_blocks_.upper_bound(start_page);
if (it != free_blocks_.begin()) {
--it;
}
// Verify the block actually contains our range.
uint32_t block_start = it->first;
uint32_t block_count = it->second;
uint32_t block_end = block_start + block_count;
assert_true(start_page >= block_start &&
start_page + page_count <= block_end);
free_blocks_.erase(it);
// Insert remnant before the allocated range.
if (block_start < start_page) {
free_blocks_[block_start] = start_page - block_start;
}
// Insert remnant after the allocated range.
uint32_t alloc_end = start_page + page_count;
if (alloc_end < block_end) {
free_blocks_[alloc_end] = block_end - alloc_end;
}
}
void BaseHeap::InsertFreeBlock(uint32_t start_page, uint32_t page_count) {
uint32_t new_start = start_page;
uint32_t new_count = page_count;
// Try to merge with block immediately after.
auto it_after = free_blocks_.find(start_page + page_count);
if (it_after != free_blocks_.end()) {
new_count += it_after->second;
free_blocks_.erase(it_after);
}
// Try to merge with block immediately before.
auto it_at = free_blocks_.lower_bound(start_page);
if (it_at != free_blocks_.begin()) {
auto it_before = std::prev(it_at);
if (it_before->first + it_before->second == start_page) {
new_start = it_before->first;
new_count += it_before->second;
free_blocks_.erase(it_before);
}
}
free_blocks_[new_start] = new_count;
}
void BaseHeap::Reset() {
// TODO(DrChat): protect pages.
std::memset(page_table_.data(), 0, sizeof(PageEntry) * page_table_.size());
unreserved_page_count_ = uint32_t(page_table_.size());
// TODO(Triang3l): Remove access callbacks from pages if this is a physical
// memory heap.
// Re-initialize free block tracker.
free_blocks_.clear();
free_blocks_[0] = uint32_t(page_table_.size());
}
bool BaseHeap::Alloc(uint32_t size, uint32_t alignment,
@@ -992,10 +1081,10 @@ bool BaseHeap::AllocFixed(uint32_t base_address, uint32_t size,
auto global_lock = global_critical_region_.Acquire();
// - If we are reserving the entire range requested must not be already
// reserved.
// - If we are reserving, the entire range must not be already reserved.
// - If we are committing it's ok for pages within the range to already be
// committed.
const bool is_pure_reserve = allocation_type == kMemoryAllocationReserve;
for (uint32_t page_number = start_page_number; page_number <= end_page_number;
++page_number) {
uint32_t state = page_table_[page_number].state;
@@ -1042,6 +1131,7 @@ bool BaseHeap::AllocFixed(uint32_t base_address, uint32_t size,
}
// Set page state.
bool had_free_pages = false;
for (uint32_t page_number = start_page_number; page_number <= end_page_number;
++page_number) {
auto& page_entry = page_table_[page_number];
@@ -1053,11 +1143,25 @@ bool BaseHeap::AllocFixed(uint32_t base_address, uint32_t size,
page_entry.allocation_protect = protect;
page_entry.current_protect = protect;
if (!(page_entry.state & kMemoryAllocationReserve)) {
had_free_pages = true;
unreserved_page_count_--;
}
page_entry.state = kMemoryAllocationReserve | allocation_type;
}
// Update free block tracker if any pages transitioned from free.
if (had_free_pages) {
if (is_pure_reserve) {
// Pure reserve: validation confirmed all pages were free, so the range
// is within a single coalesced free block.
RemoveFreeBlock(start_page_number, page_count);
} else {
// Mixed state (commit upgraded to reserve+commit): pages may span
// multiple free blocks, rebuild from page_table_.
RebuildFreeBlocks();
}
}
return true;
}
template <typename T>
@@ -1094,88 +1198,85 @@ bool BaseHeap::AllocRange(uint32_t low_address, uint32_t high_address,
auto global_lock = global_critical_region_.Acquire();
// Find a free page range.
// The base page must match the requested alignment, so we first scan for
// a free aligned page and only then check for continuous free pages.
// TODO(benvanik): optimized searching (free list buckets, bitmap, etc).
// Find a free page range using the free block tracker.
// The base page must match the requested alignment.
uint32_t start_page_number = UINT_MAX;
uint32_t end_page_number = UINT_MAX;
// chrispy:todo, page_scan_stride is probably always a power of two...
uint32_t page_scan_stride = alignment >> page_size_shift_;
high_page_number =
high_page_number - QuickMod(high_page_number, page_scan_stride);
if (top_down) {
for (int64_t base_page_number =
high_page_number - xe::round_up(page_count, page_scan_stride);
base_page_number >= low_page_number;
base_page_number -= page_scan_stride) {
if (page_table_[base_page_number].state != 0) {
// Base page not free, skip to next usable page.
continue;
}
// Check requested range to ensure free.
start_page_number = uint32_t(base_page_number);
end_page_number = uint32_t(base_page_number) + page_count - 1;
assert_true(end_page_number < page_table_.size());
bool any_taken = false;
for (uint32_t page_number = uint32_t(base_page_number);
!any_taken && page_number <= end_page_number; ++page_number) {
bool is_free = page_table_[page_number].state == 0;
if (!is_free) {
// At least one page in the range is used, skip to next.
// We know we'll be starting at least before this page.
any_taken = true;
if (page_count > page_number) {
// Not enough space left to fit entire page range. Breaks outer
// loop.
base_page_number = -1;
} else {
base_page_number = page_number - page_count;
base_page_number -= QuickMod(base_page_number, page_scan_stride);
base_page_number += page_scan_stride; // cancel out loop logic
}
break;
}
}
if (!any_taken) {
// Found our place.
// Search free blocks from high addresses downward.
// Find the first block that could overlap our range.
auto it = free_blocks_.upper_bound(high_page_number);
while (it != free_blocks_.begin()) {
--it;
uint32_t block_start = it->first;
uint32_t block_count = it->second;
uint32_t block_end = block_start + block_count;
// Block is entirely below our search range — stop.
if (block_end <= low_page_number) {
break;
}
// Skip blocks too small to possibly fit.
if (block_count < page_count) {
continue;
}
// Compute the highest aligned start within this block and range.
uint32_t usable_end = std::min(block_end, high_page_number + 1);
if (usable_end < page_count) {
continue;
}
uint32_t latest_start = usable_end - page_count;
// Align down to stride.
latest_start -= QuickMod(latest_start, page_scan_stride);
uint32_t usable_start = std::max(block_start, low_page_number);
if (latest_start >= usable_start &&
latest_start + page_count <= block_end) {
start_page_number = latest_start;
end_page_number = latest_start + page_count - 1;
break;
}
// Retry.
start_page_number = end_page_number = UINT_MAX;
}
} else {
for (uint32_t base_page_number = low_page_number;
base_page_number <= high_page_number - page_count;
base_page_number += page_scan_stride) {
if (page_table_[base_page_number].state != 0) {
// Base page not free, skip to next usable page.
continue;
// Search free blocks from low addresses upward.
auto it = free_blocks_.lower_bound(low_page_number);
// Check if the previous block extends into our range.
if (it != free_blocks_.begin()) {
auto prev = std::prev(it);
if (prev->first + prev->second > low_page_number) {
it = prev;
}
// Check requested range to ensure free.
start_page_number = base_page_number;
end_page_number = base_page_number + page_count - 1;
bool any_taken = false;
for (uint32_t page_number = base_page_number;
!any_taken && page_number <= end_page_number; ++page_number) {
bool is_free = page_table_[page_number].state == 0;
if (!is_free) {
// At least one page in the range is used, skip to next.
// We know we'll be starting at least after this page.
any_taken = true;
base_page_number = xe::round_up(page_number + 1, page_scan_stride);
base_page_number -= page_scan_stride; // cancel out loop logic
break;
}
}
if (!any_taken) {
// Found our place.
}
for (; it != free_blocks_.end(); ++it) {
uint32_t block_start = it->first;
uint32_t block_count = it->second;
uint32_t block_end = block_start + block_count;
// Block is entirely above our search range — stop.
if (block_start > high_page_number) {
break;
}
// Skip blocks too small to possibly fit.
if (block_count < page_count) {
continue;
}
// Compute the lowest aligned start within this block and range.
uint32_t earliest = std::max(block_start, low_page_number);
uint32_t aligned_start = xe::round_up(earliest, page_scan_stride, false);
if (aligned_start + page_count <= block_end &&
aligned_start + page_count - 1 <= high_page_number) {
start_page_number = aligned_start;
end_page_number = aligned_start + page_count - 1;
break;
}
// Retry.
start_page_number = end_page_number = UINT_MAX;
}
}
if (start_page_number == UINT_MAX || end_page_number == UINT_MAX) {
// Out of memory.
XELOGE("BaseHeap::Alloc failed to find contiguous range");
@@ -1183,6 +1284,9 @@ bool BaseHeap::AllocRange(uint32_t low_address, uint32_t high_address,
return false;
}
// Update free block tracker.
RemoveFreeBlock(start_page_number, page_count);
// Allocate from host.
if (allocation_type == kMemoryAllocationReserve) {
// Reserve is not needed, as we are mapped already.
@@ -1196,6 +1300,8 @@ bool BaseHeap::AllocRange(uint32_t low_address, uint32_t high_address,
page_count << page_size_shift_, alloc_type, ToPageAccess(protect));
if (!result) {
XELOGE("BaseHeap::Alloc failed to alloc range from host");
// Restore the free block since we failed.
InsertFreeBlock(start_page_number, page_count);
return false;
}
@@ -1329,6 +1435,9 @@ bool BaseHeap::Release(uint32_t base_address, uint32_t* out_region_size) {
unreserved_page_count_++;
}
// Insert freed block into tracker with coalescing.
InsertFreeBlock(base_page_number, base_page_entry.region_page_count);
return true;
}
@@ -1685,7 +1794,10 @@ bool PhysicalHeap::Alloc(uint32_t size, uint32_t alignment,
alignment, allocation_type, protect, top_down,
&parent_address)) {
XELOGE(
"PhysicalHeap::Alloc unable to alloc physical memory in parent heap");
"PhysicalHeap::Alloc unable to alloc physical memory in parent heap "
"(requested {} bytes, parent free {}/{} pages)",
size, parent_heap_->unreserved_page_count(),
parent_heap_->total_page_count());
return false;
}
@@ -1704,7 +1816,7 @@ bool PhysicalHeap::Alloc(uint32_t size, uint32_t alignment,
protect)) {
XELOGE(
"PhysicalHeap::Alloc unable to pin physical memory in physical heap");
// TODO(benvanik): don't leak parent memory.
parent_heap_->Release(parent_address);
return false;
}
*out_address = address;
@@ -1728,7 +1840,8 @@ bool PhysicalHeap::AllocFixed(uint32_t base_address, uint32_t size,
if (!parent_heap_->AllocFixed(parent_base_address, size, alignment,
allocation_type, protect)) {
XELOGE(
"PhysicalHeap::Alloc unable to alloc physical memory in parent heap");
"PhysicalHeap::AllocFixed unable to alloc physical memory in parent "
"heap");
return false;
}
@@ -1747,8 +1860,9 @@ bool PhysicalHeap::AllocFixed(uint32_t base_address, uint32_t size,
if (!BaseHeap::AllocFixed(address, size, alignment, allocation_type,
protect)) {
XELOGE(
"PhysicalHeap::Alloc unable to pin physical memory in physical heap");
// TODO(benvanik): don't leak parent memory.
"PhysicalHeap::AllocFixed unable to pin physical memory in physical "
"heap");
parent_heap_->Release(parent_base_address);
return false;
}
@@ -1777,7 +1891,10 @@ bool PhysicalHeap::AllocRange(uint32_t low_address, uint32_t high_address,
alignment, allocation_type, protect, top_down,
&parent_address)) {
XELOGE(
"PhysicalHeap::Alloc unable to alloc physical memory in parent heap");
"PhysicalHeap::AllocRange unable to alloc physical memory in parent "
"heap (requested {} bytes, parent free {}/{} pages)",
size, parent_heap_->unreserved_page_count(),
parent_heap_->total_page_count());
return false;
}
// Given the address we've reserved in the parent heap, pin that here.
@@ -1795,8 +1912,9 @@ bool PhysicalHeap::AllocRange(uint32_t low_address, uint32_t high_address,
if (!BaseHeap::AllocFixed(address, size, alignment, allocation_type,
protect)) {
XELOGE(
"PhysicalHeap::Alloc unable to pin physical memory in physical heap");
// TODO(benvanik): don't leak parent memory.
"PhysicalHeap::AllocRange unable to pin physical memory in physical "
"heap");
parent_heap_->Release(parent_address);
return false;
}
*out_address = address;