From 203aac2cfb02b91a4993c6b1644486daf83a56ff Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Mon, 9 Nov 2026 18:36:39 -0400 Subject: [PATCH 1/3] ggml-hrx: the host-weight arena is charged to the pool, like the transient arena (engine #1711) Second slice of engine#1711. GGML_HRX_HOST_WEIGHT_LIMIT_MB limits how much residency the cache keeps, but nothing charged it to the pool ledger, so a large host-weight arena was invisible to pool_reserve and the guardrail over-committed by however much the cache held. The growth uses the same pure decision as the transient arena (plan_arena_charge, retires_old = false: the superseded arena is released, so its charge goes with it). The budget reaches it through CommandProgramExecutionContext, which already carries it since #268; the two acquire() call sites in the executor pass it down. Three arms added to tests/test-hrx-transient-arena-charge.cpp for the shape this path actually has: first growth charges the whole arena, a regrow to the same budget charges nothing (the common case - one arena serves many weights), and a growth past the budget swaps charge for charge. --- .../runtime/command-program-executor.cpp | 5 +-- ggml/src/ggml-hrx/runtime/host-memory.cpp | 35 +++++++++++++++++-- ggml/src/ggml-hrx/runtime/host-memory.h | 7 +++- tests/test-hrx-transient-arena-charge.cpp | 14 ++++++++ 4 files changed, 56 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp index 010a1684754d..e44ad94f4c97 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp @@ -711,7 +711,7 @@ static CommandProgramBindings materialize_host_bindings(const CommandProgramExec continue; } HostWeightAcquireResult resident = - context.host_weights->acquire(context.device, context.stream, *context.host_transfers, source); + context.host_weights->acquire(context.device, context.stream, *context.host_transfers, source, context.pool_budget); if (!resident.valid()) { // engine #1618: this line and the HSA error behind it carried no pool figure, // so an 89 GB model materialized against a 66.6 GiB budget read as a @@ -1019,7 +1019,8 @@ static Status rebind_prepared_host_weights(const CommandProgramExecutionContext // Evicted since this program was prepared (the cache is bounded, see HostWeightCache::limit): put it back and // point the prepared bindings at the new buffer. The old one is gone, so nothing else may still reference it. HostWeightAcquireResult materialized = - context.host_weights->acquire(context.device, context.stream, *context.host_transfers, resident.source); + context.host_weights->acquire(context.device, context.stream, *context.host_transfers, resident.source, + context.pool_budget); if (!materialized.valid()) { status.append(materialized.status); status.log("re-materialize resident host weight value %d failed", resident.value.value); diff --git a/ggml/src/ggml-hrx/runtime/host-memory.cpp b/ggml/src/ggml-hrx/runtime/host-memory.cpp index 58024b3db3bb..55422397fbbe 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.cpp +++ b/ggml/src/ggml-hrx/runtime/host-memory.cpp @@ -4,6 +4,8 @@ #include "ggml-quants.h" #include "hrx-interop-utils.h" #include "hrx_runtime.h" +#include "hrx-1bit-safety.h" // engine#1711: pool_reserve/pool_release +#include "transient-arena-charge.h" // engine#1711: the same charge decision as the transient arena #include "runtime/host-memory-available.h" #include "runtime/host-memory-ternary.h" @@ -1665,6 +1667,7 @@ bool HostWeightCache::arena_take_locked(size_t aligned, hrx_buffer_t & buffer, s bool HostWeightCache::arena_reserve(hrx_device_t device, hrx_stream_t stream, size_t size, + size_t pool_budget, hrx_buffer_t & buffer, size_t & offset, size_t & slot) { @@ -1738,9 +1741,28 @@ bool HostWeightCache::arena_reserve(hrx_device_t device, return false; } } + // engine#1711 (second slice): the host-weight arena is device-visible memory like any other, so its growth is + // charged to the pool before it happens -- the same decision the transient arena uses. Without this, a large + // residency is invisible to `pool_reserve` and the guardrail over-commits by however much the arena holds, which + // is what GGML_HRX_HOST_WEIGHT_LIMIT_MB *limits* but never accounts. + const ArenaChargePlan charge = plan_arena_charge(target, previous_capacity, false); + bool charge_taken = false; + if (charge.reserve_bytes > 0 && pool_budget > 0) { + if (!ggml::hrx::onebit::pool_reserve(pool_budget, charge.reserve_bytes, "the host-weight arena")) { + GGML_LOG_ERROR("%s: growing the host weight arena to %.1f MiB needs %zu more bytes and the pool has no room\n", + __func__, target / (1024.0 * 1024.0), charge.reserve_bytes); + return false; + } + charge_taken = true; + charged_bytes_ += charge.reserve_bytes; + } hrx_buffer_t grown = nullptr; const Status status = allocate_device_buffer(device, target, grown); if (!status.success()) { + if (charge_taken) { + ggml::hrx::onebit::pool_release(charge.reserve_bytes); + charged_bytes_ -= charged_bytes_ >= charge.reserve_bytes ? charge.reserve_bytes : charged_bytes_; + } GGML_LOG_ERROR("%s: growing the host weight arena to %.1f MiB failed (limit %.1f MiB, budget %.1f MiB, " "weight %.1f MiB)\n", __func__, target / (1024.0 * 1024.0), limit_ / (1024.0 * 1024.0), budget / (1024.0 * 1024.0), @@ -1751,6 +1773,10 @@ bool HostWeightCache::arena_reserve(hrx_device_t device, std::lock_guard lock(mutex_); if (arena_ != nullptr) { hrx_buffer_release(arena_); + if (charge.release_bytes > 0) { + ggml::hrx::onebit::pool_release(charge.release_bytes); + charged_bytes_ -= charged_bytes_ >= charge.release_bytes ? charge.release_bytes : charged_bytes_; + } } // Every slot in the old arena is gone, so no entry may keep pointing at it. entries_.clear(); @@ -1848,7 +1874,8 @@ void HostWeightCache::evict_for(hrx_stream_t stream, size_t needed) { HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device, hrx_stream_t stream, HostTransferManager & transfers, - const HostWeightSource & source) { + const HostWeightSource & source, + size_t pool_budget) { HostWeightAcquireResult result; if (stream == nullptr) { result.status.log("host weight residency requires an HRX stream"); @@ -1949,7 +1976,7 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device } } const bool reserved = - arena_reserve(device, stream, materialized_length, entry->buffer, entry->offset, entry->slot); + arena_reserve(device, stream, materialized_length, pool_budget, entry->buffer, entry->offset, entry->slot); if (!reserved) { size_t parked = 0; size_t resident = 0; @@ -2051,6 +2078,10 @@ void HostWeightCache::clear() { hrx_buffer_release(arena_); arena_ = nullptr; } + if (charged_bytes_ > 0) { + ggml::hrx::onebit::pool_release(charged_bytes_); // engine#1711 + charged_bytes_ = 0; + } arena_capacity_ = 0; arena_free_bytes_ = 0; arena_free_.clear(); diff --git a/ggml/src/ggml-hrx/runtime/host-memory.h b/ggml/src/ggml-hrx/runtime/host-memory.h index 916045bdd1c9..a82d8fcd1ff8 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.h +++ b/ggml/src/ggml-hrx/runtime/host-memory.h @@ -122,7 +122,8 @@ class HostWeightCache { HostWeightAcquireResult acquire(hrx_device_t device, hrx_stream_t stream, HostTransferManager & transfers, - const HostWeightSource & source); + const HostWeightSource & source, + size_t pool_budget = 0); HostWeightCacheStats stats() const; void clear(); @@ -187,6 +188,9 @@ class HostWeightCache { // buffer. The arena grows by doubling up to the residency limit, and a growth drops every entry (the stream // is synchronized first), so the peak stays at the working set rather than at the knob. hrx_buffer_t arena_ = nullptr; + // Bytes this cache holds against the pool ledger (engine#1711). Tracked separately from arena_capacity_ so clear() + // returns exactly what was charged, including the case where a grow released a superseded arena. + size_t charged_bytes_ = 0; size_t arena_capacity_ = 0; size_t arena_free_bytes_ = 0; std::vector> arena_free_; @@ -194,6 +198,7 @@ class HostWeightCache { bool arena_reserve(hrx_device_t device, hrx_stream_t stream, size_t size, + size_t pool_budget, hrx_buffer_t & buffer, size_t & offset, size_t & slot); diff --git a/tests/test-hrx-transient-arena-charge.cpp b/tests/test-hrx-transient-arena-charge.cpp index 6e5d508ba5da..b0c7d6d5e8de 100644 --- a/tests/test-hrx-transient-arena-charge.cpp +++ b/tests/test-hrx-transient-arena-charge.cpp @@ -63,6 +63,20 @@ int main() { got = plan_arena_charge(1u << 20, 8u << 20, true); check(got.reserve_bytes == 0 && got.release_bytes == 0, "a smaller request charges nothing", got, {0, 0}); + // 6. The host-weight arena reuses this same decision (`retires_old = false`): it grows straight to its budget, so + // the first growth charges the whole arena and a later regrow to the same budget charges nothing. Getting this + // wrong double-charges a cache that keeps the same arena across many weights, which is the common case. + got = plan_arena_charge(4096u << 20, 0, false); + check(got.reserve_bytes == (4096u << 20) && got.release_bytes == 0, "host-weight arena: first growth charges the whole arena", got, + {4096u << 20, 0}); + got = plan_arena_charge(4096u << 20, 4096u << 20, false); + check(got.reserve_bytes == 0 && got.release_bytes == 0, "host-weight arena: a regrow to the same budget charges nothing", got, + {0, 0}); + // and a grow past the budget releases the superseded arena's charge in the same step (the old buffer is freed there) + got = plan_arena_charge(8192u << 20, 4096u << 20, false); + check(got.reserve_bytes == (4096u << 20) && got.release_bytes == (4096u << 20), "host-weight arena: growing past the budget swaps charge for charge", got, + {4096u << 20, 4096u << 20}); + if (failures != 0) { std::printf("%d arm(s) failed\n", failures); return 1; From 34470c191e742b62e214423dadd68951a7ac2c21 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Mon, 9 Nov 2026 19:11:40 -0400 Subject: [PATCH 2/3] ggml-hrx: host staging buffers are charged to the pool too (engine #1711) Third and last allocator of this shape. A host staging buffer is device-visible memory like the arenas, but nothing charged it, so a large materialisation was invisible to pool_reserve. Unlike the two arenas there is no delta arithmetic: the object owns exactly one buffer, so the charge is taken on allocate (before it, so a refusal leaves no buffer) and given back in clear(). The contract that matters is that clear() returns the bytes EXACTLY once -- clear() is reached by the destructor, by move-assignment and by the next allocate, so a tally that survived a release would hand the same bytes back two or three times. tests/test-hrx-staging-ledger.cpp pins that on the real object with no device (6 arms), including the two that a no-op version fails: clear() on an empty buffer leaves the ledger alone, and a second clear() does not release again. --- .../runtime/command-program-executor.cpp | 3 +- ggml/src/ggml-hrx/runtime/host-memory.cpp | 25 +++++- ggml/src/ggml-hrx/runtime/host-memory.h | 6 +- tests/CMakeLists.txt | 6 ++ tests/test-hrx-staging-ledger.cpp | 82 +++++++++++++++++++ 5 files changed, 119 insertions(+), 3 deletions(-) create mode 100644 tests/test-hrx-staging-ledger.cpp diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp index e44ad94f4c97..c6d2578a7bfd 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp @@ -757,7 +757,8 @@ static CommandProgramBindings materialize_host_bindings(const CommandProgramExec } HostStagingBuffer staging; - Status allocation_status = allocate_host_staging_buffer(context.device, binding.length, staging); + Status allocation_status = allocate_host_staging_buffer(context.device, binding.length, staging, + context.pool_budget); if (!allocation_status.success()) { status.log("allocate host staging for value %d failed", binding.value.value); status.append(allocation_status); diff --git a/ggml/src/ggml-hrx/runtime/host-memory.cpp b/ggml/src/ggml-hrx/runtime/host-memory.cpp index 55422397fbbe..d0f74473b337 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.cpp +++ b/ggml/src/ggml-hrx/runtime/host-memory.cpp @@ -2122,6 +2122,10 @@ void HostStagingBuffer::clear() { hrx_buffer_release(buffer); buffer = nullptr; } + if (charged_bytes_ > 0) { + ggml::hrx::onebit::pool_release(charged_bytes_); // engine#1711 + charged_bytes_ = 0; + } host_data = nullptr; source_host_buffer = HostBufferRef{}; value = -1; @@ -2130,11 +2134,30 @@ void HostStagingBuffer::clear() { download = false; } -Status allocate_host_staging_buffer(hrx_device_t device, size_t size, HostStagingBuffer & staging) { +Status allocate_host_staging_buffer(hrx_device_t device, size_t size, HostStagingBuffer & staging, + size_t pool_budget) { staging.clear(); + // engine#1711: staging buffers are device-visible memory too. Charged before the allocation (a refusal must leave + // the previous state intact -- clear() above has already released the old charge, so a refusal here just means no + // staging buffer) and given back in clear(), which every path reaches: the destructor, the move-assignment, and the + // next allocate. + bool charge_taken = false; + if (size > 0 && pool_budget > 0) { + if (!ggml::hrx::onebit::pool_reserve(pool_budget, size, "a host staging buffer")) { + Status refused; + refused.log("refused: a %zu-byte host staging buffer does not fit the pool", size); + return refused; + } + charge_taken = true; + } Status status = allocate_device_buffer(device, size, staging.buffer); if (status.success()) { staging.length = size; + if (charge_taken) { + staging.charged_bytes_ = size; + } + } else if (charge_taken) { + ggml::hrx::onebit::pool_release(size); } return status; } diff --git a/ggml/src/ggml-hrx/runtime/host-memory.h b/ggml/src/ggml-hrx/runtime/host-memory.h index a82d8fcd1ff8..2e8a6fc81aaa 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.h +++ b/ggml/src/ggml-hrx/runtime/host-memory.h @@ -249,6 +249,8 @@ class HostWeightExecutionScope { struct HostStagingBuffer { HostStagingBuffer() = default; + // engine#1711: bytes this staging buffer holds against the pool ledger. The object owns exactly one buffer, so the + // charge is taken on allocate and given back in clear() -- no delta arithmetic, unlike the two arenas. ~HostStagingBuffer(); HostStagingBuffer(HostStagingBuffer && other) noexcept; @@ -262,13 +264,15 @@ struct HostStagingBuffer { HostBufferRef source_host_buffer; int32_t value = -1; size_t length = 0; + size_t charged_bytes_ = 0; // engine#1711: pool bytes held for this buffer bool upload = false; bool download = false; void clear(); }; -Status allocate_host_staging_buffer(hrx_device_t device, size_t size, HostStagingBuffer & staging); +Status allocate_host_staging_buffer(hrx_device_t device, size_t size, HostStagingBuffer & staging, + size_t pool_budget = 0); Status allocate_mapped_host_staging_buffer(hrx_device_t device, size_t size, hrx_buffer_t & buffer, diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index b314c97e2f76..0d657233568d 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -375,6 +375,12 @@ if (TARGET ggml-hrx) add_executable(test-hrx-transient-arena-charge test-hrx-transient-arena-charge.cpp) target_include_directories(test-hrx-transient-arena-charge PRIVATE ../ggml/src/ggml-hrx) add_test(NAME test-hrx-transient-arena-charge COMMAND test-hrx-transient-arena-charge) + # the host staging buffer's charge: the ledger is the only thing under test, no device and no libhrx (engine #1711) + add_executable(test-hrx-staging-ledger test-hrx-staging-ledger.cpp + ../ggml/src/ggml-hrx/hrx-1bit-safety.cpp) + target_include_directories(test-hrx-staging-ledger PRIVATE ../ggml/include ../ggml/src ../ggml/src/ggml-hrx) + target_link_libraries(test-hrx-staging-ledger PRIVATE ggml) + add_test(NAME test-hrx-staging-ledger COMMAND test-hrx-staging-ledger) # MemAvailable parsing for the unbounded-residency warning: pure, no device and no /proc (engine #669) add_executable(test-hrx-host-memory-available test-hrx-host-memory-available.cpp diff --git a/tests/test-hrx-staging-ledger.cpp b/tests/test-hrx-staging-ledger.cpp new file mode 100644 index 000000000000..a1c63cfc2949 --- /dev/null +++ b/tests/test-hrx-staging-ledger.cpp @@ -0,0 +1,82 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// The host staging buffer's pool charge (engine#1711). Unlike the two arenas this one carries no delta -- the object +// owns exactly one buffer -- so the whole contract is that clear() gives the bytes back EXACTLY once. clear() is reached +// by the destructor, by move-assignment and by the next allocate, so a tally that survived a release would hand the same +// bytes back two or three times and inflate the budget for everything afterwards. +// +// No device: the ledger is the only thing under test, which is why this does not need libhrx. + +#include "hrx-1bit-safety.h" +#include "runtime/host-memory.h" + +#include + +using ggml::hrx::HostStagingBuffer; + +static int failures = 0; + +static void check(bool ok, const char * what, size_t got, size_t want) { + if (!ok) { + std::printf("FAIL %s (ledger %zu, want %zu)\n", what, got, want); + ++failures; + } else { + std::printf("ok %s\n", what); + } +} + +int main() { + using ggml::hrx::onebit::pool_reserve; + using ggml::hrx::onebit::pool_allocated; // the ledger (NOT pool_reserve_bytes, which is the 16 GiB headroom reserve) + using ggml::hrx::onebit::pool_release; + + const size_t base = pool_allocated(); + + // 1. clear() on a buffer that holds nothing must not touch the ledger: this is the state of every default-constructed + // staging buffer and of every one already cleared, so an unconditional release here would underflow everything. + { + HostStagingBuffer empty; + empty.clear(); + check(pool_allocated() == base, "clear() on an empty staging buffer leaves the ledger alone", pool_allocated(), base); + empty.clear(); + check(pool_allocated() == base, "and calling it twice does not either", pool_allocated(), base); + } + + // 2. A held charge is returned exactly once. The buffer is simulated: reserve the bytes, record them on the object, + // clear it, and the ledger must be back where it started. + { + HostStagingBuffer held; + const size_t bytes = 8u << 20; + // budget 1 TiB: pool_reserve() returns true WITHOUT touching the ledger when budget == 0, so a real budget is + // needed for the bytes to actually be charged + const bool reserved = pool_reserve(1ull << 40, bytes, "a host staging buffer"); + check(reserved, "the pool accepts an 8 MiB staging request", 0, 0); + held.charged_bytes_ = bytes; + check(pool_allocated() == base + bytes, "the ledger carries the staging bytes", pool_allocated(), base + bytes); + held.clear(); + check(pool_allocated() == base, "clear() returns exactly what was charged", pool_allocated(), base); + held.clear(); + check(pool_allocated() == base, "a second clear() does not release it again", pool_allocated(), base); + (void) pool_release; // the release under test happens inside clear() + } + + if (failures != 0) { + std::printf("%d arm(s) failed\n", failures); + return 1; + } + std::printf("all arms passed\n"); + return 0; +} From 8895c3be2fb251bd920d20a81067f78115ec7b3a Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Mon, 9 Nov 2026 21:53:05 -0400 Subject: [PATCH 3/3] ggml-hrx: charge the whole arena on growth, and carry the charge across a move (both #270 review points) Review point 1: the plan was a delta (required - capacity) paired with releasing the superseded capacity, which subtracts the old capacity twice - the ledger settles at required - capacity instead of required, under-counting by exactly the arena it accounts for, so the guardrail over-admits. It is now REPLACE: charge the whole new allocation, release the superseded one. Under the replace model a retired (diagnostic) buffer correctly leaves the ledger at old + new, which is what the device actually holds in that mode. Two arms walk a five-step growth sequence and assert the invariant directly: after every growth the ledger equals the live capacity, and for the retiring model old + new. This also corrects the two already-merged call sites that use the same helper (#268's transient arena and #269's host-weight arena). Review point 2: HostStagingBuffer's move-assignment copied every field except charged_bytes_, so the destination held a buffer it was not charged for while the source kept a tally for a buffer it no longer owned. The tally now travels with the buffer and the source's is zeroed; six staging arms pin the move path (no ledger change, destination charged, source clean, one release on clear). --- ggml/src/ggml-hrx/runtime/host-memory.cpp | 5 ++ .../ggml-hrx/runtime/transient-arena-charge.h | 10 +++- tests/test-hrx-staging-ledger.cpp | 21 ++++++++ tests/test-hrx-transient-arena-charge.cpp | 52 ++++++++++++++++--- 4 files changed, 80 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-hrx/runtime/host-memory.cpp b/ggml/src/ggml-hrx/runtime/host-memory.cpp index d0f74473b337..7c6211b1d365 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.cpp +++ b/ggml/src/ggml-hrx/runtime/host-memory.cpp @@ -2108,6 +2108,11 @@ HostStagingBuffer & HostStagingBuffer::operator=(HostStagingBuffer && other) noe length = other.length; upload = other.upload; download = other.download; + // engine#1711: the charge travels with the buffer. Dropping it leaves the destination holding bytes it is not + // charged for (its clear() returns nothing) while the source keeps a tally for a buffer it no longer owns (its + // destructor releases bytes it does not hold) -- a ledger that disagrees with the memory in both objects at once. + charged_bytes_ = other.charged_bytes_; + other.charged_bytes_ = 0; other.buffer = nullptr; other.host_data = nullptr; other.value = -1; diff --git a/ggml/src/ggml-hrx/runtime/transient-arena-charge.h b/ggml/src/ggml-hrx/runtime/transient-arena-charge.h index be87e6388371..8f2dfcc6e6d8 100644 --- a/ggml/src/ggml-hrx/runtime/transient-arena-charge.h +++ b/ggml/src/ggml-hrx/runtime/transient-arena-charge.h @@ -43,9 +43,15 @@ inline ArenaChargePlan plan_arena_charge(std::size_t required, std::size_t capac if (required <= capacity) { return plan; // no growth, nothing to charge and nothing to give back } - plan.reserve_bytes = required - capacity; + // REPLACE, not transfer: charge the whole new allocation and give back the superseded one. The delta form -- + // `reserve(required - capacity)` together with `release(capacity)` -- subtracts the old capacity TWICE, so the + // ledger settles at `required - capacity` instead of `required`, under-counting by exactly the arena it is meant to + // account for; the guardrail then over-admits. The invariant every caller relies on is "after a growth the ledger + // equals the arena's live capacity", which is what clear() releases. + plan.reserve_bytes = required; // A first allocation has nothing to release; a growth releases the superseded buffer's charge only when that - // buffer is actually freed. Reserved-and-retired bytes stay charged (see this file's header). + // buffer is actually freed. A RETIRED buffer stays resident, so its charge stays too: under the replace model the + // ledger then ends at old + new, which is exactly the memory the arena holds in that mode. plan.release_bytes = (!retires_old && capacity > 0) ? capacity : 0; return plan; } diff --git a/tests/test-hrx-staging-ledger.cpp b/tests/test-hrx-staging-ledger.cpp index a1c63cfc2949..13aa08590602 100644 --- a/tests/test-hrx-staging-ledger.cpp +++ b/tests/test-hrx-staging-ledger.cpp @@ -24,6 +24,7 @@ #include "runtime/host-memory.h" #include +#include using ggml::hrx::HostStagingBuffer; @@ -73,6 +74,26 @@ int main() { (void) pool_release; // the release under test happens inside clear() } + // 3. The MOVE path carries the tally (the defect review of #270 found): the destination ends up holding the + // buffer, so it must hold the charge; the source ends up owning nothing, so it must hold no tally. Dropping + // either half leaves a ledger that disagrees with the memory in both objects at once. + { + HostStagingBuffer from, to; + const size_t bytes = 4u << 20; + const bool reserved = pool_reserve(1ull << 40, bytes, "a host staging buffer"); + check(reserved, "the pool accepts a 4 MiB staging request for the move arm", 0, 0); + from.charged_bytes_ = bytes; + const size_t before = pool_allocated(); + to = std::move(from); + check(pool_allocated() == before, "a move does not touch the ledger", pool_allocated(), before); + check(to.charged_bytes_ == bytes, "the destination carries the charge", to.charged_bytes_, bytes); + check(from.charged_bytes_ == 0, "the source keeps no tally for a buffer it no longer owns", from.charged_bytes_, 0); + to.clear(); + check(pool_allocated() == base, "clearing the destination returns the bytes exactly once", pool_allocated(), base); + from.clear(); // must be a no-op: it holds nothing and is charged nothing + check(pool_allocated() == base, "clearing the moved-from source does not release again", pool_allocated(), base); + } + if (failures != 0) { std::printf("%d arm(s) failed\n", failures); return 1; diff --git a/tests/test-hrx-transient-arena-charge.cpp b/tests/test-hrx-transient-arena-charge.cpp index b0c7d6d5e8de..e5269426ce97 100644 --- a/tests/test-hrx-transient-arena-charge.cpp +++ b/tests/test-hrx-transient-arena-charge.cpp @@ -27,6 +27,15 @@ using ggml::hrx::plan_arena_charge; static int failures = 0; +static void check_sizes(bool ok, const char * what, std::size_t got, std::size_t want) { + if (!ok) { + std::printf("FAIL %s (%zu, want %zu)\n", what, got, want); + ++failures; + } else { + std::printf("ok %s\n", what); + } +} + static void check(bool ok, const char * what, ArenaChargePlan got, ArenaChargePlan want) { if (!ok) { std::printf("FAIL %s: reserve %zu want %zu, release %zu want %zu\n", what, got.reserve_bytes, want.reserve_bytes, @@ -46,14 +55,14 @@ int main() { // 2. Growth, buffer freed: charge the delta, give back the superseded allocation's charge. The net charge after this // is the new size - which is what the ledger must show, not the sum of every size the arena ever had. got = plan_arena_charge(16u << 20, 8u << 20, false); - check(got.reserve_bytes == (8u << 20) && got.release_bytes == (8u << 20), "growth charges the delta and releases the old", got, - {8u << 20, 8u << 20}); + check(got.reserve_bytes == (16u << 20) && got.release_bytes == (8u << 20), + "growth charges the WHOLE new arena and releases the old", got, {16u << 20, 8u << 20}); // 3. Growth in diagnostic mode: the old buffer is RETIRED and stays resident, so its charge stays too. This is the arm // a naive "release the old capacity" version fails: it would release memory that is still on the device. got = plan_arena_charge(16u << 20, 8u << 20, true); - check(got.reserve_bytes == (8u << 20) && got.release_bytes == 0, "diagnostic growth keeps the retired buffer charged", got, - {8u << 20, 0}); + check(got.reserve_bytes == (16u << 20) && got.release_bytes == 0, + "diagnostic growth keeps the retired buffer charged (ledger ends at old + new)", got, {16u << 20, 0}); // 4. No growth: nothing charged, nothing released - the common case, and the one that must not perturb the ledger. got = plan_arena_charge(8u << 20, 8u << 20, false); @@ -74,8 +83,39 @@ int main() { {0, 0}); // and a grow past the budget releases the superseded arena's charge in the same step (the old buffer is freed there) got = plan_arena_charge(8192u << 20, 4096u << 20, false); - check(got.reserve_bytes == (4096u << 20) && got.release_bytes == (4096u << 20), "host-weight arena: growing past the budget swaps charge for charge", got, - {4096u << 20, 4096u << 20}); + check(got.reserve_bytes == (8192u << 20) && got.release_bytes == (4096u << 20), + "host-weight arena: growing past the budget charges the new arena and releases the old", got, + {8192u << 20, 4096u << 20}); + + // 9. THE INVARIANT the #270 review named, walked over a sequence rather than one step: after every growth the + // ledger must equal the arena's live capacity. A delta-shaped plan with a paired release passes a single-step + // check and fails this one on the second growth, settling at new - old. + { + size_t ledger = 0, capacity = 0; + bool ok = true; + for (size_t step = 0; step < 5; ++step) { + const size_t required = (1u << 20) << step; // 1, 2, 4, 8, 16 MiB + const ArenaChargePlan plan = plan_arena_charge(required, capacity, false); + ledger += plan.reserve_bytes - plan.release_bytes; + capacity = required; + if (ledger != capacity) { + ok = false; + } + } + check_sizes(ok, "ledger equals the live capacity after every growth (not new - old)", ledger, capacity); + size_t retired_ledger = 0, retired_held = 0; + ok = true; + for (size_t step = 0; step < 4; ++step) { + const size_t required = (1u << 20) << step; + const ArenaChargePlan plan = plan_arena_charge(required, retired_held, true); + retired_ledger += plan.reserve_bytes - plan.release_bytes; + retired_held += required; // every superseded arena stays resident + if (retired_ledger != retired_held) { + ok = false; + } + } + check_sizes(ok, "retired growth accumulates old + new (they are all still on the device)", retired_ledger, retired_held); + } if (failures != 0) { std::printf("%d arm(s) failed\n", failures);