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..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; @@ -2122,6 +2127,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 +2139,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/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/CMakeLists.txt b/tests/CMakeLists.txt index f5b4c7c8cfbf..9b77d6b19d84 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -381,6 +381,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..13aa08590602 --- /dev/null +++ b/tests/test-hrx-staging-ledger.cpp @@ -0,0 +1,103 @@ +// 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 +#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() + } + + // 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; + } + std::printf("all arms passed\n"); + return 0; +} 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);